From 767522dc4e66dd26773d41d1576945187180d2b9 Mon Sep 17 00:00:00 2001
From: huangzhaowei <carlmartinmax@gmail.com>
Date: Mon, 2 Nov 2015 21:31:10 +0800
Subject: [PATCH 001/518] [SPARK-10786][SQL] Take the whole statement to
 generate the CommandProcessor

In the now implementation of `SparkSQLCLIDriver.scala`:
`val proc: CommandProcessor = CommandProcessorFactory.get(Array(tokens(0)), hconf)`
`CommandProcessorFactory` only take the first token of the statement, and this will be hard to diff the statement `delete jar xxx` and `delete from xxx`.
So maybe it's better to take the whole statement into the `CommandProcessorFactory`.

And in [HiveCommand](https://github.com/SaintBacchus/hive/blob/master/ql/src/java/org/apache/hadoop/hive/ql/processors/HiveCommand.java#L76), it already special handing these two statement.
```java
if(command.length > 1 && "from".equalsIgnoreCase(command[1])) {
  //special handling for SQL "delete from <table> where..."
  return null;
}
```

Author: huangzhaowei <carlmartinmax@gmail.com>

Closes #8895 from SaintBacchus/SPARK-10786.
---
 .../apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala
index 62e912c69abc6..6419002a2aa89 100644
--- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala
+++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala
@@ -290,7 +290,7 @@ private[hive] class SparkSQLCLIDriver extends CliDriver with Logging {
     } else {
       var ret = 0
       val hconf = conf.asInstanceOf[HiveConf]
-      val proc: CommandProcessor = CommandProcessorFactory.get(Array(tokens(0)), hconf)
+      val proc: CommandProcessor = CommandProcessorFactory.get(tokens, hconf)
 
       if (proc != null) {
         // scalastyle:off println

From 74ba95228d71a6dc4e95fef19f41dabe7c363d9e Mon Sep 17 00:00:00 2001
From: Daoyuan Wang <daoyuan.wang@intel.com>
Date: Mon, 2 Nov 2015 23:07:30 +0800
Subject: [PATCH 002/518] [SPARK-11311][SQL] spark cannot describe temporary
 functions

When describe temporary function, spark would return 'Unable to find function', this is not right.

Author: Daoyuan Wang <daoyuan.wang@intel.com>

Closes #9277 from adrian-wang/functionreg.
---
 .../scala/org/apache/spark/sql/hive/hiveUDFs.scala     |  6 +++++-
 .../spark/sql/hive/execution/HiveQuerySuite.scala      | 10 ++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
index 2ccad474b4f7a..0b5e863506142 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
@@ -119,7 +119,11 @@ private[hive] class HiveFunctionRegistry(underlying: analysis.FunctionRegistry)
           annotation.value(),
           annotation.extended()))
       } else {
-        None
+        Some(new ExpressionInfo(
+          info.getFunctionClass.getCanonicalName,
+          name,
+          null,
+          null))
       }
     }.getOrElse(None))
   }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
index b52f7d4b57899..e597d6865f67a 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
@@ -953,6 +953,16 @@ class HiveQuerySuite extends HiveComparisonTest with BeforeAndAfter {
     sql("DROP TABLE t1")
   }
 
+  test("CREATE TEMPORARY FUNCTION") {
+    val funcJar = TestHive.getHiveFile("TestUDTF.jar").getCanonicalPath
+    sql(s"ADD JAR $funcJar")
+    sql(
+      """CREATE TEMPORARY FUNCTION udtf_count2 AS
+        | 'org.apache.spark.sql.hive.execution.GenericUDTFCount2'""".stripMargin)
+    assert(sql("DESCRIBE FUNCTION udtf_count2").count > 1)
+    sql("DROP TEMPORARY FUNCTION udtf_count2")
+  }
+
   test("ADD FILE command") {
     val testFile = TestHive.getHiveFile("data/files/v1.txt").getCanonicalFile
     sql(s"ADD FILE $testFile")

From a930e624eb9feb0f7d37d99dcb8178feb9c0f177 Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Mon, 2 Nov 2015 10:23:30 -0800
Subject: [PATCH 003/518] [SPARK-9817][YARN] Improve the locality calculation
 of containers by taking pending container requests into consideraion

This is a follow-up PR to further improve the locality calculation by considering the pending container's request. Since the locality preferences of tasks may be shifted from time to time, current localities of pending container requests may not fully match the new preferences, this PR improve it by removing outdated, unmatched container requests and replace with new requests.

sryza please help to review, thanks a lot.

Author: jerryshao <sshao@hortonworks.com>

Closes #8100 from jerryshao/SPARK-9817.
---
 .../spark/deploy/yarn/ApplicationMaster.scala |  2 +-
 ...yPreferredContainerPlacementStrategy.scala | 60 +++++++++++++--
 .../spark/deploy/yarn/YarnAllocator.scala     | 73 +++++++++++++++----
 .../ContainerPlacementStrategySuite.scala     | 38 ++++++++--
 .../deploy/yarn/YarnAllocatorSuite.scala      | 26 +++----
 5 files changed, 159 insertions(+), 40 deletions(-)

diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
index 4b4d9990ce9f9..c6a6d7ac56bf3 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
@@ -375,7 +375,7 @@ private[spark] class ApplicationMaster(
             }
           }
           try {
-            val numPendingAllocate = allocator.getNumPendingAllocate
+            val numPendingAllocate = allocator.getPendingAllocate.size
             val sleepInterval =
               if (numPendingAllocate > 0) {
                 val currentAllocationInterval =
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/LocalityPreferredContainerPlacementStrategy.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/LocalityPreferredContainerPlacementStrategy.scala
index 081780204e424..2ec189de7c914 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/LocalityPreferredContainerPlacementStrategy.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/LocalityPreferredContainerPlacementStrategy.scala
@@ -18,9 +18,11 @@
 package org.apache.spark.deploy.yarn
 
 import scala.collection.mutable.{ArrayBuffer, HashMap, Set}
+import scala.collection.JavaConverters._
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.yarn.api.records.{ContainerId, Resource}
+import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest
 import org.apache.hadoop.yarn.util.RackResolver
 
 import org.apache.spark.SparkConf
@@ -30,8 +32,8 @@ private[yarn] case class ContainerLocalityPreferences(nodes: Array[String], rack
 /**
  * This strategy is calculating the optimal locality preferences of YARN containers by considering
  * the node ratio of pending tasks, number of required cores/containers and and locality of current
- * existing containers. The target of this algorithm is to maximize the number of tasks that
- * would run locally.
+ * existing and pending allocated containers. The target of this algorithm is to maximize the number
+ * of tasks that would run locally.
  *
  * Consider a situation in which we have 20 tasks that require (host1, host2, host3)
  * and 10 tasks that require (host1, host2, host4), besides each container has 2 cores
@@ -91,6 +93,11 @@ private[yarn] class LocalityPreferredContainerPlacementStrategy(
    * @param numLocalityAwareTasks number of locality required tasks
    * @param hostToLocalTaskCount a map to store the preferred hostname and possible task
    *                             numbers running on it, used as hints for container allocation
+   * @param allocatedHostToContainersMap host to allocated containers map, used to calculate the
+   *                                     expected locality preference by considering the existing
+   *                                     containers
+   * @param localityMatchedPendingAllocations A sequence of pending container request which
+   *                                          matches the localities of current required tasks.
    * @return node localities and rack localities, each locality is an array of string,
    *         the length of localities is the same as number of containers
    */
@@ -98,10 +105,12 @@ private[yarn] class LocalityPreferredContainerPlacementStrategy(
       numContainer: Int,
       numLocalityAwareTasks: Int,
       hostToLocalTaskCount: Map[String, Int],
-      allocatedHostToContainersMap: HashMap[String, Set[ContainerId]]
+      allocatedHostToContainersMap: HashMap[String, Set[ContainerId]],
+      localityMatchedPendingAllocations: Seq[ContainerRequest]
     ): Array[ContainerLocalityPreferences] = {
     val updatedHostToContainerCount = expectedHostToContainerCount(
-      numLocalityAwareTasks, hostToLocalTaskCount, allocatedHostToContainersMap)
+      numLocalityAwareTasks, hostToLocalTaskCount, allocatedHostToContainersMap,
+        localityMatchedPendingAllocations)
     val updatedLocalityAwareContainerNum = updatedHostToContainerCount.values.sum
 
     // The number of containers to allocate, divided into two groups, one with preferred locality,
@@ -158,20 +167,28 @@ private[yarn] class LocalityPreferredContainerPlacementStrategy(
    * @param localityAwareTasks number of locality aware tasks
    * @param hostToLocalTaskCount a map to store the preferred hostname and possible task
    *                             numbers running on it, used as hints for container allocation
+   * @param allocatedHostToContainersMap host to allocated containers map, used to calculate the
+   *                                     expected locality preference by considering the existing
+   *                                     containers
+   * @param localityMatchedPendingAllocations A sequence of pending container request which
+   *                                          matches the localities of current required tasks.
    * @return a map with hostname as key and required number of containers on this host as value
    */
   private def expectedHostToContainerCount(
       localityAwareTasks: Int,
       hostToLocalTaskCount: Map[String, Int],
-      allocatedHostToContainersMap: HashMap[String, Set[ContainerId]]
+      allocatedHostToContainersMap: HashMap[String, Set[ContainerId]],
+      localityMatchedPendingAllocations: Seq[ContainerRequest]
     ): Map[String, Int] = {
     val totalLocalTaskNum = hostToLocalTaskCount.values.sum
+    val pendingHostToContainersMap = pendingHostToContainerCount(localityMatchedPendingAllocations)
+
     hostToLocalTaskCount.map { case (host, count) =>
       val expectedCount =
         count.toDouble * numExecutorsPending(localityAwareTasks) / totalLocalTaskNum
-      val existedCount = allocatedHostToContainersMap.get(host)
-        .map(_.size)
-        .getOrElse(0)
+      // Take the locality of pending containers into consideration
+      val existedCount = allocatedHostToContainersMap.get(host).map(_.size).getOrElse(0) +
+        pendingHostToContainersMap.getOrElse(host, 0.0)
 
       // If existing container can not fully satisfy the expected number of container,
       // the required container number is expected count minus existed count. Otherwise the
@@ -179,4 +196,31 @@ private[yarn] class LocalityPreferredContainerPlacementStrategy(
       (host, math.max(0, (expectedCount - existedCount).ceil.toInt))
     }
   }
+
+  /**
+   * According to the locality ratio and number of container requests, calculate the host to
+   * possible number of containers for pending allocated containers.
+   *
+   * If current locality ratio of hosts is: Host1 : Host2 : Host3 = 20 : 20 : 10,
+   * and pending container requests is 3, so the possible number of containers on
+   * Host1 : Host2 : Host3 will be 1.2 : 1.2 : 0.6.
+   * @param localityMatchedPendingAllocations A sequence of pending container request which
+   *                                          matches the localities of current required tasks.
+   * @return a Map with hostname as key and possible number of containers on this host as value
+   */
+  private def pendingHostToContainerCount(
+      localityMatchedPendingAllocations: Seq[ContainerRequest]): Map[String, Double] = {
+    val pendingHostToContainerCount = new HashMap[String, Int]()
+    localityMatchedPendingAllocations.foreach { cr =>
+      cr.getNodes.asScala.foreach { n =>
+        val count = pendingHostToContainerCount.getOrElse(n, 0) + 1
+        pendingHostToContainerCount(n) = count
+      }
+    }
+
+    val possibleTotalContainerNum = pendingHostToContainerCount.values.sum
+    val localityMatchedPendingNum = localityMatchedPendingAllocations.size.toDouble
+    pendingHostToContainerCount.mapValues(_ * localityMatchedPendingNum / possibleTotalContainerNum)
+      .toMap
+  }
 }
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index 875bbd4e4e3d5..a0cf1b4aa469b 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -157,15 +157,19 @@ private[yarn] class YarnAllocator(
   def getNumExecutorsFailed: Int = numExecutorsFailed
 
   /**
-   * Number of container requests that have not yet been fulfilled.
+   * A sequence of pending container requests that have not yet been fulfilled.
    */
-  def getNumPendingAllocate: Int = getNumPendingAtLocation(ANY_HOST)
+  def getPendingAllocate: Seq[ContainerRequest] = getPendingAtLocation(ANY_HOST)
 
   /**
-   * Number of container requests at the given location that have not yet been fulfilled.
+   * A sequence of pending container requests at the given location that have not yet been
+   * fulfilled.
    */
-  private def getNumPendingAtLocation(location: String): Int =
-    amClient.getMatchingRequests(RM_REQUEST_PRIORITY, location, resource).asScala.map(_.size).sum
+  private def getPendingAtLocation(location: String): Seq[ContainerRequest] = {
+    amClient.getMatchingRequests(RM_REQUEST_PRIORITY, location, resource).asScala
+      .flatMap(_.asScala)
+      .toSeq
+  }
 
   /**
    * Request as many executors from the ResourceManager as needed to reach the desired total. If
@@ -251,20 +255,31 @@ private[yarn] class YarnAllocator(
    * Visible for testing.
    */
   def updateResourceRequests(): Unit = {
-    val numPendingAllocate = getNumPendingAllocate
+    val pendingAllocate = getPendingAllocate
+    val numPendingAllocate = pendingAllocate.size
     val missing = targetNumExecutors - numPendingAllocate - numExecutorsRunning
 
-    // TODO. Consider locality preferences of pending container requests.
-    // Since the last time we made container requests, stages have completed and been submitted,
-    // and that the localities at which we requested our pending executors
-    // no longer apply to our current needs. We should consider to remove all outstanding
-    // container requests and add requests anew each time to avoid this.
     if (missing > 0) {
       logInfo(s"Will request $missing executor containers, each with ${resource.getVirtualCores} " +
         s"cores and ${resource.getMemory} MB memory including $memoryOverhead MB overhead")
 
+      // Split the pending container request into three groups: locality matched list, locality
+      // unmatched list and non-locality list. Take the locality matched container request into
+      // consideration of container placement, treat as allocated containers.
+      // For locality unmatched and locality free container requests, cancel these container
+      // requests, since required locality preference has been changed, recalculating using
+      // container placement strategy.
+      val (localityMatched, localityUnMatched, localityFree) = splitPendingAllocationsByLocality(
+        hostToLocalTaskCounts, pendingAllocate)
+
+      // Remove the outdated container request and recalculate the requested container number
+      localityUnMatched.foreach(amClient.removeContainerRequest)
+      localityFree.foreach(amClient.removeContainerRequest)
+      val updatedNumContainer = missing + localityUnMatched.size + localityFree.size
+
       val containerLocalityPreferences = containerPlacementStrategy.localityOfRequestedContainers(
-        missing, numLocalityAwareTasks, hostToLocalTaskCounts, allocatedHostToContainersMap)
+        updatedNumContainer, numLocalityAwareTasks, hostToLocalTaskCounts,
+          allocatedHostToContainersMap, localityMatched)
 
       for (locality <- containerLocalityPreferences) {
         val request = createContainerRequest(resource, locality.nodes, locality.racks)
@@ -291,7 +306,7 @@ private[yarn] class YarnAllocator(
    * Creates a container request, handling the reflection required to use YARN features that were
    * added in recent versions.
    */
-  protected def createContainerRequest(
+  private def createContainerRequest(
       resource: Resource,
       nodes: Array[String],
       racks: Array[String]): ContainerRequest = {
@@ -535,6 +550,38 @@ private[yarn] class YarnAllocator(
 
   private[yarn] def getNumUnexpectedContainerRelease = numUnexpectedContainerRelease
 
+  /**
+   * Split the pending container requests into 3 groups based on current localities of pending
+   * tasks.
+   * @param hostToLocalTaskCount a map of preferred hostname to possible task counts to be used as
+   *                             container placement hint.
+   * @param pendingAllocations A sequence of pending allocation container request.
+   * @return A tuple of 3 sequences, first is a sequence of locality matched container
+   *         requests, second is a sequence of locality unmatched container requests, and third is a
+   *         sequence of locality free container requests.
+   */
+  private def splitPendingAllocationsByLocality(
+      hostToLocalTaskCount: Map[String, Int],
+      pendingAllocations: Seq[ContainerRequest]
+    ): (Seq[ContainerRequest], Seq[ContainerRequest], Seq[ContainerRequest]) = {
+    val localityMatched = ArrayBuffer[ContainerRequest]()
+    val localityUnMatched = ArrayBuffer[ContainerRequest]()
+    val localityFree = ArrayBuffer[ContainerRequest]()
+
+    val preferredHosts = hostToLocalTaskCount.keySet
+    pendingAllocations.foreach { cr =>
+      val nodes = cr.getNodes
+      if (nodes == null) {
+        localityFree += cr
+      } else if (nodes.asScala.toSet.intersect(preferredHosts).nonEmpty) {
+        localityMatched += cr
+      } else {
+        localityUnMatched += cr
+      }
+    }
+
+    (localityMatched.toSeq, localityUnMatched.toSeq, localityFree.toSeq)
+  }
 }
 
 private object YarnAllocator {
diff --git a/yarn/src/test/scala/org/apache/spark/deploy/yarn/ContainerPlacementStrategySuite.scala b/yarn/src/test/scala/org/apache/spark/deploy/yarn/ContainerPlacementStrategySuite.scala
index b7fe4ccc67a38..afb4b691b52de 100644
--- a/yarn/src/test/scala/org/apache/spark/deploy/yarn/ContainerPlacementStrategySuite.scala
+++ b/yarn/src/test/scala/org/apache/spark/deploy/yarn/ContainerPlacementStrategySuite.scala
@@ -17,6 +17,7 @@
 
 package org.apache.spark.deploy.yarn
 
+import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest
 import org.scalatest.{BeforeAndAfterEach, Matchers}
 
 import org.apache.spark.SparkFunSuite
@@ -26,6 +27,9 @@ class ContainerPlacementStrategySuite extends SparkFunSuite with Matchers with B
   private val yarnAllocatorSuite = new YarnAllocatorSuite
   import yarnAllocatorSuite._
 
+  def createContainerRequest(nodes: Array[String]): ContainerRequest =
+    new ContainerRequest(containerResource, nodes, null, YarnSparkHadoopUtil.RM_REQUEST_PRIORITY)
+
   override def beforeEach() {
     yarnAllocatorSuite.beforeEach()
   }
@@ -44,7 +48,8 @@ class ContainerPlacementStrategySuite extends SparkFunSuite with Matchers with B
     handler.handleAllocatedContainers(Array(createContainer("host1"), createContainer("host2")))
 
     val localities = handler.containerPlacementStrategy.localityOfRequestedContainers(
-      3, 15, Map("host3" -> 15, "host4" -> 15, "host5" -> 10), handler.allocatedHostToContainersMap)
+      3, 15, Map("host3" -> 15, "host4" -> 15, "host5" -> 10),
+        handler.allocatedHostToContainersMap, Seq.empty)
 
     assert(localities.map(_.nodes) === Array(
       Array("host3", "host4", "host5"),
@@ -66,7 +71,8 @@ class ContainerPlacementStrategySuite extends SparkFunSuite with Matchers with B
     ))
 
     val localities = handler.containerPlacementStrategy.localityOfRequestedContainers(
-      3, 15, Map("host1" -> 15, "host2" -> 15, "host3" -> 10), handler.allocatedHostToContainersMap)
+      3, 15, Map("host1" -> 15, "host2" -> 15, "host3" -> 10),
+        handler.allocatedHostToContainersMap, Seq.empty)
 
     assert(localities.map(_.nodes) ===
       Array(null, Array("host2", "host3"), Array("host2", "host3")))
@@ -86,7 +92,8 @@ class ContainerPlacementStrategySuite extends SparkFunSuite with Matchers with B
     ))
 
     val localities = handler.containerPlacementStrategy.localityOfRequestedContainers(
-      1, 15, Map("host1" -> 15, "host2" -> 15, "host3" -> 10), handler.allocatedHostToContainersMap)
+      1, 15, Map("host1" -> 15, "host2" -> 15, "host3" -> 10),
+        handler.allocatedHostToContainersMap, Seq.empty)
 
     assert(localities.map(_.nodes) === Array(Array("host2", "host3")))
   }
@@ -105,7 +112,8 @@ class ContainerPlacementStrategySuite extends SparkFunSuite with Matchers with B
     ))
 
     val localities = handler.containerPlacementStrategy.localityOfRequestedContainers(
-      3, 15, Map("host1" -> 15, "host2" -> 15, "host3" -> 10), handler.allocatedHostToContainersMap)
+      3, 15, Map("host1" -> 15, "host2" -> 15, "host3" -> 10),
+        handler.allocatedHostToContainersMap, Seq.empty)
 
     assert(localities.map(_.nodes) === Array(null, null, null))
   }
@@ -118,8 +126,28 @@ class ContainerPlacementStrategySuite extends SparkFunSuite with Matchers with B
     handler.handleAllocatedContainers(Array(createContainer("host1"), createContainer("host2")))
 
     val localities = handler.containerPlacementStrategy.localityOfRequestedContainers(
-      1, 0, Map.empty, handler.allocatedHostToContainersMap)
+      1, 0, Map.empty, handler.allocatedHostToContainersMap, Seq.empty)
 
     assert(localities.map(_.nodes) === Array(null))
   }
+
+  test("allocate locality preferred containers by considering the localities of pending requests") {
+    val handler = createAllocator(3)
+    handler.updateResourceRequests()
+    handler.handleAllocatedContainers(Array(
+      createContainer("host1"),
+      createContainer("host1"),
+      createContainer("host2")
+    ))
+
+    val pendingAllocationRequests = Seq(
+      createContainerRequest(Array("host2", "host3")),
+      createContainerRequest(Array("host1", "host4")))
+
+    val localities = handler.containerPlacementStrategy.localityOfRequestedContainers(
+      1, 15, Map("host1" -> 15, "host2" -> 15, "host3" -> 10),
+        handler.allocatedHostToContainersMap, pendingAllocationRequests)
+
+    assert(localities.map(_.nodes) === Array(Array("host3")))
+  }
 }
diff --git a/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala b/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
index 5d05f514adde3..bd80036c5cfa7 100644
--- a/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
+++ b/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnAllocatorSuite.scala
@@ -116,7 +116,7 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     val handler = createAllocator(1)
     handler.updateResourceRequests()
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (1)
+    handler.getPendingAllocate.size should be (1)
 
     val container = createContainer("host1")
     handler.handleAllocatedContainers(Array(container))
@@ -134,7 +134,7 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     val handler = createAllocator(4)
     handler.updateResourceRequests()
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (4)
+    handler.getPendingAllocate.size should be (4)
 
     val container1 = createContainer("host1")
     val container2 = createContainer("host1")
@@ -154,7 +154,7 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     val handler = createAllocator(2)
     handler.updateResourceRequests()
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (2)
+    handler.getPendingAllocate.size should be (2)
 
     val container1 = createContainer("host1")
     val container2 = createContainer("host2")
@@ -174,11 +174,11 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     val handler = createAllocator(4)
     handler.updateResourceRequests()
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (4)
+    handler.getPendingAllocate.size should be (4)
 
     handler.requestTotalExecutorsWithPreferredLocalities(3, 0, Map.empty)
     handler.updateResourceRequests()
-    handler.getNumPendingAllocate should be (3)
+    handler.getPendingAllocate.size should be (3)
 
     val container = createContainer("host1")
     handler.handleAllocatedContainers(Array(container))
@@ -189,18 +189,18 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
 
     handler.requestTotalExecutorsWithPreferredLocalities(2, 0, Map.empty)
     handler.updateResourceRequests()
-    handler.getNumPendingAllocate should be (1)
+    handler.getPendingAllocate.size should be (1)
   }
 
   test("decrease total requested executors to less than currently running") {
     val handler = createAllocator(4)
     handler.updateResourceRequests()
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (4)
+    handler.getPendingAllocate.size should be (4)
 
     handler.requestTotalExecutorsWithPreferredLocalities(3, 0, Map.empty)
     handler.updateResourceRequests()
-    handler.getNumPendingAllocate should be (3)
+    handler.getPendingAllocate.size should be (3)
 
     val container1 = createContainer("host1")
     val container2 = createContainer("host2")
@@ -210,7 +210,7 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
 
     handler.requestTotalExecutorsWithPreferredLocalities(1, 0, Map.empty)
     handler.updateResourceRequests()
-    handler.getNumPendingAllocate should be (0)
+    handler.getPendingAllocate.size should be (0)
     handler.getNumExecutorsRunning should be (2)
   }
 
@@ -218,7 +218,7 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     val handler = createAllocator(4)
     handler.updateResourceRequests()
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (4)
+    handler.getPendingAllocate.size should be (4)
 
     val container1 = createContainer("host1")
     val container2 = createContainer("host2")
@@ -233,14 +233,14 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     handler.updateResourceRequests()
     handler.processCompletedContainers(statuses.toSeq)
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (1)
+    handler.getPendingAllocate.size should be (1)
   }
 
   test("lost executor removed from backend") {
     val handler = createAllocator(4)
     handler.updateResourceRequests()
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (4)
+    handler.getPendingAllocate.size should be (4)
 
     val container1 = createContainer("host1")
     val container2 = createContainer("host2")
@@ -255,7 +255,7 @@ class YarnAllocatorSuite extends SparkFunSuite with Matchers with BeforeAndAfter
     handler.processCompletedContainers(statuses.toSeq)
     handler.updateResourceRequests()
     handler.getNumExecutorsRunning should be (0)
-    handler.getNumPendingAllocate should be (2)
+    handler.getPendingAllocate.size should be (2)
     handler.getNumExecutorsFailed should be (2)
     handler.getNumUnexpectedContainerRelease should be (2)
   }

From 71d1c907dec446db566b19f912159fd8f46deb7d Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Mon, 2 Nov 2015 10:26:36 -0800
Subject: [PATCH 004/518] [SPARK-10997][CORE] Add "client mode" to netty rpc
 env.

"Client mode" means the RPC env will not listen for incoming connections.
This allows certain processes in the Spark stack (such as Executors or
tha YARN client-mode AM) to act as pure clients when using the netty-based
RPC backend, reducing the number of sockets needed by the app and also the
number of open ports.

Client connections are also preferred when endpoints that actually have
a listening socket are involved; so, for example, if a Worker connects
to a Master and the Master needs to send a message to a Worker endpoint,
that client connection will be used, even though the Worker is also
listening for incoming connections.

With this change, the workaround for SPARK-10987 isn't necessary anymore, and
is removed. The AM connects to the driver in "client mode", and that connection
is used for all driver <-> AM communication, and so the AM is properly notified
when the connection goes down.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9210 from vanzin/SPARK-10997.
---
 .../scala/org/apache/spark/SparkEnv.scala     |   7 +-
 .../CoarseGrainedExecutorBackend.scala        |  20 +-
 .../scala/org/apache/spark/rpc/RpcEnv.scala   |   8 +-
 .../apache/spark/rpc/netty/Dispatcher.scala   |   2 +-
 .../apache/spark/rpc/netty/NettyRpcEnv.scala  | 245 ++++++++++--------
 .../org/apache/spark/rpc/netty/Outbox.scala   |  24 +-
 .../spark/rpc/netty/RpcEndpointAddress.scala  |  24 +-
 .../cluster/CoarseGrainedClusterMessage.scala |  10 +-
 .../CoarseGrainedSchedulerBackend.scala       |  18 +-
 .../cluster/YarnSchedulerBackend.scala        |   2 -
 .../org/apache/spark/rpc/RpcEnvSuite.scala    |  50 ++--
 .../spark/rpc/akka/AkkaRpcEnvSuite.scala      |  11 +-
 .../rpc/netty/NettyRpcAddressSuite.scala      |   7 +-
 .../spark/rpc/netty/NettyRpcEnvSuite.scala    |   9 +-
 .../rpc/netty/NettyRpcHandlerSuite.scala      |   8 +-
 network/yarn/pom.xml                          |   5 +
 .../spark/deploy/yarn/ApplicationMaster.scala |   6 +-
 17 files changed, 266 insertions(+), 190 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala b/core/src/main/scala/org/apache/spark/SparkEnv.scala
index 398e0936906a3..23ae9360f6a22 100644
--- a/core/src/main/scala/org/apache/spark/SparkEnv.scala
+++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala
@@ -252,7 +252,8 @@ object SparkEnv extends Logging {
 
     // Create the ActorSystem for Akka and get the port it binds to.
     val actorSystemName = if (isDriver) driverActorSystemName else executorActorSystemName
-    val rpcEnv = RpcEnv.create(actorSystemName, hostname, port, conf, securityManager)
+    val rpcEnv = RpcEnv.create(actorSystemName, hostname, port, conf, securityManager,
+      clientMode = !isDriver)
     val actorSystem: ActorSystem =
       if (rpcEnv.isInstanceOf[AkkaRpcEnv]) {
         rpcEnv.asInstanceOf[AkkaRpcEnv].actorSystem
@@ -262,9 +263,11 @@ object SparkEnv extends Logging {
       }
 
     // Figure out which port Akka actually bound to in case the original port is 0 or occupied.
+    // In the non-driver case, the RPC env's address may be null since it may not be listening
+    // for incoming connections.
     if (isDriver) {
       conf.set("spark.driver.port", rpcEnv.address.port.toString)
-    } else {
+    } else if (rpcEnv.address != null) {
       conf.set("spark.executor.port", rpcEnv.address.port.toString)
     }
 
diff --git a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala
index a9c6a05ecd434..c2ebf30596215 100644
--- a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala
+++ b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala
@@ -45,8 +45,6 @@ private[spark] class CoarseGrainedExecutorBackend(
     env: SparkEnv)
   extends ThreadSafeRpcEndpoint with ExecutorBackend with Logging {
 
-  Utils.checkHostPort(hostPort, "Expected hostport")
-
   var executor: Executor = null
   @volatile var driver: Option[RpcEndpointRef] = None
 
@@ -80,9 +78,8 @@ private[spark] class CoarseGrainedExecutorBackend(
   }
 
   override def receive: PartialFunction[Any, Unit] = {
-    case RegisteredExecutor =>
+    case RegisteredExecutor(hostname) =>
       logInfo("Successfully registered with driver")
-      val (hostname, _) = Utils.parseHostPort(hostPort)
       executor = new Executor(executorId, hostname, env, userClassPath, isLocal = false)
 
     case RegisterExecutorFailed(message) =>
@@ -163,7 +160,8 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging {
         hostname,
         port,
         executorConf,
-        new SecurityManager(executorConf))
+        new SecurityManager(executorConf),
+        clientMode = true)
       val driver = fetcher.setupEndpointRefByURI(driverUrl)
       val props = driver.askWithRetry[Seq[(String, String)]](RetrieveSparkProps) ++
         Seq[(String, String)](("spark.app.id", appId))
@@ -188,12 +186,12 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging {
       val env = SparkEnv.createExecutorEnv(
         driverConf, executorId, hostname, port, cores, isLocal = false)
 
-      // SparkEnv sets spark.driver.port so it shouldn't be 0 anymore.
-      val boundPort = env.conf.getInt("spark.executor.port", 0)
-      assert(boundPort != 0)
-
-      // Start the CoarseGrainedExecutorBackend endpoint.
-      val sparkHostPort = hostname + ":" + boundPort
+      // SparkEnv will set spark.executor.port if the rpc env is listening for incoming
+      // connections (e.g., if it's using akka). Otherwise, the executor is running in
+      // client mode only, and does not accept incoming connections.
+      val sparkHostPort = env.conf.getOption("spark.executor.port").map { port =>
+          hostname + ":" + port
+        }.orNull
       env.rpcEnv.setupEndpoint("Executor", new CoarseGrainedExecutorBackend(
         env.rpcEnv, driverUrl, executorId, sparkHostPort, cores, userClassPath, env))
       workerUrl.foreach { url =>
diff --git a/core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala
index 2c4a8b9a0a878..a560fd10cdf76 100644
--- a/core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala
@@ -43,9 +43,10 @@ private[spark] object RpcEnv {
       host: String,
       port: Int,
       conf: SparkConf,
-      securityManager: SecurityManager): RpcEnv = {
+      securityManager: SecurityManager,
+      clientMode: Boolean = false): RpcEnv = {
     // Using Reflection to create the RpcEnv to avoid to depend on Akka directly
-    val config = RpcEnvConfig(conf, name, host, port, securityManager)
+    val config = RpcEnvConfig(conf, name, host, port, securityManager, clientMode)
     getRpcEnvFactory(conf).create(config)
   }
 }
@@ -139,4 +140,5 @@ private[spark] case class RpcEnvConfig(
     name: String,
     host: String,
     port: Int,
-    securityManager: SecurityManager)
+    securityManager: SecurityManager,
+    clientMode: Boolean)
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala b/core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala
index 7bf44a6565b61..eb25d6c7b721b 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala
@@ -55,7 +55,7 @@ private[netty] class Dispatcher(nettyEnv: NettyRpcEnv) extends Logging {
   private var stopped = false
 
   def registerRpcEndpoint(name: String, endpoint: RpcEndpoint): NettyRpcEndpointRef = {
-    val addr = new RpcEndpointAddress(nettyEnv.address.host, nettyEnv.address.port, name)
+    val addr = RpcEndpointAddress(nettyEnv.address, name)
     val endpointRef = new NettyRpcEndpointRef(nettyEnv.conf, addr, nettyEnv)
     synchronized {
       if (stopped) {
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
index 284284eb805b7..09093819bb22c 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
@@ -17,10 +17,12 @@
 package org.apache.spark.rpc.netty
 
 import java.io._
+import java.lang.{Boolean => JBoolean}
 import java.net.{InetSocketAddress, URI}
 import java.nio.ByteBuffer
 import java.util.concurrent._
 import java.util.concurrent.atomic.AtomicBoolean
+import javax.annotation.Nullable;
 import javax.annotation.concurrent.GuardedBy
 
 import scala.collection.mutable
@@ -29,6 +31,7 @@ import scala.reflect.ClassTag
 import scala.util.{DynamicVariable, Failure, Success}
 import scala.util.control.NonFatal
 
+import com.google.common.base.Preconditions
 import org.apache.spark.{Logging, SecurityManager, SparkConf}
 import org.apache.spark.network.TransportContext
 import org.apache.spark.network.client._
@@ -45,15 +48,14 @@ private[netty] class NettyRpcEnv(
     host: String,
     securityManager: SecurityManager) extends RpcEnv(conf) with Logging {
 
-  // Override numConnectionsPerPeer to 1 for RPC.
   private val transportConf = SparkTransportConf.fromSparkConf(
     conf.clone.set("spark.shuffle.io.numConnectionsPerPeer", "1"),
     conf.getInt("spark.rpc.io.threads", 0))
 
   private val dispatcher: Dispatcher = new Dispatcher(this)
 
-  private val transportContext =
-    new TransportContext(transportConf, new NettyRpcHandler(dispatcher, this))
+  private val transportContext = new TransportContext(transportConf,
+    new NettyRpcHandler(dispatcher, this))
 
   private val clientFactory = {
     val bootstraps: java.util.List[TransportClientBootstrap] =
@@ -95,7 +97,7 @@ private[netty] class NettyRpcEnv(
     }
   }
 
-  def start(port: Int): Unit = {
+  def startServer(port: Int): Unit = {
     val bootstraps: java.util.List[TransportServerBootstrap] =
       if (securityManager.isAuthenticationEnabled()) {
         java.util.Arrays.asList(new SaslServerBootstrap(transportConf, securityManager))
@@ -107,9 +109,9 @@ private[netty] class NettyRpcEnv(
       RpcEndpointVerifier.NAME, new RpcEndpointVerifier(this, dispatcher))
   }
 
+  @Nullable
   override lazy val address: RpcAddress = {
-    require(server != null, "NettyRpcEnv has not yet started")
-    RpcAddress(host, server.getPort)
+    if (server != null) RpcAddress(host, server.getPort()) else null
   }
 
   override def setupEndpoint(name: String, endpoint: RpcEndpoint): RpcEndpointRef = {
@@ -120,7 +122,7 @@ private[netty] class NettyRpcEnv(
     val addr = RpcEndpointAddress(uri)
     val endpointRef = new NettyRpcEndpointRef(conf, addr, this)
     val verifier = new NettyRpcEndpointRef(
-      conf, RpcEndpointAddress(addr.host, addr.port, RpcEndpointVerifier.NAME), this)
+      conf, RpcEndpointAddress(addr.rpcAddress, RpcEndpointVerifier.NAME), this)
     verifier.ask[Boolean](RpcEndpointVerifier.CheckExistence(endpointRef.name)).flatMap { find =>
       if (find) {
         Future.successful(endpointRef)
@@ -135,28 +137,34 @@ private[netty] class NettyRpcEnv(
     dispatcher.stop(endpointRef)
   }
 
-  private def postToOutbox(address: RpcAddress, message: OutboxMessage): Unit = {
-    val targetOutbox = {
-      val outbox = outboxes.get(address)
-      if (outbox == null) {
-        val newOutbox = new Outbox(this, address)
-        val oldOutbox = outboxes.putIfAbsent(address, newOutbox)
-        if (oldOutbox == null) {
-          newOutbox
+  private def postToOutbox(receiver: NettyRpcEndpointRef, message: OutboxMessage): Unit = {
+    if (receiver.client != null) {
+      receiver.client.sendRpc(message.content, message.createCallback(receiver.client));
+    } else {
+      require(receiver.address != null,
+        "Cannot send message to client endpoint with no listen address.")
+      val targetOutbox = {
+        val outbox = outboxes.get(receiver.address)
+        if (outbox == null) {
+          val newOutbox = new Outbox(this, receiver.address)
+          val oldOutbox = outboxes.putIfAbsent(receiver.address, newOutbox)
+          if (oldOutbox == null) {
+            newOutbox
+          } else {
+            oldOutbox
+          }
         } else {
-          oldOutbox
+          outbox
         }
+      }
+      if (stopped.get) {
+        // It's possible that we put `targetOutbox` after stopping. So we need to clean it.
+        outboxes.remove(receiver.address)
+        targetOutbox.stop()
       } else {
-        outbox
+        targetOutbox.send(message)
       }
     }
-    if (stopped.get) {
-      // It's possible that we put `targetOutbox` after stopping. So we need to clean it.
-      outboxes.remove(address)
-      targetOutbox.stop()
-    } else {
-      targetOutbox.send(message)
-    }
   }
 
   private[netty] def send(message: RequestMessage): Unit = {
@@ -174,17 +182,14 @@ private[netty] class NettyRpcEnv(
       }(ThreadUtils.sameThread)
     } else {
       // Message to a remote RPC endpoint.
-      postToOutbox(remoteAddr, OutboxMessage(serialize(message), new RpcResponseCallback {
-
-        override def onFailure(e: Throwable): Unit = {
+      postToOutbox(message.receiver, OutboxMessage(serialize(message),
+        (e) => {
           logWarning(s"Exception when sending $message", e)
-        }
-
-        override def onSuccess(response: Array[Byte]): Unit = {
-          val ack = deserialize[Ack](response)
+        },
+        (client, response) => {
+          val ack = deserialize[Ack](client, response)
           logDebug(s"Receive ack from ${ack.sender}")
-        }
-      }))
+        }))
     }
   }
 
@@ -214,16 +219,14 @@ private[netty] class NettyRpcEnv(
           }
       }(ThreadUtils.sameThread)
     } else {
-      postToOutbox(remoteAddr, OutboxMessage(serialize(message), new RpcResponseCallback {
-
-        override def onFailure(e: Throwable): Unit = {
+      postToOutbox(message.receiver, OutboxMessage(serialize(message),
+        (e) => {
           if (!promise.tryFailure(e)) {
             logWarning("Ignore Exception", e)
           }
-        }
-
-        override def onSuccess(response: Array[Byte]): Unit = {
-          val reply = deserialize[AskResponse](response)
+        },
+        (client, response) => {
+          val reply = deserialize[AskResponse](client, response)
           if (reply.reply.isInstanceOf[RpcFailure]) {
             if (!promise.tryFailure(reply.reply.asInstanceOf[RpcFailure].e)) {
               logWarning(s"Ignore failure: ${reply.reply}")
@@ -231,8 +234,7 @@ private[netty] class NettyRpcEnv(
           } else if (!promise.trySuccess(reply.reply)) {
             logWarning(s"Ignore message: ${reply}")
           }
-        }
-      }))
+        }))
     }
     promise.future
   }
@@ -243,9 +245,11 @@ private[netty] class NettyRpcEnv(
       buffer.array(), buffer.arrayOffset + buffer.position, buffer.arrayOffset + buffer.limit)
   }
 
-  private[netty] def deserialize[T: ClassTag](bytes: Array[Byte]): T = {
-    deserialize { () =>
-      javaSerializerInstance.deserialize[T](ByteBuffer.wrap(bytes))
+  private[netty] def deserialize[T: ClassTag](client: TransportClient, bytes: Array[Byte]): T = {
+    NettyRpcEnv.currentClient.withValue(client) {
+      deserialize { () =>
+        javaSerializerInstance.deserialize[T](ByteBuffer.wrap(bytes))
+      }
     }
   }
 
@@ -254,7 +258,7 @@ private[netty] class NettyRpcEnv(
   }
 
   override def uriOf(systemName: String, address: RpcAddress, endpointName: String): String =
-    new RpcEndpointAddress(address.host, address.port, endpointName).toString
+    new RpcEndpointAddress(address, endpointName).toString
 
   override def shutdown(): Unit = {
     cleanup()
@@ -297,6 +301,7 @@ private[netty] class NettyRpcEnv(
       deserializationAction()
     }
   }
+
 }
 
 private[netty] object NettyRpcEnv extends Logging {
@@ -312,6 +317,13 @@ private[netty] object NettyRpcEnv extends Logging {
    * }}}
    */
   private[netty] val currentEnv = new DynamicVariable[NettyRpcEnv](null)
+
+  /**
+   * Similar to `currentEnv`, this variable references the client instance associated with an
+   * RPC, in case it's needed to find out the remote address during deserialization.
+   */
+  private[netty] val currentClient = new DynamicVariable[TransportClient](null)
+
 }
 
 private[netty] class NettyRpcEnvFactory extends RpcEnvFactory with Logging {
@@ -324,47 +336,68 @@ private[netty] class NettyRpcEnvFactory extends RpcEnvFactory with Logging {
       new JavaSerializer(sparkConf).newInstance().asInstanceOf[JavaSerializerInstance]
     val nettyEnv =
       new NettyRpcEnv(sparkConf, javaSerializerInstance, config.host, config.securityManager)
-    val startNettyRpcEnv: Int => (NettyRpcEnv, Int) = { actualPort =>
-      nettyEnv.start(actualPort)
-      (nettyEnv, actualPort)
-    }
-    try {
-      Utils.startServiceOnPort(config.port, startNettyRpcEnv, sparkConf, "NettyRpcEnv")._1
-    } catch {
-      case NonFatal(e) =>
-        nettyEnv.shutdown()
-        throw e
+    if (!config.clientMode) {
+      val startNettyRpcEnv: Int => (NettyRpcEnv, Int) = { actualPort =>
+        nettyEnv.startServer(actualPort)
+        (nettyEnv, actualPort)
+      }
+      try {
+        Utils.startServiceOnPort(config.port, startNettyRpcEnv, sparkConf, "NettyRpcEnv")._1
+      } catch {
+        case NonFatal(e) =>
+          nettyEnv.shutdown()
+          throw e
+      }
     }
+    nettyEnv
   }
 }
 
-private[netty] class NettyRpcEndpointRef(@transient private val conf: SparkConf)
+/**
+ * The NettyRpcEnv version of RpcEndpointRef.
+ *
+ * This class behaves differently depending on where it's created. On the node that "owns" the
+ * RpcEndpoint, it's a simple wrapper around the RpcEndpointAddress instance.
+ *
+ * On other machines that receive a serialized version of the reference, the behavior changes. The
+ * instance will keep track of the TransportClient that sent the reference, so that messages
+ * to the endpoint are sent over the client connection, instead of needing a new connection to
+ * be opened.
+ *
+ * The RpcAddress of this ref can be null; what that means is that the ref can only be used through
+ * a client connection, since the process hosting the endpoint is not listening for incoming
+ * connections. These refs should not be shared with 3rd parties, since they will not be able to
+ * send messages to the endpoint.
+ *
+ * @param conf Spark configuration.
+ * @param endpointAddress The address where the endpoint is listening.
+ * @param nettyEnv The RpcEnv associated with this ref.
+ * @param local Whether the referenced endpoint lives in the same process.
+ */
+private[netty] class NettyRpcEndpointRef(
+    @transient private val conf: SparkConf,
+    endpointAddress: RpcEndpointAddress,
+    @transient @volatile private var nettyEnv: NettyRpcEnv)
   extends RpcEndpointRef(conf) with Serializable with Logging {
 
-  @transient @volatile private var nettyEnv: NettyRpcEnv = _
+  @transient @volatile var client: TransportClient = _
 
-  @transient @volatile private var _address: RpcEndpointAddress = _
+  private val _address = if (endpointAddress.rpcAddress != null) endpointAddress else null
+  private val _name = endpointAddress.name
 
-  def this(conf: SparkConf, _address: RpcEndpointAddress, nettyEnv: NettyRpcEnv) {
-    this(conf)
-    this._address = _address
-    this.nettyEnv = nettyEnv
-  }
-
-  override def address: RpcAddress = _address.toRpcAddress
+  override def address: RpcAddress = if (_address != null) _address.rpcAddress else null
 
   private def readObject(in: ObjectInputStream): Unit = {
     in.defaultReadObject()
-    _address = in.readObject().asInstanceOf[RpcEndpointAddress]
     nettyEnv = NettyRpcEnv.currentEnv.value
+    client = NettyRpcEnv.currentClient.value
   }
 
   private def writeObject(out: ObjectOutputStream): Unit = {
     out.defaultWriteObject()
-    out.writeObject(_address)
   }
 
-  override def name: String = _address.name
+  override def name: String = _name
 
   override def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T] = {
     val promise = Promise[Any]()
@@ -429,41 +462,43 @@ private[netty] case class Ack(sender: NettyRpcEndpointRef) extends ResponseMessa
 private[netty] case class RpcFailure(e: Throwable)
 
 /**
- * Maintain the mapping relations between client addresses and [[RpcEnv]] addresses, broadcast
- * network events and forward messages to [[Dispatcher]].
+ * Dispatches incoming RPCs to registered endpoints.
+ *
+ * The handler keeps track of all client instances that communicate with it, so that the RpcEnv
+ * knows which `TransportClient` instance to use when sending RPCs to a client endpoint (i.e.,
+ * one that is not listening for incoming connections, but rather needs to be contacted via the
+ * client socket).
+ *
+ * Events are sent on a per-connection basis, so if a client opens multiple connections to the
+ * RpcEnv, multiple connection / disconnection events will be created for that client (albeit
+ * with different `RpcAddress` information).
  */
 private[netty] class NettyRpcHandler(
     dispatcher: Dispatcher, nettyEnv: NettyRpcEnv) extends RpcHandler with Logging {
 
-  private type ClientAddress = RpcAddress
-  private type RemoteEnvAddress = RpcAddress
-
-  // Store all client addresses and their NettyRpcEnv addresses.
-  // TODO: Is this even necessary?
-  @GuardedBy("this")
-  private val remoteAddresses = new mutable.HashMap[ClientAddress, RemoteEnvAddress]()
+  // TODO: Can we add connection callback (channel registered) to the underlying framework?
+  // A variable to track whether we should dispatch the RemoteProcessConnected message.
+  private val clients = new ConcurrentHashMap[TransportClient, JBoolean]()
 
   override def receive(
-      client: TransportClient, message: Array[Byte], callback: RpcResponseCallback): Unit = {
-    val requestMessage = nettyEnv.deserialize[RequestMessage](message)
-    val addr = client.getChannel.remoteAddress().asInstanceOf[InetSocketAddress]
+      client: TransportClient,
+      message: Array[Byte],
+      callback: RpcResponseCallback): Unit = {
+    val addr = client.getChannel().remoteAddress().asInstanceOf[InetSocketAddress]
     assert(addr != null)
-    val remoteEnvAddress = requestMessage.senderAddress
     val clientAddr = RpcAddress(addr.getHostName, addr.getPort)
-
-    // TODO: Can we add connection callback (channel registered) to the underlying framework?
-    // A variable to track whether we should dispatch the RemoteProcessConnected message.
-    var dispatchRemoteProcessConnected = false
-    synchronized {
-      if (remoteAddresses.put(clientAddr, remoteEnvAddress).isEmpty) {
-        // clientAddr connects at the first time, fire "RemoteProcessConnected"
-        dispatchRemoteProcessConnected = true
-      }
+    if (clients.putIfAbsent(client, JBoolean.TRUE) == null) {
+      dispatcher.postToAll(RemoteProcessConnected(clientAddr))
     }
-    if (dispatchRemoteProcessConnected) {
-      dispatcher.postToAll(RemoteProcessConnected(remoteEnvAddress))
-    }
-    dispatcher.postRemoteMessage(requestMessage, callback)
+    val requestMessage = nettyEnv.deserialize[RequestMessage](client, message)
+    val messageToDispatch = if (requestMessage.senderAddress == null) {
+        // Create a new message with the socket address of the client as the sender.
+        RequestMessage(clientAddr, requestMessage.receiver, requestMessage.content,
+          requestMessage.needReply)
+      } else {
+        requestMessage
+      }
+    dispatcher.postRemoteMessage(messageToDispatch, callback)
   }
 
   override def getStreamManager: StreamManager = new OneForOneStreamManager
@@ -472,15 +507,7 @@ private[netty] class NettyRpcHandler(
     val addr = client.getChannel.remoteAddress().asInstanceOf[InetSocketAddress]
     if (addr != null) {
       val clientAddr = RpcAddress(addr.getHostName, addr.getPort)
-      val broadcastMessage =
-        synchronized {
-          remoteAddresses.get(clientAddr).map(RemoteProcessConnectionError(cause, _))
-        }
-      if (broadcastMessage.isEmpty) {
-        logError(cause.getMessage, cause)
-      } else {
-        dispatcher.postToAll(broadcastMessage.get)
-      }
+      dispatcher.postToAll(RemoteProcessConnectionError(cause, clientAddr))
     } else {
       // If the channel is closed before connecting, its remoteAddress will be null.
       // See java.net.Socket.getRemoteSocketAddress
@@ -493,15 +520,9 @@ private[netty] class NettyRpcHandler(
     val addr = client.getChannel.remoteAddress().asInstanceOf[InetSocketAddress]
     if (addr != null) {
       val clientAddr = RpcAddress(addr.getHostName, addr.getPort)
+      clients.remove(client)
       nettyEnv.removeOutbox(clientAddr)
-      val messageOpt: Option[RemoteProcessDisconnected] =
-      synchronized {
-        remoteAddresses.get(clientAddr).flatMap { remoteEnvAddress =>
-          remoteAddresses -= clientAddr
-          Some(RemoteProcessDisconnected(remoteEnvAddress))
-        }
-      }
-      messageOpt.foreach(dispatcher.postToAll)
+      dispatcher.postToAll(RemoteProcessDisconnected(clientAddr))
     } else {
       // If the channel is closed before connecting, its remoteAddress will be null. In this case,
       // we can ignore it since we don't fire "Associated".
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala b/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
index 7d9d593b36241..2f6817f2eb935 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
@@ -26,7 +26,21 @@ import org.apache.spark.SparkException
 import org.apache.spark.network.client.{RpcResponseCallback, TransportClient}
 import org.apache.spark.rpc.RpcAddress
 
-private[netty] case class OutboxMessage(content: Array[Byte], callback: RpcResponseCallback)
+private[netty] case class OutboxMessage(content: Array[Byte],
+  _onFailure: (Throwable) => Unit,
+  _onSuccess: (TransportClient, Array[Byte]) => Unit) {
+
+  def createCallback(client: TransportClient): RpcResponseCallback = new RpcResponseCallback() {
+    override def onFailure(e: Throwable): Unit = {
+      _onFailure(e)
+    }
+
+    override def onSuccess(response: Array[Byte]): Unit = {
+      _onSuccess(client, response)
+    }
+  }
+
+}
 
 private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
 
@@ -68,7 +82,7 @@ private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
       }
     }
     if (dropped) {
-      message.callback.onFailure(new SparkException("Message is dropped because Outbox is stopped"))
+      message._onFailure(new SparkException("Message is dropped because Outbox is stopped"))
     } else {
       drainOutbox()
     }
@@ -108,7 +122,7 @@ private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
       try {
         val _client = synchronized { client }
         if (_client != null) {
-          _client.sendRpc(message.content, message.callback)
+          _client.sendRpc(message.content, message.createCallback(_client))
         } else {
           assert(stopped == true)
         }
@@ -181,7 +195,7 @@ private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
     // update messages and it's safe to just drain the queue.
     var message = messages.poll()
     while (message != null) {
-      message.callback.onFailure(e)
+      message._onFailure(e)
       message = messages.poll()
     }
     assert(messages.isEmpty)
@@ -215,7 +229,7 @@ private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
     // update messages and it's safe to just drain the queue.
     var message = messages.poll()
     while (message != null) {
-      message.callback.onFailure(new SparkException("Message is dropped because Outbox is stopped"))
+      message._onFailure(new SparkException("Message is dropped because Outbox is stopped"))
       message = messages.poll()
     }
   }
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/RpcEndpointAddress.scala b/core/src/main/scala/org/apache/spark/rpc/netty/RpcEndpointAddress.scala
index 87b6236936817..d2e94f943aba5 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/RpcEndpointAddress.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/RpcEndpointAddress.scala
@@ -23,15 +23,25 @@ import org.apache.spark.rpc.RpcAddress
 /**
  * An address identifier for an RPC endpoint.
  *
- * @param host host name of the remote process.
- * @param port the port the remote RPC environment binds to.
- * @param name name of the remote endpoint.
+ * The `rpcAddress` may be null, in which case the endpoint is registered via a client-only
+ * connection and can only be reached via the client that sent the endpoint reference.
+ *
+ * @param rpcAddress The socket address of the endpint.
+ * @param name Name of the endpoint.
  */
-private[netty] case class RpcEndpointAddress(host: String, port: Int, name: String) {
+private[netty] case class RpcEndpointAddress(val rpcAddress: RpcAddress, val name: String) {
+
+  require(name != null, "RpcEndpoint name must be provided.")
 
-  def toRpcAddress: RpcAddress = RpcAddress(host, port)
+  def this(host: String, port: Int, name: String) = {
+    this(RpcAddress(host, port), name)
+  }
 
-  override val toString = s"spark://$name@$host:$port"
+  override val toString = if (rpcAddress != null) {
+      s"spark://$name@${rpcAddress.host}:${rpcAddress.port}"
+    } else {
+      s"spark-client://$name"
+    }
 }
 
 private[netty] object RpcEndpointAddress {
@@ -51,7 +61,7 @@ private[netty] object RpcEndpointAddress {
           uri.getQuery != null) {
         throw new SparkException("Invalid Spark URL: " + sparkUrl)
       }
-      RpcEndpointAddress(host, port, name)
+      new RpcEndpointAddress(host, port, name)
     } catch {
       case e: java.net.URISyntaxException =>
         throw new SparkException("Invalid Spark URL: " + sparkUrl, e)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
index 8103efa7302e7..f3d0d85476772 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala
@@ -38,7 +38,7 @@ private[spark] object CoarseGrainedClusterMessages {
 
   sealed trait RegisterExecutorResponse
 
-  case object RegisteredExecutor extends CoarseGrainedClusterMessage
+  case class RegisteredExecutor(hostname: String) extends CoarseGrainedClusterMessage
     with RegisterExecutorResponse
 
   case class RegisterExecutorFailed(message: String) extends CoarseGrainedClusterMessage
@@ -51,9 +51,7 @@ private[spark] object CoarseGrainedClusterMessages {
       hostPort: String,
       cores: Int,
       logUrls: Map[String, String])
-    extends CoarseGrainedClusterMessage {
-    Utils.checkHostPort(hostPort, "Expected host port")
-  }
+    extends CoarseGrainedClusterMessage
 
   case class StatusUpdate(executorId: String, taskId: Long, state: TaskState,
     data: SerializableBuffer) extends CoarseGrainedClusterMessage
@@ -107,8 +105,4 @@ private[spark] object CoarseGrainedClusterMessages {
   // Used internally by executors to shut themselves down.
   case object Shutdown extends CoarseGrainedClusterMessage
 
-  // SPARK-10987: workaround for netty RPC issue; forces a connection from the driver back
-  // to the AM.
-  case object DriverHello extends CoarseGrainedClusterMessage
-
 }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
index 55a564b5c8eac..439a11927026b 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
@@ -131,16 +131,22 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
 
       case RegisterExecutor(executorId, executorRef, hostPort, cores, logUrls) =>
-        Utils.checkHostPort(hostPort, "Host port expected " + hostPort)
         if (executorDataMap.contains(executorId)) {
           context.reply(RegisterExecutorFailed("Duplicate executor ID: " + executorId))
         } else {
-          logInfo("Registered executor: " + executorRef + " with ID " + executorId)
-          addressToExecutorId(executorRef.address) = executorId
+          // If the executor's rpc env is not listening for incoming connections, `hostPort`
+          // will be null, and the client connection should be used to contact the executor.
+          val executorAddress = if (executorRef.address != null) {
+              executorRef.address
+            } else {
+              context.senderAddress
+            }
+          logInfo(s"Registered executor $executorRef ($executorAddress) with ID $executorId")
+          addressToExecutorId(executorAddress) = executorId
           totalCoreCount.addAndGet(cores)
           totalRegisteredExecutors.addAndGet(1)
-          val (host, _) = Utils.parseHostPort(hostPort)
-          val data = new ExecutorData(executorRef, executorRef.address, host, cores, cores, logUrls)
+          val data = new ExecutorData(executorRef, executorRef.address, executorAddress.host,
+            cores, cores, logUrls)
           // This must be synchronized because variables mutated
           // in this block are read when requesting executors
           CoarseGrainedSchedulerBackend.this.synchronized {
@@ -151,7 +157,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
             }
           }
           // Note: some tests expect the reply to come after we put the executor in the map
-          context.reply(RegisteredExecutor)
+          context.reply(RegisteredExecutor(executorAddress.host))
           listenerBus.post(
             SparkListenerExecutorAdded(System.currentTimeMillis(), executorId, data))
           makeOffers()
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
index e483688edef5f..cb24072d7d941 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
@@ -170,8 +170,6 @@ private[spark] abstract class YarnSchedulerBackend(
       case RegisterClusterManager(am) =>
         logInfo(s"ApplicationMaster registered as $am")
         amEndpoint = Option(am)
-        // See SPARK-10987.
-        am.send(DriverHello)
 
       case AddWebUIFilter(filterName, filterParams, proxyBase) =>
         addWebUIFilter(filterName, filterParams, proxyBase)
diff --git a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
index 3bead6395d384..834e4743df866 100644
--- a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
@@ -48,7 +48,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     }
   }
 
-  def createRpcEnv(conf: SparkConf, name: String, port: Int): RpcEnv
+  def createRpcEnv(conf: SparkConf, name: String, port: Int, clientMode: Boolean = false): RpcEnv
 
   test("send a message locally") {
     @volatile var message: String = null
@@ -76,7 +76,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef("local", env.address, "send-remotely")
     try {
@@ -130,7 +130,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef("local", env.address, "ask-remotely")
     try {
@@ -158,7 +158,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     val shortProp = "spark.rpc.short.timeout"
     conf.set("spark.rpc.retry.wait", "0")
     conf.set("spark.rpc.numRetries", "1")
-    val anotherEnv = createRpcEnv(conf, "remote", 13345)
+    val anotherEnv = createRpcEnv(conf, "remote", 13345, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef("local", env.address, "ask-timeout")
     try {
@@ -417,7 +417,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef("local", env.address, "sendWithReply-remotely")
     try {
@@ -457,7 +457,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef(
       "local", env.address, "sendWithReply-remotely-error")
@@ -497,26 +497,40 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
 
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef(
       "local", env.address, "network-events")
     val remoteAddress = anotherEnv.address
     rpcEndpointRef.send("hello")
     eventually(timeout(5 seconds), interval(5 millis)) {
-      assert(events === List(("onConnected", remoteAddress)))
+      // anotherEnv is connected in client mode, so the remote address may be unknown depending on
+      // the implementation. Account for that when doing checks.
+      if (remoteAddress != null) {
+        assert(events === List(("onConnected", remoteAddress)))
+      } else {
+        assert(events.size === 1)
+        assert(events(0)._1 === "onConnected")
+      }
     }
 
     anotherEnv.shutdown()
     anotherEnv.awaitTermination()
     eventually(timeout(5 seconds), interval(5 millis)) {
-      assert(events === List(
-        ("onConnected", remoteAddress),
-        ("onNetworkError", remoteAddress),
-        ("onDisconnected", remoteAddress)) ||
-        events === List(
-        ("onConnected", remoteAddress),
-        ("onDisconnected", remoteAddress)))
+      // Account for anotherEnv not having an address due to running in client mode.
+      if (remoteAddress != null) {
+        assert(events === List(
+          ("onConnected", remoteAddress),
+          ("onNetworkError", remoteAddress),
+          ("onDisconnected", remoteAddress)) ||
+          events === List(
+          ("onConnected", remoteAddress),
+          ("onDisconnected", remoteAddress)))
+      } else {
+        val eventNames = events.map(_._1)
+        assert(eventNames === List("onConnected", "onNetworkError", "onDisconnected") ||
+          eventNames === List("onConnected", "onDisconnected"))
+      }
     }
   }
 
@@ -529,7 +543,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef(
       "local", env.address, "sendWithReply-unserializable-error")
@@ -558,7 +572,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     conf.set("spark.authenticate.secret", "good")
 
     val localEnv = createRpcEnv(conf, "authentication-local", 13345)
-    val remoteEnv = createRpcEnv(conf, "authentication-remote", 14345)
+    val remoteEnv = createRpcEnv(conf, "authentication-remote", 14345, clientMode = true)
 
     try {
       @volatile var message: String = null
@@ -589,7 +603,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     conf.set("spark.authenticate.secret", "good")
 
     val localEnv = createRpcEnv(conf, "authentication-local", 13345)
-    val remoteEnv = createRpcEnv(conf, "authentication-remote", 14345)
+    val remoteEnv = createRpcEnv(conf, "authentication-remote", 14345, clientMode = true)
 
     try {
       localEnv.setupEndpoint("ask-authentication", new RpcEndpoint {
diff --git a/core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala b/core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala
index 4aa75c9230b2c..6478ab51c4da2 100644
--- a/core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala
@@ -22,9 +22,12 @@ import org.apache.spark.{SSLSampleConfigs, SecurityManager, SparkConf}
 
 class AkkaRpcEnvSuite extends RpcEnvSuite {
 
-  override def createRpcEnv(conf: SparkConf, name: String, port: Int): RpcEnv = {
+  override def createRpcEnv(conf: SparkConf,
+      name: String,
+      port: Int,
+      clientMode: Boolean = false): RpcEnv = {
     new AkkaRpcEnvFactory().create(
-      RpcEnvConfig(conf, name, "localhost", port, new SecurityManager(conf)))
+      RpcEnvConfig(conf, name, "localhost", port, new SecurityManager(conf), clientMode))
   }
 
   test("setupEndpointRef: systemName, address, endpointName") {
@@ -37,7 +40,7 @@ class AkkaRpcEnvSuite extends RpcEnvSuite {
     })
     val conf = new SparkConf()
     val newRpcEnv = new AkkaRpcEnvFactory().create(
-      RpcEnvConfig(conf, "test", "localhost", 12346, new SecurityManager(conf)))
+      RpcEnvConfig(conf, "test", "localhost", 12346, new SecurityManager(conf), false))
     try {
       val newRef = newRpcEnv.setupEndpointRef("local", ref.address, "test_endpoint")
       assert(s"akka.tcp://local@${env.address}/user/test_endpoint" ===
@@ -56,7 +59,7 @@ class AkkaRpcEnvSuite extends RpcEnvSuite {
     val conf = SSLSampleConfigs.sparkSSLConfig()
     val securityManager = new SecurityManager(conf)
     val rpcEnv = new AkkaRpcEnvFactory().create(
-      RpcEnvConfig(conf, "test", "localhost", 12346, securityManager))
+      RpcEnvConfig(conf, "test", "localhost", 12346, securityManager, false))
     try {
       val uri = rpcEnv.uriOf("local", RpcAddress("1.2.3.4", 12345), "test_endpoint")
       assert("akka.ssl.tcp://local@1.2.3.4:12345/user/test_endpoint" === uri)
diff --git a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcAddressSuite.scala b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcAddressSuite.scala
index 973a07a0bde3a..56743ba650b41 100644
--- a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcAddressSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcAddressSuite.scala
@@ -22,8 +22,13 @@ import org.apache.spark.SparkFunSuite
 class NettyRpcAddressSuite extends SparkFunSuite {
 
   test("toString") {
-    val addr = RpcEndpointAddress("localhost", 12345, "test")
+    val addr = new RpcEndpointAddress("localhost", 12345, "test")
     assert(addr.toString === "spark://test@localhost:12345")
   }
 
+  test("toString for client mode") {
+    val addr = RpcEndpointAddress(null, "test")
+    assert(addr.toString === "spark-client://test")
+  }
+
 }
diff --git a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcEnvSuite.scala b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcEnvSuite.scala
index be19668e17c04..ce83087ec04d6 100644
--- a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcEnvSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcEnvSuite.scala
@@ -22,8 +22,13 @@ import org.apache.spark.rpc._
 
 class NettyRpcEnvSuite extends RpcEnvSuite {
 
-  override def createRpcEnv(conf: SparkConf, name: String, port: Int): RpcEnv = {
-    val config = RpcEnvConfig(conf, "test", "localhost", port, new SecurityManager(conf))
+  override def createRpcEnv(
+      conf: SparkConf,
+      name: String,
+      port: Int,
+      clientMode: Boolean = false): RpcEnv = {
+    val config = RpcEnvConfig(conf, "test", "localhost", port, new SecurityManager(conf),
+      clientMode)
     new NettyRpcEnvFactory().create(config)
   }
 
diff --git a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
index 5430e4c0c4d6c..f9d8e80c98b66 100644
--- a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
@@ -30,7 +30,7 @@ import org.apache.spark.rpc._
 class NettyRpcHandlerSuite extends SparkFunSuite {
 
   val env = mock(classOf[NettyRpcEnv])
-  when(env.deserialize(any(classOf[Array[Byte]]))(any())).
+  when(env.deserialize(any(classOf[TransportClient]), any(classOf[Array[Byte]]))(any())).
     thenReturn(RequestMessage(RpcAddress("localhost", 12345), null, null, false))
 
   test("receive") {
@@ -42,7 +42,7 @@ class NettyRpcHandlerSuite extends SparkFunSuite {
     when(channel.remoteAddress()).thenReturn(new InetSocketAddress("localhost", 40000))
     nettyRpcHandler.receive(client, null, null)
 
-    verify(dispatcher, times(1)).postToAll(RemoteProcessConnected(RpcAddress("localhost", 12345)))
+    verify(dispatcher, times(1)).postToAll(RemoteProcessConnected(RpcAddress("localhost", 40000)))
   }
 
   test("connectionTerminated") {
@@ -57,9 +57,9 @@ class NettyRpcHandlerSuite extends SparkFunSuite {
     when(channel.remoteAddress()).thenReturn(new InetSocketAddress("localhost", 40000))
     nettyRpcHandler.connectionTerminated(client)
 
-    verify(dispatcher, times(1)).postToAll(RemoteProcessConnected(RpcAddress("localhost", 12345)))
+    verify(dispatcher, times(1)).postToAll(RemoteProcessConnected(RpcAddress("localhost", 40000)))
     verify(dispatcher, times(1)).postToAll(
-      RemoteProcessDisconnected(RpcAddress("localhost", 12345)))
+      RemoteProcessDisconnected(RpcAddress("localhost", 40000)))
   }
 
 }
diff --git a/network/yarn/pom.xml b/network/yarn/pom.xml
index 541ed9a8d0ab6..e2360eff5cfe1 100644
--- a/network/yarn/pom.xml
+++ b/network/yarn/pom.xml
@@ -54,6 +54,11 @@
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-client</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.slf4j</groupId>
+      <artifactId>slf4j-api</artifactId>
+      <scope>provided</scope>
+    </dependency>
   </dependencies>
 
   <build>
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
index c6a6d7ac56bf3..12ae350e4cef6 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
@@ -321,7 +321,8 @@ private[spark] class ApplicationMaster(
 
   private def runExecutorLauncher(securityMgr: SecurityManager): Unit = {
     val port = sparkConf.getInt("spark.yarn.am.port", 0)
-    rpcEnv = RpcEnv.create("sparkYarnAM", Utils.localHostName, port, sparkConf, securityMgr)
+    rpcEnv = RpcEnv.create("sparkYarnAM", Utils.localHostName, port, sparkConf, securityMgr,
+      clientMode = true)
     val driverRef = waitForSparkDriver()
     addAmIpFilter()
     registerAM(rpcEnv, driverRef, sparkConf.get("spark.driver.appUIAddress", ""), securityMgr)
@@ -574,9 +575,6 @@ private[spark] class ApplicationMaster(
       case x: AddWebUIFilter =>
         logInfo(s"Add WebUI Filter. $x")
         driver.send(x)
-
-      case DriverHello =>
-        // SPARK-10987: no action needed for this message.
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {

From f92f334ca47c03b980b06cf300aa652d0ffa1880 Mon Sep 17 00:00:00 2001
From: Jason White <jason.white@shopify.com>
Date: Mon, 2 Nov 2015 10:49:06 -0800
Subject: [PATCH 005/518] [SPARK-11437] [PYSPARK] Don't .take when converting
 RDD to DataFrame with provided schema

When creating a DataFrame from an RDD in PySpark, `createDataFrame` calls `.take(10)` to verify the first 10 rows of the RDD match the provided schema. Similar to https://issues.apache.org/jira/browse/SPARK-8070, but that issue affected cases where a schema was not provided.

Verifying the first 10 rows is of limited utility and causes the DAG to be executed non-lazily. If necessary, I believe this verification should be done lazily on all rows. However, since the caller is providing a schema to follow, I think it's acceptable to simply fail if the schema is incorrect.

marmbrus We chatted about this at SparkSummitEU. davies you made a similar change for the infer-schema path in https://github.com/apache/spark/pull/6606

Author: Jason White <jason.white@shopify.com>

Closes #9392 from JasonMWhite/createDataFrame_without_take.
---
 python/pyspark/sql/context.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/python/pyspark/sql/context.py b/python/pyspark/sql/context.py
index 79453658a167a..924bb6433de0e 100644
--- a/python/pyspark/sql/context.py
+++ b/python/pyspark/sql/context.py
@@ -318,13 +318,7 @@ def _createFromRDD(self, rdd, schema, samplingRatio):
                     struct.names[i] = name
             schema = struct
 
-        elif isinstance(schema, StructType):
-            # take the first few rows to verify schema
-            rows = rdd.take(10)
-            for row in rows:
-                _verify_type(row, schema)
-
-        else:
+        elif not isinstance(schema, StructType):
             raise TypeError("schema should be StructType or list or None, but got: %s" % schema)
 
         # convert python objects to sql data

From b3aedca6b55c678e40a5961e2fd3af4cb8c52bba Mon Sep 17 00:00:00 2001
From: Sean Owen <sowen@cloudera.com>
Date: Mon, 2 Nov 2015 14:36:37 -0600
Subject: [PATCH 006/518] [SPARK-11456][TESTS] Remove deprecated
 junit.framework in Java tests

Replace use of `junit.framework` with `org.junit`, and touch up tests in question

Author: Sean Owen <sowen@cloudera.com>

Closes #9411 from srowen/SPARK-11456.
---
 .../spark/unsafe/bitset/BitSetSuite.java      |  4 +-
 .../unsafe/hash/Murmur3_x86_32Suite.java      | 11 +--
 .../unsafe/types/CalendarIntervalSuite.java   | 78 +++++++++----------
 .../spark/unsafe/types/UTF8StringSuite.java   | 74 +++++++++---------
 4 files changed, 84 insertions(+), 83 deletions(-)

diff --git a/unsafe/src/test/java/org/apache/spark/unsafe/bitset/BitSetSuite.java b/unsafe/src/test/java/org/apache/spark/unsafe/bitset/BitSetSuite.java
index a93fc0ee297c4..14e38683df4ab 100644
--- a/unsafe/src/test/java/org/apache/spark/unsafe/bitset/BitSetSuite.java
+++ b/unsafe/src/test/java/org/apache/spark/unsafe/bitset/BitSetSuite.java
@@ -17,7 +17,7 @@
 
 package org.apache.spark.unsafe.bitset;
 
-import junit.framework.Assert;
+import org.junit.Assert;
 import org.junit.Test;
 
 import org.apache.spark.unsafe.memory.MemoryBlock;
@@ -25,7 +25,7 @@
 public class BitSetSuite {
 
   private static BitSet createBitSet(int capacity) {
-    assert capacity % 64 == 0;
+    Assert.assertEquals(0, capacity % 64);
     return new BitSet(MemoryBlock.fromLongArray(new long[capacity / 64]));
   }
 
diff --git a/unsafe/src/test/java/org/apache/spark/unsafe/hash/Murmur3_x86_32Suite.java b/unsafe/src/test/java/org/apache/spark/unsafe/hash/Murmur3_x86_32Suite.java
index 2f8cb132ac8b4..e759cb33b3e6a 100644
--- a/unsafe/src/test/java/org/apache/spark/unsafe/hash/Murmur3_x86_32Suite.java
+++ b/unsafe/src/test/java/org/apache/spark/unsafe/hash/Murmur3_x86_32Suite.java
@@ -17,12 +17,13 @@
 
 package org.apache.spark.unsafe.hash;
 
+import java.nio.charset.StandardCharsets;
 import java.util.HashSet;
 import java.util.Random;
 import java.util.Set;
 
-import junit.framework.Assert;
 import org.apache.spark.unsafe.Platform;
+import org.junit.Assert;
 import org.junit.Test;
 
 /**
@@ -56,7 +57,7 @@ public void randomizedStressTest() {
     Random rand = new Random();
 
     // A set used to track collision rate.
-    Set<Integer> hashcodes = new HashSet<Integer>();
+    Set<Integer> hashcodes = new HashSet<>();
     for (int i = 0; i < size; i++) {
       int vint = rand.nextInt();
       long lint = rand.nextLong();
@@ -76,7 +77,7 @@ public void randomizedStressTestBytes() {
     Random rand = new Random();
 
     // A set used to track collision rate.
-    Set<Integer> hashcodes = new HashSet<Integer>();
+    Set<Integer> hashcodes = new HashSet<>();
     for (int i = 0; i < size; i++) {
       int byteArrSize = rand.nextInt(100) * 8;
       byte[] bytes = new byte[byteArrSize];
@@ -98,10 +99,10 @@ public void randomizedStressTestBytes() {
   public void randomizedStressTestPaddedStrings() {
     int size = 64000;
     // A set used to track collision rate.
-    Set<Integer> hashcodes = new HashSet<Integer>();
+    Set<Integer> hashcodes = new HashSet<>();
     for (int i = 0; i < size; i++) {
       int byteArrSize = 8;
-      byte[] strBytes = ("" + i).getBytes();
+      byte[] strBytes = String.valueOf(i).getBytes(StandardCharsets.UTF_8);
       byte[] paddedBytes = new byte[byteArrSize];
       System.arraycopy(strBytes, 0, paddedBytes, 0, strBytes.length);
 
diff --git a/unsafe/src/test/java/org/apache/spark/unsafe/types/CalendarIntervalSuite.java b/unsafe/src/test/java/org/apache/spark/unsafe/types/CalendarIntervalSuite.java
index 80d4982c4b576..9e69e264ff287 100644
--- a/unsafe/src/test/java/org/apache/spark/unsafe/types/CalendarIntervalSuite.java
+++ b/unsafe/src/test/java/org/apache/spark/unsafe/types/CalendarIntervalSuite.java
@@ -19,7 +19,7 @@
 
 import org.junit.Test;
 
-import static junit.framework.Assert.*;
+import static org.junit.Assert.*;
 import static org.apache.spark.unsafe.types.CalendarInterval.*;
 
 public class CalendarIntervalSuite {
@@ -42,19 +42,19 @@ public void toStringTest() {
     CalendarInterval i;
 
     i = new CalendarInterval(34, 0);
-    assertEquals(i.toString(), "interval 2 years 10 months");
+    assertEquals("interval 2 years 10 months", i.toString());
 
     i = new CalendarInterval(-34, 0);
-    assertEquals(i.toString(), "interval -2 years -10 months");
+    assertEquals("interval -2 years -10 months", i.toString());
 
     i = new CalendarInterval(0, 3 * MICROS_PER_WEEK + 13 * MICROS_PER_HOUR + 123);
-    assertEquals(i.toString(), "interval 3 weeks 13 hours 123 microseconds");
+    assertEquals("interval 3 weeks 13 hours 123 microseconds", i.toString());
 
     i = new CalendarInterval(0, -3 * MICROS_PER_WEEK - 13 * MICROS_PER_HOUR - 123);
-    assertEquals(i.toString(), "interval -3 weeks -13 hours -123 microseconds");
+    assertEquals("interval -3 weeks -13 hours -123 microseconds", i.toString());
 
     i = new CalendarInterval(34, 3 * MICROS_PER_WEEK + 13 * MICROS_PER_HOUR + 123);
-    assertEquals(i.toString(), "interval 2 years 10 months 3 weeks 13 hours 123 microseconds");
+    assertEquals("interval 2 years 10 months 3 weeks 13 hours 123 microseconds", i.toString());
   }
 
   @Test
@@ -73,32 +73,32 @@ public void fromStringTest() {
 
     input = "interval   -5  years  23   month";
     CalendarInterval result = new CalendarInterval(-5 * 12 + 23, 0);
-    assertEquals(CalendarInterval.fromString(input), result);
+    assertEquals(fromString(input), result);
 
     input = "interval   -5  years  23   month   ";
-    assertEquals(CalendarInterval.fromString(input), result);
+    assertEquals(fromString(input), result);
 
     input = "  interval   -5  years  23   month   ";
-    assertEquals(CalendarInterval.fromString(input), result);
+    assertEquals(fromString(input), result);
 
     // Error cases
     input = "interval   3month 1 hour";
-    assertEquals(CalendarInterval.fromString(input), null);
+    assertNull(fromString(input));
 
     input = "interval 3 moth 1 hour";
-    assertEquals(CalendarInterval.fromString(input), null);
+    assertNull(fromString(input));
 
     input = "interval";
-    assertEquals(CalendarInterval.fromString(input), null);
+    assertNull(fromString(input));
 
     input = "int";
-    assertEquals(CalendarInterval.fromString(input), null);
+    assertNull(fromString(input));
 
     input = "";
-    assertEquals(CalendarInterval.fromString(input), null);
+    assertNull(fromString(input));
 
     input = null;
-    assertEquals(CalendarInterval.fromString(input), null);
+    assertNull(fromString(input));
   }
 
   @Test
@@ -108,15 +108,15 @@ public void fromYearMonthStringTest() {
 
     input = "99-10";
     i = new CalendarInterval(99 * 12 + 10, 0L);
-    assertEquals(CalendarInterval.fromYearMonthString(input), i);
+    assertEquals(fromYearMonthString(input), i);
 
     input = "-8-10";
     i = new CalendarInterval(-8 * 12 - 10, 0L);
-    assertEquals(CalendarInterval.fromYearMonthString(input), i);
+    assertEquals(fromYearMonthString(input), i);
 
     try {
       input = "99-15";
-      CalendarInterval.fromYearMonthString(input);
+      fromYearMonthString(input);
       fail("Expected to throw an exception for the invalid input");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("month 15 outside range"));
@@ -131,19 +131,19 @@ public void fromDayTimeStringTest() {
     input = "5 12:40:30.999999999";
     i = new CalendarInterval(0, 5 * MICROS_PER_DAY + 12 * MICROS_PER_HOUR +
       40 * MICROS_PER_MINUTE + 30 * MICROS_PER_SECOND + 999999L);
-    assertEquals(CalendarInterval.fromDayTimeString(input), i);
+    assertEquals(fromDayTimeString(input), i);
 
     input = "10 0:12:0.888";
     i = new CalendarInterval(0, 10 * MICROS_PER_DAY + 12 * MICROS_PER_MINUTE);
-    assertEquals(CalendarInterval.fromDayTimeString(input), i);
+    assertEquals(fromDayTimeString(input), i);
 
     input = "-3 0:0:0";
     i = new CalendarInterval(0, -3 * MICROS_PER_DAY);
-    assertEquals(CalendarInterval.fromDayTimeString(input), i);
+    assertEquals(fromDayTimeString(input), i);
 
     try {
       input = "5 30:12:20";
-      CalendarInterval.fromDayTimeString(input);
+      fromDayTimeString(input);
       fail("Expected to throw an exception for the invalid input");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("hour 30 outside range"));
@@ -151,7 +151,7 @@ public void fromDayTimeStringTest() {
 
     try {
       input = "5 30-12";
-      CalendarInterval.fromDayTimeString(input);
+      fromDayTimeString(input);
       fail("Expected to throw an exception for the invalid input");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("not match day-time format"));
@@ -165,19 +165,19 @@ public void fromSingleUnitStringTest() {
 
     input = "12";
     i = new CalendarInterval(12 * 12, 0L);
-    assertEquals(CalendarInterval.fromSingleUnitString("year", input), i);
+    assertEquals(fromSingleUnitString("year", input), i);
 
     input = "100";
     i = new CalendarInterval(0, 100 * MICROS_PER_DAY);
-    assertEquals(CalendarInterval.fromSingleUnitString("day", input), i);
+    assertEquals(fromSingleUnitString("day", input), i);
 
     input = "1999.38888";
     i = new CalendarInterval(0, 1999 * MICROS_PER_SECOND + 38);
-    assertEquals(CalendarInterval.fromSingleUnitString("second", input), i);
+    assertEquals(fromSingleUnitString("second", input), i);
 
     try {
       input = String.valueOf(Integer.MAX_VALUE);
-      CalendarInterval.fromSingleUnitString("year", input);
+      fromSingleUnitString("year", input);
       fail("Expected to throw an exception for the invalid input");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("outside range"));
@@ -185,7 +185,7 @@ public void fromSingleUnitStringTest() {
 
     try {
       input = String.valueOf(Long.MAX_VALUE / MICROS_PER_HOUR + 1);
-      CalendarInterval.fromSingleUnitString("hour", input);
+      fromSingleUnitString("hour", input);
       fail("Expected to throw an exception for the invalid input");
     } catch (IllegalArgumentException e) {
       assertTrue(e.getMessage().contains("outside range"));
@@ -197,16 +197,16 @@ public void addTest() {
     String input = "interval 3 month 1 hour";
     String input2 = "interval 2 month 100 hour";
 
-    CalendarInterval interval = CalendarInterval.fromString(input);
-    CalendarInterval interval2 = CalendarInterval.fromString(input2);
+    CalendarInterval interval = fromString(input);
+    CalendarInterval interval2 = fromString(input2);
 
     assertEquals(interval.add(interval2), new CalendarInterval(5, 101 * MICROS_PER_HOUR));
 
     input = "interval -10 month -81 hour";
     input2 = "interval 75 month 200 hour";
 
-    interval = CalendarInterval.fromString(input);
-    interval2 = CalendarInterval.fromString(input2);
+    interval = fromString(input);
+    interval2 = fromString(input2);
 
     assertEquals(interval.add(interval2), new CalendarInterval(65, 119 * MICROS_PER_HOUR));
   }
@@ -216,25 +216,25 @@ public void subtractTest() {
     String input = "interval 3 month 1 hour";
     String input2 = "interval 2 month 100 hour";
 
-    CalendarInterval interval = CalendarInterval.fromString(input);
-    CalendarInterval interval2 = CalendarInterval.fromString(input2);
+    CalendarInterval interval = fromString(input);
+    CalendarInterval interval2 = fromString(input2);
 
     assertEquals(interval.subtract(interval2), new CalendarInterval(1, -99 * MICROS_PER_HOUR));
 
     input = "interval -10 month -81 hour";
     input2 = "interval 75 month 200 hour";
 
-    interval = CalendarInterval.fromString(input);
-    interval2 = CalendarInterval.fromString(input2);
+    interval = fromString(input);
+    interval2 = fromString(input2);
 
     assertEquals(interval.subtract(interval2), new CalendarInterval(-85, -281 * MICROS_PER_HOUR));
   }
 
-  private void testSingleUnit(String unit, int number, int months, long microseconds) {
+  private static void testSingleUnit(String unit, int number, int months, long microseconds) {
     String input1 = "interval " + number + " " + unit;
     String input2 = "interval " + number + " " + unit + "s";
     CalendarInterval result = new CalendarInterval(months, microseconds);
-    assertEquals(CalendarInterval.fromString(input1), result);
-    assertEquals(CalendarInterval.fromString(input2), result);
+    assertEquals(fromString(input1), result);
+    assertEquals(fromString(input2), result);
   }
 }
diff --git a/unsafe/src/test/java/org/apache/spark/unsafe/types/UTF8StringSuite.java b/unsafe/src/test/java/org/apache/spark/unsafe/types/UTF8StringSuite.java
index 98aa8a2469a75..e21ffdcff9abf 100644
--- a/unsafe/src/test/java/org/apache/spark/unsafe/types/UTF8StringSuite.java
+++ b/unsafe/src/test/java/org/apache/spark/unsafe/types/UTF8StringSuite.java
@@ -24,13 +24,13 @@
 import com.google.common.collect.ImmutableMap;
 import org.junit.Test;
 
-import static junit.framework.Assert.*;
+import static org.junit.Assert.*;
 
 import static org.apache.spark.unsafe.types.UTF8String.*;
 
 public class UTF8StringSuite {
 
-  private void checkBasic(String str, int len) throws UnsupportedEncodingException {
+  private static void checkBasic(String str, int len) throws UnsupportedEncodingException {
     UTF8String s1 = fromString(str);
     UTF8String s2 = fromBytes(str.getBytes("utf8"));
     assertEquals(s1.numChars(), len);
@@ -42,12 +42,12 @@ private void checkBasic(String str, int len) throws UnsupportedEncodingException
 
     assertEquals(s1.hashCode(), s2.hashCode());
 
-    assertEquals(s1.compareTo(s2), 0);
+    assertEquals(0, s1.compareTo(s2));
 
-    assertEquals(s1.contains(s2), true);
-    assertEquals(s2.contains(s1), true);
-    assertEquals(s1.startsWith(s1), true);
-    assertEquals(s1.endsWith(s1), true);
+    assertTrue(s1.contains(s2));
+    assertTrue(s2.contains(s1));
+    assertTrue(s1.startsWith(s1));
+    assertTrue(s1.endsWith(s1));
   }
 
   @Test
@@ -59,8 +59,8 @@ public void basicTest() throws UnsupportedEncodingException {
 
   @Test
   public void emptyStringTest() {
-    assertEquals(fromString(""), EMPTY_UTF8);
-    assertEquals(fromBytes(new byte[0]), EMPTY_UTF8);
+    assertEquals(EMPTY_UTF8, fromString(""));
+    assertEquals(EMPTY_UTF8, fromBytes(new byte[0]));
     assertEquals(0, EMPTY_UTF8.numChars());
     assertEquals(0, EMPTY_UTF8.numBytes());
   }
@@ -76,9 +76,9 @@ public void prefix() {
 
     byte[] buf1 = {1, 2, 3, 4, 5, 6, 7, 8, 9};
     byte[] buf2 = {1, 2, 3};
-    UTF8String str1 = UTF8String.fromBytes(buf1, 0, 3);
-    UTF8String str2 = UTF8String.fromBytes(buf1, 0, 8);
-    UTF8String str3 = UTF8String.fromBytes(buf2);
+    UTF8String str1 = fromBytes(buf1, 0, 3);
+    UTF8String str2 = fromBytes(buf1, 0, 8);
+    UTF8String str3 = fromBytes(buf2);
     assertTrue(str1.getPrefix() - str2.getPrefix() < 0);
     assertEquals(str1.getPrefix(), str3.getPrefix());
   }
@@ -98,7 +98,7 @@ public void compareTo() {
     assertTrue(fromString("你好123").compareTo(fromString("你好122")) > 0);
   }
 
-  protected void testUpperandLower(String upper, String lower) {
+  protected static void testUpperandLower(String upper, String lower) {
     UTF8String us = fromString(upper);
     UTF8String ls = fromString(lower);
     assertEquals(ls, us.toLowerCase());
@@ -127,22 +127,22 @@ public void titleCase() {
   @Test
   public void concatTest() {
     assertEquals(EMPTY_UTF8, concat());
-    assertEquals(null, concat((UTF8String) null));
+    assertNull(concat((UTF8String) null));
     assertEquals(EMPTY_UTF8, concat(EMPTY_UTF8));
     assertEquals(fromString("ab"), concat(fromString("ab")));
     assertEquals(fromString("ab"), concat(fromString("a"), fromString("b")));
     assertEquals(fromString("abc"), concat(fromString("a"), fromString("b"), fromString("c")));
-    assertEquals(null, concat(fromString("a"), null, fromString("c")));
-    assertEquals(null, concat(fromString("a"), null, null));
-    assertEquals(null, concat(null, null, null));
+    assertNull(concat(fromString("a"), null, fromString("c")));
+    assertNull(concat(fromString("a"), null, null));
+    assertNull(concat(null, null, null));
     assertEquals(fromString("数据砖头"), concat(fromString("数据"), fromString("砖头")));
   }
 
   @Test
   public void concatWsTest() {
     // Returns null if the separator is null
-    assertEquals(null, concatWs(null, (UTF8String)null));
-    assertEquals(null, concatWs(null, fromString("a")));
+    assertNull(concatWs(null, (UTF8String) null));
+    assertNull(concatWs(null, fromString("a")));
 
     // If separator is null, concatWs should skip all null inputs and never return null.
     UTF8String sep = fromString("哈哈");
@@ -381,16 +381,16 @@ public void split() {
   
   @Test
   public void levenshteinDistance() {
-    assertEquals(EMPTY_UTF8.levenshteinDistance(EMPTY_UTF8), 0);
-    assertEquals(EMPTY_UTF8.levenshteinDistance(fromString("a")), 1);
-    assertEquals(fromString("aaapppp").levenshteinDistance(EMPTY_UTF8), 7);
-    assertEquals(fromString("frog").levenshteinDistance(fromString("fog")), 1);
-    assertEquals(fromString("fly").levenshteinDistance(fromString("ant")),3);
-    assertEquals(fromString("elephant").levenshteinDistance(fromString("hippo")), 7);
-    assertEquals(fromString("hippo").levenshteinDistance(fromString("elephant")), 7);
-    assertEquals(fromString("hippo").levenshteinDistance(fromString("zzzzzzzz")), 8);
-    assertEquals(fromString("hello").levenshteinDistance(fromString("hallo")),1);
-    assertEquals(fromString("世界千世").levenshteinDistance(fromString("千a世b")),4);
+    assertEquals(0, EMPTY_UTF8.levenshteinDistance(EMPTY_UTF8));
+    assertEquals(1, EMPTY_UTF8.levenshteinDistance(fromString("a")));
+    assertEquals(7, fromString("aaapppp").levenshteinDistance(EMPTY_UTF8));
+    assertEquals(1, fromString("frog").levenshteinDistance(fromString("fog")));
+    assertEquals(3, fromString("fly").levenshteinDistance(fromString("ant")));
+    assertEquals(7, fromString("elephant").levenshteinDistance(fromString("hippo")));
+    assertEquals(7, fromString("hippo").levenshteinDistance(fromString("elephant")));
+    assertEquals(8, fromString("hippo").levenshteinDistance(fromString("zzzzzzzz")));
+    assertEquals(1, fromString("hello").levenshteinDistance(fromString("hallo")));
+    assertEquals(4, fromString("世界千世").levenshteinDistance(fromString("千a世b")));
   }
 
   @Test
@@ -432,14 +432,14 @@ public void createBlankString() {
 
   @Test
   public void findInSet() {
-    assertEquals(fromString("ab").findInSet(fromString("ab")), 1);
-    assertEquals(fromString("a,b").findInSet(fromString("b")), 2);
-    assertEquals(fromString("abc,b,ab,c,def").findInSet(fromString("ab")), 3);
-    assertEquals(fromString("ab,abc,b,ab,c,def").findInSet(fromString("ab")), 1);
-    assertEquals(fromString(",,,ab,abc,b,ab,c,def").findInSet(fromString("ab")), 4);
-    assertEquals(fromString(",ab,abc,b,ab,c,def").findInSet(fromString("")), 1);
-    assertEquals(fromString("数据砖头,abc,b,ab,c,def").findInSet(fromString("ab")), 4);
-    assertEquals(fromString("数据砖头,abc,b,ab,c,def").findInSet(fromString("def")), 6);
+    assertEquals(1, fromString("ab").findInSet(fromString("ab")));
+    assertEquals(2, fromString("a,b").findInSet(fromString("b")));
+    assertEquals(3, fromString("abc,b,ab,c,def").findInSet(fromString("ab")));
+    assertEquals(1, fromString("ab,abc,b,ab,c,def").findInSet(fromString("ab")));
+    assertEquals(4, fromString(",,,ab,abc,b,ab,c,def").findInSet(fromString("ab")));
+    assertEquals(1, fromString(",ab,abc,b,ab,c,def").findInSet(fromString("")));
+    assertEquals(4, fromString("数据砖头,abc,b,ab,c,def").findInSet(fromString("ab")));
+    assertEquals(6, fromString("数据砖头,abc,b,ab,c,def").findInSet(fromString("def")));
   }
 
   @Test

From 33ae7a35daa86c34f1f9f72f997e0c2d4cd8abec Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Mon, 2 Nov 2015 13:42:16 -0800
Subject: [PATCH 007/518] [SPARK-11358][MLLIB] deprecate runs in k-means

This PR deprecates `runs` in k-means. `runs` introduces extra complexity and overhead in MLlib's k-means implementation. I haven't seen much usage with `runs` not equal to `1`. We don't have a unit test for it either. We can deprecate this method in 1.6, and void it in 1.7. It helps us simplify the implementation.

cc: srowen

Author: Xiangrui Meng <meng@databricks.com>

Closes #9322 from mengxr/SPARK-11358.
---
 .../main/scala/org/apache/spark/mllib/clustering/KMeans.scala | 4 ++--
 python/pyspark/mllib/clustering.py                            | 4 ++++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala
index 7168aac32c997..2895db7c9061b 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala
@@ -107,7 +107,7 @@ class KMeans private (
    * Number of runs of the algorithm to execute in parallel.
    */
   @Since("1.4.0")
-  @Experimental
+  @deprecated("Support for runs is deprecated. This param will have no effect in 1.7.0.", "1.6.0")
   def getRuns: Int = runs
 
   /**
@@ -117,7 +117,7 @@ class KMeans private (
    * return the best clustering found over any run. Default: 1.
    */
   @Since("0.8.0")
-  @Experimental
+  @deprecated("Support for runs is deprecated. This param will have no effect in 1.7.0.", "1.6.0")
   def setRuns(runs: Int): this.type = {
     if (runs <= 0) {
       throw new IllegalArgumentException("Number of runs must be positive")
diff --git a/python/pyspark/mllib/clustering.py b/python/pyspark/mllib/clustering.py
index d1c3755a785f2..8629aa5a17164 100644
--- a/python/pyspark/mllib/clustering.py
+++ b/python/pyspark/mllib/clustering.py
@@ -17,6 +17,7 @@
 
 import sys
 import array as pyarray
+import warnings
 
 if sys.version > '3':
     xrange = range
@@ -170,6 +171,9 @@ class KMeans(object):
     def train(cls, rdd, k, maxIterations=100, runs=1, initializationMode="k-means||",
               seed=None, initializationSteps=5, epsilon=1e-4, initialModel=None):
         """Train a k-means clustering model."""
+        if runs != 1:
+            warnings.warn(
+                "Support for runs is deprecated in 1.6.0. This param will have no effect in 1.7.0.")
         clusterInitialModel = []
         if initialModel is not None:
             if not isinstance(initialModel, KMeansModel):

From db11ee5e56e5fac59895c772a9a87c5ac86888ef Mon Sep 17 00:00:00 2001
From: tedyu <yuzhihong@gmail.com>
Date: Mon, 2 Nov 2015 13:51:53 -0800
Subject: [PATCH 008/518] [SPARK-11371] Make "mean" an alias for "avg" operator

From Reynold in the thread 'Exception when using some aggregate operators' (http://search-hadoop.com/m/q3RTt0xFr22nXB4/):

I don't think these are bugs. The SQL standard for average is "avg", not "mean". Similarly, a distinct count is supposed to be written as "count(distinct col)", not "countDistinct(col)".
We can, however, make "mean" an alias for "avg" to improve compatibility between DataFrame and SQL.

Author: tedyu <yuzhihong@gmail.com>

Closes #9332 from ted-yu/master.
---
 .../spark/sql/catalyst/analysis/FunctionRegistry.scala   | 1 +
 .../spark/sql/hive/execution/AggregationQuerySuite.scala | 9 +++++++++
 2 files changed, 10 insertions(+)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index 5f3ec74ac0d92..24c1a7b7ac5af 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -185,6 +185,7 @@ object FunctionRegistry {
     expression[Last]("last"),
     expression[Last]("last_value"),
     expression[Max]("max"),
+    expression[Average]("mean"),
     expression[Min]("min"),
     expression[Stddev]("stddev"),
     expression[StddevPop]("stddev_pop"),
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 0cf0e0aab9eb2..74061db0f28af 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -298,6 +298,15 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
         """.stripMargin),
       Row(1, 20.0) :: Row(2, -0.5) :: Row(3, null) :: Row(null, 10.0) :: Nil)
 
+    checkAnswer(
+      sqlContext.sql(
+        """
+          |SELECT key, mean(value)
+          |FROM agg1
+          |GROUP BY key
+        """.stripMargin),
+      Row(1, 20.0) :: Row(2, -0.5) :: Row(3, null) :: Row(null, 10.0) :: Nil)
+
     checkAnswer(
       sqlContext.sql(
         """

From 2804674a7af8f11eeb1280459bc9145815398eed Mon Sep 17 00:00:00 2001
From: Rishabh Bhardwaj <rbnext29@gmail.com>
Date: Mon, 2 Nov 2015 14:03:50 -0800
Subject: [PATCH 009/518] [SPARK-11383][DOCS] Replaced example code in
 mllib-naive-bayes.md/mllib-isotonic-regression.md using include_example

I have made the required changes in mllib-naive-bayes.md/mllib-isotonic-regression.md and also verified them.
Kindle Review it.

Author: Rishabh Bhardwaj <rbnext29@gmail.com>

Closes #9353 from rishabhbhardwaj/SPARK-11383.
---
 docs/mllib-isotonic-regression.md             | 124 +-----------------
 docs/mllib-naive-bayes.md                     |  89 +------------
 .../mllib/JavaIsotonicRegressionExample.java  |  86 ++++++++++++
 .../examples/mllib/JavaNaiveBayesExample.java |  64 +++++++++
 .../mllib/isotonic_regression_example.py      |  56 ++++++++
 .../main/python/mllib/naive_bayes_example.py  |  56 ++++++++
 .../mllib/IsotonicRegressionExample.scala     |  66 ++++++++++
 .../examples/mllib/NaiveBayesExample.scala    |  57 ++++++++
 8 files changed, 391 insertions(+), 207 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaIsotonicRegressionExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaNaiveBayesExample.java
 create mode 100644 examples/src/main/python/mllib/isotonic_regression_example.py
 create mode 100644 examples/src/main/python/mllib/naive_bayes_example.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/IsotonicRegressionExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/NaiveBayesExample.scala

diff --git a/docs/mllib-isotonic-regression.md b/docs/mllib-isotonic-regression.md
index f91a697b31891..85f9226b43416 100644
--- a/docs/mllib-isotonic-regression.md
+++ b/docs/mllib-isotonic-regression.md
@@ -61,42 +61,8 @@ labels and real labels in the test set.
 
 Refer to the [`IsotonicRegression` Scala docs](api/scala/index.html#org.apache.spark.mllib.regression.IsotonicRegression) and [`IsotonicRegressionModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.regression.IsotonicRegressionModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.regression.{IsotonicRegression, IsotonicRegressionModel}
-
-val data = sc.textFile("data/mllib/sample_isotonic_regression_data.txt")
-
-// Create label, feature, weight tuples from input data with weight set to default value 1.0.
-val parsedData = data.map { line =>
-  val parts = line.split(',').map(_.toDouble)
-  (parts(0), parts(1), 1.0)
-}
-
-// Split data into training (60%) and test (40%) sets.
-val splits = parsedData.randomSplit(Array(0.6, 0.4), seed = 11L)
-val training = splits(0)
-val test = splits(1)
-
-// Create isotonic regression model from training data.
-// Isotonic parameter defaults to true so it is only shown for demonstration
-val model = new IsotonicRegression().setIsotonic(true).run(training)
-
-// Create tuples of predicted and real labels.
-val predictionAndLabel = test.map { point =>
-  val predictedLabel = model.predict(point._2)
-  (predictedLabel, point._1)
-}
-
-// Calculate mean squared error between predicted and real labels.
-val meanSquaredError = predictionAndLabel.map{case(p, l) => math.pow((p - l), 2)}.mean()
-println("Mean Squared Error = " + meanSquaredError)
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = IsotonicRegressionModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/IsotonicRegressionExample.scala %}
 </div>
-
 <div data-lang="java" markdown="1">
 Data are read from a file where each line has a format label,feature
 i.e. 4710.28,500.00. The data are split to training and testing set.
@@ -105,66 +71,8 @@ labels and real labels in the test set.
 
 Refer to the [`IsotonicRegression` Java docs](api/java/org/apache/spark/mllib/regression/IsotonicRegression.html) and [`IsotonicRegressionModel` Java docs](api/java/org/apache/spark/mllib/regression/IsotonicRegressionModel.html) for details on the API.
 
-{% highlight java %}
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.JavaDoubleRDD;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.IsotonicRegressionModel;
-import scala.Tuple2;
-import scala.Tuple3;
-
-JavaRDD<String> data = sc.textFile("data/mllib/sample_isotonic_regression_data.txt");
-
-// Create label, feature, weight tuples from input data with weight set to default value 1.0.
-JavaRDD<Tuple3<Double, Double, Double>> parsedData = data.map(
-  new Function<String, Tuple3<Double, Double, Double>>() {
-    public Tuple3<Double, Double, Double> call(String line) {
-      String[] parts = line.split(",");
-      return new Tuple3<>(new Double(parts[0]), new Double(parts[1]), 1.0);
-    }
-  }
-);
-
-// Split data into training (60%) and test (40%) sets.
-JavaRDD<Tuple3<Double, Double, Double>>[] splits = parsedData.randomSplit(new double[] {0.6, 0.4}, 11L);
-JavaRDD<Tuple3<Double, Double, Double>> training = splits[0];
-JavaRDD<Tuple3<Double, Double, Double>> test = splits[1];
-
-// Create isotonic regression model from training data.
-// Isotonic parameter defaults to true so it is only shown for demonstration
-IsotonicRegressionModel model = new IsotonicRegression().setIsotonic(true).run(training);
-
-// Create tuples of predicted and real labels.
-JavaPairRDD<Double, Double> predictionAndLabel = test.mapToPair(
-  new PairFunction<Tuple3<Double, Double, Double>, Double, Double>() {
-    @Override public Tuple2<Double, Double> call(Tuple3<Double, Double, Double> point) {
-      Double predictedLabel = model.predict(point._2());
-      return new Tuple2<Double, Double>(predictedLabel, point._1());
-    }
-  }
-);
-
-// Calculate mean squared error between predicted and real labels.
-Double meanSquaredError = new JavaDoubleRDD(predictionAndLabel.map(
-  new Function<Tuple2<Double, Double>, Object>() {
-    @Override public Object call(Tuple2<Double, Double> pl) {
-      return Math.pow(pl._1() - pl._2(), 2);
-    }
-  }
-).rdd()).mean();
-
-System.out.println("Mean Squared Error = " + meanSquaredError);
-
-// Save and load model
-model.save(sc.sc(), "myModelPath");
-IsotonicRegressionModel sameModel = IsotonicRegressionModel.load(sc.sc(), "myModelPath");
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaIsotonicRegressionExample.java %}
 </div>
-
 <div data-lang="python" markdown="1">
 Data are read from a file where each line has a format label,feature
 i.e. 4710.28,500.00. The data are split to training and testing set.
@@ -173,32 +81,6 @@ labels and real labels in the test set.
 
 Refer to the [`IsotonicRegression` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.regression.IsotonicRegression) and [`IsotonicRegressionModel` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.regression.IsotonicRegressionModel) for more details on the API.
 
-{% highlight python %}
-import math
-from pyspark.mllib.regression import IsotonicRegression, IsotonicRegressionModel
-
-data = sc.textFile("data/mllib/sample_isotonic_regression_data.txt")
-
-# Create label, feature, weight tuples from input data with weight set to default value 1.0.
-parsedData = data.map(lambda line: tuple([float(x) for x in line.split(',')]) + (1.0,))
-
-# Split data into training (60%) and test (40%) sets.
-training, test = parsedData.randomSplit([0.6, 0.4], 11)
-
-# Create isotonic regression model from training data.
-# Isotonic parameter defaults to true so it is only shown for demonstration
-model = IsotonicRegression.train(training)
-
-# Create tuples of predicted and real labels.
-predictionAndLabel = test.map(lambda p: (model.predict(p[1]), p[0]))
-
-# Calculate mean squared error between predicted and real labels.
-meanSquaredError = predictionAndLabel.map(lambda pl: math.pow((pl[0] - pl[1]), 2)).mean()
-print("Mean Squared Error = " + str(meanSquaredError))
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = IsotonicRegressionModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example python/mllib/isotonic_regression_example.py %}
 </div>
 </div>
diff --git a/docs/mllib-naive-bayes.md b/docs/mllib-naive-bayes.md
index f4f6a10c8299e..60ac6c7e5bb1a 100644
--- a/docs/mllib-naive-bayes.md
+++ b/docs/mllib-naive-bayes.md
@@ -40,32 +40,8 @@ can be used for evaluation and prediction.
 
 Refer to the [`NaiveBayes` Scala docs](api/scala/index.html#org.apache.spark.mllib.classification.NaiveBayes) and [`NaiveBayesModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.classification.NaiveBayesModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.classification.{NaiveBayes, NaiveBayesModel}
-import org.apache.spark.mllib.linalg.Vectors
-import org.apache.spark.mllib.regression.LabeledPoint
-
-val data = sc.textFile("data/mllib/sample_naive_bayes_data.txt")
-val parsedData = data.map { line =>
-  val parts = line.split(',')
-  LabeledPoint(parts(0).toDouble, Vectors.dense(parts(1).split(' ').map(_.toDouble)))
-}
-// Split data into training (60%) and test (40%).
-val splits = parsedData.randomSplit(Array(0.6, 0.4), seed = 11L)
-val training = splits(0)
-val test = splits(1)
-
-val model = NaiveBayes.train(training, lambda = 1.0, modelType = "multinomial")
-
-val predictionAndLabel = test.map(p => (model.predict(p.features), p.label))
-val accuracy = 1.0 * predictionAndLabel.filter(x => x._1 == x._2).count() / test.count()
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = NaiveBayesModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/NaiveBayesExample.scala %}
 </div>
-
 <div data-lang="java" markdown="1">
 
 [NaiveBayes](api/java/org/apache/spark/mllib/classification/NaiveBayes.html) implements
@@ -77,40 +53,8 @@ can be used for evaluation and prediction.
 
 Refer to the [`NaiveBayes` Java docs](api/java/org/apache/spark/mllib/classification/NaiveBayes.html) and [`NaiveBayesModel` Java docs](api/java/org/apache/spark/mllib/classification/NaiveBayesModel.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.classification.NaiveBayes;
-import org.apache.spark.mllib.classification.NaiveBayesModel;
-import org.apache.spark.mllib.regression.LabeledPoint;
-
-JavaRDD<LabeledPoint> training = ... // training set
-JavaRDD<LabeledPoint> test = ... // test set
-
-final NaiveBayesModel model = NaiveBayes.train(training.rdd(), 1.0);
-
-JavaPairRDD<Double, Double> predictionAndLabel = 
-  test.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-    @Override public Tuple2<Double, Double> call(LabeledPoint p) {
-      return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-    }
-  });
-double accuracy = predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
-    @Override public Boolean call(Tuple2<Double, Double> pl) {
-      return pl._1().equals(pl._2());
-    }
-  }).count() / (double) test.count();
-
-// Save and load model
-model.save(sc.sc(), "myModelPath");
-NaiveBayesModel sameModel = NaiveBayesModel.load(sc.sc(), "myModelPath");
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaNaiveBayesExample.java %}
 </div>
-
 <div data-lang="python" markdown="1">
 
 [NaiveBayes](api/python/pyspark.mllib.html#pyspark.mllib.classification.NaiveBayes) implements multinomial
@@ -124,33 +68,6 @@ Note that the Python API does not yet support model save/load but will in the fu
 
 Refer to the [`NaiveBayes` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.classification.NaiveBayes) and [`NaiveBayesModel` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.classification.NaiveBayesModel) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.classification import NaiveBayes, NaiveBayesModel
-from pyspark.mllib.linalg import Vectors
-from pyspark.mllib.regression import LabeledPoint
-
-def parseLine(line):
-    parts = line.split(',')
-    label = float(parts[0])
-    features = Vectors.dense([float(x) for x in parts[1].split(' ')])
-    return LabeledPoint(label, features)
-
-data = sc.textFile('data/mllib/sample_naive_bayes_data.txt').map(parseLine)
-
-# Split data aproximately into training (60%) and test (40%)
-training, test = data.randomSplit([0.6, 0.4], seed = 0)
-
-# Train a naive Bayes model.
-model = NaiveBayes.train(training, 1.0)
-
-# Make prediction and test accuracy.
-predictionAndLabel = test.map(lambda p : (model.predict(p.features), p.label))
-accuracy = 1.0 * predictionAndLabel.filter(lambda (x, v): x == v).count() / test.count()
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = NaiveBayesModel.load(sc, "myModelPath")
-{% endhighlight %}
-
+{% include_example python/mllib/naive_bayes_example.py %}
 </div>
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaIsotonicRegressionExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaIsotonicRegressionExample.java
new file mode 100644
index 0000000000000..37e709b4cbc03
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaIsotonicRegressionExample.java
@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import scala.Tuple2;
+import scala.Tuple3;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.api.java.JavaDoubleRDD;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.mllib.regression.IsotonicRegression;
+import org.apache.spark.mllib.regression.IsotonicRegressionModel;
+// $example off$
+import org.apache.spark.SparkConf;
+
+public class JavaIsotonicRegressionExample {
+  public static void main(String[] args) {
+    SparkConf sparkConf = new SparkConf().setAppName("JavaIsotonicRegressionExample");
+    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
+    // $example on$
+    JavaRDD<String> data = jsc.textFile("data/mllib/sample_isotonic_regression_data.txt");
+
+    // Create label, feature, weight tuples from input data with weight set to default value 1.0.
+    JavaRDD<Tuple3<Double, Double, Double>> parsedData = data.map(
+      new Function<String, Tuple3<Double, Double, Double>>() {
+        public Tuple3<Double, Double, Double> call(String line) {
+          String[] parts = line.split(",");
+          return new Tuple3<>(new Double(parts[0]), new Double(parts[1]), 1.0);
+        }
+      }
+    );
+
+    // Split data into training (60%) and test (40%) sets.
+    JavaRDD<Tuple3<Double, Double, Double>>[] splits = parsedData.randomSplit(new double[]{0.6, 0.4}, 11L);
+    JavaRDD<Tuple3<Double, Double, Double>> training = splits[0];
+    JavaRDD<Tuple3<Double, Double, Double>> test = splits[1];
+
+    // Create isotonic regression model from training data.
+    // Isotonic parameter defaults to true so it is only shown for demonstration
+    final IsotonicRegressionModel model = new IsotonicRegression().setIsotonic(true).run(training);
+
+    // Create tuples of predicted and real labels.
+    JavaPairRDD<Double, Double> predictionAndLabel = test.mapToPair(
+      new PairFunction<Tuple3<Double, Double, Double>, Double, Double>() {
+        @Override
+        public Tuple2<Double, Double> call(Tuple3<Double, Double, Double> point) {
+          Double predictedLabel = model.predict(point._2());
+          return new Tuple2<Double, Double>(predictedLabel, point._1());
+        }
+      }
+    );
+
+    // Calculate mean squared error between predicted and real labels.
+    Double meanSquaredError = new JavaDoubleRDD(predictionAndLabel.map(
+      new Function<Tuple2<Double, Double>, Object>() {
+        @Override
+        public Object call(Tuple2<Double, Double> pl) {
+          return Math.pow(pl._1() - pl._2(), 2);
+        }
+      }
+    ).rdd()).mean();
+    System.out.println("Mean Squared Error = " + meanSquaredError);
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myIsotonicRegressionModel");
+    IsotonicRegressionModel sameModel = IsotonicRegressionModel.load(jsc.sc(), "target/tmp/myIsotonicRegressionModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaNaiveBayesExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaNaiveBayesExample.java
new file mode 100644
index 0000000000000..e6a5904bd71f0
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaNaiveBayesExample.java
@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import scala.Tuple2;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.mllib.classification.NaiveBayes;
+import org.apache.spark.mllib.classification.NaiveBayesModel;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.util.MLUtils;
+// $example off$
+import org.apache.spark.SparkConf;
+
+public class JavaNaiveBayesExample {
+  public static void main(String[] args) {
+    SparkConf sparkConf = new SparkConf().setAppName("JavaNaiveBayesExample");
+    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
+    // $example on$
+    String path = "data/mllib/sample_naive_bayes_data.txt";
+    JavaRDD<LabeledPoint> inputData = MLUtils.loadLibSVMFile(jsc.sc(), path).toJavaRDD();
+    JavaRDD<LabeledPoint>[] tmp = inputData.randomSplit(new double[]{0.6, 0.4}, 12345);
+    JavaRDD<LabeledPoint> training = tmp[0]; // training set
+    JavaRDD<LabeledPoint> test = tmp[1]; // test set
+    final NaiveBayesModel model = NaiveBayes.train(training.rdd(), 1.0);
+    JavaPairRDD<Double, Double> predictionAndLabel =
+      test.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
+        @Override
+        public Tuple2<Double, Double> call(LabeledPoint p) {
+          return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
+        }
+      });
+    double accuracy = predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
+      @Override
+      public Boolean call(Tuple2<Double, Double> pl) {
+        return pl._1().equals(pl._2());
+      }
+    }).count() / (double) test.count();
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myNaiveBayesModel");
+    NaiveBayesModel sameModel = NaiveBayesModel.load(jsc.sc(), "target/tmp/myNaiveBayesModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/python/mllib/isotonic_regression_example.py b/examples/src/main/python/mllib/isotonic_regression_example.py
new file mode 100644
index 0000000000000..89dc9f4b6611a
--- /dev/null
+++ b/examples/src/main/python/mllib/isotonic_regression_example.py
@@ -0,0 +1,56 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Isotonic Regression Example.
+"""
+from __future__ import print_function
+
+from pyspark import SparkContext
+# $example on$
+import math
+from pyspark.mllib.regression import IsotonicRegression, IsotonicRegressionModel
+# $example off$
+
+if __name__ == "__main__":
+
+    sc = SparkContext(appName="PythonIsotonicRegressionExample")
+
+    # $example on$
+    data = sc.textFile("data/mllib/sample_isotonic_regression_data.txt")
+
+    # Create label, feature, weight tuples from input data with weight set to default value 1.0.
+    parsedData = data.map(lambda line: tuple([float(x) for x in line.split(',')]) + (1.0,))
+
+    # Split data into training (60%) and test (40%) sets.
+    training, test = parsedData.randomSplit([0.6, 0.4], 11)
+
+    # Create isotonic regression model from training data.
+    # Isotonic parameter defaults to true so it is only shown for demonstration
+    model = IsotonicRegression.train(training)
+
+    # Create tuples of predicted and real labels.
+    predictionAndLabel = test.map(lambda p: (model.predict(p[1]), p[0]))
+
+    # Calculate mean squared error between predicted and real labels.
+    meanSquaredError = predictionAndLabel.map(lambda pl: math.pow((pl[0] - pl[1]), 2)).mean()
+    print("Mean Squared Error = " + str(meanSquaredError))
+
+    # Save and load model
+    model.save(sc, "target/tmp/myIsotonicRegressionModel")
+    sameModel = IsotonicRegressionModel.load(sc, "target/tmp/myIsotonicRegressionModel")
+    # $example off$
diff --git a/examples/src/main/python/mllib/naive_bayes_example.py b/examples/src/main/python/mllib/naive_bayes_example.py
new file mode 100644
index 0000000000000..a2e7dacf25491
--- /dev/null
+++ b/examples/src/main/python/mllib/naive_bayes_example.py
@@ -0,0 +1,56 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+NaiveBayes Example.
+"""
+from __future__ import print_function
+
+# $example on$
+from pyspark.mllib.classification import NaiveBayes, NaiveBayesModel
+from pyspark.mllib.linalg import Vectors
+from pyspark.mllib.regression import LabeledPoint
+
+
+def parseLine(line):
+    parts = line.split(',')
+    label = float(parts[0])
+    features = Vectors.dense([float(x) for x in parts[1].split(' ')])
+    return LabeledPoint(label, features)
+# $example off$
+
+if __name__ == "__main__":
+
+    sc = SparkContext(appName="PythonNaiveBayesExample")
+
+    # $example on$
+    data = sc.textFile('data/mllib/sample_naive_bayes_data.txt').map(parseLine)
+
+    # Split data aproximately into training (60%) and test (40%)
+    training, test = data.randomSplit([0.6, 0.4], seed=0)
+
+    # Train a naive Bayes model.
+    model = NaiveBayes.train(training, 1.0)
+
+    # Make prediction and test accuracy.
+    predictionAndLabel = test.map(lambda p: (model.predict(p.features), p.label))
+    accuracy = 1.0 * predictionAndLabel.filter(lambda (x, v): x == v).count() / test.count()
+
+    # Save and load model
+    model.save(sc, "target/tmp/myNaiveBayesModel")
+    sameModel = NaiveBayesModel.load(sc, "target/tmp/myNaiveBayesModel")
+    # $example off$
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/IsotonicRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/IsotonicRegressionExample.scala
new file mode 100644
index 0000000000000..52ac9ae7dd2d0
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/IsotonicRegressionExample.scala
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.regression.{IsotonicRegression, IsotonicRegressionModel}
+// $example off$
+import org.apache.spark.{SparkConf, SparkContext}
+
+object IsotonicRegressionExample {
+
+  def main(args: Array[String]) : Unit = {
+
+    val conf = new SparkConf().setAppName("IsotonicRegressionExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    val data = sc.textFile("data/mllib/sample_isotonic_regression_data.txt")
+
+    // Create label, feature, weight tuples from input data with weight set to default value 1.0.
+    val parsedData = data.map { line =>
+      val parts = line.split(',').map(_.toDouble)
+      (parts(0), parts(1), 1.0)
+    }
+
+    // Split data into training (60%) and test (40%) sets.
+    val splits = parsedData.randomSplit(Array(0.6, 0.4), seed = 11L)
+    val training = splits(0)
+    val test = splits(1)
+
+    // Create isotonic regression model from training data.
+    // Isotonic parameter defaults to true so it is only shown for demonstration
+    val model = new IsotonicRegression().setIsotonic(true).run(training)
+
+    // Create tuples of predicted and real labels.
+    val predictionAndLabel = test.map { point =>
+      val predictedLabel = model.predict(point._2)
+      (predictedLabel, point._1)
+    }
+
+    // Calculate mean squared error between predicted and real labels.
+    val meanSquaredError = predictionAndLabel.map { case (p, l) => math.pow((p - l), 2) }.mean()
+    println("Mean Squared Error = " + meanSquaredError)
+
+    // Save and load model
+    model.save(sc, "target/tmp/myIsotonicRegressionModel")
+    val sameModel = IsotonicRegressionModel.load(sc, "target/tmp/myIsotonicRegressionModel")
+    // $example off$
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/NaiveBayesExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/NaiveBayesExample.scala
new file mode 100644
index 0000000000000..a7a47c2a3556a
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/NaiveBayesExample.scala
@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.classification.{NaiveBayes, NaiveBayesModel}
+import org.apache.spark.mllib.linalg.Vectors
+import org.apache.spark.mllib.regression.LabeledPoint
+// $example off$
+import org.apache.spark.{SparkConf, SparkContext}
+
+object NaiveBayesExample {
+
+  def main(args: Array[String]) : Unit = {
+    val conf = new SparkConf().setAppName("NaiveBayesExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    val data = sc.textFile("data/mllib/sample_naive_bayes_data.txt")
+    val parsedData = data.map { line =>
+      val parts = line.split(',')
+      LabeledPoint(parts(0).toDouble, Vectors.dense(parts(1).split(' ').map(_.toDouble)))
+    }
+
+    // Split data into training (60%) and test (40%).
+    val splits = parsedData.randomSplit(Array(0.6, 0.4), seed = 11L)
+    val training = splits(0)
+    val test = splits(1)
+
+    val model = NaiveBayes.train(training, lambda = 1.0, modelType = "multinomial")
+
+    val predictionAndLabel = test.map(p => (model.predict(p.features), p.label))
+    val accuracy = 1.0 * predictionAndLabel.filter(x => x._1 == x._2).count() / test.count()
+
+    // Save and load model
+    model.save(sc, "target/tmp/myNaiveBayesModel")
+    val sameModel = NaiveBayesModel.load(sc, "target/tmp/myNaiveBayesModel")
+    // $example off$
+  }
+}
+
+// scalastyle:on println

From ecfb3e73fd0a99f0be96034710974e78b6f9d624 Mon Sep 17 00:00:00 2001
From: lihao <lihaowhu@gmail.com>
Date: Mon, 2 Nov 2015 16:09:22 -0800
Subject: [PATCH 010/518] [SPARK-10286][ML][PYSPARK][DOCS] Add @since
 annotation to pyspark.ml.param and pyspark.ml.*

Author: lihao <lihaowhu@gmail.com>

Closes #9275 from lidinghao/SPARK-10286.
---
 python/pyspark/ml/evaluation.py     |  20 ++++
 python/pyspark/ml/feature.py        | 164 ++++++++++++++++++++++++++++
 python/pyspark/ml/param/__init__.py |  16 +++
 python/pyspark/ml/pipeline.py       |  30 +++++
 4 files changed, 230 insertions(+)

diff --git a/python/pyspark/ml/evaluation.py b/python/pyspark/ml/evaluation.py
index cb3b07947e488..dcc1738ec518b 100644
--- a/python/pyspark/ml/evaluation.py
+++ b/python/pyspark/ml/evaluation.py
@@ -17,6 +17,7 @@
 
 from abc import abstractmethod, ABCMeta
 
+from pyspark import since
 from pyspark.ml.wrapper import JavaWrapper
 from pyspark.ml.param import Param, Params
 from pyspark.ml.param.shared import HasLabelCol, HasPredictionCol, HasRawPredictionCol
@@ -31,6 +32,8 @@
 class Evaluator(Params):
     """
     Base class for evaluators that compute metrics from predictions.
+
+    .. versionadded:: 1.4.0
     """
 
     __metaclass__ = ABCMeta
@@ -46,6 +49,7 @@ def _evaluate(self, dataset):
         """
         raise NotImplementedError()
 
+    @since("1.4.0")
     def evaluate(self, dataset, params=None):
         """
         Evaluates the output with optional parameters.
@@ -66,6 +70,7 @@ def evaluate(self, dataset, params=None):
         else:
             raise ValueError("Params must be a param map but got %s." % type(params))
 
+    @since("1.5.0")
     def isLargerBetter(self):
         """
         Indicates whether the metric returned by :py:meth:`evaluate` should be maximized
@@ -114,6 +119,8 @@ class BinaryClassificationEvaluator(JavaEvaluator, HasLabelCol, HasRawPrediction
     0.70...
     >>> evaluator.evaluate(dataset, {evaluator.metricName: "areaUnderPR"})
     0.83...
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -138,6 +145,7 @@ def __init__(self, rawPredictionCol="rawPrediction", labelCol="label",
         kwargs = self.__init__._input_kwargs
         self._set(**kwargs)
 
+    @since("1.4.0")
     def setMetricName(self, value):
         """
         Sets the value of :py:attr:`metricName`.
@@ -145,6 +153,7 @@ def setMetricName(self, value):
         self._paramMap[self.metricName] = value
         return self
 
+    @since("1.4.0")
     def getMetricName(self):
         """
         Gets the value of metricName or its default value.
@@ -152,6 +161,7 @@ def getMetricName(self):
         return self.getOrDefault(self.metricName)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, rawPredictionCol="rawPrediction", labelCol="label",
                   metricName="areaUnderROC"):
         """
@@ -180,6 +190,8 @@ class RegressionEvaluator(JavaEvaluator, HasLabelCol, HasPredictionCol):
     0.993...
     >>> evaluator.evaluate(dataset, {evaluator.metricName: "mae"})
     2.649...
+
+    .. versionadded:: 1.4.0
     """
     # Because we will maximize evaluation value (ref: `CrossValidator`),
     # when we evaluate a metric that is needed to minimize (e.g., `"rmse"`, `"mse"`, `"mae"`),
@@ -205,6 +217,7 @@ def __init__(self, predictionCol="prediction", labelCol="label",
         kwargs = self.__init__._input_kwargs
         self._set(**kwargs)
 
+    @since("1.4.0")
     def setMetricName(self, value):
         """
         Sets the value of :py:attr:`metricName`.
@@ -212,6 +225,7 @@ def setMetricName(self, value):
         self._paramMap[self.metricName] = value
         return self
 
+    @since("1.4.0")
     def getMetricName(self):
         """
         Gets the value of metricName or its default value.
@@ -219,6 +233,7 @@ def getMetricName(self):
         return self.getOrDefault(self.metricName)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, predictionCol="prediction", labelCol="label",
                   metricName="rmse"):
         """
@@ -246,6 +261,8 @@ class MulticlassClassificationEvaluator(JavaEvaluator, HasLabelCol, HasPredictio
     0.66...
     >>> evaluator.evaluate(dataset, {evaluator.metricName: "recall"})
     0.66...
+
+    .. versionadded:: 1.5.0
     """
     # a placeholder to make it appear in the generated doc
     metricName = Param(Params._dummy(), "metricName",
@@ -271,6 +288,7 @@ def __init__(self, predictionCol="prediction", labelCol="label",
         kwargs = self.__init__._input_kwargs
         self._set(**kwargs)
 
+    @since("1.5.0")
     def setMetricName(self, value):
         """
         Sets the value of :py:attr:`metricName`.
@@ -278,6 +296,7 @@ def setMetricName(self, value):
         self._paramMap[self.metricName] = value
         return self
 
+    @since("1.5.0")
     def getMetricName(self):
         """
         Gets the value of metricName or its default value.
@@ -285,6 +304,7 @@ def getMetricName(self):
         return self.getOrDefault(self.metricName)
 
     @keyword_only
+    @since("1.5.0")
     def setParams(self, predictionCol="prediction", labelCol="label",
                   metricName="f1"):
         """
diff --git a/python/pyspark/ml/feature.py b/python/pyspark/ml/feature.py
index 55bde6d0ea4fb..c7b6dd926c3e8 100644
--- a/python/pyspark/ml/feature.py
+++ b/python/pyspark/ml/feature.py
@@ -19,6 +19,7 @@
 if sys.version > '3':
     basestring = str
 
+from pyspark import since
 from pyspark.rdd import ignore_unicode_prefix
 from pyspark.ml.param.shared import *
 from pyspark.ml.util import keyword_only
@@ -51,6 +52,8 @@ class Binarizer(JavaTransformer, HasInputCol, HasOutputCol):
     >>> params = {binarizer.threshold: -0.5, binarizer.outputCol: "vector"}
     >>> binarizer.transform(df, params).head().vector
     1.0
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -71,6 +74,7 @@ def __init__(self, threshold=0.0, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, threshold=0.0, inputCol=None, outputCol=None):
         """
         setParams(self, threshold=0.0, inputCol=None, outputCol=None)
@@ -79,6 +83,7 @@ def setParams(self, threshold=0.0, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setThreshold(self, value):
         """
         Sets the value of :py:attr:`threshold`.
@@ -86,6 +91,7 @@ def setThreshold(self, value):
         self._paramMap[self.threshold] = value
         return self
 
+    @since("1.4.0")
     def getThreshold(self):
         """
         Gets the value of threshold or its default value.
@@ -114,6 +120,8 @@ class Bucketizer(JavaTransformer, HasInputCol, HasOutputCol):
     2.0
     >>> bucketizer.setParams(outputCol="b").transform(df).head().b
     0.0
+
+    .. versionadded:: 1.3.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -150,6 +158,7 @@ def __init__(self, splits=None, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, splits=None, inputCol=None, outputCol=None):
         """
         setParams(self, splits=None, inputCol=None, outputCol=None)
@@ -158,6 +167,7 @@ def setParams(self, splits=None, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setSplits(self, value):
         """
         Sets the value of :py:attr:`splits`.
@@ -165,6 +175,7 @@ def setSplits(self, value):
         self._paramMap[self.splits] = value
         return self
 
+    @since("1.4.0")
     def getSplits(self):
         """
         Gets the value of threshold or its default value.
@@ -194,6 +205,8 @@ class CountVectorizer(JavaEstimator, HasInputCol, HasOutputCol):
     ...
     >>> sorted(map(str, model.vocabulary))
     ['a', 'b', 'c']
+
+    .. versionadded:: 1.6.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -242,6 +255,7 @@ def __init__(self, minTF=1.0, minDF=1.0, vocabSize=1 << 18, inputCol=None, outpu
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.6.0")
     def setParams(self, minTF=1.0, minDF=1.0, vocabSize=1 << 18, inputCol=None, outputCol=None):
         """
         setParams(self, minTF=1.0, minDF=1.0, vocabSize=1 << 18, inputCol=None, outputCol=None)
@@ -250,6 +264,7 @@ def setParams(self, minTF=1.0, minDF=1.0, vocabSize=1 << 18, inputCol=None, outp
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.6.0")
     def setMinTF(self, value):
         """
         Sets the value of :py:attr:`minTF`.
@@ -257,12 +272,14 @@ def setMinTF(self, value):
         self._paramMap[self.minTF] = value
         return self
 
+    @since("1.6.0")
     def getMinTF(self):
         """
         Gets the value of minTF or its default value.
         """
         return self.getOrDefault(self.minTF)
 
+    @since("1.6.0")
     def setMinDF(self, value):
         """
         Sets the value of :py:attr:`minDF`.
@@ -270,12 +287,14 @@ def setMinDF(self, value):
         self._paramMap[self.minDF] = value
         return self
 
+    @since("1.6.0")
     def getMinDF(self):
         """
         Gets the value of minDF or its default value.
         """
         return self.getOrDefault(self.minDF)
 
+    @since("1.6.0")
     def setVocabSize(self, value):
         """
         Sets the value of :py:attr:`vocabSize`.
@@ -283,6 +302,7 @@ def setVocabSize(self, value):
         self._paramMap[self.vocabSize] = value
         return self
 
+    @since("1.6.0")
     def getVocabSize(self):
         """
         Gets the value of vocabSize or its default value.
@@ -298,9 +318,12 @@ class CountVectorizerModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by CountVectorizer.
+
+    .. versionadded:: 1.6.0
     """
 
     @property
+    @since("1.6.0")
     def vocabulary(self):
         """
         An array of terms in the vocabulary.
@@ -331,6 +354,8 @@ class DCT(JavaTransformer, HasInputCol, HasOutputCol):
     >>> df3 = DCT(inverse=True, inputCol="resultVec", outputCol="origVec").transform(df2)
     >>> df3.head().origVec
     DenseVector([5.0, 8.0, 6.0])
+
+    .. versionadded:: 1.6.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -351,6 +376,7 @@ def __init__(self, inverse=False, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.6.0")
     def setParams(self, inverse=False, inputCol=None, outputCol=None):
         """
         setParams(self, inverse=False, inputCol=None, outputCol=None)
@@ -359,6 +385,7 @@ def setParams(self, inverse=False, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.6.0")
     def setInverse(self, value):
         """
         Sets the value of :py:attr:`inverse`.
@@ -366,6 +393,7 @@ def setInverse(self, value):
         self._paramMap[self.inverse] = value
         return self
 
+    @since("1.6.0")
     def getInverse(self):
         """
         Gets the value of inverse or its default value.
@@ -390,6 +418,8 @@ class ElementwiseProduct(JavaTransformer, HasInputCol, HasOutputCol):
     DenseVector([2.0, 2.0, 9.0])
     >>> ep.setParams(scalingVec=Vectors.dense([2.0, 3.0, 5.0])).transform(df).head().eprod
     DenseVector([4.0, 3.0, 15.0])
+
+    .. versionadded:: 1.5.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -410,6 +440,7 @@ def __init__(self, scalingVec=None, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.5.0")
     def setParams(self, scalingVec=None, inputCol=None, outputCol=None):
         """
         setParams(self, scalingVec=None, inputCol=None, outputCol=None)
@@ -418,6 +449,7 @@ def setParams(self, scalingVec=None, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.5.0")
     def setScalingVec(self, value):
         """
         Sets the value of :py:attr:`scalingVec`.
@@ -425,6 +457,7 @@ def setScalingVec(self, value):
         self._paramMap[self.scalingVec] = value
         return self
 
+    @since("1.5.0")
     def getScalingVec(self):
         """
         Gets the value of scalingVec or its default value.
@@ -449,6 +482,8 @@ class HashingTF(JavaTransformer, HasInputCol, HasOutputCol, HasNumFeatures):
     >>> params = {hashingTF.numFeatures: 5, hashingTF.outputCol: "vector"}
     >>> hashingTF.transform(df, params).head().vector
     SparseVector(5, {2: 1.0, 3: 1.0, 4: 1.0})
+
+    .. versionadded:: 1.3.0
     """
 
     @keyword_only
@@ -463,6 +498,7 @@ def __init__(self, numFeatures=1 << 18, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.3.0")
     def setParams(self, numFeatures=1 << 18, inputCol=None, outputCol=None):
         """
         setParams(self, numFeatures=1 << 18, inputCol=None, outputCol=None)
@@ -490,6 +526,8 @@ class IDF(JavaEstimator, HasInputCol, HasOutputCol):
     >>> params = {idf.minDocFreq: 1, idf.outputCol: "vector"}
     >>> idf.fit(df, params).transform(df).head().vector
     DenseVector([0.2877, 0.0])
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -510,6 +548,7 @@ def __init__(self, minDocFreq=0, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, minDocFreq=0, inputCol=None, outputCol=None):
         """
         setParams(self, minDocFreq=0, inputCol=None, outputCol=None)
@@ -518,6 +557,7 @@ def setParams(self, minDocFreq=0, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setMinDocFreq(self, value):
         """
         Sets the value of :py:attr:`minDocFreq`.
@@ -525,6 +565,7 @@ def setMinDocFreq(self, value):
         self._paramMap[self.minDocFreq] = value
         return self
 
+    @since("1.4.0")
     def getMinDocFreq(self):
         """
         Gets the value of minDocFreq or its default value.
@@ -540,6 +581,8 @@ class IDFModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by IDF.
+
+    .. versionadded:: 1.4.0
     """
 
 
@@ -571,6 +614,8 @@ class MinMaxScaler(JavaEstimator, HasInputCol, HasOutputCol):
     |[2.0]| [1.0]|
     +-----+------+
     ...
+
+    .. versionadded:: 1.6.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -591,6 +636,7 @@ def __init__(self, min=0.0, max=1.0, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.6.0")
     def setParams(self, min=0.0, max=1.0, inputCol=None, outputCol=None):
         """
         setParams(self, min=0.0, max=1.0, inputCol=None, outputCol=None)
@@ -599,6 +645,7 @@ def setParams(self, min=0.0, max=1.0, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.6.0")
     def setMin(self, value):
         """
         Sets the value of :py:attr:`min`.
@@ -606,12 +653,14 @@ def setMin(self, value):
         self._paramMap[self.min] = value
         return self
 
+    @since("1.6.0")
     def getMin(self):
         """
         Gets the value of min or its default value.
         """
         return self.getOrDefault(self.min)
 
+    @since("1.6.0")
     def setMax(self, value):
         """
         Sets the value of :py:attr:`max`.
@@ -619,6 +668,7 @@ def setMax(self, value):
         self._paramMap[self.max] = value
         return self
 
+    @since("1.6.0")
     def getMax(self):
         """
         Gets the value of max or its default value.
@@ -634,6 +684,8 @@ class MinMaxScalerModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by :py:class:`MinMaxScaler`.
+
+    .. versionadded:: 1.6.0
     """
 
 
@@ -668,6 +720,8 @@ class NGram(JavaTransformer, HasInputCol, HasOutputCol):
     Traceback (most recent call last):
         ...
     TypeError: Method setParams forces keyword arguments.
+
+    .. versionadded:: 1.5.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -686,6 +740,7 @@ def __init__(self, n=2, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.5.0")
     def setParams(self, n=2, inputCol=None, outputCol=None):
         """
         setParams(self, n=2, inputCol=None, outputCol=None)
@@ -694,6 +749,7 @@ def setParams(self, n=2, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.5.0")
     def setN(self, value):
         """
         Sets the value of :py:attr:`n`.
@@ -701,6 +757,7 @@ def setN(self, value):
         self._paramMap[self.n] = value
         return self
 
+    @since("1.5.0")
     def getN(self):
         """
         Gets the value of n or its default value.
@@ -726,6 +783,8 @@ class Normalizer(JavaTransformer, HasInputCol, HasOutputCol):
     >>> params = {normalizer.p: 1.0, normalizer.inputCol: "dense", normalizer.outputCol: "vector"}
     >>> normalizer.transform(df, params).head().vector
     DenseVector([0.4286, -0.5714])
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -744,6 +803,7 @@ def __init__(self, p=2.0, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, p=2.0, inputCol=None, outputCol=None):
         """
         setParams(self, p=2.0, inputCol=None, outputCol=None)
@@ -752,6 +812,7 @@ def setParams(self, p=2.0, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setP(self, value):
         """
         Sets the value of :py:attr:`p`.
@@ -759,6 +820,7 @@ def setP(self, value):
         self._paramMap[self.p] = value
         return self
 
+    @since("1.4.0")
     def getP(self):
         """
         Gets the value of p or its default value.
@@ -800,6 +862,8 @@ class OneHotEncoder(JavaTransformer, HasInputCol, HasOutputCol):
     >>> params = {encoder.dropLast: False, encoder.outputCol: "test"}
     >>> encoder.transform(td, params).head().test
     SparseVector(3, {0: 1.0})
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -818,6 +882,7 @@ def __init__(self, dropLast=True, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, dropLast=True, inputCol=None, outputCol=None):
         """
         setParams(self, dropLast=True, inputCol=None, outputCol=None)
@@ -826,6 +891,7 @@ def setParams(self, dropLast=True, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setDropLast(self, value):
         """
         Sets the value of :py:attr:`dropLast`.
@@ -833,6 +899,7 @@ def setDropLast(self, value):
         self._paramMap[self.dropLast] = value
         return self
 
+    @since("1.4.0")
     def getDropLast(self):
         """
         Gets the value of dropLast or its default value.
@@ -858,6 +925,8 @@ class PolynomialExpansion(JavaTransformer, HasInputCol, HasOutputCol):
     DenseVector([0.5, 0.25, 2.0, 1.0, 4.0])
     >>> px.setParams(outputCol="test").transform(df).head().test
     DenseVector([0.5, 0.25, 2.0, 1.0, 4.0])
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -877,6 +946,7 @@ def __init__(self, degree=2, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, degree=2, inputCol=None, outputCol=None):
         """
         setParams(self, degree=2, inputCol=None, outputCol=None)
@@ -885,6 +955,7 @@ def setParams(self, degree=2, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setDegree(self, value):
         """
         Sets the value of :py:attr:`degree`.
@@ -892,6 +963,7 @@ def setDegree(self, value):
         self._paramMap[self.degree] = value
         return self
 
+    @since("1.4.0")
     def getDegree(self):
         """
         Gets the value of degree or its default value.
@@ -929,6 +1001,8 @@ class RegexTokenizer(JavaTransformer, HasInputCol, HasOutputCol):
     Traceback (most recent call last):
         ...
     TypeError: Method setParams forces keyword arguments.
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -951,6 +1025,7 @@ def __init__(self, minTokenLength=1, gaps=True, pattern="\\s+", inputCol=None, o
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, minTokenLength=1, gaps=True, pattern="\\s+", inputCol=None, outputCol=None):
         """
         setParams(self, minTokenLength=1, gaps=True, pattern="\\s+", inputCol=None, outputCol=None)
@@ -959,6 +1034,7 @@ def setParams(self, minTokenLength=1, gaps=True, pattern="\\s+", inputCol=None,
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setMinTokenLength(self, value):
         """
         Sets the value of :py:attr:`minTokenLength`.
@@ -966,12 +1042,14 @@ def setMinTokenLength(self, value):
         self._paramMap[self.minTokenLength] = value
         return self
 
+    @since("1.4.0")
     def getMinTokenLength(self):
         """
         Gets the value of minTokenLength or its default value.
         """
         return self.getOrDefault(self.minTokenLength)
 
+    @since("1.4.0")
     def setGaps(self, value):
         """
         Sets the value of :py:attr:`gaps`.
@@ -979,12 +1057,14 @@ def setGaps(self, value):
         self._paramMap[self.gaps] = value
         return self
 
+    @since("1.4.0")
     def getGaps(self):
         """
         Gets the value of gaps or its default value.
         """
         return self.getOrDefault(self.gaps)
 
+    @since("1.4.0")
     def setPattern(self, value):
         """
         Sets the value of :py:attr:`pattern`.
@@ -992,6 +1072,7 @@ def setPattern(self, value):
         self._paramMap[self.pattern] = value
         return self
 
+    @since("1.4.0")
     def getPattern(self):
         """
         Gets the value of pattern or its default value.
@@ -1013,6 +1094,8 @@ class SQLTransformer(JavaTransformer):
     ...     statement="SELECT *, (v1 + v2) AS v3, (v1 * v2) AS v4 FROM __THIS__")
     >>> sqlTrans.transform(df).head()
     Row(id=0, v1=1.0, v2=3.0, v3=4.0, v4=3.0)
+
+    .. versionadded:: 1.6.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -1030,6 +1113,7 @@ def __init__(self, statement=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.6.0")
     def setParams(self, statement=None):
         """
         setParams(self, statement=None)
@@ -1038,6 +1122,7 @@ def setParams(self, statement=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.6.0")
     def setStatement(self, value):
         """
         Sets the value of :py:attr:`statement`.
@@ -1045,6 +1130,7 @@ def setStatement(self, value):
         self._paramMap[self.statement] = value
         return self
 
+    @since("1.6.0")
     def getStatement(self):
         """
         Gets the value of statement or its default value.
@@ -1070,6 +1156,8 @@ class StandardScaler(JavaEstimator, HasInputCol, HasOutputCol):
     DenseVector([1.4142])
     >>> model.transform(df).collect()[1].scaled
     DenseVector([1.4142])
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -1090,6 +1178,7 @@ def __init__(self, withMean=False, withStd=True, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, withMean=False, withStd=True, inputCol=None, outputCol=None):
         """
         setParams(self, withMean=False, withStd=True, inputCol=None, outputCol=None)
@@ -1098,6 +1187,7 @@ def setParams(self, withMean=False, withStd=True, inputCol=None, outputCol=None)
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setWithMean(self, value):
         """
         Sets the value of :py:attr:`withMean`.
@@ -1105,12 +1195,14 @@ def setWithMean(self, value):
         self._paramMap[self.withMean] = value
         return self
 
+    @since("1.4.0")
     def getWithMean(self):
         """
         Gets the value of withMean or its default value.
         """
         return self.getOrDefault(self.withMean)
 
+    @since("1.4.0")
     def setWithStd(self, value):
         """
         Sets the value of :py:attr:`withStd`.
@@ -1118,6 +1210,7 @@ def setWithStd(self, value):
         self._paramMap[self.withStd] = value
         return self
 
+    @since("1.4.0")
     def getWithStd(self):
         """
         Gets the value of withStd or its default value.
@@ -1133,9 +1226,12 @@ class StandardScalerModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by StandardScaler.
+
+    .. versionadded:: 1.4.0
     """
 
     @property
+    @since("1.5.0")
     def std(self):
         """
         Standard deviation of the StandardScalerModel.
@@ -1143,6 +1239,7 @@ def std(self):
         return self._call_java("std")
 
     @property
+    @since("1.5.0")
     def mean(self):
         """
         Mean of the StandardScalerModel.
@@ -1171,6 +1268,8 @@ class StringIndexer(JavaEstimator, HasInputCol, HasOutputCol, HasHandleInvalid):
     >>> sorted(set([(i[0], str(i[1])) for i in itd.select(itd.id, itd.label2).collect()]),
     ...     key=lambda x: x[0])
     [(0, 'a'), (1, 'b'), (2, 'c'), (3, 'a'), (4, 'a'), (5, 'c')]
+
+    .. versionadded:: 1.4.0
     """
 
     @keyword_only
@@ -1185,6 +1284,7 @@ def __init__(self, inputCol=None, outputCol=None, handleInvalid="error"):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, inputCol=None, outputCol=None, handleInvalid="error"):
         """
         setParams(self, inputCol=None, outputCol=None, handleInvalid="error")
@@ -1202,8 +1302,11 @@ class StringIndexerModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by StringIndexer.
+
+    .. versionadded:: 1.4.0
     """
     @property
+    @since("1.5.0")
     def labels(self):
         """
         Ordered list of labels, corresponding to indices to be assigned.
@@ -1221,6 +1324,8 @@ class IndexToString(JavaTransformer, HasInputCol, HasOutputCol):
     The index-string mapping is either from the ML attributes of the input column,
     or from user-supplied labels (which take precedence over ML attributes).
     See L{StringIndexer} for converting strings into indices.
+
+    .. versionadded:: 1.6.0
     """
 
     # a placeholder to make the labels show up in generated doc
@@ -1243,6 +1348,7 @@ def __init__(self, inputCol=None, outputCol=None, labels=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.6.0")
     def setParams(self, inputCol=None, outputCol=None, labels=None):
         """
         setParams(self, inputCol=None, outputCol=None, labels=None)
@@ -1251,6 +1357,7 @@ def setParams(self, inputCol=None, outputCol=None, labels=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.6.0")
     def setLabels(self, value):
         """
         Sets the value of :py:attr:`labels`.
@@ -1258,6 +1365,7 @@ def setLabels(self, value):
         self._paramMap[self.labels] = value
         return self
 
+    @since("1.6.0")
     def getLabels(self):
         """
         Gets the value of :py:attr:`labels` or its default value.
@@ -1271,6 +1379,8 @@ class StopWordsRemover(JavaTransformer, HasInputCol, HasOutputCol):
 
     A feature transformer that filters out stop words from input.
     Note: null values from input array are preserved unless adding null to stopWords explicitly.
+
+    .. versionadded:: 1.6.0
     """
     # a placeholder to make the stopwords show up in generated doc
     stopWords = Param(Params._dummy(), "stopWords", "The words to be filtered out")
@@ -1297,6 +1407,7 @@ def __init__(self, inputCol=None, outputCol=None, stopWords=None,
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.6.0")
     def setParams(self, inputCol=None, outputCol=None, stopWords=None,
                   caseSensitive=False):
         """
@@ -1307,6 +1418,7 @@ def setParams(self, inputCol=None, outputCol=None, stopWords=None,
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.6.0")
     def setStopWords(self, value):
         """
         Specify the stopwords to be filtered.
@@ -1314,12 +1426,14 @@ def setStopWords(self, value):
         self._paramMap[self.stopWords] = value
         return self
 
+    @since("1.6.0")
     def getStopWords(self):
         """
         Get the stopwords.
         """
         return self.getOrDefault(self.stopWords)
 
+    @since("1.6.0")
     def setCaseSensitive(self, value):
         """
         Set whether to do a case sensitive comparison over the stop words
@@ -1327,6 +1441,7 @@ def setCaseSensitive(self, value):
         self._paramMap[self.caseSensitive] = value
         return self
 
+    @since("1.6.0")
     def getCaseSensitive(self):
         """
         Get whether to do a case sensitive comparison over the stop words.
@@ -1360,6 +1475,8 @@ class Tokenizer(JavaTransformer, HasInputCol, HasOutputCol):
     Traceback (most recent call last):
         ...
     TypeError: Method setParams forces keyword arguments.
+
+    .. versionadded:: 1.3.0
     """
 
     @keyword_only
@@ -1373,6 +1490,7 @@ def __init__(self, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.3.0")
     def setParams(self, inputCol=None, outputCol=None):
         """
         setParams(self, inputCol="input", outputCol="output")
@@ -1398,6 +1516,8 @@ class VectorAssembler(JavaTransformer, HasInputCols, HasOutputCol):
     >>> params = {vecAssembler.inputCols: ["b", "a"], vecAssembler.outputCol: "vector"}
     >>> vecAssembler.transform(df, params).head().vector
     DenseVector([0.0, 1.0])
+
+    .. versionadded:: 1.4.0
     """
 
     @keyword_only
@@ -1411,6 +1531,7 @@ def __init__(self, inputCols=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, inputCols=None, outputCol=None):
         """
         setParams(self, inputCols=None, outputCol=None)
@@ -1477,6 +1598,8 @@ class VectorIndexer(JavaEstimator, HasInputCol, HasOutputCol):
     >>> model2 = indexer.fit(df, params)
     >>> model2.transform(df).head().vector
     DenseVector([1.0, 0.0])
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -1501,6 +1624,7 @@ def __init__(self, maxCategories=20, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, maxCategories=20, inputCol=None, outputCol=None):
         """
         setParams(self, maxCategories=20, inputCol=None, outputCol=None)
@@ -1509,6 +1633,7 @@ def setParams(self, maxCategories=20, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setMaxCategories(self, value):
         """
         Sets the value of :py:attr:`maxCategories`.
@@ -1516,6 +1641,7 @@ def setMaxCategories(self, value):
         self._paramMap[self.maxCategories] = value
         return self
 
+    @since("1.4.0")
     def getMaxCategories(self):
         """
         Gets the value of maxCategories or its default value.
@@ -1531,9 +1657,12 @@ class VectorIndexerModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by VectorIndexer.
+
+    .. versionadded:: 1.4.0
     """
 
     @property
+    @since("1.4.0")
     def numFeatures(self):
         """
         Number of features, i.e., length of Vectors which this transforms.
@@ -1541,6 +1670,7 @@ def numFeatures(self):
         return self._call_java("numFeatures")
 
     @property
+    @since("1.4.0")
     def categoryMaps(self):
         """
         Feature value index.  Keys are categorical feature indices (column indices).
@@ -1573,6 +1703,8 @@ class VectorSlicer(JavaTransformer, HasInputCol, HasOutputCol):
     >>> vs = VectorSlicer(inputCol="features", outputCol="sliced", indices=[1, 4])
     >>> vs.transform(df).head().sliced
     DenseVector([2.3, 1.0])
+
+    .. versionadded:: 1.6.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -1600,6 +1732,7 @@ def __init__(self, inputCol=None, outputCol=None, indices=None, names=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.6.0")
     def setParams(self, inputCol=None, outputCol=None, indices=None, names=None):
         """
         setParams(self, inputCol=None, outputCol=None, indices=None, names=None):
@@ -1608,6 +1741,7 @@ def setParams(self, inputCol=None, outputCol=None, indices=None, names=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.6.0")
     def setIndices(self, value):
         """
         Sets the value of :py:attr:`indices`.
@@ -1615,12 +1749,14 @@ def setIndices(self, value):
         self._paramMap[self.indices] = value
         return self
 
+    @since("1.6.0")
     def getIndices(self):
         """
         Gets the value of indices or its default value.
         """
         return self.getOrDefault(self.indices)
 
+    @since("1.6.0")
     def setNames(self, value):
         """
         Sets the value of :py:attr:`names`.
@@ -1628,6 +1764,7 @@ def setNames(self, value):
         self._paramMap[self.names] = value
         return self
 
+    @since("1.6.0")
     def getNames(self):
         """
         Gets the value of names or its default value.
@@ -1666,6 +1803,8 @@ class Word2Vec(JavaEstimator, HasStepSize, HasMaxIter, HasSeed, HasInputCol, Has
     ...
     >>> model.transform(doc).head().model
     DenseVector([-0.0422, -0.5138, -0.2546, 0.6885, 0.276])
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -1699,6 +1838,7 @@ def __init__(self, vectorSize=100, minCount=5, numPartitions=1, stepSize=0.025,
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, vectorSize=100, minCount=5, numPartitions=1, stepSize=0.025, maxIter=1,
                   seed=None, inputCol=None, outputCol=None):
         """
@@ -1709,6 +1849,7 @@ def setParams(self, vectorSize=100, minCount=5, numPartitions=1, stepSize=0.025,
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.4.0")
     def setVectorSize(self, value):
         """
         Sets the value of :py:attr:`vectorSize`.
@@ -1716,12 +1857,14 @@ def setVectorSize(self, value):
         self._paramMap[self.vectorSize] = value
         return self
 
+    @since("1.4.0")
     def getVectorSize(self):
         """
         Gets the value of vectorSize or its default value.
         """
         return self.getOrDefault(self.vectorSize)
 
+    @since("1.4.0")
     def setNumPartitions(self, value):
         """
         Sets the value of :py:attr:`numPartitions`.
@@ -1729,12 +1872,14 @@ def setNumPartitions(self, value):
         self._paramMap[self.numPartitions] = value
         return self
 
+    @since("1.4.0")
     def getNumPartitions(self):
         """
         Gets the value of numPartitions or its default value.
         """
         return self.getOrDefault(self.numPartitions)
 
+    @since("1.4.0")
     def setMinCount(self, value):
         """
         Sets the value of :py:attr:`minCount`.
@@ -1742,6 +1887,7 @@ def setMinCount(self, value):
         self._paramMap[self.minCount] = value
         return self
 
+    @since("1.4.0")
     def getMinCount(self):
         """
         Gets the value of minCount or its default value.
@@ -1757,8 +1903,11 @@ class Word2VecModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by Word2Vec.
+
+    .. versionadded:: 1.4.0
     """
 
+    @since("1.5.0")
     def getVectors(self):
         """
         Returns the vector representation of the words as a dataframe
@@ -1766,6 +1915,7 @@ def getVectors(self):
         """
         return self._call_java("getVectors")
 
+    @since("1.5.0")
     def findSynonyms(self, word, num):
         """
         Find "num" number of words closest in similarity to "word".
@@ -1794,6 +1944,8 @@ class PCA(JavaEstimator, HasInputCol, HasOutputCol):
     >>> model = pca.fit(df)
     >>> model.transform(df).collect()[0].pca_features
     DenseVector([1.648..., -4.013...])
+
+    .. versionadded:: 1.5.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -1811,6 +1963,7 @@ def __init__(self, k=None, inputCol=None, outputCol=None):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.5.0")
     def setParams(self, k=None, inputCol=None, outputCol=None):
         """
         setParams(self, k=None, inputCol=None, outputCol=None)
@@ -1819,6 +1972,7 @@ def setParams(self, k=None, inputCol=None, outputCol=None):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.5.0")
     def setK(self, value):
         """
         Sets the value of :py:attr:`k`.
@@ -1826,6 +1980,7 @@ def setK(self, value):
         self._paramMap[self.k] = value
         return self
 
+    @since("1.5.0")
     def getK(self):
         """
         Gets the value of k or its default value.
@@ -1841,6 +1996,8 @@ class PCAModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by PCA.
+
+    .. versionadded:: 1.5.0
     """
 
 
@@ -1879,6 +2036,8 @@ class RFormula(JavaEstimator, HasFeaturesCol, HasLabelCol):
     |0.0|0.0|  a|   [0.0]|  0.0|
     +---+---+---+--------+-----+
     ...
+
+    .. versionadded:: 1.5.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -1896,6 +2055,7 @@ def __init__(self, formula=None, featuresCol="features", labelCol="label"):
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.5.0")
     def setParams(self, formula=None, featuresCol="features", labelCol="label"):
         """
         setParams(self, formula=None, featuresCol="features", labelCol="label")
@@ -1904,6 +2064,7 @@ def setParams(self, formula=None, featuresCol="features", labelCol="label"):
         kwargs = self.setParams._input_kwargs
         return self._set(**kwargs)
 
+    @since("1.5.0")
     def setFormula(self, value):
         """
         Sets the value of :py:attr:`formula`.
@@ -1911,6 +2072,7 @@ def setFormula(self, value):
         self._paramMap[self.formula] = value
         return self
 
+    @since("1.5.0")
     def getFormula(self):
         """
         Gets the value of :py:attr:`formula`.
@@ -1926,6 +2088,8 @@ class RFormulaModel(JavaModel):
     .. note:: Experimental
 
     Model fitted by :py:class:`RFormula`.
+
+    .. versionadded:: 1.5.0
     """
 
 
diff --git a/python/pyspark/ml/param/__init__.py b/python/pyspark/ml/param/__init__.py
index 2e0c63cb47b17..35c9b776a3d5e 100644
--- a/python/pyspark/ml/param/__init__.py
+++ b/python/pyspark/ml/param/__init__.py
@@ -18,6 +18,7 @@
 from abc import ABCMeta
 import copy
 
+from pyspark import since
 from pyspark.ml.util import Identifiable
 
 
@@ -27,6 +28,8 @@
 class Param(object):
     """
     A param with self-contained documentation.
+
+    .. versionadded:: 1.3.0
     """
 
     def __init__(self, parent, name, doc):
@@ -56,6 +59,8 @@ class Params(Identifiable):
     """
     Components that take parameters. This also provides an internal
     param map to store parameter values attached to the instance.
+
+    .. versionadded:: 1.3.0
     """
 
     __metaclass__ = ABCMeta
@@ -72,6 +77,7 @@ def __init__(self):
         self._params = None
 
     @property
+    @since("1.3.0")
     def params(self):
         """
         Returns all params ordered by name. The default implementation
@@ -83,6 +89,7 @@ def params(self):
                                        [getattr(self, x) for x in dir(self) if x != "params"]))
         return self._params
 
+    @since("1.4.0")
     def explainParam(self, param):
         """
         Explains a single param and returns its name, doc, and optional
@@ -100,6 +107,7 @@ def explainParam(self, param):
         valueStr = "(" + ", ".join(values) + ")"
         return "%s: %s %s" % (param.name, param.doc, valueStr)
 
+    @since("1.4.0")
     def explainParams(self):
         """
         Returns the documentation of all params with their optionally
@@ -107,6 +115,7 @@ def explainParams(self):
         """
         return "\n".join([self.explainParam(param) for param in self.params])
 
+    @since("1.4.0")
     def getParam(self, paramName):
         """
         Gets a param by its name.
@@ -117,6 +126,7 @@ def getParam(self, paramName):
         else:
             raise ValueError("Cannot find param with name %s." % paramName)
 
+    @since("1.4.0")
     def isSet(self, param):
         """
         Checks whether a param is explicitly set by user.
@@ -124,6 +134,7 @@ def isSet(self, param):
         param = self._resolveParam(param)
         return param in self._paramMap
 
+    @since("1.4.0")
     def hasDefault(self, param):
         """
         Checks whether a param has a default value.
@@ -131,6 +142,7 @@ def hasDefault(self, param):
         param = self._resolveParam(param)
         return param in self._defaultParamMap
 
+    @since("1.4.0")
     def isDefined(self, param):
         """
         Checks whether a param is explicitly set by user or has
@@ -138,6 +150,7 @@ def isDefined(self, param):
         """
         return self.isSet(param) or self.hasDefault(param)
 
+    @since("1.4.0")
     def hasParam(self, paramName):
         """
         Tests whether this instance contains a param with a given
@@ -146,6 +159,7 @@ def hasParam(self, paramName):
         param = self._resolveParam(paramName)
         return param in self.params
 
+    @since("1.4.0")
     def getOrDefault(self, param):
         """
         Gets the value of a param in the user-supplied param map or its
@@ -157,6 +171,7 @@ def getOrDefault(self, param):
         else:
             return self._defaultParamMap[param]
 
+    @since("1.4.0")
     def extractParamMap(self, extra=None):
         """
         Extracts the embedded default param values and user-supplied
@@ -175,6 +190,7 @@ def extractParamMap(self, extra=None):
         paramMap.update(extra)
         return paramMap
 
+    @since("1.4.0")
     def copy(self, extra=None):
         """
         Creates a copy of this instance with the same uid and some
diff --git a/python/pyspark/ml/pipeline.py b/python/pyspark/ml/pipeline.py
index 312a8502b3a2c..4475451edb781 100644
--- a/python/pyspark/ml/pipeline.py
+++ b/python/pyspark/ml/pipeline.py
@@ -17,6 +17,7 @@
 
 from abc import ABCMeta, abstractmethod
 
+from pyspark import since
 from pyspark.ml.param import Param, Params
 from pyspark.ml.util import keyword_only
 from pyspark.mllib.common import inherit_doc
@@ -26,6 +27,8 @@
 class Estimator(Params):
     """
     Abstract class for estimators that fit models to data.
+
+    .. versionadded:: 1.3.0
     """
 
     __metaclass__ = ABCMeta
@@ -42,6 +45,7 @@ def _fit(self, dataset):
         """
         raise NotImplementedError()
 
+    @since("1.3.0")
     def fit(self, dataset, params=None):
         """
         Fits a model to the input dataset with optional parameters.
@@ -73,6 +77,8 @@ class Transformer(Params):
     """
     Abstract class for transformers that transform one dataset into
     another.
+
+    .. versionadded:: 1.3.0
     """
 
     __metaclass__ = ABCMeta
@@ -88,6 +94,7 @@ def _transform(self, dataset):
         """
         raise NotImplementedError()
 
+    @since("1.3.0")
     def transform(self, dataset, params=None):
         """
         Transforms the input dataset with optional parameters.
@@ -113,6 +120,8 @@ def transform(self, dataset, params=None):
 class Model(Transformer):
     """
     Abstract class for models that are fitted by estimators.
+
+    .. versionadded:: 1.4.0
     """
 
     __metaclass__ = ABCMeta
@@ -136,6 +145,8 @@ class Pipeline(Estimator):
     consists of fitted models and transformers, corresponding to the
     pipeline stages. If there are no stages, the pipeline acts as an
     identity transformer.
+
+    .. versionadded:: 1.3.0
     """
 
     @keyword_only
@@ -151,6 +162,7 @@ def __init__(self, stages=None):
         kwargs = self.__init__._input_kwargs
         self.setParams(**kwargs)
 
+    @since("1.3.0")
     def setStages(self, value):
         """
         Set pipeline stages.
@@ -161,6 +173,7 @@ def setStages(self, value):
         self._paramMap[self.stages] = value
         return self
 
+    @since("1.3.0")
     def getStages(self):
         """
         Get pipeline stages.
@@ -169,6 +182,7 @@ def getStages(self):
             return self._paramMap[self.stages]
 
     @keyword_only
+    @since("1.3.0")
     def setParams(self, stages=None):
         """
         setParams(self, stages=None)
@@ -204,7 +218,14 @@ def _fit(self, dataset):
                 transformers.append(stage)
         return PipelineModel(transformers)
 
+    @since("1.4.0")
     def copy(self, extra=None):
+        """
+        Creates a copy of this instance.
+
+        :param extra: extra parameters
+        :returns: new instance
+        """
         if extra is None:
             extra = dict()
         that = Params.copy(self, extra)
@@ -216,6 +237,8 @@ def copy(self, extra=None):
 class PipelineModel(Model):
     """
     Represents a compiled pipeline with transformers and fitted models.
+
+    .. versionadded:: 1.3.0
     """
 
     def __init__(self, stages):
@@ -227,7 +250,14 @@ def _transform(self, dataset):
             dataset = t.transform(dataset)
         return dataset
 
+    @since("1.4.0")
     def copy(self, extra=None):
+        """
+        Creates a copy of this instance.
+
+        :param extra: extra parameters
+        :returns: new instance
+        """
         if extra is None:
             extra = dict()
         stages = [stage.copy(extra) for stage in self.stages]

From ec03866a7ef2d0826520755d47c8c9480148a76c Mon Sep 17 00:00:00 2001
From: Dominik Dahlem <dominik.dahlem@gmail.combination>
Date: Mon, 2 Nov 2015 16:11:42 -0800
Subject: [PATCH 011/518] [SPARK-11343][ML] Allow float and double
 prediction/label columns in RegressionEvaluator

mengxr, felixcheung

This pull request just relaxes the type of the prediction/label columns to be float and double. Internally, these columns are casted to double. The other evaluators might need to be changed also.

Author: Dominik Dahlem <dominik.dahlem@gmail.combination>

Closes #9296 from dahlem/ddahlem_regression_evaluator_double_predictions_27102015.
---
 .../spark/ml/evaluation/RegressionEvaluator.scala    | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/evaluation/RegressionEvaluator.scala b/mllib/src/main/scala/org/apache/spark/ml/evaluation/RegressionEvaluator.scala
index 3fd34d8571017..ba012f444d3e0 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/evaluation/RegressionEvaluator.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/evaluation/RegressionEvaluator.scala
@@ -23,7 +23,8 @@ import org.apache.spark.ml.param.shared.{HasLabelCol, HasPredictionCol}
 import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
 import org.apache.spark.mllib.evaluation.RegressionMetrics
 import org.apache.spark.sql.{DataFrame, Row}
-import org.apache.spark.sql.types.DoubleType
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.types.{DoubleType, FloatType}
 
 /**
  * :: Experimental ::
@@ -72,10 +73,13 @@ final class RegressionEvaluator @Since("1.4.0") (@Since("1.4.0") override val ui
   @Since("1.4.0")
   override def evaluate(dataset: DataFrame): Double = {
     val schema = dataset.schema
-    SchemaUtils.checkColumnType(schema, $(predictionCol), DoubleType)
-    SchemaUtils.checkColumnType(schema, $(labelCol), DoubleType)
+    val predictionType = schema($(predictionCol)).dataType
+    require(predictionType == FloatType || predictionType == DoubleType)
+    val labelType = schema($(labelCol)).dataType
+    require(labelType == FloatType || labelType == DoubleType)
 
-    val predictionAndLabels = dataset.select($(predictionCol), $(labelCol))
+    val predictionAndLabels = dataset
+      .select(col($(predictionCol)).cast(DoubleType), col($(labelCol)).cast(DoubleType))
       .map { case Row(prediction: Double, label: Double) =>
         (prediction, label)
       }

From c020f7d9d43548d27ae4a9564ba38981fd530cb1 Mon Sep 17 00:00:00 2001
From: vectorijk <jiangkai@gmail.com>
Date: Mon, 2 Nov 2015 16:12:04 -0800
Subject: [PATCH 012/518] [SPARK-10592] [ML] [PySpark] Deprecate weights and
 use coefficients instead in ML models

Deprecated in `LogisticRegression` and `LinearRegression`

Author: vectorijk <jiangkai@gmail.com>

Closes #9311 from vectorijk/spark-10592.
---
 R/pkg/R/mllib.R                               |   6 +-
 .../classification/LogisticRegression.scala   |  11 +-
 .../apache/spark/ml/r/SparkRWrappers.scala    |  15 +-
 .../ml/regression/AFTSurvivalRegression.scala |  32 +--
 .../ml/regression/IsotonicRegression.scala    |   4 +-
 .../ml/regression/LinearRegression.scala      |  15 +-
 .../ml/classification/JavaOneVsRestSuite.java |   6 +-
 .../LogisticRegressionSuite.scala             | 152 ++++++++-------
 .../MultilayerPerceptronClassifierSuite.scala |   6 +-
 .../ml/classification/OneVsRestSuite.scala    |   6 +-
 .../AFTSurvivalRegressionSuite.scala          |  12 +-
 .../ml/regression/LinearRegressionSuite.scala | 184 +++++++++---------
 python/pyspark/ml/classification.py           |  13 ++
 python/pyspark/ml/regression.py               |  12 ++
 14 files changed, 263 insertions(+), 211 deletions(-)

diff --git a/R/pkg/R/mllib.R b/R/pkg/R/mllib.R
index aadd5b8da5e3b..60bfadb8e7503 100644
--- a/R/pkg/R/mllib.R
+++ b/R/pkg/R/mllib.R
@@ -92,9 +92,9 @@ setMethod("summary", signature(x = "PipelineModel"),
           function(x, ...) {
             features <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
                                    "getModelFeatures", x@model)
-            weights <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
-                                   "getModelWeights", x@model)
-            coefficients <- as.matrix(unlist(weights))
+            coefficients <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
+                                   "getModelCoefficients", x@model)
+            coefficients <- as.matrix(unlist(coefficients))
             colnames(coefficients) <- c("Estimate")
             rownames(coefficients) <- unlist(features)
             return(list(coefficients = coefficients))
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
index 6f839ff4d7cd8..a1335e7a1bde8 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
@@ -392,11 +392,14 @@ class LogisticRegression(override val uid: String)
 @Experimental
 class LogisticRegressionModel private[ml] (
     override val uid: String,
-    val weights: Vector,
+    val coefficients: Vector,
     val intercept: Double)
   extends ProbabilisticClassificationModel[Vector, LogisticRegressionModel]
   with LogisticRegressionParams {
 
+  @deprecated("Use coefficients instead.", "1.6.0")
+  def weights: Vector = coefficients
+
   override def setThreshold(value: Double): this.type = super.setThreshold(value)
 
   override def getThreshold: Double = super.getThreshold
@@ -407,7 +410,7 @@ class LogisticRegressionModel private[ml] (
 
   /** Margin (rawPrediction) for class label 1.  For binary classification only. */
   private val margin: Vector => Double = (features) => {
-    BLAS.dot(features, weights) + intercept
+    BLAS.dot(features, coefficients) + intercept
   }
 
   /** Score (probability) for class label 1.  For binary classification only. */
@@ -416,7 +419,7 @@ class LogisticRegressionModel private[ml] (
     1.0 / (1.0 + math.exp(-m))
   }
 
-  override val numFeatures: Int = weights.size
+  override val numFeatures: Int = coefficients.size
 
   override val numClasses: Int = 2
 
@@ -483,7 +486,7 @@ class LogisticRegressionModel private[ml] (
   }
 
   override def copy(extra: ParamMap): LogisticRegressionModel = {
-    val newModel = copyValues(new LogisticRegressionModel(uid, weights, intercept), extra)
+    val newModel = copyValues(new LogisticRegressionModel(uid, coefficients, intercept), extra)
     if (trainingSummary.isDefined) newModel.setSummary(trainingSummary.get)
     newModel.setParent(parent)
   }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala b/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
index 21ebf6d916db7..9162ec0e4e153 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
@@ -51,13 +51,22 @@ private[r] object SparkRWrappers {
     pipeline.fit(df)
   }
 
+  @deprecated("Use getModelCoefficients instead.", "1.6.0")
   def getModelWeights(model: PipelineModel): Array[Double] = {
     model.stages.last match {
       case m: LinearRegressionModel =>
         Array(m.intercept) ++ m.weights.toArray
-      case _: LogisticRegressionModel =>
-        throw new UnsupportedOperationException(
-          "No weights available for LogisticRegressionModel")  // SPARK-9492
+      case m: LogisticRegressionModel =>
+        Array(m.intercept) ++ m.weights.toArray
+    }
+  }
+
+  def getModelCoefficients(model: PipelineModel): Array[Double] = {
+    model.stages.last match {
+      case m: LinearRegressionModel =>
+        Array(m.intercept) ++ m.coefficients.toArray
+      case m: LogisticRegressionModel =>
+        Array(m.intercept) ++ m.coefficients.toArray
     }
   }
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
index ac2c3d825f13c..4dbbc7d39931b 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
@@ -200,17 +200,17 @@ class AFTSurvivalRegression @Since("1.6.0") (@Since("1.6.0") override val uid: S
 
     val numFeatures = dataset.select($(featuresCol)).take(1)(0).getAs[Vector](0).size
     /*
-       The weights vector has three parts:
+       The coefficients vector has three parts:
        the first element: Double, log(sigma), the log of scale parameter
        the second element: Double, intercept of the beta parameter
        the third to the end elements: Doubles, regression coefficients vector of the beta parameter
      */
-    val initialWeights = Vectors.zeros(numFeatures + 2)
+    val initialCoefficients = Vectors.zeros(numFeatures + 2)
 
     val states = optimizer.iterations(new CachedDiffFunction(costFun),
-      initialWeights.toBreeze.toDenseVector)
+      initialCoefficients.toBreeze.toDenseVector)
 
-    val weights = {
+    val coefficients = {
       val arrayBuilder = mutable.ArrayBuilder.make[Double]
       var state: optimizer.State = null
       while (states.hasNext) {
@@ -227,10 +227,10 @@ class AFTSurvivalRegression @Since("1.6.0") (@Since("1.6.0") override val uid: S
 
     if (handlePersistence) instances.unpersist()
 
-    val coefficients = Vectors.dense(weights.slice(2, weights.length))
-    val intercept = weights(1)
-    val scale = math.exp(weights(0))
-    val model = new AFTSurvivalRegressionModel(uid, coefficients, intercept, scale)
+    val regressionCoefficients = Vectors.dense(coefficients.slice(2, coefficients.length))
+    val intercept = coefficients(1)
+    val scale = math.exp(coefficients(0))
+    val model = new AFTSurvivalRegressionModel(uid, regressionCoefficients, intercept, scale)
     copyValues(model.setParent(this))
   }
 
@@ -251,7 +251,7 @@ class AFTSurvivalRegression @Since("1.6.0") (@Since("1.6.0") override val uid: S
 @Since("1.6.0")
 class AFTSurvivalRegressionModel private[ml] (
     @Since("1.6.0") override val uid: String,
-    @Since("1.6.0") val coefficients: Vector,
+    @Since("1.6.0") val regressionCoefficients: Vector,
     @Since("1.6.0") val intercept: Double,
     @Since("1.6.0") val scale: Double)
   extends Model[AFTSurvivalRegressionModel] with AFTSurvivalRegressionParams {
@@ -275,7 +275,7 @@ class AFTSurvivalRegressionModel private[ml] (
   @Since("1.6.0")
   def predictQuantiles(features: Vector): Vector = {
     // scale parameter for the Weibull distribution of lifetime
-    val lambda = math.exp(BLAS.dot(coefficients, features) + intercept)
+    val lambda = math.exp(BLAS.dot(regressionCoefficients, features) + intercept)
     // shape parameter for the Weibull distribution of lifetime
     val k = 1 / scale
     val quantiles = $(quantileProbabilities).map {
@@ -286,7 +286,7 @@ class AFTSurvivalRegressionModel private[ml] (
 
   @Since("1.6.0")
   def predict(features: Vector): Double = {
-    math.exp(BLAS.dot(coefficients, features) + intercept)
+    math.exp(BLAS.dot(regressionCoefficients, features) + intercept)
   }
 
   @Since("1.6.0")
@@ -309,7 +309,7 @@ class AFTSurvivalRegressionModel private[ml] (
 
   @Since("1.6.0")
   override def copy(extra: ParamMap): AFTSurvivalRegressionModel = {
-    copyValues(new AFTSurvivalRegressionModel(uid, coefficients, intercept, scale), extra)
+    copyValues(new AFTSurvivalRegressionModel(uid, regressionCoefficients, intercept, scale), extra)
       .setParent(parent)
   }
 }
@@ -369,17 +369,17 @@ class AFTSurvivalRegressionModel private[ml] (
  *   \frac{\partial (-\iota)}{\partial (\log\sigma)}=
  *   \sum_{i=1}^{n}[\delta_{i}+(\delta_{i}-e^{\epsilon_{i}})\epsilon_{i}]
  * }}}
- * @param weights The log of scale parameter, the intercept and
+ * @param coefficients including three part: The log of scale parameter, the intercept and
  *                regression coefficients corresponding to the features.
  * @param fitIntercept Whether to fit an intercept term.
  */
-private class AFTAggregator(weights: BDV[Double], fitIntercept: Boolean)
+private class AFTAggregator(coefficients: BDV[Double], fitIntercept: Boolean)
   extends Serializable {
 
   // beta is the intercept and regression coefficients to the covariates
-  private val beta = weights.slice(1, weights.length)
+  private val beta = coefficients.slice(1, coefficients.length)
   // sigma is the scale parameter of the AFT model
-  private val sigma = math.exp(weights(0))
+  private val sigma = math.exp(coefficients(0))
 
   private var totalCnt: Long = 0L
   private var lossSum = 0.0
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
index 2ff500f291abc..f4a17c8f9a582 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
@@ -87,8 +87,8 @@ private[regression] trait IsotonicRegressionBase extends Params with HasFeatures
       lit(1.0)
     }
     dataset.select(col($(labelCol)), f, w)
-      .map { case Row(label: Double, feature: Double, weights: Double) =>
-      (label, feature, weights)
+      .map { case Row(label: Double, feature: Double, weight: Double) =>
+      (label, feature, weight)
     }
   }
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index f663b9bd9ac73..6e9c7442b8110 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -203,7 +203,7 @@ class LinearRegression(override val uid: String)
     val yMean = ySummarizer.mean(0)
     val yStd = math.sqrt(ySummarizer.variance(0))
 
-    // If the yStd is zero, then the intercept is yMean with zero weights;
+    // If the yStd is zero, then the intercept is yMean with zero coefficient;
     // as a result, training is not needed.
     if (yStd == 0.0) {
       logWarning(s"The standard deviation of the label is zero, so the coefficients will be " +
@@ -331,14 +331,17 @@ class LinearRegression(override val uid: String)
 @Experimental
 class LinearRegressionModel private[ml] (
     override val uid: String,
-    val weights: Vector,
+    val coefficients: Vector,
     val intercept: Double)
   extends RegressionModel[Vector, LinearRegressionModel]
   with LinearRegressionParams {
 
   private var trainingSummary: Option[LinearRegressionTrainingSummary] = None
 
-  override val numFeatures: Int = weights.size
+  @deprecated("Use coefficients instead.", "1.6.0")
+  def weights: Vector = coefficients
+
+  override val numFeatures: Int = coefficients.size
 
   /**
    * Gets summary (e.g. residuals, mse, r-squared ) of model on training set. An exception is
@@ -387,11 +390,11 @@ class LinearRegressionModel private[ml] (
 
 
   override protected def predict(features: Vector): Double = {
-    dot(features, weights) + intercept
+    dot(features, coefficients) + intercept
   }
 
   override def copy(extra: ParamMap): LinearRegressionModel = {
-    val newModel = copyValues(new LinearRegressionModel(uid, weights, intercept), extra)
+    val newModel = copyValues(new LinearRegressionModel(uid, coefficients, intercept), extra)
     if (trainingSummary.isDefined) newModel.setSummary(trainingSummary.get)
     newModel.setParent(parent)
   }
@@ -400,7 +403,7 @@ class LinearRegressionModel private[ml] (
 /**
  * :: Experimental ::
  * Linear regression training results. Currently, the training summary ignores the
- * training weights except for the objective trace.
+ * training coefficients except for the objective trace.
  * @param predictions predictions outputted by the model's `transform` method.
  * @param objectiveHistory objective function (scaled loss + regularization) at each iteration.
  */
diff --git a/mllib/src/test/java/org/apache/spark/ml/classification/JavaOneVsRestSuite.java b/mllib/src/test/java/org/apache/spark/ml/classification/JavaOneVsRestSuite.java
index 253cabf0133d0..cbabafe1b541d 100644
--- a/mllib/src/test/java/org/apache/spark/ml/classification/JavaOneVsRestSuite.java
+++ b/mllib/src/test/java/org/apache/spark/ml/classification/JavaOneVsRestSuite.java
@@ -47,16 +47,16 @@ public void setUp() {
         jsql = new SQLContext(jsc);
         int nPoints = 3;
 
-        // The following weights and xMean/xVariance are computed from iris dataset with lambda=0.2.
+        // The following coefficients and xMean/xVariance are computed from iris dataset with lambda=0.2.
         // As a result, we are drawing samples from probability distribution of an actual model.
-        double[] weights = {
+        double[] coefficients = {
                 -0.57997, 0.912083, -0.371077, -0.819866, 2.688191,
                 -0.16624, -0.84355, -0.048509, -0.301789, 4.170682 };
 
         double[] xMean = {5.843, 3.057, 3.758, 1.199};
         double[] xVariance = {0.6856, 0.1899, 3.116, 0.581};
         List<LabeledPoint> points = JavaConverters.seqAsJavaListConverter(
-            generateMultinomialLogisticInput(weights, xMean, xVariance, true, nPoints, 42)
+            generateMultinomialLogisticInput(coefficients, xMean, xVariance, true, nPoints, 42)
         ).asJava();
         datasetRDD = jsc.parallelize(points, 2);
         dataset = jsql.createDataFrame(datasetRDD, LabeledPoint.class);
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
index e0a795e5e0b00..325faf37e8eea 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
@@ -48,21 +48,22 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
 
        import org.apache.spark.mllib.classification.LogisticRegressionSuite
        val nPoints = 10000
-       val weights = Array(-0.57997, 0.912083, -0.371077, -0.819866, 2.688191)
+       val coefficients = Array(-0.57997, 0.912083, -0.371077, -0.819866, 2.688191)
        val xMean = Array(5.843, 3.057, 3.758, 1.199)
        val xVariance = Array(0.6856, 0.1899, 3.116, 0.581)
        val data = sc.parallelize(LogisticRegressionSuite.generateMultinomialLogisticInput(
-         weights, xMean, xVariance, true, nPoints, 42), 1)
+         coefficients, xMean, xVariance, true, nPoints, 42), 1)
        data.map(x=> x.label + ", " + x.features(0) + ", " + x.features(1) + ", "
          + x.features(2) + ", " + x.features(3)).saveAsTextFile("path")
      */
     binaryDataset = {
       val nPoints = 10000
-      val weights = Array(-0.57997, 0.912083, -0.371077, -0.819866, 2.688191)
+      val coefficients = Array(-0.57997, 0.912083, -0.371077, -0.819866, 2.688191)
       val xMean = Array(5.843, 3.057, 3.758, 1.199)
       val xVariance = Array(0.6856, 0.1899, 3.116, 0.581)
 
-      val testData = generateMultinomialLogisticInput(weights, xMean, xVariance, true, nPoints, 42)
+      val testData =
+        generateMultinomialLogisticInput(coefficients, xMean, xVariance, true, nPoints, 42)
 
       sqlContext.createDataFrame(sc.parallelize(testData, 4))
     }
@@ -296,8 +297,8 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 0))
-       weights
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 0))
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                            s0
@@ -308,14 +309,14 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.7996864
      */
     val interceptR = 2.8366423
-    val weightsR = Vectors.dense(-0.5895848, 0.8931147, -0.3925051, -0.7996864)
+    val coefficientsR = Vectors.dense(-0.5895848, 0.8931147, -0.3925051, -0.7996864)
 
     assert(model1.intercept ~== interceptR relTol 1E-3)
-    assert(model1.weights ~= weightsR relTol 1E-3)
+    assert(model1.coefficients ~= coefficientsR relTol 1E-3)
 
     // Without regularization, with or without standardization will converge to the same solution.
     assert(model2.intercept ~== interceptR relTol 1E-3)
-    assert(model2.weights ~= weightsR relTol 1E-3)
+    assert(model2.coefficients ~= coefficientsR relTol 1E-3)
   }
 
   test("binary logistic regression without intercept without regularization") {
@@ -332,9 +333,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights =
+       coefficients =
            coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 0, intercept=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                            s0
@@ -345,14 +346,14 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.7407946
      */
     val interceptR = 0.0
-    val weightsR = Vectors.dense(-0.3534996, 1.2964482, -0.3571741, -0.7407946)
+    val coefficientsR = Vectors.dense(-0.3534996, 1.2964482, -0.3571741, -0.7407946)
 
     assert(model1.intercept ~== interceptR relTol 1E-3)
-    assert(model1.weights ~= weightsR relTol 1E-2)
+    assert(model1.coefficients ~= coefficientsR relTol 1E-2)
 
     // Without regularization, with or without standardization should converge to the same solution.
     assert(model2.intercept ~== interceptR relTol 1E-3)
-    assert(model2.weights ~= weightsR relTol 1E-2)
+    assert(model2.coefficients ~= coefficientsR relTol 1E-2)
   }
 
   test("binary logistic regression with intercept with L1 regularization") {
@@ -371,8 +372,8 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 1, lambda = 0.12))
-       weights
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 1, lambda = 0.12))
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -383,10 +384,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.02481551
      */
     val interceptR1 = -0.05627428
-    val weightsR1 = Vectors.dense(0.0, 0.0, -0.04325749, -0.02481551)
+    val coefficientsR1 = Vectors.dense(0.0, 0.0, -0.04325749, -0.02481551)
 
     assert(model1.intercept ~== interceptR1 relTol 1E-2)
-    assert(model1.weights ~= weightsR1 absTol 2E-2)
+    assert(model1.coefficients ~= coefficientsR1 absTol 2E-2)
 
     /*
        Using the following R code to load the data and train the model using glmnet package.
@@ -395,9 +396,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 1, lambda = 0.12,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 1, lambda = 0.12,
            standardize=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                            s0
@@ -408,10 +409,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5       .
      */
     val interceptR2 = 0.3722152
-    val weightsR2 = Vectors.dense(0.0, 0.0, -0.1665453, 0.0)
+    val coefficientsR2 = Vectors.dense(0.0, 0.0, -0.1665453, 0.0)
 
     assert(model2.intercept ~== interceptR2 relTol 1E-2)
-    assert(model2.weights ~= weightsR2 absTol 1E-3)
+    assert(model2.coefficients ~= coefficientsR2 absTol 1E-3)
   }
 
   test("binary logistic regression without intercept with L1 regularization") {
@@ -430,9 +431,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 1, lambda = 0.12,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 1, lambda = 0.12,
            intercept=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -443,10 +444,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.03891782
      */
     val interceptR1 = 0.0
-    val weightsR1 = Vectors.dense(0.0, 0.0, -0.05189203, -0.03891782)
+    val coefficientsR1 = Vectors.dense(0.0, 0.0, -0.05189203, -0.03891782)
 
     assert(model1.intercept ~== interceptR1 relTol 1E-3)
-    assert(model1.weights ~= weightsR1 absTol 1E-3)
+    assert(model1.coefficients ~= coefficientsR1 absTol 1E-3)
 
     /*
        Using the following R code to load the data and train the model using glmnet package.
@@ -455,9 +456,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 1, lambda = 0.12,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 1, lambda = 0.12,
            intercept=FALSE, standardize=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -468,10 +469,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5       .
      */
     val interceptR2 = 0.0
-    val weightsR2 = Vectors.dense(0.0, 0.0, -0.08420782, 0.0)
+    val coefficientsR2 = Vectors.dense(0.0, 0.0, -0.08420782, 0.0)
 
     assert(model2.intercept ~== interceptR2 absTol 1E-3)
-    assert(model2.weights ~= weightsR2 absTol 1E-3)
+    assert(model2.coefficients ~= coefficientsR2 absTol 1E-3)
   }
 
   test("binary logistic regression with intercept with L2 regularization") {
@@ -490,8 +491,8 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 1.37))
-       weights
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 1.37))
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -502,10 +503,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.10062872
      */
     val interceptR1 = 0.15021751
-    val weightsR1 = Vectors.dense(-0.07251837, 0.10724191, -0.04865309, -0.10062872)
+    val coefficientsR1 = Vectors.dense(-0.07251837, 0.10724191, -0.04865309, -0.10062872)
 
     assert(model1.intercept ~== interceptR1 relTol 1E-3)
-    assert(model1.weights ~= weightsR1 relTol 1E-3)
+    assert(model1.coefficients ~= coefficientsR1 relTol 1E-3)
 
     /*
        Using the following R code to load the data and train the model using glmnet package.
@@ -514,9 +515,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 1.37,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 1.37,
            standardize=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -527,10 +528,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.06266838
      */
     val interceptR2 = 0.48657516
-    val weightsR2 = Vectors.dense(-0.05155371, 0.02301057, -0.11482896, -0.06266838)
+    val coefficientsR2 = Vectors.dense(-0.05155371, 0.02301057, -0.11482896, -0.06266838)
 
     assert(model2.intercept ~== interceptR2 relTol 1E-3)
-    assert(model2.weights ~= weightsR2 relTol 1E-3)
+    assert(model2.coefficients ~= coefficientsR2 relTol 1E-3)
   }
 
   test("binary logistic regression without intercept with L2 regularization") {
@@ -549,9 +550,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 1.37,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 1.37,
            intercept=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -562,10 +563,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.09799775
      */
     val interceptR1 = 0.0
-    val weightsR1 = Vectors.dense(-0.06099165, 0.12857058, -0.04708770, -0.09799775)
+    val coefficientsR1 = Vectors.dense(-0.06099165, 0.12857058, -0.04708770, -0.09799775)
 
     assert(model1.intercept ~== interceptR1 absTol 1E-3)
-    assert(model1.weights ~= weightsR1 relTol 1E-2)
+    assert(model1.coefficients ~= coefficientsR1 relTol 1E-2)
 
     /*
        Using the following R code to load the data and train the model using glmnet package.
@@ -574,9 +575,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 1.37,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0, lambda = 1.37,
            intercept=FALSE, standardize=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                              s0
@@ -587,10 +588,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.053314311
      */
     val interceptR2 = 0.0
-    val weightsR2 = Vectors.dense(-0.005679651, 0.048967094, -0.093714016, -0.053314311)
+    val coefficientsR2 = Vectors.dense(-0.005679651, 0.048967094, -0.093714016, -0.053314311)
 
     assert(model2.intercept ~== interceptR2 absTol 1E-3)
-    assert(model2.weights ~= weightsR2 relTol 1E-2)
+    assert(model2.coefficients ~= coefficientsR2 relTol 1E-2)
   }
 
   test("binary logistic regression with intercept with ElasticNet regularization") {
@@ -609,8 +610,8 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0.38, lambda = 0.21))
-       weights
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0.38, lambda = 0.21))
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -621,10 +622,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.15458796
      */
     val interceptR1 = 0.57734851
-    val weightsR1 = Vectors.dense(-0.05310287, 0.0, -0.08849250, -0.15458796)
+    val coefficientsR1 = Vectors.dense(-0.05310287, 0.0, -0.08849250, -0.15458796)
 
     assert(model1.intercept ~== interceptR1 relTol 6E-3)
-    assert(model1.weights ~== weightsR1 absTol 5E-3)
+    assert(model1.coefficients ~== coefficientsR1 absTol 5E-3)
 
     /*
        Using the following R code to load the data and train the model using glmnet package.
@@ -633,9 +634,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0.38, lambda = 0.21,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0.38, lambda = 0.21,
            standardize=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -646,10 +647,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.05350074
      */
     val interceptR2 = 0.51555993
-    val weightsR2 = Vectors.dense(0.0, 0.0, -0.18807395, -0.05350074)
+    val coefficientsR2 = Vectors.dense(0.0, 0.0, -0.18807395, -0.05350074)
 
     assert(model2.intercept ~== interceptR2 relTol 6E-3)
-    assert(model2.weights ~= weightsR2 absTol 1E-3)
+    assert(model2.coefficients ~= coefficientsR2 absTol 1E-3)
   }
 
   test("binary logistic regression without intercept with ElasticNet regularization") {
@@ -668,9 +669,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0.38, lambda = 0.21,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0.38, lambda = 0.21,
            intercept=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -681,10 +682,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5     -0.142534158
      */
     val interceptR1 = 0.0
-    val weightsR1 = Vectors.dense(-0.001005743, 0.072577857, -0.081203769, -0.142534158)
+    val coefficientsR1 = Vectors.dense(-0.001005743, 0.072577857, -0.081203769, -0.142534158)
 
     assert(model1.intercept ~== interceptR1 relTol 1E-3)
-    assert(model1.weights ~= weightsR1 absTol 1E-2)
+    assert(model1.coefficients ~= coefficientsR1 absTol 1E-2)
 
     /*
        Using the following R code to load the data and train the model using glmnet package.
@@ -693,9 +694,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 0.38, lambda = 0.21,
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 0.38, lambda = 0.21,
            intercept=FALSE, standardize=FALSE))
-       weights
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -706,10 +707,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5       .
      */
     val interceptR2 = 0.0
-    val weightsR2 = Vectors.dense(0.0, 0.03345223, -0.11304532, 0.0)
+    val coefficientsR2 = Vectors.dense(0.0, 0.03345223, -0.11304532, 0.0)
 
     assert(model2.intercept ~== interceptR2 absTol 1E-3)
-    assert(model2.weights ~= weightsR2 absTol 1E-3)
+    assert(model2.coefficients ~= coefficientsR2 absTol 1E-3)
   }
 
   test("binary logistic regression with intercept with strong L1 regularization") {
@@ -732,8 +733,8 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         }).histogram
 
     /*
-       For binary logistic regression with strong L1 regularization, all the weights will be zeros.
-       As a result,
+       For binary logistic regression with strong L1 regularization, all the coefficients
+       will be zeros. As a result,
        {{{
        P(0) = 1 / (1 + \exp(b)), and
        P(1) = \exp(b) / (1 + \exp(b))
@@ -743,13 +744,13 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        }}}
      */
     val interceptTheory = math.log(histogram(1) / histogram(0))
-    val weightsTheory = Vectors.dense(0.0, 0.0, 0.0, 0.0)
+    val coefficientsTheory = Vectors.dense(0.0, 0.0, 0.0, 0.0)
 
     assert(model1.intercept ~== interceptTheory relTol 1E-5)
-    assert(model1.weights ~= weightsTheory absTol 1E-6)
+    assert(model1.coefficients ~= coefficientsTheory absTol 1E-6)
 
     assert(model2.intercept ~== interceptTheory relTol 1E-5)
-    assert(model2.weights ~= weightsTheory absTol 1E-6)
+    assert(model2.coefficients ~= coefficientsTheory absTol 1E-6)
 
     /*
        Using the following R code to load the data and train the model using glmnet package.
@@ -758,8 +759,8 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data <- read.csv("path", header=FALSE)
        label = factor(data$V1)
        features = as.matrix(data.frame(data$V2, data$V3, data$V4, data$V5))
-       weights = coef(glmnet(features,label, family="binomial", alpha = 1.0, lambda = 6.0))
-       weights
+       coefficients = coef(glmnet(features,label, family="binomial", alpha = 1.0, lambda = 6.0))
+       coefficients
 
        5 x 1 sparse Matrix of class "dgCMatrix"
                             s0
@@ -770,10 +771,10 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
        data.V5       .
      */
     val interceptR = -0.248065
-    val weightsR = Vectors.dense(0.0, 0.0, 0.0, 0.0)
+    val coefficientsR = Vectors.dense(0.0, 0.0, 0.0, 0.0)
 
     assert(model1.intercept ~== interceptR relTol 1E-5)
-    assert(model1.weights ~== weightsR absTol 1E-6)
+    assert(model1.coefficients ~== coefficientsR absTol 1E-6)
   }
 
   test("evaluate on test set") {
@@ -814,10 +815,11 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
   test("binary logistic regression with weighted samples") {
     val (dataset, weightedDataset) = {
       val nPoints = 1000
-      val weights = Array(-0.57997, 0.912083, -0.371077, -0.819866, 2.688191)
+      val coefficients = Array(-0.57997, 0.912083, -0.371077, -0.819866, 2.688191)
       val xMean = Array(5.843, 3.057, 3.758, 1.199)
       val xVariance = Array(0.6856, 0.1899, 3.116, 0.581)
-      val testData = generateMultinomialLogisticInput(weights, xMean, xVariance, true, nPoints, 42)
+      val testData =
+        generateMultinomialLogisticInput(coefficients, xMean, xVariance, true, nPoints, 42)
 
       // Let's over-sample the positive samples twice.
       val data1 = testData.flatMap { case labeledPoint: LabeledPoint =>
@@ -863,9 +865,9 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
     val model1a0 = trainer1a.fit(dataset)
     val model1a1 = trainer1a.fit(weightedDataset)
     val model1b = trainer1b.fit(weightedDataset)
-    assert(model1a0.weights !~= model1a1.weights absTol 1E-3)
+    assert(model1a0.coefficients !~= model1a1.coefficients absTol 1E-3)
     assert(model1a0.intercept !~= model1a1.intercept absTol 1E-3)
-    assert(model1a0.weights ~== model1b.weights absTol 1E-3)
+    assert(model1a0.coefficients ~== model1b.coefficients absTol 1E-3)
     assert(model1a0.intercept ~== model1b.intercept absTol 1E-3)
 
   }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifierSuite.scala
index 2d1df9b2b82e8..17db8c44777d4 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifierSuite.scala
@@ -53,16 +53,16 @@ class MultilayerPerceptronClassifierSuite extends SparkFunSuite with MLlibTestSp
   test("3 class classification with 2 hidden layers") {
     val nPoints = 1000
 
-    // The following weights are taken from OneVsRestSuite.scala
+    // The following coefficients are taken from OneVsRestSuite.scala
     // they represent 3-class iris dataset
-    val weights = Array(
+    val coefficients = Array(
       -0.57997, 0.912083, -0.371077, -0.819866, 2.688191,
       -0.16624, -0.84355, -0.048509, -0.301789, 4.170682)
 
     val xMean = Array(5.843, 3.057, 3.758, 1.199)
     val xVariance = Array(0.6856, 0.1899, 3.116, 0.581)
     val rdd = sc.parallelize(generateMultinomialLogisticInput(
-      weights, xMean, xVariance, true, nPoints, 42), 2)
+      coefficients, xMean, xVariance, true, nPoints, 42), 2)
     val dataFrame = sqlContext.createDataFrame(rdd).toDF("label", "features")
     val numClasses = 3
     val numIterations = 100
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala
index 977f0e0b70c1a..5ea71c5317b7a 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala
@@ -43,16 +43,16 @@ class OneVsRestSuite extends SparkFunSuite with MLlibTestSparkContext {
 
     val nPoints = 1000
 
-    // The following weights and xMean/xVariance are computed from iris dataset with lambda=0.2.
+    // The following coefficients and xMean/xVariance are computed from iris dataset with lambda=0.2
     // As a result, we are drawing samples from probability distribution of an actual model.
-    val weights = Array(
+    val coefficients = Array(
       -0.57997, 0.912083, -0.371077, -0.819866, 2.688191,
       -0.16624, -0.84355, -0.048509, -0.301789, 4.170682)
 
     val xMean = Array(5.843, 3.057, 3.758, 1.199)
     val xVariance = Array(0.6856, 0.1899, 3.116, 0.581)
     rdd = sc.parallelize(generateMultinomialLogisticInput(
-      weights, xMean, xVariance, true, nPoints, 42), 2)
+      coefficients, xMean, xVariance, true, nPoints, 42), 2)
     dataset = sqlContext.createDataFrame(rdd)
   }
 
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
index 359f31027172b..c0f791bce13d1 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
@@ -141,12 +141,12 @@ class AFTSurvivalRegressionSuite extends SparkFunSuite with MLlibTestSparkContex
        Number of Newton-Raphson Iterations: 5
        n= 1000
      */
-    val coefficientsR = Vectors.dense(-0.039)
+    val regressionCoefficientsR = Vectors.dense(-0.039)
     val interceptR = 1.759
     val scaleR = 1.41
 
     assert(model.intercept ~== interceptR relTol 1E-3)
-    assert(model.coefficients ~== coefficientsR relTol 1E-3)
+    assert(model.regressionCoefficients ~== regressionCoefficientsR relTol 1E-3)
     assert(model.scale ~== scaleR relTol 1E-3)
 
     /*
@@ -212,12 +212,12 @@ class AFTSurvivalRegressionSuite extends SparkFunSuite with MLlibTestSparkContex
        Number of Newton-Raphson Iterations: 5
        n= 1000
      */
-    val coefficientsR = Vectors.dense(-0.0844, 0.0677)
+    val regressionCoefficientsR = Vectors.dense(-0.0844, 0.0677)
     val interceptR = 1.9206
     val scaleR = 0.977
 
     assert(model.intercept ~== interceptR relTol 1E-3)
-    assert(model.coefficients ~== coefficientsR relTol 1E-3)
+    assert(model.regressionCoefficients ~== regressionCoefficientsR relTol 1E-3)
     assert(model.scale ~== scaleR relTol 1E-3)
 
     /*
@@ -282,12 +282,12 @@ class AFTSurvivalRegressionSuite extends SparkFunSuite with MLlibTestSparkContex
        Number of Newton-Raphson Iterations: 6
        n= 1000
      */
-    val coefficientsR = Vectors.dense(0.896, -0.709)
+    val regressionCoefficientsR = Vectors.dense(0.896, -0.709)
     val interceptR = 0.0
     val scaleR = 1.52
 
     assert(model.intercept === interceptR)
-    assert(model.coefficients ~== coefficientsR relTol 1E-3)
+    assert(model.regressionCoefficients ~== regressionCoefficientsR relTol 1E-3)
     assert(model.scale ~== scaleR relTol 1E-3)
 
     /*
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
index a2a5c0bbdcb90..235c796d785a6 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
@@ -122,8 +122,8 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
          data <- read.csv("path", header=FALSE, stringsAsFactors=FALSE)
          features <- as.matrix(data.frame(as.numeric(data$V2), as.numeric(data$V3)))
          label <- as.numeric(data$V1)
-         weights <- coef(glmnet(features, label, family="gaussian", alpha = 0, lambda = 0))
-         > weights
+         coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0, lambda = 0))
+         > coefficients
           3 x 1 sparse Matrix of class "dgCMatrix"
                                    s0
          (Intercept)         6.298698
@@ -131,17 +131,18 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
          as.numeric.data.V3. 7.199082
        */
       val interceptR = 6.298698
-      val weightsR = Vectors.dense(4.700706, 7.199082)
+      val coefficientsR = Vectors.dense(4.700706, 7.199082)
 
       assert(model1.intercept ~== interceptR relTol 1E-3)
-      assert(model1.weights ~= weightsR relTol 1E-3)
+      assert(model1.coefficients ~= coefficientsR relTol 1E-3)
       assert(model2.intercept ~== interceptR relTol 1E-3)
-      assert(model2.weights ~= weightsR relTol 1E-3)
+      assert(model2.coefficients ~= coefficientsR relTol 1E-3)
 
       model1.transform(datasetWithDenseFeature).select("features", "prediction").collect().foreach {
         case Row(features: DenseVector, prediction1: Double) =>
           val prediction2 =
-            features(0) * model1.weights(0) + features(1) * model1.weights(1) + model1.intercept
+            features(0) * model1.coefficients(0) + features(1) * model1.coefficients(1) +
+              model1.intercept
           assert(prediction1 ~== prediction2 relTol 1E-5)
       }
     }
@@ -159,37 +160,37 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
       val modelWithoutIntercept2 = trainer2.fit(datasetWithDenseFeatureWithoutIntercept)
 
       /*
-         weights <- coef(glmnet(features, label, family="gaussian", alpha = 0, lambda = 0,
+         coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0, lambda = 0,
            intercept = FALSE))
-         > weights
+         > coefficients
           3 x 1 sparse Matrix of class "dgCMatrix"
                                    s0
          (Intercept)         .
          as.numeric.data.V2. 6.995908
          as.numeric.data.V3. 5.275131
        */
-      val weightsR = Vectors.dense(6.995908, 5.275131)
+      val coefficientsR = Vectors.dense(6.995908, 5.275131)
 
       assert(model1.intercept ~== 0 absTol 1E-3)
-      assert(model1.weights ~= weightsR relTol 1E-3)
+      assert(model1.coefficients ~= coefficientsR relTol 1E-3)
       assert(model2.intercept ~== 0 absTol 1E-3)
-      assert(model2.weights ~= weightsR relTol 1E-3)
+      assert(model2.coefficients ~= coefficientsR relTol 1E-3)
 
       /*
          Then again with the data with no intercept:
-         > weightsWithoutIntercept
+         > coefficientsWithourIntercept
           3 x 1 sparse Matrix of class "dgCMatrix"
                                    s0
          (Intercept)           .
          as.numeric.data3.V2. 4.70011
          as.numeric.data3.V3. 7.19943
        */
-      val weightsWithoutInterceptR = Vectors.dense(4.70011, 7.19943)
+      val coefficientsWithourInterceptR = Vectors.dense(4.70011, 7.19943)
 
       assert(modelWithoutIntercept1.intercept ~== 0 absTol 1E-3)
-      assert(modelWithoutIntercept1.weights ~= weightsWithoutInterceptR relTol 1E-3)
+      assert(modelWithoutIntercept1.coefficients ~= coefficientsWithourInterceptR relTol 1E-3)
       assert(modelWithoutIntercept2.intercept ~== 0 absTol 1E-3)
-      assert(modelWithoutIntercept2.weights ~= weightsWithoutInterceptR relTol 1E-3)
+      assert(modelWithoutIntercept2.coefficients ~= coefficientsWithourInterceptR relTol 1E-3)
     }
   }
 
@@ -211,8 +212,9 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         val model2 = trainer2.fit(datasetWithDenseFeature)
 
         /*
-           weights <- coef(glmnet(features, label, family="gaussian", alpha = 1.0, lambda = 0.57))
-           > weights
+           coefficients <- coef(glmnet(features, label, family="gaussian",
+             alpha = 1.0, lambda = 0.57 ))
+           > coefficients
             3 x 1 sparse Matrix of class "dgCMatrix"
                                     s0
            (Intercept)         6.24300
@@ -220,14 +222,14 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
            as.numeric.data.V3. 6.679841
          */
         val interceptR1 = 6.24300
-        val weightsR1 = Vectors.dense(4.024821, 6.679841)
+        val coefficientsR1 = Vectors.dense(4.024821, 6.679841)
         assert(model1.intercept ~== interceptR1 relTol 1E-3)
-        assert(model1.weights ~= weightsR1 relTol 1E-3)
+        assert(model1.coefficients ~= coefficientsR1 relTol 1E-3)
 
         /*
-           weights <- coef(glmnet(features, label, family="gaussian", alpha = 1.0, lambda = 0.57,
-             standardize=FALSE))
-           > weights
+           coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 1.0,
+             lambda = 0.57, standardize=FALSE ))
+           > coefficients
             3 x 1 sparse Matrix of class "dgCMatrix"
                                     s0
            (Intercept)         6.416948
@@ -235,16 +237,17 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
            as.numeric.data.V3. 6.724286
          */
         val interceptR2 = 6.416948
-        val weightsR2 = Vectors.dense(3.893869, 6.724286)
+        val coefficientsR2 = Vectors.dense(3.893869, 6.724286)
 
         assert(model2.intercept ~== interceptR2 relTol 1E-3)
-        assert(model2.weights ~= weightsR2 relTol 1E-3)
+        assert(model2.coefficients ~= coefficientsR2 relTol 1E-3)
 
         model1.transform(datasetWithDenseFeature).select("features", "prediction")
           .collect().foreach {
             case Row(features: DenseVector, prediction1: Double) =>
-              val prediction2 = features(0) * model1.weights(0) + features(1) * model1.weights(1) +
-                model1.intercept
+              val prediction2 =
+                features(0) * model1.coefficients(0) + features(1) * model1.coefficients(1) +
+                  model1.intercept
               assert(prediction1 ~== prediction2 relTol 1E-5)
         }
       }
@@ -269,9 +272,9 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         val model2 = trainer2.fit(datasetWithDenseFeature)
 
         /*
-           weights <- coef(glmnet(features, label, family="gaussian", alpha = 1.0, lambda = 0.57,
-             intercept=FALSE))
-           > weights
+           coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 1.0,
+             lambda = 0.57, intercept=FALSE ))
+           > coefficients
             3 x 1 sparse Matrix of class "dgCMatrix"
                                      s0
            (Intercept)          .
@@ -279,15 +282,15 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
            as.numeric.data.V3. 4.772913
          */
         val interceptR1 = 0.0
-        val weightsR1 = Vectors.dense(6.299752, 4.772913)
+        val coefficientsR1 = Vectors.dense(6.299752, 4.772913)
 
         assert(model1.intercept ~== interceptR1 absTol 1E-3)
-        assert(model1.weights ~= weightsR1 relTol 1E-3)
+        assert(model1.coefficients ~= coefficientsR1 relTol 1E-3)
 
         /*
-           weights <- coef(glmnet(features, label, family="gaussian", alpha = 1.0, lambda = 0.57,
-             intercept=FALSE, standardize=FALSE))
-           > weights
+           coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 1.0,
+             lambda = 0.57, intercept=FALSE, standardize=FALSE ))
+           > coefficients
             3 x 1 sparse Matrix of class "dgCMatrix"
                                      s0
            (Intercept)         .
@@ -295,16 +298,17 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
            as.numeric.data.V3. 4.764229
          */
         val interceptR2 = 0.0
-        val weightsR2 = Vectors.dense(6.232193, 4.764229)
+        val coefficientsR2 = Vectors.dense(6.232193, 4.764229)
 
         assert(model2.intercept ~== interceptR2 absTol 1E-3)
-        assert(model2.weights ~= weightsR2 relTol 1E-3)
+        assert(model2.coefficients ~= coefficientsR2 relTol 1E-3)
 
         model1.transform(datasetWithDenseFeature).select("features", "prediction")
           .collect().foreach {
             case Row(features: DenseVector, prediction1: Double) =>
-              val prediction2 = features(0) * model1.weights(0) + features(1) * model1.weights(1) +
-                model1.intercept
+              val prediction2 =
+                features(0) * model1.coefficients(0) + features(1) * model1.coefficients(1) +
+                  model1.intercept
               assert(prediction1 ~== prediction2 relTol 1E-5)
         }
       }
@@ -321,8 +325,8 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
       val model2 = trainer2.fit(datasetWithDenseFeature)
 
       /*
-         weights <- coef(glmnet(features, label, family="gaussian", alpha = 0.0, lambda = 2.3))
-         > weights
+         coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0.0, lambda = 2.3))
+         > coefficients
           3 x 1 sparse Matrix of class "dgCMatrix"
                                    s0
          (Intercept)         5.269376
@@ -330,15 +334,15 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
          as.numeric.data.V3. 5.712356)
        */
       val interceptR1 = 5.269376
-      val weightsR1 = Vectors.dense(3.736216, 5.712356)
+      val coefficientsR1 = Vectors.dense(3.736216, 5.712356)
 
       assert(model1.intercept ~== interceptR1 relTol 1E-3)
-      assert(model1.weights ~= weightsR1 relTol 1E-3)
+      assert(model1.coefficients ~= coefficientsR1 relTol 1E-3)
 
       /*
-         weights <- coef(glmnet(features, label, family="gaussian", alpha = 0.0, lambda = 2.3,
+         coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0.0, lambda = 2.3,
            standardize=FALSE))
-         > weights
+         > coefficients
           3 x 1 sparse Matrix of class "dgCMatrix"
                                    s0
          (Intercept)         5.791109
@@ -346,15 +350,16 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
          as.numeric.data.V3. 5.910406
        */
       val interceptR2 = 5.791109
-      val weightsR2 = Vectors.dense(3.435466, 5.910406)
+      val coefficientsR2 = Vectors.dense(3.435466, 5.910406)
 
       assert(model2.intercept ~== interceptR2 relTol 1E-3)
-      assert(model2.weights ~= weightsR2 relTol 1E-3)
+      assert(model2.coefficients ~= coefficientsR2 relTol 1E-3)
 
       model1.transform(datasetWithDenseFeature).select("features", "prediction").collect().foreach {
         case Row(features: DenseVector, prediction1: Double) =>
           val prediction2 =
-            features(0) * model1.weights(0) + features(1) * model1.weights(1) + model1.intercept
+            features(0) * model1.coefficients(0) + features(1) * model1.coefficients(1) +
+              model1.intercept
           assert(prediction1 ~== prediction2 relTol 1E-5)
       }
     }
@@ -370,9 +375,9 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
       val model2 = trainer2.fit(datasetWithDenseFeature)
 
       /*
-         weights <- coef(glmnet(features, label, family="gaussian", alpha = 0.0, lambda = 2.3,
+         coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0.0, lambda = 2.3,
            intercept = FALSE))
-         > weights
+         > coefficients
           3 x 1 sparse Matrix of class "dgCMatrix"
                                    s0
          (Intercept)         .
@@ -380,15 +385,15 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
          as.numeric.data.V3. 4.214502
        */
       val interceptR1 = 0.0
-      val weightsR1 = Vectors.dense(5.522875, 4.214502)
+      val coefficientsR1 = Vectors.dense(5.522875, 4.214502)
 
       assert(model1.intercept ~== interceptR1 absTol 1E-3)
-      assert(model1.weights ~= weightsR1 relTol 1E-3)
+      assert(model1.coefficients ~= coefficientsR1 relTol 1E-3)
 
       /*
-         weights <- coef(glmnet(features, label, family="gaussian", alpha = 0.0, lambda = 2.3,
+         coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0.0, lambda = 2.3,
            intercept = FALSE, standardize=FALSE))
-         > weights
+         > coefficients
           3 x 1 sparse Matrix of class "dgCMatrix"
                                    s0
          (Intercept)         .
@@ -396,15 +401,16 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
          as.numeric.data.V3. 4.187419
        */
       val interceptR2 = 0.0
-      val weightsR2 = Vectors.dense(5.263704, 4.187419)
+      val coefficientsR2 = Vectors.dense(5.263704, 4.187419)
 
       assert(model2.intercept ~== interceptR2 absTol 1E-3)
-      assert(model2.weights ~= weightsR2 relTol 1E-3)
+      assert(model2.coefficients ~= coefficientsR2 relTol 1E-3)
 
       model1.transform(datasetWithDenseFeature).select("features", "prediction").collect().foreach {
         case Row(features: DenseVector, prediction1: Double) =>
           val prediction2 =
-            features(0) * model1.weights(0) + features(1) * model1.weights(1) + model1.intercept
+            features(0) * model1.coefficients(0) + features(1) * model1.coefficients(1) +
+              model1.intercept
           assert(prediction1 ~== prediction2 relTol 1E-5)
       }
     }
@@ -428,8 +434,9 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         val model2 = trainer2.fit(datasetWithDenseFeature)
 
         /*
-           weights <- coef(glmnet(features, label, family="gaussian", alpha = 0.3, lambda = 1.6))
-           > weights
+           coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0.3,
+             lambda = 1.6 ))
+           > coefficients
             3 x 1 sparse Matrix of class "dgCMatrix"
                                      s0
            (Intercept)         6.324108
@@ -437,15 +444,15 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
            as.numeric.data.V3. 5.200403
          */
         val interceptR1 = 5.696056
-        val weightsR1 = Vectors.dense(3.670489, 6.001122)
+        val coefficientsR1 = Vectors.dense(3.670489, 6.001122)
 
         assert(model1.intercept ~== interceptR1 relTol 1E-3)
-        assert(model1.weights ~= weightsR1 relTol 1E-3)
+        assert(model1.coefficients ~= coefficientsR1 relTol 1E-3)
 
         /*
-           weights <- coef(glmnet(features, label, family="gaussian", alpha = 0.3, lambda = 1.6
+           coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0.3, lambda = 1.6
              standardize=FALSE))
-           > weights
+           > coefficients
             3 x 1 sparse Matrix of class "dgCMatrix"
                                      s0
            (Intercept)         6.114723
@@ -453,16 +460,17 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
            as.numeric.data.V3. 6.146531
          */
         val interceptR2 = 6.114723
-        val weightsR2 = Vectors.dense(3.409937, 6.146531)
+        val coefficientsR2 = Vectors.dense(3.409937, 6.146531)
 
         assert(model2.intercept ~== interceptR2 relTol 1E-3)
-        assert(model2.weights ~= weightsR2 relTol 1E-3)
+        assert(model2.coefficients ~= coefficientsR2 relTol 1E-3)
 
         model1.transform(datasetWithDenseFeature).select("features", "prediction")
           .collect().foreach {
           case Row(features: DenseVector, prediction1: Double) =>
-            val prediction2 = features(0) * model1.weights(0) + features(1) * model1.weights(1) +
-              model1.intercept
+            val prediction2 =
+              features(0) * model1.coefficients(0) + features(1) * model1.coefficients(1) +
+                model1.intercept
             assert(prediction1 ~== prediction2 relTol 1E-5)
         }
       }
@@ -487,9 +495,9 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         val model2 = trainer2.fit(datasetWithDenseFeature)
 
         /*
-           weights <- coef(glmnet(features, label, family="gaussian", alpha = 0.3, lambda = 1.6,
-             intercept=FALSE))
-           > weights
+           coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0.3,
+             lambda = 1.6, intercept=FALSE ))
+           > coefficients
             3 x 1 sparse Matrix of class "dgCMatrix"
                                       s0
            (Intercept)         .
@@ -497,15 +505,15 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
            as.numeric.dataM.V3. 4.322251
          */
         val interceptR1 = 0.0
-        val weightsR1 = Vectors.dense(5.673348, 4.322251)
+        val coefficientsR1 = Vectors.dense(5.673348, 4.322251)
 
         assert(model1.intercept ~== interceptR1 absTol 1E-3)
-        assert(model1.weights ~= weightsR1 relTol 1E-3)
+        assert(model1.coefficients ~= coefficientsR1 relTol 1E-3)
 
         /*
-           weights <- coef(glmnet(features, label, family="gaussian", alpha = 0.3, lambda = 1.6,
-             intercept=FALSE, standardize=FALSE))
-           > weights
+           coefficients <- coef(glmnet(features, label, family="gaussian", alpha = 0.3,
+             lambda = 1.6, intercept=FALSE, standardize=FALSE ))
+           > coefficients
             3 x 1 sparse Matrix of class "dgCMatrix"
                                      s0
            (Intercept)         .
@@ -513,16 +521,17 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
            as.numeric.data.V3. 4.297622
          */
         val interceptR2 = 0.0
-        val weightsR2 = Vectors.dense(5.477988, 4.297622)
+        val coefficientsR2 = Vectors.dense(5.477988, 4.297622)
 
         assert(model2.intercept ~== interceptR2 absTol 1E-3)
-        assert(model2.weights ~= weightsR2 relTol 1E-3)
+        assert(model2.coefficients ~= coefficientsR2 relTol 1E-3)
 
         model1.transform(datasetWithDenseFeature).select("features", "prediction")
           .collect().foreach {
           case Row(features: DenseVector, prediction1: Double) =>
-            val prediction2 = features(0) * model1.weights(0) + features(1) * model1.weights(1) +
-              model1.intercept
+            val prediction2 =
+              features(0) * model1.coefficients(0) + features(1) * model1.coefficients(1) +
+                model1.intercept
             assert(prediction1 ~== prediction2 relTol 1E-5)
         }
       }
@@ -554,7 +563,8 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
       val expectedResiduals = datasetWithDenseFeature.select("features", "label")
         .map { case Row(features: DenseVector, label: Double) =>
         val prediction =
-          features(0) * model.weights(0) + features(1) * model.weights(1) + model.intercept
+          features(0) * model.coefficients(0) + features(1) * model.coefficients(1) +
+            model.intercept
         label - prediction
       }
         .zip(model.summary.residuals.map(_.getDouble(0)))
@@ -663,9 +673,9 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
       val model1a1 = trainer1a.fit(weightedData)
       val model1b = trainer1b.fit(weightedData)
 
-      assert(model1a0.weights !~= model1a1.weights absTol 1E-3)
+      assert(model1a0.coefficients !~= model1a1.coefficients absTol 1E-3)
       assert(model1a0.intercept !~= model1a1.intercept absTol 1E-3)
-      assert(model1a0.weights ~== model1b.weights absTol 1E-3)
+      assert(model1a0.coefficients ~== model1b.coefficients absTol 1E-3)
       assert(model1a0.intercept ~== model1b.intercept absTol 1E-3)
 
       val trainer2a = (new LinearRegression).setFitIntercept(true)
@@ -675,9 +685,9 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
       val model2a0 = trainer2a.fit(data)
       val model2a1 = trainer2a.fit(weightedData)
       val model2b = trainer2b.fit(weightedData)
-      assert(model2a0.weights !~= model2a1.weights absTol 1E-3)
+      assert(model2a0.coefficients !~= model2a1.coefficients absTol 1E-3)
       assert(model2a0.intercept !~= model2a1.intercept absTol 1E-3)
-      assert(model2a0.weights ~== model2b.weights absTol 1E-3)
+      assert(model2a0.coefficients ~== model2b.coefficients absTol 1E-3)
       assert(model2a0.intercept ~== model2b.intercept absTol 1E-3)
 
       val trainer3a = (new LinearRegression).setFitIntercept(false)
@@ -687,8 +697,8 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
       val model3a0 = trainer3a.fit(data)
       val model3a1 = trainer3a.fit(weightedData)
       val model3b = trainer3b.fit(weightedData)
-      assert(model3a0.weights !~= model3a1.weights absTol 1E-3)
-      assert(model3a0.weights ~== model3b.weights absTol 1E-3)
+      assert(model3a0.coefficients !~= model3a1.coefficients absTol 1E-3)
+      assert(model3a0.coefficients ~== model3b.coefficients absTol 1E-3)
 
       val trainer4a = (new LinearRegression).setFitIntercept(false)
         .setElasticNetParam(0.0).setRegParam(0.21).setStandardization(false).setSolver(solver)
@@ -697,8 +707,8 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
       val model4a0 = trainer4a.fit(data)
       val model4a1 = trainer4a.fit(weightedData)
       val model4b = trainer4b.fit(weightedData)
-      assert(model4a0.weights !~= model4a1.weights absTol 1E-3)
-      assert(model4a0.weights ~== model4b.weights absTol 1E-3)
+      assert(model4a0.coefficients !~= model4a1.coefficients absTol 1E-3)
+      assert(model4a0.coefficients ~== model4b.coefficients absTol 1E-3)
     }
   }
 
diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py
index 4cbe7fbd482da..2e468f67b8987 100644
--- a/python/pyspark/ml/classification.py
+++ b/python/pyspark/ml/classification.py
@@ -15,6 +15,9 @@
 # limitations under the License.
 #
 
+import warnings
+
+from pyspark import since
 from pyspark.ml.util import keyword_only
 from pyspark.ml.wrapper import JavaEstimator, JavaModel
 from pyspark.ml.param.shared import *
@@ -189,8 +192,18 @@ def weights(self):
         """
         Model weights.
         """
+
+        warnings.warn("weights is deprecated. Use coefficients instead.")
         return self._call_java("weights")
 
+    @property
+    @since("1.6.0")
+    def coefficients(self):
+        """
+        Model coefficients.
+        """
+        return self._call_java("coefficients")
+
     @property
     def intercept(self):
         """
diff --git a/python/pyspark/ml/regression.py b/python/pyspark/ml/regression.py
index dc68815556d4e..ab26616f4a01d 100644
--- a/python/pyspark/ml/regression.py
+++ b/python/pyspark/ml/regression.py
@@ -15,6 +15,8 @@
 # limitations under the License.
 #
 
+import warnings
+
 from pyspark import since
 from pyspark.ml.util import keyword_only
 from pyspark.ml.wrapper import JavaEstimator, JavaModel
@@ -117,8 +119,18 @@ def weights(self):
         """
         Model weights.
         """
+
+        warnings.warn("weights is deprecated. Use coefficients instead.")
         return self._call_java("weights")
 
+    @property
+    @since("1.6.0")
+    def coefficients(self):
+        """
+        Model coefficients.
+        """
+        return self._call_java("coefficients")
+
     @property
     @since("1.4.0")
     def intercept(self):

From 476f4348e2ea57ea05f4b470abfe76d97eeb20ce Mon Sep 17 00:00:00 2001
From: Calvin Jia <jia.calvin@gmail.com>
Date: Mon, 2 Nov 2015 17:02:31 -0800
Subject: [PATCH 013/518] [SPARK-11236] [TEST-MAVEN] [TEST-HADOOP1.0] [CORE]
 Update Tachyon dependency 0.7.1 -> 0.8.1

This is a reopening of #9204 which failed hadoop1 sbt tests.

With the original PR, a classpath issue would occur due to the MIMA plugin pulling in hadoop-2.2 dependencies regardless of the hadoop version when building the `oldDeps` project. These affect the hadoop1 sbt build because they are placed in `lib_managed` and Tachyon 0.8.0's default hadoop version is 2.2.

Author: Calvin Jia <jia.calvin@gmail.com>

Closes #9395 from calvinjia/spark-11236.
---
 core/pom.xml         | 6 +-----
 make-distribution.sh | 8 ++++----
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/core/pom.xml b/core/pom.xml
index 1b6b13517bd56..570a25cf325a2 100644
--- a/core/pom.xml
+++ b/core/pom.xml
@@ -262,7 +262,7 @@
     <dependency>
       <groupId>org.tachyonproject</groupId>
       <artifactId>tachyon-client</artifactId>
-      <version>0.7.1</version>
+      <version>0.8.1</version>
       <exclusions>
         <exclusion>
           <groupId>org.apache.hadoop</groupId>
@@ -284,10 +284,6 @@
           <groupId>org.tachyonproject</groupId>
           <artifactId>tachyon-underfs-glusterfs</artifactId>
         </exclusion>
-        <exclusion>
-          <groupId>org.tachyonproject</groupId>
-          <artifactId>tachyon-underfs-s3</artifactId>
-        </exclusion>
       </exclusions>
     </dependency>
     <dependency>
diff --git a/make-distribution.sh b/make-distribution.sh
index 24418ace26270..e1c2afdbc6d87 100755
--- a/make-distribution.sh
+++ b/make-distribution.sh
@@ -33,9 +33,9 @@ SPARK_HOME="$(cd "`dirname "$0"`"; pwd)"
 DISTDIR="$SPARK_HOME/dist"
 
 SPARK_TACHYON=false
-TACHYON_VERSION="0.7.1"
+TACHYON_VERSION="0.8.1"
 TACHYON_TGZ="tachyon-${TACHYON_VERSION}-bin.tar.gz"
-TACHYON_URL="https://github.com/amplab/tachyon/releases/download/v${TACHYON_VERSION}/${TACHYON_TGZ}"
+TACHYON_URL="http://tachyon-project.org/downloads/files/${TACHYON_VERSION}/${TACHYON_TGZ}"
 
 MAKE_TGZ=false
 NAME=none
@@ -240,10 +240,10 @@ if [ "$SPARK_TACHYON" == "true" ]; then
   fi
 
   tar xzf "${TACHYON_TGZ}"
-  cp "tachyon-${TACHYON_VERSION}/core/target/tachyon-${TACHYON_VERSION}-jar-with-dependencies.jar" "$DISTDIR/lib"
+  cp "tachyon-${TACHYON_VERSION}/assembly/target/tachyon-assemblies-${TACHYON_VERSION}-jar-with-dependencies.jar" "$DISTDIR/lib"
   mkdir -p "$DISTDIR/tachyon/src/main/java/tachyon/web"
   cp -r "tachyon-${TACHYON_VERSION}"/{bin,conf,libexec} "$DISTDIR/tachyon"
-  cp -r "tachyon-${TACHYON_VERSION}"/core/src/main/java/tachyon/web "$DISTDIR/tachyon/src/main/java/tachyon/web"
+  cp -r "tachyon-${TACHYON_VERSION}"/servers/src/main/java/tachyon/web "$DISTDIR/tachyon/src/main/java/tachyon/web"
 
   if [[ `uname -a` == Darwin* ]]; then
     # need to run sed differently on osx

From 21ad846238a9a79564e2e99a1def89fd31a0870d Mon Sep 17 00:00:00 2001
From: DB Tsai <dbt@netflix.com>
Date: Mon, 2 Nov 2015 19:07:31 -0800
Subject: [PATCH 014/518] [MINOR][ML] removed the old `getModelWeights`
 function

Removed the old `getModelWeights` function which was private and renamed into `getModelCoefficients`

Author: DB Tsai <dbt@netflix.com>

Closes #9426 from dbtsai/feature-minor.
---
 .../scala/org/apache/spark/ml/r/SparkRWrappers.scala   | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala b/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
index 9162ec0e4e153..24f76de806d8f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
@@ -51,16 +51,6 @@ private[r] object SparkRWrappers {
     pipeline.fit(df)
   }
 
-  @deprecated("Use getModelCoefficients instead.", "1.6.0")
-  def getModelWeights(model: PipelineModel): Array[Double] = {
-    model.stages.last match {
-      case m: LinearRegressionModel =>
-        Array(m.intercept) ++ m.weights.toArray
-      case m: LogisticRegressionModel =>
-        Array(m.intercept) ++ m.weights.toArray
-    }
-  }
-
   def getModelCoefficients(model: PipelineModel): Array[Double] = {
     model.stages.last match {
       case m: LinearRegressionModel =>

From 2cef1bb0b560a03aa7308f694b0c66347b90c9ea Mon Sep 17 00:00:00 2001
From: Nong Li <nongli@gmail.com>
Date: Mon, 2 Nov 2015 19:18:45 -0800
Subject: [PATCH 015/518] =?UTF-8?q?[SPARK-5354][SQL]=20Cached=20tables=20s?=
 =?UTF-8?q?hould=20preserve=20partitioning=20and=20ord=E2=80=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…ering.

For cached tables, we can just maintain the partitioning and ordering from the
source relation.

Author: Nong Li <nongli@gmail.com>

Closes #9404 from nongli/spark-5354.
---
 .../columnar/InMemoryColumnarTableScan.scala  |  7 +++
 .../apache/spark/sql/execution/Exchange.scala | 40 ++++++++++---
 .../apache/spark/sql/CachedTableSuite.scala   | 59 +++++++++++++++++++
 3 files changed, 97 insertions(+), 9 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
index b4607b12fcefa..7eb1ad7cd8198 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
@@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Statistics}
+import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.execution.{ConvertToUnsafe, LeafNode, SparkPlan}
 import org.apache.spark.sql.types.UserDefinedType
 import org.apache.spark.storage.StorageLevel
@@ -209,6 +210,12 @@ private[sql] case class InMemoryColumnarTableScan(
 
   override def output: Seq[Attribute] = attributes
 
+  // The cached version does not change the outputPartitioning of the original SparkPlan.
+  override def outputPartitioning: Partitioning = relation.child.outputPartitioning
+
+  // The cached version does not change the outputOrdering of the original SparkPlan.
+  override def outputOrdering: Seq[SortOrder] = relation.child.outputOrdering
+
   override def outputsUnsafeRows: Boolean = true
 
   private def statsFor(a: Attribute) = relation.partitionStatistics.forAttribute(a)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
index 7f60c8f5eaa95..e81108b7884d1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
@@ -194,12 +194,13 @@ case class Exchange(newPartitioning: Partitioning, child: SparkPlan) extends Una
  */
 private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[SparkPlan] {
   // TODO: Determine the number of partitions.
-  private def numPartitions: Int = sqlContext.conf.numShufflePartitions
+  private def defaultPartitions: Int = sqlContext.conf.numShufflePartitions
 
   /**
    * Given a required distribution, returns a partitioning that satisfies that distribution.
    */
-  private def canonicalPartitioning(requiredDistribution: Distribution): Partitioning = {
+  private def createPartitioning(requiredDistribution: Distribution,
+      numPartitions: Int): Partitioning = {
     requiredDistribution match {
       case AllTuples => SinglePartition
       case ClusteredDistribution(clustering) => HashPartitioning(clustering, numPartitions)
@@ -220,7 +221,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       if (child.outputPartitioning.satisfies(distribution)) {
         child
       } else {
-        Exchange(canonicalPartitioning(distribution), child)
+        Exchange(createPartitioning(distribution, defaultPartitions), child)
       }
     }
 
@@ -229,12 +230,33 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
     if (children.length > 1
         && requiredChildDistributions.toSet != Set(UnspecifiedDistribution)
         && !Partitioning.allCompatible(children.map(_.outputPartitioning))) {
-      children = children.zip(requiredChildDistributions).map { case (child, distribution) =>
-        val targetPartitioning = canonicalPartitioning(distribution)
-        if (child.outputPartitioning.guarantees(targetPartitioning)) {
-          child
-        } else {
-          Exchange(targetPartitioning, child)
+
+      // First check if the existing partitions of the children all match. This means they are
+      // partitioned by the same partitioning into the same number of partitions. In that case,
+      // don't try to make them match `defaultPartitions`, just use the existing partitioning.
+      // TODO: this should be a cost based descision. For example, a big relation should probably
+      // maintain its existing number of partitions and smaller partitions should be shuffled.
+      // defaultPartitions is arbitrary.
+      val numPartitions = children.head.outputPartitioning.numPartitions
+      val useExistingPartitioning = children.zip(requiredChildDistributions).forall {
+        case (child, distribution) => {
+          child.outputPartitioning.guarantees(
+            createPartitioning(distribution, numPartitions))
+        }
+      }
+
+      children = if (useExistingPartitioning) {
+        children
+      } else {
+        children.zip(requiredChildDistributions).map {
+          case (child, distribution) => {
+            val targetPartitioning = createPartitioning(distribution, defaultPartitions)
+            if (child.outputPartitioning.guarantees(targetPartitioning)) {
+              child
+            } else {
+              Exchange(targetPartitioning, child)
+            }
+          }
         }
       }
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
index fd566c8276bc1..605954b105d1e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql
 
 import org.apache.spark.sql.catalyst.analysis.NoSuchTableException
+import org.apache.spark.sql.execution.Exchange
 import org.apache.spark.sql.execution.PhysicalRDD
 
 import scala.concurrent.duration._
@@ -353,4 +354,62 @@ class CachedTableSuite extends QueryTest with SharedSQLContext {
     assert(sparkPlan.collect { case e: InMemoryColumnarTableScan => e }.size === 3)
     assert(sparkPlan.collect { case e: PhysicalRDD => e }.size === 0)
   }
+
+  /**
+   * Verifies that the plan for `df` contains `expected` number of Exchange operators.
+   */
+  private def verifyNumExchanges(df: DataFrame, expected: Int): Unit = {
+    assert(df.queryExecution.executedPlan.collect { case e: Exchange => e }.size == expected)
+  }
+
+  test("A cached table preserves the partitioning and ordering of its cached SparkPlan") {
+    val table3x = testData.unionAll(testData).unionAll(testData)
+    table3x.registerTempTable("testData3x")
+
+    sql("SELECT key, value FROM testData3x ORDER BY key").registerTempTable("orderedTable")
+    sqlContext.cacheTable("orderedTable")
+    assertCached(sqlContext.table("orderedTable"))
+    // Should not have an exchange as the query is already sorted on the group by key.
+    verifyNumExchanges(sql("SELECT key, count(*) FROM orderedTable GROUP BY key"), 0)
+    checkAnswer(
+      sql("SELECT key, count(*) FROM orderedTable GROUP BY key ORDER BY key"),
+      sql("SELECT key, count(*) FROM testData3x GROUP BY key ORDER BY key").collect())
+    sqlContext.uncacheTable("orderedTable")
+
+    // Set up two tables distributed in the same way. Try this with the data distributed into
+    // different number of partitions.
+    for (numPartitions <- 1 until 10 by 4) {
+      testData.distributeBy(Column("key") :: Nil, numPartitions).registerTempTable("t1")
+      testData2.distributeBy(Column("a") :: Nil, numPartitions).registerTempTable("t2")
+      sqlContext.cacheTable("t1")
+      sqlContext.cacheTable("t2")
+
+      // Joining them should result in no exchanges.
+      verifyNumExchanges(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"), 0)
+      checkAnswer(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"),
+        sql("SELECT * FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a"))
+
+      // Grouping on the partition key should result in no exchanges
+      verifyNumExchanges(sql("SELECT count(*) FROM t1 GROUP BY key"), 0)
+      checkAnswer(sql("SELECT count(*) FROM t1 GROUP BY key"),
+        sql("SELECT count(*) FROM testData GROUP BY key"))
+
+      sqlContext.uncacheTable("t1")
+      sqlContext.uncacheTable("t2")
+      sqlContext.dropTempTable("t1")
+      sqlContext.dropTempTable("t2")
+    }
+
+    // Distribute the tables into non-matching number of partitions. Need to shuffle.
+    testData.distributeBy(Column("key") :: Nil, 6).registerTempTable("t1")
+    testData2.distributeBy(Column("a") :: Nil, 3).registerTempTable("t2")
+    sqlContext.cacheTable("t1")
+    sqlContext.cacheTable("t2")
+
+    verifyNumExchanges(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"), 2)
+    sqlContext.uncacheTable("t1")
+    sqlContext.uncacheTable("t2")
+    sqlContext.dropTempTable("t1")
+    sqlContext.dropTempTable("t2")
+  }
 }

From 9cb5c731dadff9539126362827a258d6b65754bb Mon Sep 17 00:00:00 2001
From: Nong Li <nongli@gmail.com>
Date: Mon, 2 Nov 2015 20:32:08 -0800
Subject: [PATCH 016/518] [SPARK-11329][SQL] Support star expansion for
 structs.

1. Supporting expanding structs in Projections. i.e.
  "SELECT s.*" where s is a struct type.
  This is fixed by allowing the expand function to handle structs in addition to tables.

2. Supporting expanding * inside aggregate functions of structs.
   "SELECT max(struct(col1, structCol.*))"
   This requires recursively expanding the expressions. In this case, it it the aggregate
   expression "max(...)" and we need to recursively expand its children inputs.

Author: Nong Li <nongli@gmail.com>

Closes #9343 from nongli/spark-11329.
---
 .../apache/spark/sql/catalyst/SqlParser.scala |   6 +-
 .../sql/catalyst/analysis/Analyzer.scala      |  46 ++++--
 .../sql/catalyst/analysis/unresolved.scala    |  78 +++++++---
 .../scala/org/apache/spark/sql/Column.scala   |   3 +-
 .../org/apache/spark/sql/SQLQuerySuite.scala  | 133 ++++++++++++++++++
 .../org/apache/spark/sql/hive/HiveQl.scala    |   2 +-
 6 files changed, 230 insertions(+), 38 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
index 0fef04302714e..d7567e8613e3c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
@@ -466,9 +466,9 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
 
   protected lazy val baseExpression: Parser[Expression] =
     ( "*" ^^^ UnresolvedStar(None)
-    | ident <~ "." ~ "*" ^^ { case tableName => UnresolvedStar(Option(tableName)) }
-    | primary
-    )
+    | (ident <~ "."). + <~ "*" ^^ { case target => { UnresolvedStar(Option(target)) }
+    } | primary
+   )
 
   protected lazy val signedPrimary: Parser[Expression] =
     sign ~ primary ^^ { case s ~ e => if (s == "-") UnaryMinus(e) else e }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index beabacfc88e32..912c967b95f08 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -279,6 +279,24 @@ class Analyzer(
    * a logical plan node's children.
    */
   object ResolveReferences extends Rule[LogicalPlan] {
+    /**
+     * Foreach expression, expands the matching attribute.*'s in `child`'s input for the subtree
+     * rooted at each expression.
+     */
+    def expandStarExpressions(exprs: Seq[Expression], child: LogicalPlan): Seq[Expression] = {
+      exprs.flatMap {
+        case s: Star => s.expand(child, resolver)
+        case e =>
+          e.transformDown {
+            case f1: UnresolvedFunction if containsStar(f1.children) =>
+              f1.copy(children = f1.children.flatMap {
+                case s: Star => s.expand(child, resolver)
+                case o => o :: Nil
+              })
+          } :: Nil
+      }
+    }
+
     def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
       case p: LogicalPlan if !p.childrenResolved => p
 
@@ -286,44 +304,42 @@ class Analyzer(
       case p @ Project(projectList, child) if containsStar(projectList) =>
         Project(
           projectList.flatMap {
-            case s: Star => s.expand(child.output, resolver)
+            case s: Star => s.expand(child, resolver)
             case UnresolvedAlias(f @ UnresolvedFunction(_, args, _)) if containsStar(args) =>
-              val expandedArgs = args.flatMap {
-                case s: Star => s.expand(child.output, resolver)
-                case o => o :: Nil
-              }
-              UnresolvedAlias(child = f.copy(children = expandedArgs)) :: Nil
+              val newChildren = expandStarExpressions(args, child)
+              UnresolvedAlias(child = f.copy(children = newChildren)) :: Nil
+            case Alias(f @ UnresolvedFunction(_, args, _), name) if containsStar(args) =>
+              val newChildren = expandStarExpressions(args, child)
+              Alias(child = f.copy(children = newChildren), name)() :: Nil
             case UnresolvedAlias(c @ CreateArray(args)) if containsStar(args) =>
               val expandedArgs = args.flatMap {
-                case s: Star => s.expand(child.output, resolver)
+                case s: Star => s.expand(child, resolver)
                 case o => o :: Nil
               }
               UnresolvedAlias(c.copy(children = expandedArgs)) :: Nil
             case UnresolvedAlias(c @ CreateStruct(args)) if containsStar(args) =>
               val expandedArgs = args.flatMap {
-                case s: Star => s.expand(child.output, resolver)
+                case s: Star => s.expand(child, resolver)
                 case o => o :: Nil
               }
               UnresolvedAlias(c.copy(children = expandedArgs)) :: Nil
             case o => o :: Nil
           },
           child)
+
       case t: ScriptTransformation if containsStar(t.input) =>
         t.copy(
           input = t.input.flatMap {
-            case s: Star => s.expand(t.child.output, resolver)
+            case s: Star => s.expand(t.child, resolver)
             case o => o :: Nil
           }
         )
 
       // If the aggregate function argument contains Stars, expand it.
       case a: Aggregate if containsStar(a.aggregateExpressions) =>
-        a.copy(
-          aggregateExpressions = a.aggregateExpressions.flatMap {
-            case s: Star => s.expand(a.child.output, resolver)
-            case o => o :: Nil
-          }
-        )
+        val expanded = expandStarExpressions(a.aggregateExpressions, a.child)
+            .map(_.asInstanceOf[NamedExpression])
+        a.copy(aggregateExpressions = expanded)
 
       // Special handling for cases when self-join introduce duplicate expression ids.
       case j @ Join(left, right, _, _) if !j.selfJoinResolved =>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
index c97365003935e..6975662e2b738 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
@@ -18,12 +18,12 @@
 package org.apache.spark.sql.catalyst.analysis
 
 import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
-import org.apache.spark.sql.catalyst.{TableIdentifier, errors}
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.LeafNode
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, LeafNode}
 import org.apache.spark.sql.catalyst.trees.TreeNode
-import org.apache.spark.sql.types.DataType
+import org.apache.spark.sql.catalyst.{TableIdentifier, errors}
+import org.apache.spark.sql.types.{DataType, StructType}
 
 /**
  * Thrown when an invalid attempt is made to access a property of a tree that has yet to be fully
@@ -158,7 +158,7 @@ abstract class Star extends LeafExpression with NamedExpression {
   override def toAttribute: Attribute = throw new UnresolvedException(this, "toAttribute")
   override lazy val resolved = false
 
-  def expand(input: Seq[Attribute], resolver: Resolver): Seq[NamedExpression]
+  def expand(input: LogicalPlan, resolver: Resolver): Seq[NamedExpression]
 }
 
 
@@ -166,26 +166,68 @@ abstract class Star extends LeafExpression with NamedExpression {
  * Represents all of the input attributes to a given relational operator, for example in
  * "SELECT * FROM ...".
  *
- * @param table an optional table that should be the target of the expansion.  If omitted all
- *              tables' columns are produced.
+ * This is also used to expand structs. For example:
+ * "SELECT record.* from (SELECT struct(a,b,c) as record ...)
+ *
+ * @param target an optional name that should be the target of the expansion.  If omitted all
+ *              targets' columns are produced. This can either be a table name or struct name. This
+ *              is a list of identifiers that is the path of the expansion.
  */
-case class UnresolvedStar(table: Option[String]) extends Star with Unevaluable {
+case class UnresolvedStar(target: Option[Seq[String]]) extends Star with Unevaluable {
+
+  override def expand(input: LogicalPlan, resolver: Resolver): Seq[NamedExpression] = {
 
-  override def expand(input: Seq[Attribute], resolver: Resolver): Seq[NamedExpression] = {
-    val expandedAttributes: Seq[Attribute] = table match {
+    // First try to expand assuming it is table.*.
+    val expandedAttributes: Seq[Attribute] = target match {
       // If there is no table specified, use all input attributes.
-      case None => input
+      case None => input.output
       // If there is a table, pick out attributes that are part of this table.
-      case Some(t) => input.filter(_.qualifiers.filter(resolver(_, t)).nonEmpty)
+      case Some(t) => if (t.size == 1) {
+        input.output.filter(_.qualifiers.filter(resolver(_, t.head)).nonEmpty)
+      } else {
+        List()
+      }
     }
-    expandedAttributes.zip(input).map {
-      case (n: NamedExpression, _) => n
-      case (e, originalAttribute) =>
-        Alias(e, originalAttribute.name)(qualifiers = originalAttribute.qualifiers)
+    if (!expandedAttributes.isEmpty) {
+      if (expandedAttributes.forall(_.isInstanceOf[NamedExpression])) {
+        return expandedAttributes
+      } else {
+        require(expandedAttributes.size == input.output.size)
+        expandedAttributes.zip(input.output).map {
+          case (e, originalAttribute) =>
+            Alias(e, originalAttribute.name)(qualifiers = originalAttribute.qualifiers)
+        }
+      }
+      return expandedAttributes
+    }
+
+    require(target.isDefined)
+
+    // Try to resolve it as a struct expansion. If there is a conflict and both are possible,
+    // (i.e. [name].* is both a table and a struct), the struct path can always be qualified.
+    val attribute = input.resolve(target.get, resolver)
+    if (attribute.isDefined) {
+      // This target resolved to an attribute in child. It must be a struct. Expand it.
+      attribute.get.dataType match {
+        case s: StructType => {
+          s.fields.map( f => {
+            val extract = GetStructField(attribute.get, f, s.getFieldIndex(f.name).get)
+            Alias(extract, target.get + "." + f.name)()
+          })
+        }
+        case _ => {
+          throw new AnalysisException("Can only star expand struct data types. Attribute: `" +
+            target.get + "`")
+        }
+      }
+    } else {
+      val from = input.inputSet.map(_.name).mkString(", ")
+      val targetString = target.get.mkString(".")
+      throw new AnalysisException(s"cannot resolve '$targetString.*' give input columns '$from'")
     }
   }
 
-  override def toString: String = table.map(_ + ".").getOrElse("") + "*"
+  override def toString: String = target.map(_ + ".").getOrElse("") + "*"
 }
 
 /**
@@ -225,7 +267,7 @@ case class MultiAlias(child: Expression, names: Seq[String])
  * @param expressions Expressions to expand.
  */
 case class ResolvedStar(expressions: Seq[NamedExpression]) extends Star with Unevaluable {
-  override def expand(input: Seq[Attribute], resolver: Resolver): Seq[NamedExpression] = expressions
+  override def expand(input: LogicalPlan, resolver: Resolver): Seq[NamedExpression] = expressions
   override def toString: String = expressions.mkString("ResolvedStar(", ", ", ")")
 }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index e4f4cf1533ac4..3cde9d6cb4708 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -60,7 +60,8 @@ class Column(protected[sql] val expr: Expression) extends Logging {
 
   def this(name: String) = this(name match {
     case "*" => UnresolvedStar(None)
-    case _ if name.endsWith(".*") => UnresolvedStar(Some(name.substring(0, name.length - 2)))
+    case _ if name.endsWith(".*") => UnresolvedStar(Some(UnresolvedAttribute.parseAttributeName(
+      name.substring(0, name.length - 2))))
     case _ => UnresolvedAttribute.quotedString(name)
   })
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index 5413ef1287da1..ee54bff24b196 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -1932,4 +1932,137 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
       assert(sampled.count() == sampledOdd.count() + sampledEven.count())
     }
   }
+
+  test("Struct Star Expansion") {
+    val structDf = testData2.select("a", "b").as("record")
+
+    checkAnswer(
+      structDf.select($"record.a", $"record.b"),
+      Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
+
+    checkAnswer(
+      structDf.select($"record.*"),
+      Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
+
+    checkAnswer(
+      structDf.select($"record.*", $"record.*"),
+      Row(1, 1, 1, 1) :: Row(1, 2, 1, 2) :: Row(2, 1, 2, 1) :: Row(2, 2, 2, 2) ::
+        Row(3, 1, 3, 1) :: Row(3, 2, 3, 2) :: Nil)
+
+    checkAnswer(
+      sql("select struct(a, b) as r1, struct(b, a) as r2 from testData2").select($"r1.*", $"r2.*"),
+      Row(1, 1, 1, 1) :: Row(1, 2, 2, 1) :: Row(2, 1, 1, 2) :: Row(2, 2, 2, 2) ::
+        Row(3, 1, 1, 3) :: Row(3, 2, 2, 3) :: Nil)
+
+    // Try with a registered table.
+    sql("select struct(a, b) as record from testData2").registerTempTable("structTable")
+    checkAnswer(sql("SELECT record.* FROM structTable"),
+      Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
+
+    checkAnswer(sql(
+      """
+        | SELECT min(struct(record.*)) FROM
+        |   (select struct(a,b) as record from testData2) tmp
+      """.stripMargin),
+      Row(Row(1, 1)) :: Nil)
+
+    // Try with an alias on the select list
+    checkAnswer(sql(
+      """
+        | SELECT max(struct(record.*)) as r FROM
+        |   (select struct(a,b) as record from testData2) tmp
+      """.stripMargin).select($"r.*"),
+      Row(3, 2) :: Nil)
+
+    // With GROUP BY
+    checkAnswer(sql(
+      """
+        | SELECT min(struct(record.*)) FROM
+        |   (select a as a, struct(a,b) as record from testData2) tmp
+        | GROUP BY a
+      """.stripMargin),
+      Row(Row(1, 1)) :: Row(Row(2, 1)) :: Row(Row(3, 1)) :: Nil)
+
+    // With GROUP BY and alias
+    checkAnswer(sql(
+      """
+        | SELECT max(struct(record.*)) as r FROM
+        |   (select a as a, struct(a,b) as record from testData2) tmp
+        | GROUP BY a
+      """.stripMargin).select($"r.*"),
+      Row(1, 2) :: Row(2, 2) :: Row(3, 2) :: Nil)
+
+    // With GROUP BY and alias and additional fields in the struct
+    checkAnswer(sql(
+      """
+        | SELECT max(struct(a, record.*, b)) as r FROM
+        |   (select a as a, b as b, struct(a,b) as record from testData2) tmp
+        | GROUP BY a
+      """.stripMargin).select($"r.*"),
+      Row(1, 1, 2, 2) :: Row(2, 2, 2, 2) :: Row(3, 3, 2, 2) :: Nil)
+
+    // Create a data set that contains nested structs.
+    val nestedStructData = sql(
+      """
+        | SELECT struct(r1, r2) as record FROM
+        |   (SELECT struct(a, b) as r1, struct(b, a) as r2 FROM testData2) tmp
+      """.stripMargin)
+
+    checkAnswer(nestedStructData.select($"record.*"),
+      Row(Row(1, 1), Row(1, 1)) :: Row(Row(1, 2), Row(2, 1)) :: Row(Row(2, 1), Row(1, 2)) ::
+        Row(Row(2, 2), Row(2, 2)) :: Row(Row(3, 1), Row(1, 3)) :: Row(Row(3, 2), Row(2, 3)) :: Nil)
+    checkAnswer(nestedStructData.select($"record.r1"),
+      Row(Row(1, 1)) :: Row(Row(1, 2)) :: Row(Row(2, 1)) :: Row(Row(2, 2)) ::
+        Row(Row(3, 1)) :: Row(Row(3, 2)) :: Nil)
+    checkAnswer(
+      nestedStructData.select($"record.r1.*"),
+      Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
+
+    // Try with a registered table
+    nestedStructData.registerTempTable("nestedStructTable")
+    checkAnswer(sql("SELECT record.* FROM nestedStructTable"),
+      nestedStructData.select($"record.*"))
+    checkAnswer(sql("SELECT record.r1 FROM nestedStructTable"),
+      nestedStructData.select($"record.r1"))
+    checkAnswer(sql("SELECT record.r1.* FROM nestedStructTable"),
+      nestedStructData.select($"record.r1.*"))
+
+    // Create paths with unusual characters.
+    val specialCharacterPath = sql(
+      """
+        | SELECT struct(`col$.a_`, `a.b.c.`) as `r&&b.c` FROM
+        |   (SELECT struct(a, b) as `col$.a_`, struct(b, a) as `a.b.c.` FROM testData2) tmp
+      """.stripMargin)
+    specialCharacterPath.registerTempTable("specialCharacterTable")
+    checkAnswer(specialCharacterPath.select($"`r&&b.c`.*"),
+      nestedStructData.select($"record.*"))
+    checkAnswer(sql("SELECT `r&&b.c`.`col$.a_` FROM specialCharacterTable"),
+      nestedStructData.select($"record.r1"))
+    checkAnswer(sql("SELECT `r&&b.c`.`a.b.c.` FROM specialCharacterTable"),
+      nestedStructData.select($"record.r2"))
+    checkAnswer(sql("SELECT `r&&b.c`.`col$.a_`.* FROM specialCharacterTable"),
+      nestedStructData.select($"record.r1.*"))
+
+    // Try star expanding a scalar. This should fail.
+    assert(intercept[AnalysisException](sql("select a.* from testData2")).getMessage.contains(
+      "Can only star expand struct data types."))
+
+    // Try resolving something not there.
+    assert(intercept[AnalysisException](sql("SELECT abc.* FROM nestedStructTable"))
+      .getMessage.contains("cannot resolve"))
+  }
+
+
+  test("Struct Star Expansion - Name conflict") {
+    // Create a data set that contains a naming conflict
+    val nameConflict = sql("SELECT struct(a, b) as nameConflict, a as a FROM testData2")
+    nameConflict.registerTempTable("nameConflict")
+    // Unqualified should resolve to table.
+    checkAnswer(sql("SELECT nameConflict.* FROM nameConflict"),
+      Row(Row(1, 1), 1) :: Row(Row(1, 2), 1) :: Row(Row(2, 1), 2) :: Row(Row(2, 2), 2) ::
+        Row(Row(3, 1), 3) :: Row(Row(3, 2), 3) :: Nil)
+    // Qualify the struct type with the table name.
+    checkAnswer(sql("SELECT nameConflict.nameConflict.* FROM nameConflict"),
+      Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
+  }
 }
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
index 3697761f20c28..ab88c1e68fd72 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
@@ -1505,7 +1505,7 @@ https://cwiki.apache.org/confluence/display/Hive/Enhanced+Aggregation%2C+Cube%2C
     // The format of dbName.tableName.* cannot be parsed by HiveParser. TOK_TABNAME will only
     // has a single child which is tableName.
     case Token("TOK_ALLCOLREF", Token("TOK_TABNAME", Token(name, Nil) :: Nil) :: Nil) =>
-      UnresolvedStar(Some(name))
+      UnresolvedStar(Some(UnresolvedAttribute.parseAttributeName(name)))
 
     /* Aggregate Functions */
     case Token("TOK_FUNCTIONSTAR", Token(COUNT(), Nil) :: Nil) => Count(Literal(1))

From efaa4721b511a1d29229facde6457a6dcda18966 Mon Sep 17 00:00:00 2001
From: Yves Raimond <yraimond@netflix.com>
Date: Mon, 2 Nov 2015 20:35:59 -0800
Subject: [PATCH 017/518] [SPARK-11432][GRAPHX] Personalized PageRank shouldn't
 use uniform initialization

Changes the personalized pagerank initialization to be non-uniform.

Author: Yves Raimond <yraimond@netflix.com>

Closes #9386 from moustaki/personalized-pagerank-init.
---
 .../apache/spark/graphx/lib/PageRank.scala    | 29 ++++++++++++-------
 .../spark/graphx/lib/PageRankSuite.scala      | 13 ++++++---
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/graphx/src/main/scala/org/apache/spark/graphx/lib/PageRank.scala b/graphx/src/main/scala/org/apache/spark/graphx/lib/PageRank.scala
index 8c0a461e99fa4..52b237fc15093 100644
--- a/graphx/src/main/scala/org/apache/spark/graphx/lib/PageRank.scala
+++ b/graphx/src/main/scala/org/apache/spark/graphx/lib/PageRank.scala
@@ -104,18 +104,23 @@ object PageRank extends Logging {
       graph: Graph[VD, ED], numIter: Int, resetProb: Double = 0.15,
       srcId: Option[VertexId] = None): Graph[Double, Double] =
   {
+    val personalized = srcId isDefined
+    val src: VertexId = srcId.getOrElse(-1L)
+
     // Initialize the PageRank graph with each edge attribute having
-    // weight 1/outDegree and each vertex with attribute 1.0.
+    // weight 1/outDegree and each vertex with attribute resetProb.
+    // When running personalized pagerank, only the source vertex
+    // has an attribute resetProb. All others are set to 0.
     var rankGraph: Graph[Double, Double] = graph
       // Associate the degree with each vertex
       .outerJoinVertices(graph.outDegrees) { (vid, vdata, deg) => deg.getOrElse(0) }
       // Set the weight on the edges based on the degree
       .mapTriplets( e => 1.0 / e.srcAttr, TripletFields.Src )
       // Set the vertex attributes to the initial pagerank values
-      .mapVertices( (id, attr) => resetProb )
+      .mapVertices { (id, attr) =>
+        if (!(id != src && personalized)) resetProb else 0.0
+      }
 
-    val personalized = srcId isDefined
-    val src: VertexId = srcId.getOrElse(-1L)
     def delta(u: VertexId, v: VertexId): Double = { if (u == v) 1.0 else 0.0 }
 
     var iteration = 0
@@ -192,6 +197,9 @@ object PageRank extends Logging {
       graph: Graph[VD, ED], tol: Double, resetProb: Double = 0.15,
       srcId: Option[VertexId] = None): Graph[Double, Double] =
   {
+    val personalized = srcId.isDefined
+    val src: VertexId = srcId.getOrElse(-1L)
+
     // Initialize the pagerankGraph with each edge attribute
     // having weight 1/outDegree and each vertex with attribute 1.0.
     val pagerankGraph: Graph[(Double, Double), Double] = graph
@@ -202,13 +210,11 @@ object PageRank extends Logging {
       // Set the weight on the edges based on the degree
       .mapTriplets( e => 1.0 / e.srcAttr )
       // Set the vertex attributes to (initalPR, delta = 0)
-      .mapVertices( (id, attr) => (0.0, 0.0) )
+      .mapVertices { (id, attr) =>
+        if (id == src) (resetProb, Double.NegativeInfinity) else (0.0, 0.0)
+      }
       .cache()
 
-    val personalized = srcId.isDefined
-    val src: VertexId = srcId.getOrElse(-1L)
-
-
     // Define the three functions needed to implement PageRank in the GraphX
     // version of Pregel
     def vertexProgram(id: VertexId, attr: (Double, Double), msgSum: Double): (Double, Double) = {
@@ -225,7 +231,8 @@ object PageRank extends Logging {
       teleport = oldPR*delta
 
       val newPR = teleport + (1.0 - resetProb) * msgSum
-      (newPR, newPR - oldPR)
+      val newDelta = if (lastDelta == Double.NegativeInfinity) newPR else newPR - oldPR
+      (newPR, newDelta)
     }
 
     def sendMessage(edge: EdgeTriplet[(Double, Double), Double]) = {
@@ -239,7 +246,7 @@ object PageRank extends Logging {
     def messageCombiner(a: Double, b: Double): Double = a + b
 
     // The initial message received by all vertices in PageRank
-    val initialMessage = resetProb / (1.0 - resetProb)
+    val initialMessage = if (personalized) 0.0 else resetProb / (1.0 - resetProb)
 
     // Execute a dynamic version of Pregel.
     val vp = if (personalized) {
diff --git a/graphx/src/test/scala/org/apache/spark/graphx/lib/PageRankSuite.scala b/graphx/src/test/scala/org/apache/spark/graphx/lib/PageRankSuite.scala
index 45f1e3011035e..bdff31446f8ee 100644
--- a/graphx/src/test/scala/org/apache/spark/graphx/lib/PageRankSuite.scala
+++ b/graphx/src/test/scala/org/apache/spark/graphx/lib/PageRankSuite.scala
@@ -109,17 +109,22 @@ class PageRankSuite extends SparkFunSuite with LocalSparkContext {
       assert(notMatching === 0)
 
       val staticErrors = staticRanks2.map { case (vid, pr) =>
-        val correct = (vid > 0 && pr == resetProb) ||
-          (vid == 0 && math.abs(pr - (resetProb + (1.0 - resetProb) * (resetProb *
-            (nVertices - 1)) )) < 1.0E-5)
+        val correct = (vid > 0 && pr == 0.0) ||
+          (vid == 0 && pr == resetProb)
         if (!correct) 1 else 0
       }
       assert(staticErrors.sum === 0)
 
       val dynamicRanks = starGraph.personalizedPageRank(0, 0, resetProb).vertices.cache()
       assert(compareRanks(staticRanks2, dynamicRanks) < errorTol)
+
+      // We have one outbound edge from 1 to 0
+      val otherStaticRanks2 = starGraph.staticPersonalizedPageRank(1, numIter = 2, resetProb)
+        .vertices.cache()
+      val otherDynamicRanks = starGraph.personalizedPageRank(1, 0, resetProb).vertices.cache()
+      assert(compareRanks(otherDynamicRanks, otherStaticRanks2) < errorTol)
     }
-  } // end of test Star PageRank
+  } // end of test Star PersonalPageRank
 
   test("Grid PageRank") {
     withSpark { sc =>

From 9cf56c96b7d02a14175d40b336da14c2e1c88339 Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Mon, 2 Nov 2015 21:18:38 -0800
Subject: [PATCH 018/518] [SPARK-11469][SQL] Allow users to define
 nondeterministic udfs.

This is the first task (https://issues.apache.org/jira/browse/SPARK-11469) of https://issues.apache.org/jira/browse/SPARK-11438

Author: Yin Huai <yhuai@databricks.com>

Closes #9393 from yhuai/udfNondeterministic.
---
 project/MimaExcludes.scala                    |  47 +++++
 .../sql/catalyst/expressions/ScalaUDF.scala   |   7 +-
 .../apache/spark/sql/UDFRegistration.scala    | 164 ++++++++++--------
 .../spark/sql/UserDefinedFunction.scala       |  13 +-
 .../scala/org/apache/spark/sql/UDFSuite.scala | 105 +++++++++++
 .../datasources/parquet/ParquetIOSuite.scala  |   4 +-
 6 files changed, 262 insertions(+), 78 deletions(-)

diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 8282f7ea62400..ec0e44b7f2d66 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -112,6 +112,53 @@ object MimaExcludes {
           "org.apache.spark.rdd.MapPartitionsWithPreparationRDD"),
         ProblemFilters.exclude[MissingClassProblem](
           "org.apache.spark.rdd.MapPartitionsWithPreparationRDD$")
+      ) ++ Seq(
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$2"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$3"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$4"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$5"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$6"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$7"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$8"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$9"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$10"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$11"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$12"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$13"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$14"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$15"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$16"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$17"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$18"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$19"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$20"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$21"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$22"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$23"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$24")
       )
     case v if v.startsWith("1.5") =>
       Seq(
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
index 11c7950c0613b..a04af7f1dd877 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
@@ -30,13 +30,18 @@ case class ScalaUDF(
     function: AnyRef,
     dataType: DataType,
     children: Seq[Expression],
-    inputTypes: Seq[DataType] = Nil)
+    inputTypes: Seq[DataType] = Nil,
+    isDeterministic: Boolean = true)
   extends Expression with ImplicitCastInputTypes with CodegenFallback {
 
   override def nullable: Boolean = true
 
   override def toString: String = s"UDF(${children.mkString(",")})"
 
+  override def foldable: Boolean = deterministic && children.forall(_.foldable)
+
+  override def deterministic: Boolean = isDeterministic && children.forall(_.deterministic)
+
   // scalastyle:off
 
   /** This method has been generated by this script
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala b/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
index fc4d0938c533a..f5b95e13e47bc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
@@ -58,8 +58,10 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
    * Register a user-defined aggregate function (UDAF).
    *
    * @param name the name of the UDAF.
-   * @param udaf the UDAF needs to be registered.
+   * @param udaf the UDAF that needs to be registered.
    * @return the registered UDAF.
+   *
+   * @since 1.5.0
    */
   def register(
       name: String,
@@ -69,6 +71,22 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     udaf
   }
 
+  /**
+   * Register a user-defined function (UDF).
+   *
+   * @param name the name of the UDF.
+   * @param udf the UDF that needs to be registered.
+   * @return the registered UDF.
+   *
+   * @since 1.6.0
+   */
+  def register(
+      name: String,
+      udf: UserDefinedFunction): UserDefinedFunction = {
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
+  }
+
   // scalastyle:off
 
   /* register 0-22 were generated by this script
@@ -86,9 +104,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
         def register[$typeTags](name: String, func: Function$x[$types]): UserDefinedFunction = {
           val dataType = ScalaReflection.schemaFor[RT].dataType
           val inputTypes = Try($inputTypes).getOrElse(Nil)
-          def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-          functionRegistry.registerFunction(name, builder)
-          UserDefinedFunction(func, dataType)
+          val udf = UserDefinedFunction(func, dataType, inputTypes)
+          functionRegistry.registerFunction(name, udf.builder)
+          udf
         }""")
     }
 
@@ -118,9 +136,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag](name: String, func: Function0[RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -131,9 +149,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag](name: String, func: Function1[A1, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -144,9 +162,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag](name: String, func: Function2[A1, A2, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -157,9 +175,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag](name: String, func: Function3[A1, A2, A3, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -170,9 +188,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag](name: String, func: Function4[A1, A2, A3, A4, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -183,9 +201,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag](name: String, func: Function5[A1, A2, A3, A4, A5, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -196,9 +214,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag](name: String, func: Function6[A1, A2, A3, A4, A5, A6, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -209,9 +227,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag](name: String, func: Function7[A1, A2, A3, A4, A5, A6, A7, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -222,9 +240,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag](name: String, func: Function8[A1, A2, A3, A4, A5, A6, A7, A8, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -235,9 +253,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag](name: String, func: Function9[A1, A2, A3, A4, A5, A6, A7, A8, A9, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -248,9 +266,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag](name: String, func: Function10[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -261,9 +279,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag](name: String, func: Function11[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -274,9 +292,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag](name: String, func: Function12[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -287,9 +305,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag](name: String, func: Function13[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -300,9 +318,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag](name: String, func: Function14[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -313,9 +331,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag](name: String, func: Function15[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -326,9 +344,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag](name: String, func: Function16[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -339,9 +357,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag](name: String, func: Function17[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -352,9 +370,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag](name: String, func: Function18[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -365,9 +383,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag, A19: TypeTag](name: String, func: Function19[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, A19, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -378,9 +396,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag, A19: TypeTag, A20: TypeTag](name: String, func: Function20[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, A19, A20, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -391,9 +409,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag, A19: TypeTag, A20: TypeTag, A21: TypeTag](name: String, func: Function21[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, A19, A20, A21, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: ScalaReflection.schemaFor[A21].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   /**
@@ -404,9 +422,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag, A19: TypeTag, A20: TypeTag, A21: TypeTag, A22: TypeTag](name: String, func: Function22[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, A19, A20, A21, A22, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: ScalaReflection.schemaFor[A21].dataType :: ScalaReflection.schemaFor[A22].dataType :: Nil).getOrElse(Nil)
-    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
-    functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    val udf = UserDefinedFunction(func, dataType, inputTypes)
+    functionRegistry.registerFunction(name, udf.builder)
+    udf
   }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/UserDefinedFunction.scala b/sql/core/src/main/scala/org/apache/spark/sql/UserDefinedFunction.scala
index 0f8cd280b5acb..1319391db5375 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/UserDefinedFunction.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/UserDefinedFunction.scala
@@ -44,11 +44,20 @@ import org.apache.spark.sql.types.DataType
 case class UserDefinedFunction protected[sql] (
     f: AnyRef,
     dataType: DataType,
-    inputTypes: Seq[DataType] = Nil) {
+    inputTypes: Seq[DataType] = Nil,
+    deterministic: Boolean = true) {
 
   def apply(exprs: Column*): Column = {
-    Column(ScalaUDF(f, dataType, exprs.map(_.expr), inputTypes))
+    Column(ScalaUDF(f, dataType, exprs.map(_.expr), inputTypes, deterministic))
   }
+
+  protected[sql] def builder: Seq[Expression] => ScalaUDF = {
+    (exprs: Seq[Expression]) =>
+      ScalaUDF(f, dataType, exprs, inputTypes, deterministic)
+  }
+
+  def nondeterministic: UserDefinedFunction =
+    UserDefinedFunction(f, dataType, inputTypes, deterministic = false)
 }
 
 /**
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
index e0435a0dba6ad..6e510f0b8aff4 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql
 
+import org.apache.spark.sql.catalyst.expressions.ScalaUDF
+import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.test.SQLTestData._
 
@@ -191,4 +193,107 @@ class UDFSuite extends QueryTest with SharedSQLContext {
     // pass a decimal to intExpected.
     assert(sql("SELECT intExpected(1.0)").head().getInt(0) === 1)
   }
+
+  private def checkNumUDFs(df: DataFrame, expectedNumUDFs: Int): Unit = {
+    val udfs = df.queryExecution.optimizedPlan.collect {
+      case p: logical.Project => p.projectList.flatMap {
+        case e => e.collect {
+          case udf: ScalaUDF => udf
+        }
+      }
+    }.flatten
+    assert(udfs.length === expectedNumUDFs)
+  }
+
+  test("foldable udf") {
+    import org.apache.spark.sql.functions._
+
+    val myUDF = udf((x: Int) => x + 1)
+
+    {
+      val df = sql("SELECT 1 as a")
+        .select(col("a"), myUDF(col("a")).as("b"))
+        .select(col("a"), col("b"), myUDF(col("b")).as("c"))
+      checkNumUDFs(df, 0)
+      checkAnswer(df, Row(1, 2, 3))
+    }
+  }
+
+  test("nondeterministic udf: using UDFRegistration") {
+    import org.apache.spark.sql.functions._
+
+    val myUDF = sqlContext.udf.register("plusOne1", (x: Int) => x + 1)
+    sqlContext.udf.register("plusOne2", myUDF.nondeterministic)
+
+    {
+      val df = sqlContext.range(1, 2).select(col("id").as("a"))
+        .select(col("a"), myUDF(col("a")).as("b"))
+        .select(col("a"), col("b"), myUDF(col("b")).as("c"))
+      checkNumUDFs(df, 3)
+      checkAnswer(df, Row(1, 2, 3))
+    }
+
+    {
+      val df = sqlContext.range(1, 2).select(col("id").as("a"))
+        .select(col("a"), callUDF("plusOne1", col("a")).as("b"))
+        .select(col("a"), col("b"), callUDF("plusOne1", col("b")).as("c"))
+      checkNumUDFs(df, 3)
+      checkAnswer(df, Row(1, 2, 3))
+    }
+
+    {
+      val df = sqlContext.range(1, 2).select(col("id").as("a"))
+        .select(col("a"), myUDF.nondeterministic(col("a")).as("b"))
+        .select(col("a"), col("b"), myUDF.nondeterministic(col("b")).as("c"))
+      checkNumUDFs(df, 2)
+      checkAnswer(df, Row(1, 2, 3))
+    }
+
+    {
+      val df = sqlContext.range(1, 2).select(col("id").as("a"))
+        .select(col("a"), callUDF("plusOne2", col("a")).as("b"))
+        .select(col("a"), col("b"), callUDF("plusOne2", col("b")).as("c"))
+      checkNumUDFs(df, 2)
+      checkAnswer(df, Row(1, 2, 3))
+    }
+  }
+
+  test("nondeterministic udf: using udf function") {
+    import org.apache.spark.sql.functions._
+
+    val myUDF = udf((x: Int) => x + 1)
+
+    {
+      val df = sqlContext.range(1, 2).select(col("id").as("a"))
+        .select(col("a"), myUDF(col("a")).as("b"))
+        .select(col("a"), col("b"), myUDF(col("b")).as("c"))
+      checkNumUDFs(df, 3)
+      checkAnswer(df, Row(1, 2, 3))
+    }
+
+    {
+      val df = sqlContext.range(1, 2).select(col("id").as("a"))
+        .select(col("a"), myUDF.nondeterministic(col("a")).as("b"))
+        .select(col("a"), col("b"), myUDF.nondeterministic(col("b")).as("c"))
+      checkNumUDFs(df, 2)
+      checkAnswer(df, Row(1, 2, 3))
+    }
+
+    {
+      // nondeterministicUDF will not be foldable.
+      val df = sql("SELECT 1 as a")
+        .select(col("a"), myUDF.nondeterministic(col("a")).as("b"))
+        .select(col("a"), col("b"), myUDF.nondeterministic(col("b")).as("c"))
+      checkNumUDFs(df, 2)
+      checkAnswer(df, Row(1, 2, 3))
+    }
+  }
+
+  test("override a registered udf") {
+    sqlContext.udf.register("intExpected", (x: Int) => x)
+    assert(sql("SELECT intExpected(1.0)").head().getInt(0) === 1)
+
+    sqlContext.udf.register("intExpected", (x: Int) => x + 1)
+    assert(sql("SELECT intExpected(1.0)").head().getInt(0) === 2)
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index 72744799897be..f14b2886a9ecb 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -381,7 +381,7 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
       sqlContext.udf.register("div0", (x: Int) => x / 0)
       withTempPath { dir =>
         intercept[org.apache.spark.SparkException] {
-          sqlContext.sql("select div0(1)").write.parquet(dir.getCanonicalPath)
+          sqlContext.range(1, 2).selectExpr("div0(id) as a").write.parquet(dir.getCanonicalPath)
         }
         val path = new Path(dir.getCanonicalPath, "_temporary")
         val fs = path.getFileSystem(hadoopConfiguration)
@@ -405,7 +405,7 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
       sqlContext.udf.register("div0", (x: Int) => x / 0)
       withTempPath { dir =>
         intercept[org.apache.spark.SparkException] {
-          sqlContext.sql("select div0(1)").write.parquet(dir.getCanonicalPath)
+          sqlContext.range(1, 2).selectExpr("div0(id) as a").write.parquet(dir.getCanonicalPath)
         }
         val path = new Path(dir.getCanonicalPath, "_temporary")
         val fs = path.getFileSystem(hadoopConfiguration)

From c34c27fe9244939d8c905cd689536dfb81c74d7d Mon Sep 17 00:00:00 2001
From: "navis.ryu" <navis@apache.org>
Date: Mon, 2 Nov 2015 23:52:36 -0800
Subject: [PATCH 019/518] [SPARK-9034][SQL] Reflect field names defined in
 GenericUDTF

Hive GenericUDTF#initialize() defines field names in a returned schema though,
the current HiveGenericUDTF drops these names.
We might need to reflect these in a logical plan tree.

Author: navis.ryu <navis@apache.org>

Closes #8456 from navis/SPARK-9034.
---
 .../spark/sql/catalyst/analysis/Analyzer.scala       | 11 +++++------
 .../spark/sql/catalyst/expressions/generators.scala  | 12 +++++++-----
 .../main/scala/org/apache/spark/sql/DataFrame.scala  | 10 +++++-----
 .../sql/hive/execution/HiveCompatibilitySuite.scala  |  1 +
 .../scala/org/apache/spark/sql/hive/hiveUDFs.scala   |  2 +-
 ...GenericUDTF #1-0-ff502d8c06f4b32f57aa45057b7fab0e |  1 +
 ...GenericUDTF #2-0-d6d0def30a7fad5f90fd835361820c30 |  1 +
 ...l_view_noalias-0-50131c0ba7b7a6b65c789a5a8497bada |  1 +
 ...l_view_noalias-1-72509f06e1f7c5d5ccc292f775f8eea7 |  0
 ...l_view_noalias-2-6d5806dd1d2511911a5de1e205523f42 |  2 ++
 ...l_view_noalias-3-155b3cc2f5054725a9c2acca3c38c00a |  0
 ...l_view_noalias-4-3b7045ace234af8e5e86d8ac23ccee56 |  2 ++
 ...al_view_noalias-5-e1eca4e08216897d090259d4fd1e3fe |  0
 ...l_view_noalias-6-16d227442dd775615c6ecfceedc6c612 |  0
 ...l_view_noalias-7-66cb5ab20690dd85b2ed95bbfb9481d3 |  2 ++
 .../spark/sql/hive/execution/HiveQuerySuite.scala    |  6 ++++++
 16 files changed, 34 insertions(+), 17 deletions(-)
 create mode 100644 sql/hive/src/test/resources/golden/SPARK-9034 Reflect field names defined in GenericUDTF #1-0-ff502d8c06f4b32f57aa45057b7fab0e
 create mode 100644 sql/hive/src/test/resources/golden/SPARK-9034 Reflect field names defined in GenericUDTF #2-0-d6d0def30a7fad5f90fd835361820c30
 create mode 100644 sql/hive/src/test/resources/golden/lateral_view_noalias-0-50131c0ba7b7a6b65c789a5a8497bada
 create mode 100644 sql/hive/src/test/resources/golden/lateral_view_noalias-1-72509f06e1f7c5d5ccc292f775f8eea7
 create mode 100644 sql/hive/src/test/resources/golden/lateral_view_noalias-2-6d5806dd1d2511911a5de1e205523f42
 create mode 100644 sql/hive/src/test/resources/golden/lateral_view_noalias-3-155b3cc2f5054725a9c2acca3c38c00a
 create mode 100644 sql/hive/src/test/resources/golden/lateral_view_noalias-4-3b7045ace234af8e5e86d8ac23ccee56
 create mode 100644 sql/hive/src/test/resources/golden/lateral_view_noalias-5-e1eca4e08216897d090259d4fd1e3fe
 create mode 100644 sql/hive/src/test/resources/golden/lateral_view_noalias-6-16d227442dd775615c6ecfceedc6c612
 create mode 100644 sql/hive/src/test/resources/golden/lateral_view_noalias-7-66cb5ab20690dd85b2ed95bbfb9481d3

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index 912c967b95f08..899ee67352df4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -147,7 +147,7 @@ class Analyzer(
             case u @ UnresolvedAlias(child) => child match {
               case ne: NamedExpression => ne
               case e if !e.resolved => u
-              case g: Generator if g.elementTypes.size > 1 => MultiAlias(g, Nil)
+              case g: Generator => MultiAlias(g, Nil)
               case c @ Cast(ne: NamedExpression, _) => Alias(c, ne.name)()
               case other => Alias(other, s"_c$i")()
             }
@@ -722,7 +722,7 @@ class Analyzer(
 
     /**
      * Construct the output attributes for a [[Generator]], given a list of names.  If the list of
-     * names is empty names are assigned by ordinal (i.e., _c0, _c1, ...) to match Hive's defaults.
+     * names is empty names are assigned from field names in generator.
      */
     private def makeGeneratorOutput(
         generator: Generator,
@@ -731,13 +731,12 @@ class Analyzer(
 
       if (names.length == elementTypes.length) {
         names.zip(elementTypes).map {
-          case (name, (t, nullable)) =>
+          case (name, (t, nullable, _)) =>
             AttributeReference(name, t, nullable)()
         }
       } else if (names.isEmpty) {
-        elementTypes.zipWithIndex.map {
-          // keep the default column names as Hive does _c0, _c1, _cN
-          case ((t, nullable), i) => AttributeReference(s"_c$i", t, nullable)()
+        elementTypes.map {
+          case (t, nullable, name) => AttributeReference(name, t, nullable)()
         }
       } else {
         failAnalysis(
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala
index 1a2092c909c56..894a0730d1c2a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala
@@ -53,7 +53,7 @@ trait Generator extends Expression {
    * The output element data types in structure of Seq[(DataType, Nullable)]
    * TODO we probably need to add more information like metadata etc.
    */
-  def elementTypes: Seq[(DataType, Boolean)]
+  def elementTypes: Seq[(DataType, Boolean, String)]
 
   /** Should be implemented by child classes to perform specific Generators. */
   override def eval(input: InternalRow): TraversableOnce[InternalRow]
@@ -69,7 +69,7 @@ trait Generator extends Expression {
  * A generator that produces its output using the provided lambda function.
  */
 case class UserDefinedGenerator(
-    elementTypes: Seq[(DataType, Boolean)],
+    elementTypes: Seq[(DataType, Boolean, String)],
     function: Row => TraversableOnce[InternalRow],
     children: Seq[Expression])
   extends Generator with CodegenFallback {
@@ -112,9 +112,11 @@ case class Explode(child: Expression) extends UnaryExpression with Generator wit
     }
   }
 
-  override def elementTypes: Seq[(DataType, Boolean)] = child.dataType match {
-    case ArrayType(et, containsNull) => (et, containsNull) :: Nil
-    case MapType(kt, vt, valueContainsNull) => (kt, false) :: (vt, valueContainsNull) :: Nil
+  // hive-compatible default alias for explode function ("col" for array, "key", "value" for map)
+  override def elementTypes: Seq[(DataType, Boolean, String)] = child.dataType match {
+    case ArrayType(et, containsNull) => (et, containsNull, "col") :: Nil
+    case MapType(kt, vt, valueContainsNull) =>
+      (kt, false, "key") :: (vt, valueContainsNull, "value") :: Nil
   }
 
   override def eval(input: InternalRow): TraversableOnce[InternalRow] = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 53ad3c0266cdb..fc0ab632f9930 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -1175,7 +1175,8 @@ class DataFrame private[sql](
   def explode[A <: Product : TypeTag](input: Column*)(f: Row => TraversableOnce[A]): DataFrame = {
     val schema = ScalaReflection.schemaFor[A].dataType.asInstanceOf[StructType]
 
-    val elementTypes = schema.toAttributes.map { attr => (attr.dataType, attr.nullable) }
+    val elementTypes = schema.toAttributes.map {
+      attr => (attr.dataType, attr.nullable, attr.name) }
     val names = schema.toAttributes.map(_.name)
     val convert = CatalystTypeConverters.createToCatalystConverter(schema)
 
@@ -1184,7 +1185,7 @@ class DataFrame private[sql](
     val generator = UserDefinedGenerator(elementTypes, rowFunction, input.map(_.expr))
 
     Generate(generator, join = true, outer = false,
-      qualifier = None, names.map(UnresolvedAttribute(_)), logicalPlan)
+      qualifier = None, generatorOutput = Nil, logicalPlan)
   }
 
   /**
@@ -1203,8 +1204,7 @@ class DataFrame private[sql](
     val dataType = ScalaReflection.schemaFor[B].dataType
     val attributes = AttributeReference(outputColumn, dataType)() :: Nil
     // TODO handle the metadata?
-    val elementTypes = attributes.map { attr => (attr.dataType, attr.nullable) }
-    val names = attributes.map(_.name)
+    val elementTypes = attributes.map { attr => (attr.dataType, attr.nullable, attr.name) }
 
     def rowFunction(row: Row): TraversableOnce[InternalRow] = {
       val convert = CatalystTypeConverters.createToCatalystConverter(dataType)
@@ -1213,7 +1213,7 @@ class DataFrame private[sql](
     val generator = UserDefinedGenerator(elementTypes, rowFunction, apply(inputColumn).expr :: Nil)
 
     Generate(generator, join = true, outer = false,
-      qualifier = None, names.map(UnresolvedAttribute(_)), logicalPlan)
+      qualifier = None, generatorOutput = Nil, logicalPlan)
   }
 
   /////////////////////////////////////////////////////////////////////////////
diff --git a/sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala b/sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala
index 6ed40b03975d0..2d0d7b8af3581 100644
--- a/sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala
+++ b/sql/hive/compatibility/src/test/scala/org/apache/spark/sql/hive/execution/HiveCompatibilitySuite.scala
@@ -661,6 +661,7 @@ class HiveCompatibilitySuite extends HiveQueryFileTest with BeforeAndAfter {
     "join_star",
     "lateral_view",
     "lateral_view_cp",
+    "lateral_view_noalias",
     "lateral_view_ppd",
     "leftsemijoin",
     "leftsemijoin_mr",
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
index 0b5e863506142..a9db70119d011 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
@@ -511,7 +511,7 @@ private[hive] case class HiveGenericUDTF(
   protected lazy val collector = new UDTFCollector
 
   override lazy val elementTypes = outputInspector.getAllStructFieldRefs.asScala.map {
-    field => (inspectorToDataType(field.getFieldObjectInspector), true)
+    field => (inspectorToDataType(field.getFieldObjectInspector), true, field.getFieldName)
   }
 
   @transient
diff --git a/sql/hive/src/test/resources/golden/SPARK-9034 Reflect field names defined in GenericUDTF #1-0-ff502d8c06f4b32f57aa45057b7fab0e b/sql/hive/src/test/resources/golden/SPARK-9034 Reflect field names defined in GenericUDTF #1-0-ff502d8c06f4b32f57aa45057b7fab0e
new file mode 100644
index 0000000000000..1cf253f92c055
--- /dev/null
+++ b/sql/hive/src/test/resources/golden/SPARK-9034 Reflect field names defined in GenericUDTF #1-0-ff502d8c06f4b32f57aa45057b7fab0e	
@@ -0,0 +1 @@
+238
diff --git a/sql/hive/src/test/resources/golden/SPARK-9034 Reflect field names defined in GenericUDTF #2-0-d6d0def30a7fad5f90fd835361820c30 b/sql/hive/src/test/resources/golden/SPARK-9034 Reflect field names defined in GenericUDTF #2-0-d6d0def30a7fad5f90fd835361820c30
new file mode 100644
index 0000000000000..60878ffb77064
--- /dev/null
+++ b/sql/hive/src/test/resources/golden/SPARK-9034 Reflect field names defined in GenericUDTF #2-0-d6d0def30a7fad5f90fd835361820c30	
@@ -0,0 +1 @@
+238	val_238
diff --git a/sql/hive/src/test/resources/golden/lateral_view_noalias-0-50131c0ba7b7a6b65c789a5a8497bada b/sql/hive/src/test/resources/golden/lateral_view_noalias-0-50131c0ba7b7a6b65c789a5a8497bada
new file mode 100644
index 0000000000000..573541ac9702d
--- /dev/null
+++ b/sql/hive/src/test/resources/golden/lateral_view_noalias-0-50131c0ba7b7a6b65c789a5a8497bada
@@ -0,0 +1 @@
+0
diff --git a/sql/hive/src/test/resources/golden/lateral_view_noalias-1-72509f06e1f7c5d5ccc292f775f8eea7 b/sql/hive/src/test/resources/golden/lateral_view_noalias-1-72509f06e1f7c5d5ccc292f775f8eea7
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/sql/hive/src/test/resources/golden/lateral_view_noalias-2-6d5806dd1d2511911a5de1e205523f42 b/sql/hive/src/test/resources/golden/lateral_view_noalias-2-6d5806dd1d2511911a5de1e205523f42
new file mode 100644
index 0000000000000..0da0d93886e01
--- /dev/null
+++ b/sql/hive/src/test/resources/golden/lateral_view_noalias-2-6d5806dd1d2511911a5de1e205523f42
@@ -0,0 +1,2 @@
+key1	100
+key2	200
diff --git a/sql/hive/src/test/resources/golden/lateral_view_noalias-3-155b3cc2f5054725a9c2acca3c38c00a b/sql/hive/src/test/resources/golden/lateral_view_noalias-3-155b3cc2f5054725a9c2acca3c38c00a
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/sql/hive/src/test/resources/golden/lateral_view_noalias-4-3b7045ace234af8e5e86d8ac23ccee56 b/sql/hive/src/test/resources/golden/lateral_view_noalias-4-3b7045ace234af8e5e86d8ac23ccee56
new file mode 100644
index 0000000000000..0da0d93886e01
--- /dev/null
+++ b/sql/hive/src/test/resources/golden/lateral_view_noalias-4-3b7045ace234af8e5e86d8ac23ccee56
@@ -0,0 +1,2 @@
+key1	100
+key2	200
diff --git a/sql/hive/src/test/resources/golden/lateral_view_noalias-5-e1eca4e08216897d090259d4fd1e3fe b/sql/hive/src/test/resources/golden/lateral_view_noalias-5-e1eca4e08216897d090259d4fd1e3fe
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/sql/hive/src/test/resources/golden/lateral_view_noalias-6-16d227442dd775615c6ecfceedc6c612 b/sql/hive/src/test/resources/golden/lateral_view_noalias-6-16d227442dd775615c6ecfceedc6c612
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/sql/hive/src/test/resources/golden/lateral_view_noalias-7-66cb5ab20690dd85b2ed95bbfb9481d3 b/sql/hive/src/test/resources/golden/lateral_view_noalias-7-66cb5ab20690dd85b2ed95bbfb9481d3
new file mode 100644
index 0000000000000..4ba46bbda5b04
--- /dev/null
+++ b/sql/hive/src/test/resources/golden/lateral_view_noalias-7-66cb5ab20690dd85b2ed95bbfb9481d3
@@ -0,0 +1,2 @@
+key1	100	key1	100
+key2	200	key2	200
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
index e597d6865f67a..fc72e3c7dc6aa 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
@@ -563,6 +563,12 @@ class HiveQuerySuite extends HiveComparisonTest with BeforeAndAfter {
   createQueryTest("Specify the udtf output",
     "SELECT d FROM (SELECT explode(array(1,1)) d FROM src LIMIT 1) t")
 
+  createQueryTest("SPARK-9034 Reflect field names defined in GenericUDTF #1",
+    "SELECT col FROM (SELECT explode(array(key,value)) FROM src LIMIT 1) t")
+
+  createQueryTest("SPARK-9034 Reflect field names defined in GenericUDTF #2",
+    "SELECT key,value FROM (SELECT explode(map(key,value)) FROM src LIMIT 1) t")
+
   test("sampling") {
     sql("SELECT * FROM src TABLESAMPLE(0.1 PERCENT) s")
     sql("SELECT * FROM src TABLESAMPLE(100 PERCENT) s")

From d728d5c98658c44ed2949b55d36edeaa46f8c980 Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Tue, 3 Nov 2015 00:12:49 -0800
Subject: [PATCH 020/518] [SPARK-9858][SPARK-9859][SPARK-9861][SQL] Add an
 ExchangeCoordinator to estimate the number of post-shuffle partitions for
 aggregates and joins

https://issues.apache.org/jira/browse/SPARK-9858
https://issues.apache.org/jira/browse/SPARK-9859
https://issues.apache.org/jira/browse/SPARK-9861

Author: Yin Huai <yhuai@databricks.com>

Closes #9276 from yhuai/numReducer.
---
 .../plans/physical/partitioning.scala         |   8 +
 .../scala/org/apache/spark/sql/SQLConf.scala  |  27 +
 .../apache/spark/sql/execution/Exchange.scala | 217 +++++++-
 .../sql/execution/ExchangeCoordinator.scala   | 260 ++++++++++
 .../spark/sql/execution/ShuffledRowRDD.scala  | 134 ++++-
 .../execution/ExchangeCoordinatorSuite.scala  | 479 ++++++++++++++++++
 .../spark/sql/execution/PlannerSuite.scala    |   8 +-
 .../execution/UnsafeRowSerializerSuite.scala  |   7 +-
 .../sql/execution/joins/InnerJoinSuite.scala  |  19 +-
 9 files changed, 1115 insertions(+), 44 deletions(-)
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/ExchangeCoordinator.scala
 create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index 86b9417477ba3..9312c8123e92e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -165,6 +165,11 @@ sealed trait Partitioning {
    * produced by `A` could have also been produced by `B`.
    */
   def guarantees(other: Partitioning): Boolean = this == other
+
+  def withNumPartitions(newNumPartitions: Int): Partitioning = {
+    throw new IllegalStateException(
+      s"It is not allowed to call withNumPartitions method of a ${this.getClass.getSimpleName}")
+  }
 }
 
 object Partitioning {
@@ -249,6 +254,9 @@ case class HashPartitioning(expressions: Seq[Expression], numPartitions: Int)
     case _ => false
   }
 
+  override def withNumPartitions(newNumPartitions: Int): HashPartitioning = {
+    HashPartitioning(expressions, newNumPartitions)
+  }
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
index 6f2892085a8f8..ed8b634ad5630 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
@@ -233,6 +233,25 @@ private[spark] object SQLConf {
     defaultValue = Some(200),
     doc = "The default number of partitions to use when shuffling data for joins or aggregations.")
 
+  val SHUFFLE_TARGET_POSTSHUFFLE_INPUT_SIZE =
+    longConf("spark.sql.adaptive.shuffle.targetPostShuffleInputSize",
+      defaultValue = Some(64 * 1024 * 1024),
+      doc = "The target post-shuffle input size in bytes of a task.")
+
+  val ADAPTIVE_EXECUTION_ENABLED = booleanConf("spark.sql.adaptive.enabled",
+    defaultValue = Some(false),
+    doc = "When true, enable adaptive query execution.")
+
+  val SHUFFLE_MIN_NUM_POSTSHUFFLE_PARTITIONS =
+    intConf("spark.sql.adaptive.minNumPostShufflePartitions",
+      defaultValue = Some(-1),
+      doc = "The advisory minimal number of post-shuffle partitions provided to " +
+        "ExchangeCoordinator. This setting is used in our test to make sure we " +
+        "have enough parallelism to expose issues that will not be exposed with a " +
+        "single partition. When the value is a non-positive value, this setting will" +
+        "not be provided to ExchangeCoordinator.",
+      isPublic = false)
+
   val TUNGSTEN_ENABLED = booleanConf("spark.sql.tungsten.enabled",
     defaultValue = Some(true),
     doc = "When true, use the optimized Tungsten physical execution backend which explicitly " +
@@ -487,6 +506,14 @@ private[sql] class SQLConf extends Serializable with CatalystConf {
 
   private[spark] def numShufflePartitions: Int = getConf(SHUFFLE_PARTITIONS)
 
+  private[spark] def targetPostShuffleInputSize: Long =
+    getConf(SHUFFLE_TARGET_POSTSHUFFLE_INPUT_SIZE)
+
+  private[spark] def adaptiveExecutionEnabled: Boolean = getConf(ADAPTIVE_EXECUTION_ENABLED)
+
+  private[spark] def minNumPostShufflePartitions: Int =
+    getConf(SHUFFLE_MIN_NUM_POSTSHUFFLE_PARTITIONS)
+
   private[spark] def parquetFilterPushDown: Boolean = getConf(PARQUET_FILTER_PUSHDOWN_ENABLED)
 
   private[spark] def orcFilterPushDown: Boolean = getConf(ORC_FILTER_PUSHDOWN_ENABLED)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
index e81108b7884d1..0f72ec6cc107a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
@@ -36,9 +36,23 @@ import org.apache.spark.util.MutablePair
 /**
  * Performs a shuffle that will result in the desired `newPartitioning`.
  */
-case class Exchange(newPartitioning: Partitioning, child: SparkPlan) extends UnaryNode {
+case class Exchange(
+    var newPartitioning: Partitioning,
+    child: SparkPlan,
+    @transient coordinator: Option[ExchangeCoordinator]) extends UnaryNode {
 
-  override def nodeName: String = if (tungstenMode) "TungstenExchange" else "Exchange"
+  override def nodeName: String = {
+    val extraInfo = coordinator match {
+      case Some(exchangeCoordinator) if exchangeCoordinator.isEstimated =>
+        "Shuffle"
+      case Some(exchangeCoordinator) if !exchangeCoordinator.isEstimated =>
+        "May shuffle"
+      case None => "Shuffle without coordinator"
+    }
+
+    val simpleNodeName = if (tungstenMode) "TungstenExchange" else "Exchange"
+    s"$simpleNodeName($extraInfo)"
+  }
 
   /**
    * Returns true iff we can support the data type, and we are not doing range partitioning.
@@ -129,7 +143,27 @@ case class Exchange(newPartitioning: Partitioning, child: SparkPlan) extends Una
     }
   }
 
-  protected override def doExecute(): RDD[InternalRow] = attachTree(this , "execute") {
+  override protected def doPrepare(): Unit = {
+    // If an ExchangeCoordinator is needed, we register this Exchange operator
+    // to the coordinator when we do prepare. It is important to make sure
+    // we register this operator right before the execution instead of register it
+    // in the constructor because it is possible that we create new instances of
+    // Exchange operators when we transform the physical plan
+    // (then the ExchangeCoordinator will hold references of unneeded Exchanges).
+    // So, we should only call registerExchange just before we start to execute
+    // the plan.
+    coordinator match {
+      case Some(exchangeCoordinator) => exchangeCoordinator.registerExchange(this)
+      case None =>
+    }
+  }
+
+  /**
+   * Returns a [[ShuffleDependency]] that will partition rows of its child based on
+   * the partitioning scheme defined in `newPartitioning`. Those partitions of
+   * the returned ShuffleDependency will be the input of shuffle.
+   */
+  private[sql] def prepareShuffleDependency(): ShuffleDependency[Int, InternalRow, InternalRow] = {
     val rdd = child.execute()
     val part: Partitioner = newPartitioning match {
       case RoundRobinPartitioning(numPartitions) => new HashPartitioner(numPartitions)
@@ -181,7 +215,54 @@ case class Exchange(newPartitioning: Partitioning, child: SparkPlan) extends Una
         }
       }
     }
-    new ShuffledRowRDD(rddWithPartitionIds, serializer, part.numPartitions)
+
+    // Now, we manually create a ShuffleDependency. Because pairs in rddWithPartitionIds
+    // are in the form of (partitionId, row) and every partitionId is in the expected range
+    // [0, part.numPartitions - 1]. The partitioner of this is a PartitionIdPassthrough.
+    val dependency =
+      new ShuffleDependency[Int, InternalRow, InternalRow](
+        rddWithPartitionIds,
+        new PartitionIdPassthrough(part.numPartitions),
+        Some(serializer))
+
+    dependency
+  }
+
+  /**
+   * Returns a [[ShuffledRowRDD]] that represents the post-shuffle dataset.
+   * This [[ShuffledRowRDD]] is created based on a given [[ShuffleDependency]] and an optional
+   * partition start indices array. If this optional array is defined, the returned
+   * [[ShuffledRowRDD]] will fetch pre-shuffle partitions based on indices of this array.
+   */
+  private[sql] def preparePostShuffleRDD(
+      shuffleDependency: ShuffleDependency[Int, InternalRow, InternalRow],
+      specifiedPartitionStartIndices: Option[Array[Int]] = None): ShuffledRowRDD = {
+    // If an array of partition start indices is provided, we need to use this array
+    // to create the ShuffledRowRDD. Also, we need to update newPartitioning to
+    // update the number of post-shuffle partitions.
+    specifiedPartitionStartIndices.foreach { indices =>
+      assert(newPartitioning.isInstanceOf[HashPartitioning])
+      newPartitioning = newPartitioning.withNumPartitions(indices.length)
+    }
+    new ShuffledRowRDD(shuffleDependency, specifiedPartitionStartIndices)
+  }
+
+  protected override def doExecute(): RDD[InternalRow] = attachTree(this , "execute") {
+    coordinator match {
+      case Some(exchangeCoordinator) =>
+        val shuffleRDD = exchangeCoordinator.postShuffleRDD(this)
+        assert(shuffleRDD.partitions.length == newPartitioning.numPartitions)
+        shuffleRDD
+      case None =>
+        val shuffleDependency = prepareShuffleDependency()
+        preparePostShuffleRDD(shuffleDependency)
+    }
+  }
+}
+
+object Exchange {
+  def apply(newPartitioning: Partitioning, child: SparkPlan): Exchange = {
+    Exchange(newPartitioning, child, None: Option[ExchangeCoordinator])
   }
 }
 
@@ -193,13 +274,22 @@ case class Exchange(newPartitioning: Partitioning, child: SparkPlan) extends Una
  * input partition ordering requirements are met.
  */
 private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[SparkPlan] {
-  // TODO: Determine the number of partitions.
-  private def defaultPartitions: Int = sqlContext.conf.numShufflePartitions
+  private def defaultNumPreShufflePartitions: Int = sqlContext.conf.numShufflePartitions
+
+  private def targetPostShuffleInputSize: Long = sqlContext.conf.targetPostShuffleInputSize
+
+  private def adaptiveExecutionEnabled: Boolean = sqlContext.conf.adaptiveExecutionEnabled
+
+  private def minNumPostShufflePartitions: Option[Int] = {
+    val minNumPostShufflePartitions = sqlContext.conf.minNumPostShufflePartitions
+    if (minNumPostShufflePartitions > 0) Some(minNumPostShufflePartitions) else None
+  }
 
   /**
    * Given a required distribution, returns a partitioning that satisfies that distribution.
    */
-  private def createPartitioning(requiredDistribution: Distribution,
+  private def createPartitioning(
+      requiredDistribution: Distribution,
       numPartitions: Int): Partitioning = {
     requiredDistribution match {
       case AllTuples => SinglePartition
@@ -209,6 +299,98 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
     }
   }
 
+  /**
+   * Adds [[ExchangeCoordinator]] to [[Exchange]]s if adaptive query execution is enabled
+   * and partitioning schemes of these [[Exchange]]s support [[ExchangeCoordinator]].
+   */
+  private def withExchangeCoordinator(
+      children: Seq[SparkPlan],
+      requiredChildDistributions: Seq[Distribution]): Seq[SparkPlan] = {
+    val supportsCoordinator =
+      if (children.exists(_.isInstanceOf[Exchange])) {
+        // Right now, ExchangeCoordinator only support HashPartitionings.
+        children.forall {
+          case e @ Exchange(hash: HashPartitioning, _, _) => true
+          case child =>
+            child.outputPartitioning match {
+              case hash: HashPartitioning => true
+              case collection: PartitioningCollection =>
+                collection.partitionings.exists(_.isInstanceOf[HashPartitioning])
+              case _ => false
+            }
+        }
+      } else {
+        // In this case, although we do not have Exchange operators, we may still need to
+        // shuffle data when we have more than one children because data generated by
+        // these children may not be partitioned in the same way.
+        // Please see the comment in withCoordinator for more details.
+        val supportsDistribution =
+          requiredChildDistributions.forall(_.isInstanceOf[ClusteredDistribution])
+        children.length > 1 && supportsDistribution
+      }
+
+    val withCoordinator =
+      if (adaptiveExecutionEnabled && supportsCoordinator) {
+        val coordinator =
+          new ExchangeCoordinator(
+            children.length,
+            targetPostShuffleInputSize,
+            minNumPostShufflePartitions)
+        children.zip(requiredChildDistributions).map {
+          case (e: Exchange, _) =>
+            // This child is an Exchange, we need to add the coordinator.
+            e.copy(coordinator = Some(coordinator))
+          case (child, distribution) =>
+            // If this child is not an Exchange, we need to add an Exchange for now.
+            // Ideally, we can try to avoid this Exchange. However, when we reach here,
+            // there are at least two children operators (because if there is a single child
+            // and we can avoid Exchange, supportsCoordinator will be false and we
+            // will not reach here.). Although we can make two children have the same number of
+            // post-shuffle partitions. Their numbers of pre-shuffle partitions may be different.
+            // For example, let's say we have the following plan
+            //         Join
+            //         /  \
+            //       Agg  Exchange
+            //       /      \
+            //    Exchange  t2
+            //      /
+            //     t1
+            // In this case, because a post-shuffle partition can include multiple pre-shuffle
+            // partitions, a HashPartitioning will not be strictly partitioned by the hashcodes
+            // after shuffle. So, even we can use the child Exchange operator of the Join to
+            // have a number of post-shuffle partitions that matches the number of partitions of
+            // Agg, we cannot say these two children are partitioned in the same way.
+            // Here is another case
+            //         Join
+            //         /  \
+            //       Agg1  Agg2
+            //       /      \
+            //   Exchange1  Exchange2
+            //       /       \
+            //      t1       t2
+            // In this case, two Aggs shuffle data with the same column of the join condition.
+            // After we use ExchangeCoordinator, these two Aggs may not be partitioned in the same
+            // way. Let's say that Agg1 and Agg2 both have 5 pre-shuffle partitions and 2
+            // post-shuffle partitions. It is possible that Agg1 fetches those pre-shuffle
+            // partitions by using a partitionStartIndices [0, 3]. However, Agg2 may fetch its
+            // pre-shuffle partitions by using another partitionStartIndices [0, 4].
+            // So, Agg1 and Agg2 are actually not co-partitioned.
+            //
+            // It will be great to introduce a new Partitioning to represent the post-shuffle
+            // partitions when one post-shuffle partition includes multiple pre-shuffle partitions.
+            val targetPartitioning =
+              createPartitioning(distribution, defaultNumPreShufflePartitions)
+            assert(targetPartitioning.isInstanceOf[HashPartitioning])
+            Exchange(targetPartitioning, child, Some(coordinator))
+        }
+      } else {
+        // If we do not need ExchangeCoordinator, the original children are returned.
+        children
+      }
+
+    withCoordinator
+  }
+
   private def ensureDistributionAndOrdering(operator: SparkPlan): SparkPlan = {
     val requiredChildDistributions: Seq[Distribution] = operator.requiredChildDistribution
     val requiredChildOrderings: Seq[Seq[SortOrder]] = operator.requiredChildOrdering
@@ -221,7 +403,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       if (child.outputPartitioning.satisfies(distribution)) {
         child
       } else {
-        Exchange(createPartitioning(distribution, defaultPartitions), child)
+        Exchange(createPartitioning(distribution, defaultNumPreShufflePartitions), child)
       }
     }
 
@@ -234,7 +416,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       // First check if the existing partitions of the children all match. This means they are
       // partitioned by the same partitioning into the same number of partitions. In that case,
       // don't try to make them match `defaultPartitions`, just use the existing partitioning.
-      // TODO: this should be a cost based descision. For example, a big relation should probably
+      // TODO: this should be a cost based decision. For example, a big relation should probably
       // maintain its existing number of partitions and smaller partitions should be shuffled.
       // defaultPartitions is arbitrary.
       val numPartitions = children.head.outputPartitioning.numPartitions
@@ -250,7 +432,8 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       } else {
         children.zip(requiredChildDistributions).map {
           case (child, distribution) => {
-            val targetPartitioning = createPartitioning(distribution, defaultPartitions)
+            val targetPartitioning =
+              createPartitioning(distribution, defaultNumPreShufflePartitions)
             if (child.outputPartitioning.guarantees(targetPartitioning)) {
               child
             } else {
@@ -261,12 +444,24 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       }
     }
 
+    // Now, we need to add ExchangeCoordinator if necessary.
+    // Actually, it is not a good idea to add ExchangeCoordinators while we are adding Exchanges.
+    // However, with the way that we plan the query, we do not have a place where we have a
+    // global picture of all shuffle dependencies of a post-shuffle stage. So, we add coordinator
+    // at here for now.
+    // Once we finish https://issues.apache.org/jira/browse/SPARK-10665,
+    // we can first add Exchanges and then add coordinator once we have a DAG of query fragments.
+    children = withExchangeCoordinator(children, requiredChildDistributions)
+
     // Now that we've performed any necessary shuffles, add sorts to guarantee output orderings:
     children = children.zip(requiredChildOrderings).map { case (child, requiredOrdering) =>
       if (requiredOrdering.nonEmpty) {
         // If child.outputOrdering is [a, b] and requiredOrdering is [a], we do not need to sort.
         if (requiredOrdering != child.outputOrdering.take(requiredOrdering.length)) {
-          sqlContext.planner.BasicOperators.getSortOperator(requiredOrdering, global = false, child)
+          sqlContext.planner.BasicOperators.getSortOperator(
+            requiredOrdering,
+            global = false,
+            child)
         } else {
           child
         }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExchangeCoordinator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExchangeCoordinator.scala
new file mode 100644
index 0000000000000..8dbd69e1f44b8
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExchangeCoordinator.scala
@@ -0,0 +1,260 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import java.util.{Map => JMap, HashMap => JHashMap}
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.spark.{Logging, SimpleFutureAction, ShuffleDependency, MapOutputStatistics}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+
+/**
+ * A coordinator used to determines how we shuffle data between stages generated by Spark SQL.
+ * Right now, the work of this coordinator is to determine the number of post-shuffle partitions
+ * for a stage that needs to fetch shuffle data from one or multiple stages.
+ *
+ * A coordinator is constructed with three parameters, `numExchanges`,
+ * `targetPostShuffleInputSize`, and `minNumPostShufflePartitions`.
+ *  - `numExchanges` is used to indicated that how many [[Exchange]]s that will be registered to
+ *    this coordinator. So, when we start to do any actual work, we have a way to make sure that
+ *    we have got expected number of [[Exchange]]s.
+ *  - `targetPostShuffleInputSize` is the targeted size of a post-shuffle partition's
+ *    input data size. With this parameter, we can estimate the number of post-shuffle partitions.
+ *    This parameter is configured through
+ *    `spark.sql.adaptive.shuffle.targetPostShuffleInputSize`.
+ *  - `minNumPostShufflePartitions` is an optional parameter. If it is defined, this coordinator
+ *    will try to make sure that there are at least `minNumPostShufflePartitions` post-shuffle
+ *    partitions.
+ *
+ * The workflow of this coordinator is described as follows:
+ *  - Before the execution of a [[SparkPlan]], for an [[Exchange]] operator,
+ *    if an [[ExchangeCoordinator]] is assigned to it, it registers itself to this coordinator.
+ *    This happens in the `doPrepare` method.
+ *  - Once we start to execute a physical plan, an [[Exchange]] registered to this coordinator will
+ *    call `postShuffleRDD` to get its corresponding post-shuffle [[ShuffledRowRDD]].
+ *    If this coordinator has made the decision on how to shuffle data, this [[Exchange]] will
+ *    immediately get its corresponding post-shuffle [[ShuffledRowRDD]].
+ *  - If this coordinator has not made the decision on how to shuffle data, it will ask those
+ *    registered [[Exchange]]s to submit their pre-shuffle stages. Then, based on the the size
+ *    statistics of pre-shuffle partitions, this coordinator will determine the number of
+ *    post-shuffle partitions and pack multiple pre-shuffle partitions with continuous indices
+ *    to a single post-shuffle partition whenever necessary.
+ *  - Finally, this coordinator will create post-shuffle [[ShuffledRowRDD]]s for all registered
+ *    [[Exchange]]s. So, when an [[Exchange]] calls `postShuffleRDD`, this coordinator can
+ *    lookup the corresponding [[RDD]].
+ *
+ * The strategy used to determine the number of post-shuffle partitions is described as follows.
+ * To determine the number of post-shuffle partitions, we have a target input size for a
+ * post-shuffle partition. Once we have size statistics of pre-shuffle partitions from stages
+ * corresponding to the registered [[Exchange]]s, we will do a pass of those statistics and
+ * pack pre-shuffle partitions with continuous indices to a single post-shuffle partition until
+ * the size of a post-shuffle partition is equal or greater than the target size.
+ * For example, we have two stages with the following pre-shuffle partition size statistics:
+ * stage 1: [100 MB, 20 MB, 100 MB, 10MB, 30 MB]
+ * stage 2: [10 MB,  10 MB, 70 MB,  5 MB, 5 MB]
+ * assuming the target input size is 128 MB, we will have three post-shuffle partitions,
+ * which are:
+ *  - post-shuffle partition 0: pre-shuffle partition 0 and 1
+ *  - post-shuffle partition 1: pre-shuffle partition 2
+ *  - post-shuffle partition 2: pre-shuffle partition 3 and 4
+ */
+private[sql] class ExchangeCoordinator(
+    numExchanges: Int,
+    advisoryTargetPostShuffleInputSize: Long,
+    minNumPostShufflePartitions: Option[Int] = None)
+  extends Logging {
+
+  // The registered Exchange operators.
+  private[this] val exchanges = ArrayBuffer[Exchange]()
+
+  // This map is used to lookup the post-shuffle ShuffledRowRDD for an Exchange operator.
+  private[this] val postShuffleRDDs: JMap[Exchange, ShuffledRowRDD] =
+    new JHashMap[Exchange, ShuffledRowRDD](numExchanges)
+
+  // A boolean that indicates if this coordinator has made decision on how to shuffle data.
+  // This variable will only be updated by doEstimationIfNecessary, which is protected by
+  // synchronized.
+  @volatile private[this] var estimated: Boolean = false
+
+  /**
+   * Registers an [[Exchange]] operator to this coordinator. This method is only allowed to be
+   * called in the `doPrepare` method of an [[Exchange]] operator.
+   */
+  def registerExchange(exchange: Exchange): Unit = synchronized {
+    exchanges += exchange
+  }
+
+  def isEstimated: Boolean = estimated
+
+  /**
+   * Estimates partition start indices for post-shuffle partitions based on
+   * mapOutputStatistics provided by all pre-shuffle stages.
+   */
+  private[sql] def estimatePartitionStartIndices(
+      mapOutputStatistics: Array[MapOutputStatistics]): Array[Int] = {
+    // If we have mapOutputStatistics.length <= numExchange, it is because we do not submit
+    // a stage when the number of partitions of this dependency is 0.
+    assert(mapOutputStatistics.length <= numExchanges)
+
+    // If minNumPostShufflePartitions is defined, it is possible that we need to use a
+    // value less than advisoryTargetPostShuffleInputSize as the target input size of
+    // a post shuffle task.
+    val targetPostShuffleInputSize = minNumPostShufflePartitions match {
+      case Some(numPartitions) =>
+        val totalPostShuffleInputSize = mapOutputStatistics.map(_.bytesByPartitionId.sum).sum
+        // The max at here is to make sure that when we have an empty table, we
+        // only have a single post-shuffle partition.
+        val maxPostShuffleInputSize =
+          math.max(math.ceil(totalPostShuffleInputSize / numPartitions.toDouble).toLong, 16)
+        math.min(maxPostShuffleInputSize, advisoryTargetPostShuffleInputSize)
+
+      case None => advisoryTargetPostShuffleInputSize
+    }
+
+    logInfo(
+      s"advisoryTargetPostShuffleInputSize: $advisoryTargetPostShuffleInputSize, " +
+      s"targetPostShuffleInputSize $targetPostShuffleInputSize.")
+
+    // Make sure we do get the same number of pre-shuffle partitions for those stages.
+    val distinctNumPreShufflePartitions =
+      mapOutputStatistics.map(stats => stats.bytesByPartitionId.length).distinct
+    assert(
+      distinctNumPreShufflePartitions.length == 1,
+      "There should be only one distinct value of the number pre-shuffle partitions " +
+        "among registered Exchange operator.")
+    val numPreShufflePartitions = distinctNumPreShufflePartitions.head
+
+    val partitionStartIndices = ArrayBuffer[Int]()
+    // The first element of partitionStartIndices is always 0.
+    partitionStartIndices += 0
+
+    var postShuffleInputSize = 0L
+
+    var i = 0
+    while (i < numPreShufflePartitions) {
+      // We calculate the total size of ith pre-shuffle partitions from all pre-shuffle stages.
+      // Then, we add the total size to postShuffleInputSize.
+      var j = 0
+      while (j < mapOutputStatistics.length) {
+        postShuffleInputSize += mapOutputStatistics(j).bytesByPartitionId(i)
+        j += 1
+      }
+
+      // If the current postShuffleInputSize is equal or greater than the
+      // targetPostShuffleInputSize, We need to add a new element in partitionStartIndices.
+      if (postShuffleInputSize >= targetPostShuffleInputSize) {
+        if (i < numPreShufflePartitions - 1) {
+          // Next start index.
+          partitionStartIndices += i + 1
+        } else {
+          // This is the last element. So, we do not need to append the next start index to
+          // partitionStartIndices.
+        }
+        // reset postShuffleInputSize.
+        postShuffleInputSize = 0L
+      }
+
+      i += 1
+    }
+
+    partitionStartIndices.toArray
+  }
+
+  private def doEstimationIfNecessary(): Unit = synchronized {
+    // It is unlikely that this method will be called from multiple threads
+    // (when multiple threads trigger the execution of THIS physical)
+    // because in common use cases, we will create new physical plan after
+    // users apply operations (e.g. projection) to an existing DataFrame.
+    // However, if it happens, we have synchronized to make sure only one
+    // thread will trigger the job submission.
+    if (!estimated) {
+      // Make sure we have the expected number of registered Exchange operators.
+      assert(exchanges.length == numExchanges)
+
+      val newPostShuffleRDDs = new JHashMap[Exchange, ShuffledRowRDD](numExchanges)
+
+      // Submit all map stages
+      val shuffleDependencies = ArrayBuffer[ShuffleDependency[Int, InternalRow, InternalRow]]()
+      val submittedStageFutures = ArrayBuffer[SimpleFutureAction[MapOutputStatistics]]()
+      var i = 0
+      while (i < numExchanges) {
+        val exchange = exchanges(i)
+        val shuffleDependency = exchange.prepareShuffleDependency()
+        shuffleDependencies += shuffleDependency
+        if (shuffleDependency.rdd.partitions.length != 0) {
+          // submitMapStage does not accept RDD with 0 partition.
+          // So, we will not submit this dependency.
+          submittedStageFutures +=
+            exchange.sqlContext.sparkContext.submitMapStage(shuffleDependency)
+        }
+        i += 1
+      }
+
+      // Wait for the finishes of those submitted map stages.
+      val mapOutputStatistics = new Array[MapOutputStatistics](submittedStageFutures.length)
+      i = 0
+      while (i < submittedStageFutures.length) {
+        // This call is a blocking call. If the stage has not finished, we will wait at here.
+        mapOutputStatistics(i) = submittedStageFutures(i).get()
+        i += 1
+      }
+
+      // Now, we estimate partitionStartIndices. partitionStartIndices.length will be the
+      // number of post-shuffle partitions.
+      val partitionStartIndices =
+        if (mapOutputStatistics.length == 0) {
+          None
+        } else {
+          Some(estimatePartitionStartIndices(mapOutputStatistics))
+        }
+
+      i = 0
+      while (i < numExchanges) {
+        val exchange = exchanges(i)
+        val rdd =
+          exchange.preparePostShuffleRDD(shuffleDependencies(i), partitionStartIndices)
+        newPostShuffleRDDs.put(exchange, rdd)
+
+        i += 1
+      }
+
+      // Finally, we set postShuffleRDDs and estimated.
+      assert(postShuffleRDDs.isEmpty)
+      assert(newPostShuffleRDDs.size() == numExchanges)
+      postShuffleRDDs.putAll(newPostShuffleRDDs)
+      estimated = true
+    }
+  }
+
+  def postShuffleRDD(exchange: Exchange): ShuffledRowRDD = {
+    doEstimationIfNecessary()
+
+    if (!postShuffleRDDs.containsKey(exchange)) {
+      throw new IllegalStateException(
+        s"The given $exchange is not registered in this coordinator.")
+    }
+
+    postShuffleRDDs.get(exchange)
+  }
+
+  override def toString: String = {
+    s"coordinator[target post-shuffle partition size: $advisoryTargetPostShuffleInputSize]"
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala
index fb338b90bf79b..42891287a3006 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ShuffledRowRDD.scala
@@ -17,14 +17,23 @@
 
 package org.apache.spark.sql.execution
 
+import java.util.Arrays
+
 import org.apache.spark._
 import org.apache.spark.rdd.RDD
-import org.apache.spark.serializer.Serializer
 import org.apache.spark.sql.catalyst.InternalRow
 
-private class ShuffledRowRDDPartition(val idx: Int) extends Partition {
-  override val index: Int = idx
-  override def hashCode(): Int = idx
+/**
+ * The [[Partition]] used by [[ShuffledRowRDD]]. A post-shuffle partition
+ * (identified by `postShufflePartitionIndex`) contains a range of pre-shuffle partitions
+ * (`startPreShufflePartitionIndex` to `endPreShufflePartitionIndex - 1`, inclusive).
+ */
+private final class ShuffledRowRDDPartition(
+    val postShufflePartitionIndex: Int,
+    val startPreShufflePartitionIndex: Int,
+    val endPreShufflePartitionIndex: Int) extends Partition {
+  override val index: Int = postShufflePartitionIndex
+  override def hashCode(): Int = postShufflePartitionIndex
 }
 
 /**
@@ -35,33 +44,107 @@ private class PartitionIdPassthrough(override val numPartitions: Int) extends Pa
   override def getPartition(key: Any): Int = key.asInstanceOf[Int]
 }
 
+/**
+ * A Partitioner that might group together one or more partitions from the parent.
+ *
+ * @param parent a parent partitioner
+ * @param partitionStartIndices indices of partitions in parent that should create new partitions
+ *   in child (this should be an array of increasing partition IDs). For example, if we have a
+ *   parent with 5 partitions, and partitionStartIndices is [0, 2, 4], we get three output
+ *   partitions, corresponding to partition ranges [0, 1], [2, 3] and [4] of the parent partitioner.
+ */
+class CoalescedPartitioner(val parent: Partitioner, val partitionStartIndices: Array[Int])
+  extends Partitioner {
+
+  @transient private lazy val parentPartitionMapping: Array[Int] = {
+    val n = parent.numPartitions
+    val result = new Array[Int](n)
+    for (i <- 0 until partitionStartIndices.length) {
+      val start = partitionStartIndices(i)
+      val end = if (i < partitionStartIndices.length - 1) partitionStartIndices(i + 1) else n
+      for (j <- start until end) {
+        result(j) = i
+      }
+    }
+    result
+  }
+
+  override def numPartitions: Int = partitionStartIndices.length
+
+  override def getPartition(key: Any): Int = {
+    parentPartitionMapping(parent.getPartition(key))
+  }
+
+  override def equals(other: Any): Boolean = other match {
+    case c: CoalescedPartitioner =>
+      c.parent == parent && Arrays.equals(c.partitionStartIndices, partitionStartIndices)
+    case _ =>
+      false
+  }
+
+  override def hashCode(): Int = 31 * parent.hashCode() + Arrays.hashCode(partitionStartIndices)
+}
+
 /**
  * This is a specialized version of [[org.apache.spark.rdd.ShuffledRDD]] that is optimized for
  * shuffling rows instead of Java key-value pairs. Note that something like this should eventually
  * be implemented in Spark core, but that is blocked by some more general refactorings to shuffle
  * interfaces / internals.
  *
- * @param prev the RDD being shuffled. Elements of this RDD are (partitionId, Row) pairs.
- *             Partition ids should be in the range [0, numPartitions - 1].
- * @param serializer the serializer used during the shuffle.
- * @param numPartitions the number of post-shuffle partitions.
+ * This RDD takes a [[ShuffleDependency]] (`dependency`),
+ * and a optional array of partition start indices as input arguments
+ * (`specifiedPartitionStartIndices`).
+ *
+ * The `dependency` has the parent RDD of this RDD, which represents the dataset before shuffle
+ * (i.e. map output). Elements of this RDD are (partitionId, Row) pairs.
+ * Partition ids should be in the range [0, numPartitions - 1].
+ * `dependency.partitioner` is the original partitioner used to partition
+ * map output, and `dependency.partitioner.numPartitions` is the number of pre-shuffle partitions
+ * (i.e. the number of partitions of the map output).
+ *
+ * When `specifiedPartitionStartIndices` is defined, `specifiedPartitionStartIndices.length`
+ * will be the number of post-shuffle partitions. For this case, the `i`th post-shuffle
+ * partition includes `specifiedPartitionStartIndices[i]` to
+ * `specifiedPartitionStartIndices[i+1] - 1` (inclusive).
+ *
+ * When `specifiedPartitionStartIndices` is not defined, there will be
+ * `dependency.partitioner.numPartitions` post-shuffle partitions. For this case,
+ * a post-shuffle partition is created for every pre-shuffle partition.
  */
 class ShuffledRowRDD(
-    @transient var prev: RDD[Product2[Int, InternalRow]],
-    serializer: Serializer,
-    numPartitions: Int)
-  extends RDD[InternalRow](prev.context, Nil) {
+    var dependency: ShuffleDependency[Int, InternalRow, InternalRow],
+    specifiedPartitionStartIndices: Option[Array[Int]] = None)
+  extends RDD[InternalRow](dependency.rdd.context, Nil) {
 
-  private val part: Partitioner = new PartitionIdPassthrough(numPartitions)
+  private[this] val numPreShufflePartitions = dependency.partitioner.numPartitions
 
-  override def getDependencies: Seq[Dependency[_]] = {
-    List(new ShuffleDependency[Int, InternalRow, InternalRow](prev, part, Some(serializer)))
+  private[this] val partitionStartIndices: Array[Int] = specifiedPartitionStartIndices match {
+    case Some(indices) => indices
+    case None =>
+      // When specifiedPartitionStartIndices is not defined, every post-shuffle partition
+      // corresponds to a pre-shuffle partition.
+      (0 until numPreShufflePartitions).toArray
   }
 
-  override val partitioner = Some(part)
+  private[this] val part: Partitioner =
+    new CoalescedPartitioner(dependency.partitioner, partitionStartIndices)
+
+  override def getDependencies: Seq[Dependency[_]] = List(dependency)
+
+  override val partitioner: Option[Partitioner] = Some(part)
 
   override def getPartitions: Array[Partition] = {
-    Array.tabulate[Partition](part.numPartitions)(i => new ShuffledRowRDDPartition(i))
+    assert(partitionStartIndices.length == part.numPartitions)
+    Array.tabulate[Partition](partitionStartIndices.length) { i =>
+      val startIndex = partitionStartIndices(i)
+      val endIndex =
+        if (i < partitionStartIndices.length - 1) {
+          partitionStartIndices(i + 1)
+        } else {
+          numPreShufflePartitions
+        }
+      new ShuffledRowRDDPartition(i, startIndex, endIndex)
+    }
   }
 
   override def getPreferredLocations(partition: Partition): Seq[String] = {
@@ -71,15 +154,20 @@ class ShuffledRowRDD(
   }
 
   override def compute(split: Partition, context: TaskContext): Iterator[InternalRow] = {
-    val dep = dependencies.head.asInstanceOf[ShuffleDependency[Int, InternalRow, InternalRow]]
-    SparkEnv.get.shuffleManager.getReader(dep.shuffleHandle, split.index, split.index + 1, context)
-      .read()
-      .asInstanceOf[Iterator[Product2[Int, InternalRow]]]
-      .map(_._2)
+    val shuffledRowPartition = split.asInstanceOf[ShuffledRowRDDPartition]
+    // The range of pre-shuffle partitions that we are fetching at here is
+    // [startPreShufflePartitionIndex, endPreShufflePartitionIndex - 1].
+    val reader =
+      SparkEnv.get.shuffleManager.getReader(
+        dependency.shuffleHandle,
+        shuffledRowPartition.startPreShufflePartitionIndex,
+        shuffledRowPartition.endPreShufflePartitionIndex,
+        context)
+    reader.read().asInstanceOf[Iterator[Product2[Int, InternalRow]]].map(_._2)
   }
 
   override def clearDependencies() {
     super.clearDependencies()
-    prev = null
+    dependency = null
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala
new file mode 100644
index 0000000000000..25f2f5caeed15
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala
@@ -0,0 +1,479 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.scalatest.BeforeAndAfterAll
+
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.test.TestSQLContext
+import org.apache.spark.sql._
+import org.apache.spark.{SparkFunSuite, SparkContext, SparkConf, MapOutputStatistics}
+
+class ExchangeCoordinatorSuite extends SparkFunSuite with BeforeAndAfterAll {
+
+  private var originalActiveSQLContext: Option[SQLContext] = _
+  private var originalInstantiatedSQLContext: Option[SQLContext] = _
+
+  override protected def beforeAll(): Unit = {
+    originalActiveSQLContext = SQLContext.getActiveContextOption()
+    originalInstantiatedSQLContext = SQLContext.getInstantiatedContextOption()
+
+    SQLContext.clearActive()
+    originalInstantiatedSQLContext.foreach(ctx => SQLContext.clearInstantiatedContext(ctx))
+  }
+
+  override protected def afterAll(): Unit = {
+    // Set these states back.
+    originalActiveSQLContext.foreach(ctx => SQLContext.setActive(ctx))
+    originalInstantiatedSQLContext.foreach(ctx => SQLContext.setInstantiatedContext(ctx))
+  }
+
+  private def checkEstimation(
+      coordinator: ExchangeCoordinator,
+      bytesByPartitionIdArray: Array[Array[Long]],
+      expectedPartitionStartIndices: Array[Int]): Unit = {
+    val mapOutputStatistics = bytesByPartitionIdArray.zipWithIndex.map {
+      case (bytesByPartitionId, index) =>
+        new MapOutputStatistics(index, bytesByPartitionId)
+    }
+    val estimatedPartitionStartIndices =
+      coordinator.estimatePartitionStartIndices(mapOutputStatistics)
+    assert(estimatedPartitionStartIndices === expectedPartitionStartIndices)
+  }
+
+  test("test estimatePartitionStartIndices - 1 Exchange") {
+    val coordinator = new ExchangeCoordinator(1, 100L)
+
+    {
+      // All bytes per partition are 0.
+      val bytesByPartitionId = Array[Long](0, 0, 0, 0, 0)
+      val expectedPartitionStartIndices = Array[Int](0)
+      checkEstimation(coordinator, Array(bytesByPartitionId), expectedPartitionStartIndices)
+    }
+
+    {
+      // Some bytes per partition are 0 and total size is less than the target size.
+      // 1 post-shuffle partition is needed.
+      val bytesByPartitionId = Array[Long](10, 0, 20, 0, 0)
+      val expectedPartitionStartIndices = Array[Int](0)
+      checkEstimation(coordinator, Array(bytesByPartitionId), expectedPartitionStartIndices)
+    }
+
+    {
+      // 2 post-shuffle partitions are needed.
+      val bytesByPartitionId = Array[Long](10, 0, 90, 20, 0)
+      val expectedPartitionStartIndices = Array[Int](0, 3)
+      checkEstimation(coordinator, Array(bytesByPartitionId), expectedPartitionStartIndices)
+    }
+
+    {
+      // There are a few large pre-shuffle partitions.
+      val bytesByPartitionId = Array[Long](110, 10, 100, 110, 0)
+      val expectedPartitionStartIndices = Array[Int](0, 1, 3, 4)
+      checkEstimation(coordinator, Array(bytesByPartitionId), expectedPartitionStartIndices)
+    }
+
+    {
+      // All pre-shuffle partitions are larger than the targeted size.
+      val bytesByPartitionId = Array[Long](100, 110, 100, 110, 110)
+      val expectedPartitionStartIndices = Array[Int](0, 1, 2, 3, 4)
+      checkEstimation(coordinator, Array(bytesByPartitionId), expectedPartitionStartIndices)
+    }
+
+    {
+      // The last pre-shuffle partition is in a single post-shuffle partition.
+      val bytesByPartitionId = Array[Long](30, 30, 0, 40, 110)
+      val expectedPartitionStartIndices = Array[Int](0, 4)
+      checkEstimation(coordinator, Array(bytesByPartitionId), expectedPartitionStartIndices)
+    }
+  }
+
+  test("test estimatePartitionStartIndices - 2 Exchanges") {
+    val coordinator = new ExchangeCoordinator(2, 100L)
+
+    {
+      // If there are multiple values of the number of pre-shuffle partitions,
+      // we should see an assertion error.
+      val bytesByPartitionId1 = Array[Long](0, 0, 0, 0, 0)
+      val bytesByPartitionId2 = Array[Long](0, 0, 0, 0, 0, 0)
+      val mapOutputStatistics =
+        Array(
+          new MapOutputStatistics(0, bytesByPartitionId1),
+          new MapOutputStatistics(1, bytesByPartitionId2))
+      intercept[AssertionError](coordinator.estimatePartitionStartIndices(mapOutputStatistics))
+    }
+
+    {
+      // All bytes per partition are 0.
+      val bytesByPartitionId1 = Array[Long](0, 0, 0, 0, 0)
+      val bytesByPartitionId2 = Array[Long](0, 0, 0, 0, 0)
+      val expectedPartitionStartIndices = Array[Int](0)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+
+    {
+      // Some bytes per partition are 0.
+      // 1 post-shuffle partition is needed.
+      val bytesByPartitionId1 = Array[Long](0, 10, 0, 20, 0)
+      val bytesByPartitionId2 = Array[Long](30, 0, 20, 0, 20)
+      val expectedPartitionStartIndices = Array[Int](0)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+
+    {
+      // 2 post-shuffle partition are needed.
+      val bytesByPartitionId1 = Array[Long](0, 10, 0, 20, 0)
+      val bytesByPartitionId2 = Array[Long](30, 0, 70, 0, 30)
+      val expectedPartitionStartIndices = Array[Int](0, 3)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+
+    {
+      // 2 post-shuffle partition are needed.
+      val bytesByPartitionId1 = Array[Long](0, 99, 0, 20, 0)
+      val bytesByPartitionId2 = Array[Long](30, 0, 70, 0, 30)
+      val expectedPartitionStartIndices = Array[Int](0, 2)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+
+    {
+      // 2 post-shuffle partition are needed.
+      val bytesByPartitionId1 = Array[Long](0, 100, 0, 30, 0)
+      val bytesByPartitionId2 = Array[Long](30, 0, 70, 0, 30)
+      val expectedPartitionStartIndices = Array[Int](0, 2, 4)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+
+    {
+      // There are a few large pre-shuffle partitions.
+      val bytesByPartitionId1 = Array[Long](0, 100, 40, 30, 0)
+      val bytesByPartitionId2 = Array[Long](30, 0, 60, 0, 110)
+      val expectedPartitionStartIndices = Array[Int](0, 2, 3)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+
+    {
+      // All pairs of pre-shuffle partitions are larger than the targeted size.
+      val bytesByPartitionId1 = Array[Long](100, 100, 40, 30, 0)
+      val bytesByPartitionId2 = Array[Long](30, 0, 60, 70, 110)
+      val expectedPartitionStartIndices = Array[Int](0, 1, 2, 3, 4)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+  }
+
+  test("test estimatePartitionStartIndices and enforce minimal number of reducers") {
+    val coordinator = new ExchangeCoordinator(2, 100L, Some(2))
+
+    {
+      // The minimal number of post-shuffle partitions is not enforced because
+      // the size of data is 0.
+      val bytesByPartitionId1 = Array[Long](0, 0, 0, 0, 0)
+      val bytesByPartitionId2 = Array[Long](0, 0, 0, 0, 0)
+      val expectedPartitionStartIndices = Array[Int](0)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+
+    {
+      // The minimal number of post-shuffle partitions is enforced.
+      val bytesByPartitionId1 = Array[Long](10, 5, 5, 0, 20)
+      val bytesByPartitionId2 = Array[Long](5, 10, 0, 10, 5)
+      val expectedPartitionStartIndices = Array[Int](0, 3)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+
+    {
+      // The number of post-shuffle partitions is determined by the coordinator.
+      val bytesByPartitionId1 = Array[Long](10, 50, 20, 80, 20)
+      val bytesByPartitionId2 = Array[Long](40, 10, 0, 10, 30)
+      val expectedPartitionStartIndices = Array[Int](0, 2, 4)
+      checkEstimation(
+        coordinator,
+        Array(bytesByPartitionId1, bytesByPartitionId2),
+        expectedPartitionStartIndices)
+    }
+  }
+
+  ///////////////////////////////////////////////////////////////////////////
+  // Query tests
+  ///////////////////////////////////////////////////////////////////////////
+
+  val numInputPartitions: Int = 10
+
+  def checkAnswer(actual: => DataFrame, expectedAnswer: Seq[Row]): Unit = {
+    QueryTest.checkAnswer(actual, expectedAnswer) match {
+      case Some(errorMessage) => fail(errorMessage)
+      case None =>
+    }
+  }
+
+  def withSQLContext(
+      f: SQLContext => Unit,
+      targetNumPostShufflePartitions: Int,
+      minNumPostShufflePartitions: Option[Int]): Unit = {
+    val sparkConf =
+      new SparkConf(false)
+        .setMaster("local[*]")
+        .setAppName("test")
+        .set("spark.ui.enabled", "false")
+        .set("spark.driver.allowMultipleContexts", "true")
+        .set(SQLConf.SHUFFLE_PARTITIONS.key, "5")
+        .set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true")
+        .set(
+          SQLConf.SHUFFLE_TARGET_POSTSHUFFLE_INPUT_SIZE.key,
+          targetNumPostShufflePartitions.toString)
+    minNumPostShufflePartitions match {
+      case Some(numPartitions) =>
+        sparkConf.set(SQLConf.SHUFFLE_MIN_NUM_POSTSHUFFLE_PARTITIONS.key, numPartitions.toString)
+      case None =>
+        sparkConf.set(SQLConf.SHUFFLE_MIN_NUM_POSTSHUFFLE_PARTITIONS.key, "-1")
+    }
+    val sparkContext = new SparkContext(sparkConf)
+    val sqlContext = new TestSQLContext(sparkContext)
+    try f(sqlContext) finally sparkContext.stop()
+  }
+
+  Seq(Some(3), None).foreach { minNumPostShufflePartitions =>
+    val testNameNote = minNumPostShufflePartitions match {
+      case Some(numPartitions) => "(minNumPostShufflePartitions: 3)"
+      case None => ""
+    }
+
+    test(s"determining the number of reducers: aggregate operator$testNameNote") {
+      val test = { sqlContext: SQLContext =>
+        val df =
+          sqlContext
+            .range(0, 1000, 1, numInputPartitions)
+            .selectExpr("id % 20 as key", "id as value")
+        val agg = df.groupBy("key").count
+
+        // Check the answer first.
+        checkAnswer(
+          agg,
+          sqlContext.range(0, 20).selectExpr("id", "50 as cnt").collect())
+
+        // Then, let's look at the number of post-shuffle partitions estimated
+        // by the ExchangeCoordinator.
+        val exchanges = agg.queryExecution.executedPlan.collect {
+          case e: Exchange => e
+        }
+        assert(exchanges.length === 1)
+        minNumPostShufflePartitions match {
+          case Some(numPartitions) =>
+            exchanges.foreach {
+              case e: Exchange =>
+                assert(e.coordinator.isDefined)
+                assert(e.outputPartitioning.numPartitions === 3)
+              case o =>
+            }
+
+          case None =>
+            exchanges.foreach {
+              case e: Exchange =>
+                assert(e.coordinator.isDefined)
+                assert(e.outputPartitioning.numPartitions === 2)
+              case o =>
+            }
+        }
+      }
+
+      withSQLContext(test, 1536, minNumPostShufflePartitions)
+    }
+
+    test(s"determining the number of reducers: join operator$testNameNote") {
+      val test = { sqlContext: SQLContext =>
+        val df1 =
+          sqlContext
+            .range(0, 1000, 1, numInputPartitions)
+            .selectExpr("id % 500 as key1", "id as value1")
+        val df2 =
+          sqlContext
+            .range(0, 1000, 1, numInputPartitions)
+            .selectExpr("id % 500 as key2", "id as value2")
+
+        val join = df1.join(df2, col("key1") === col("key2")).select(col("key1"), col("value2"))
+
+        // Check the answer first.
+        val expectedAnswer =
+          sqlContext
+            .range(0, 1000)
+            .selectExpr("id % 500 as key", "id as value")
+            .unionAll(sqlContext.range(0, 1000).selectExpr("id % 500 as key", "id as value"))
+        checkAnswer(
+          join,
+          expectedAnswer.collect())
+
+        // Then, let's look at the number of post-shuffle partitions estimated
+        // by the ExchangeCoordinator.
+        val exchanges = join.queryExecution.executedPlan.collect {
+          case e: Exchange => e
+        }
+        assert(exchanges.length === 2)
+        minNumPostShufflePartitions match {
+          case Some(numPartitions) =>
+            exchanges.foreach {
+              case e: Exchange =>
+                assert(e.coordinator.isDefined)
+                assert(e.outputPartitioning.numPartitions === 3)
+              case o =>
+            }
+
+          case None =>
+            exchanges.foreach {
+              case e: Exchange =>
+                assert(e.coordinator.isDefined)
+                assert(e.outputPartitioning.numPartitions === 2)
+              case o =>
+            }
+        }
+      }
+
+      withSQLContext(test, 16384, minNumPostShufflePartitions)
+    }
+
+    test(s"determining the number of reducers: complex query 1$testNameNote") {
+      val test = { sqlContext: SQLContext =>
+        val df1 =
+          sqlContext
+            .range(0, 1000, 1, numInputPartitions)
+            .selectExpr("id % 500 as key1", "id as value1")
+            .groupBy("key1")
+            .count
+            .toDF("key1", "cnt1")
+        val df2 =
+          sqlContext
+            .range(0, 1000, 1, numInputPartitions)
+            .selectExpr("id % 500 as key2", "id as value2")
+            .groupBy("key2")
+            .count
+            .toDF("key2", "cnt2")
+
+        val join = df1.join(df2, col("key1") === col("key2")).select(col("key1"), col("cnt2"))
+
+        // Check the answer first.
+        val expectedAnswer =
+          sqlContext
+            .range(0, 500)
+            .selectExpr("id", "2 as cnt")
+        checkAnswer(
+          join,
+          expectedAnswer.collect())
+
+        // Then, let's look at the number of post-shuffle partitions estimated
+        // by the ExchangeCoordinator.
+        val exchanges = join.queryExecution.executedPlan.collect {
+          case e: Exchange => e
+        }
+        assert(exchanges.length === 4)
+        minNumPostShufflePartitions match {
+          case Some(numPartitions) =>
+            exchanges.foreach {
+              case e: Exchange =>
+                assert(e.coordinator.isDefined)
+                assert(e.outputPartitioning.numPartitions === 3)
+              case o =>
+            }
+
+          case None =>
+            assert(exchanges.forall(_.coordinator.isDefined))
+            assert(exchanges.map(_.outputPartitioning.numPartitions).toSeq.toSet === Set(1, 2))
+        }
+      }
+
+      withSQLContext(test, 6144, minNumPostShufflePartitions)
+    }
+
+    test(s"determining the number of reducers: complex query 2$testNameNote") {
+      val test = { sqlContext: SQLContext =>
+        val df1 =
+          sqlContext
+            .range(0, 1000, 1, numInputPartitions)
+            .selectExpr("id % 500 as key1", "id as value1")
+            .groupBy("key1")
+            .count
+            .toDF("key1", "cnt1")
+        val df2 =
+          sqlContext
+            .range(0, 1000, 1, numInputPartitions)
+            .selectExpr("id % 500 as key2", "id as value2")
+
+        val join =
+          df1
+            .join(df2, col("key1") === col("key2"))
+            .select(col("key1"), col("cnt1"), col("value2"))
+
+        // Check the answer first.
+        val expectedAnswer =
+          sqlContext
+            .range(0, 1000)
+            .selectExpr("id % 500 as key", "2 as cnt", "id as value")
+        checkAnswer(
+          join,
+          expectedAnswer.collect())
+
+        // Then, let's look at the number of post-shuffle partitions estimated
+        // by the ExchangeCoordinator.
+        val exchanges = join.queryExecution.executedPlan.collect {
+          case e: Exchange => e
+        }
+        assert(exchanges.length === 3)
+        minNumPostShufflePartitions match {
+          case Some(numPartitions) =>
+            exchanges.foreach {
+              case e: Exchange =>
+                assert(e.coordinator.isDefined)
+                assert(e.outputPartitioning.numPartitions === 3)
+              case o =>
+            }
+
+          case None =>
+            assert(exchanges.forall(_.coordinator.isDefined))
+            assert(exchanges.map(_.outputPartitioning.numPartitions).toSeq.toSet === Set(2, 3))
+        }
+      }
+
+      withSQLContext(test, 6144, minNumPostShufflePartitions)
+    }
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
index ebdab1c26d7bd..2076c573b56c1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -268,7 +268,7 @@ class PlannerSuite extends SharedSQLContext {
     )
     val outputPlan = EnsureRequirements(sqlContext).apply(inputPlan)
     assertDistributionRequirementsAreSatisfied(outputPlan)
-    if (outputPlan.collect { case Exchange(_, _) => true }.isEmpty) {
+    if (outputPlan.collect { case e: Exchange => true }.isEmpty) {
       fail(s"Exchange should have been added:\n$outputPlan")
     }
   }
@@ -306,7 +306,7 @@ class PlannerSuite extends SharedSQLContext {
     )
     val outputPlan = EnsureRequirements(sqlContext).apply(inputPlan)
     assertDistributionRequirementsAreSatisfied(outputPlan)
-    if (outputPlan.collect { case Exchange(_, _) => true }.isEmpty) {
+    if (outputPlan.collect { case e: Exchange => true }.isEmpty) {
       fail(s"Exchange should have been added:\n$outputPlan")
     }
   }
@@ -326,7 +326,7 @@ class PlannerSuite extends SharedSQLContext {
     )
     val outputPlan = EnsureRequirements(sqlContext).apply(inputPlan)
     assertDistributionRequirementsAreSatisfied(outputPlan)
-    if (outputPlan.collect { case Exchange(_, _) => true }.nonEmpty) {
+    if (outputPlan.collect { case e: Exchange => true }.nonEmpty) {
       fail(s"Exchange should not have been added:\n$outputPlan")
     }
   }
@@ -349,7 +349,7 @@ class PlannerSuite extends SharedSQLContext {
     )
     val outputPlan = EnsureRequirements(sqlContext).apply(inputPlan)
     assertDistributionRequirementsAreSatisfied(outputPlan)
-    if (outputPlan.collect { case Exchange(_, _) => true }.nonEmpty) {
+    if (outputPlan.collect { case e: Exchange => true }.nonEmpty) {
       fail(s"No Exchanges should have been added:\n$outputPlan")
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeRowSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeRowSerializerSuite.scala
index d32572b54b8a8..09e258299de5a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeRowSerializerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeRowSerializerSuite.scala
@@ -152,7 +152,12 @@ class UnsafeRowSerializerSuite extends SparkFunSuite with LocalSparkContext {
     val unsafeRow = toUnsafeRow(row, Array(StringType, IntegerType))
     val rowsRDD = sc.parallelize(Seq((0, unsafeRow), (1, unsafeRow), (0, unsafeRow)))
       .asInstanceOf[RDD[Product2[Int, InternalRow]]]
-    val shuffled = new ShuffledRowRDD(rowsRDD, new UnsafeRowSerializer(2), 2)
+    val dependency =
+      new ShuffleDependency[Int, InternalRow, InternalRow](
+        rowsRDD,
+        new PartitionIdPassthrough(2),
+        Some(new UnsafeRowSerializer(2)))
+    val shuffled = new ShuffledRowRDD(dependency)
     shuffled.count()
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/InnerJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/InnerJoinSuite.scala
index da58e96f3e6f7..066c16e535c76 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/InnerJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/InnerJoinSuite.scala
@@ -49,7 +49,16 @@ class InnerJoinSuite extends SparkPlanTest with SharedSQLContext {
       Row(null, "e")
     )), new StructType().add("n", IntegerType).add("l", StringType))
 
-  private lazy val myTestData = Seq(
+  private lazy val myTestData1 = Seq(
+    (1, 1),
+    (1, 2),
+    (2, 1),
+    (2, 2),
+    (3, 1),
+    (3, 2)
+  ).toDF("a", "b")
+
+  private lazy val myTestData2 = Seq(
     (1, 1),
     (1, 2),
     (2, 1),
@@ -184,8 +193,8 @@ class InnerJoinSuite extends SparkPlanTest with SharedSQLContext {
   )
 
   {
-    lazy val left = myTestData.where("a = 1")
-    lazy val right = myTestData.where("a = 1")
+    lazy val left = myTestData1.where("a = 1")
+    lazy val right = myTestData2.where("a = 1")
     testInnerJoin(
       "inner join, multiple matches",
       left,
@@ -201,8 +210,8 @@ class InnerJoinSuite extends SparkPlanTest with SharedSQLContext {
   }
 
   {
-    lazy val left = myTestData.where("a = 1")
-    lazy val right = myTestData.where("a = 2")
+    lazy val left = myTestData1.where("a = 1")
+    lazy val right = myTestData2.where("a = 2")
     testInnerJoin(
       "inner join, no matches",
       left,

From 67e23b39ac3cdee06668fa9131951278b9731e29 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Tue, 3 Nov 2015 11:42:08 +0100
Subject: [PATCH 021/518] [SPARK-10429] [SQL] make mutableProjection atomic

Right now, SQL's mutable projection updates every value of the mutable project after it evaluates the corresponding expression. This makes the behavior of MutableProjection confusing and complicate the implementation of common aggregate functions like stddev because developers need to be aware that when evaluating {{i+1}}th expression of a mutable projection, {{i}}th slot of the mutable row has already been updated.

This PR make the MutableProjection atomic, by generating all the results of expressions first, then copy them into mutableRow.

Had run a mircro-benchmark, there is no notable performance difference between using class members and local variables.

cc yhuai

Author: Davies Liu <davies@databricks.com>

Closes #9422 from davies/atomic_mutable and squashes the following commits:

bbc1758 [Davies Liu] support wide table
8a0ae14 [Davies Liu] fix bug
bec07da [Davies Liu] refactor
2891628 [Davies Liu] make mutableProjection atomic
---
 .../sql/catalyst/expressions/Projection.scala |  13 +-
 .../expressions/aggregate/functions.scala     | 154 ++++++++----------
 .../codegen/GenerateMutableProjection.scala   |  28 +++-
 3 files changed, 97 insertions(+), 98 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
index afe52e6a667eb..a6fe730f6dad4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
@@ -19,8 +19,7 @@ package org.apache.spark.sql.catalyst.expressions
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateSafeProjection, GenerateUnsafeProjection}
-import org.apache.spark.sql.types.{DataType, Decimal, StructType, _}
-import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
+import org.apache.spark.sql.types.{DataType, StructType}
 
 /**
  * A [[Projection]] that is calculated by calling the `eval` of each of the specified expressions.
@@ -62,6 +61,8 @@ case class InterpretedMutableProjection(expressions: Seq[Expression]) extends Mu
   def this(expressions: Seq[Expression], inputSchema: Seq[Attribute]) =
     this(expressions.map(BindReferences.bindReference(_, inputSchema)))
 
+  private[this] val buffer = new Array[Any](expressions.size)
+
   expressions.foreach(_.foreach {
     case n: Nondeterministic => n.setInitialValues()
     case _ =>
@@ -79,7 +80,13 @@ case class InterpretedMutableProjection(expressions: Seq[Expression]) extends Mu
   override def apply(input: InternalRow): InternalRow = {
     var i = 0
     while (i < exprArray.length) {
-      mutableRow(i) = exprArray(i).eval(input)
+      // Store the result into buffer first, to make the projection atomic (needed by aggregation)
+      buffer(i) = exprArray(i).eval(input)
+      i += 1
+    }
+    i = 0
+    while (i < exprArray.length) {
+      mutableRow(i) = buffer(i)
       i += 1
     }
     mutableRow
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala
index 5d2eb7b017ab9..f2c3eca095115 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala
@@ -57,37 +57,37 @@ case class Average(child: Expression) extends DeclarativeAggregate {
     case _ => DoubleType
   }
 
-  private val currentSum = AttributeReference("currentSum", sumDataType)()
-  private val currentCount = AttributeReference("currentCount", LongType)()
+  private val sum = AttributeReference("sum", sumDataType)()
+  private val count = AttributeReference("count", LongType)()
 
-  override val aggBufferAttributes = currentSum :: currentCount :: Nil
+  override val aggBufferAttributes = sum :: count :: Nil
 
   override val initialValues = Seq(
-    /* currentSum = */ Cast(Literal(0), sumDataType),
-    /* currentCount = */ Literal(0L)
+    /* sum = */ Cast(Literal(0), sumDataType),
+    /* count = */ Literal(0L)
   )
 
   override val updateExpressions = Seq(
-    /* currentSum = */
+    /* sum = */
     Add(
-      currentSum,
+      sum,
       Coalesce(Cast(child, sumDataType) :: Cast(Literal(0), sumDataType) :: Nil)),
-    /* currentCount = */ If(IsNull(child), currentCount, currentCount + 1L)
+    /* count = */ If(IsNull(child), count, count + 1L)
   )
 
   override val mergeExpressions = Seq(
-    /* currentSum = */ currentSum.left + currentSum.right,
-    /* currentCount = */ currentCount.left + currentCount.right
+    /* sum = */ sum.left + sum.right,
+    /* count = */ count.left + count.right
   )
 
-  // If all input are nulls, currentCount will be 0 and we will get null after the division.
+  // If all input are nulls, count will be 0 and we will get null after the division.
   override val evaluateExpression = child.dataType match {
     case DecimalType.Fixed(p, s) =>
       // increase the precision and scale to prevent precision loss
       val dt = DecimalType.bounded(p + 14, s + 4)
-      Cast(Cast(currentSum, dt) / Cast(currentCount, dt), resultType)
+      Cast(Cast(sum, dt) / Cast(count, dt), resultType)
     case _ =>
-      Cast(currentSum, resultType) / Cast(currentCount, resultType)
+      Cast(sum, resultType) / Cast(count, resultType)
   }
 }
 
@@ -102,23 +102,23 @@ case class Count(child: Expression) extends DeclarativeAggregate {
   // Expected input data type.
   override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
 
-  private val currentCount = AttributeReference("currentCount", LongType)()
+  private val count = AttributeReference("count", LongType)()
 
-  override val aggBufferAttributes = currentCount :: Nil
+  override val aggBufferAttributes = count :: Nil
 
   override val initialValues = Seq(
-    /* currentCount = */ Literal(0L)
+    /* count = */ Literal(0L)
   )
 
   override val updateExpressions = Seq(
-    /* currentCount = */ If(IsNull(child), currentCount, currentCount + 1L)
+    /* count = */ If(IsNull(child), count, count + 1L)
   )
 
   override val mergeExpressions = Seq(
-    /* currentCount = */ currentCount.left + currentCount.right
+    /* count = */ count.left + count.right
   )
 
-  override val evaluateExpression = Cast(currentCount, LongType)
+  override val evaluateExpression = Cast(count, LongType)
 }
 
 /**
@@ -372,101 +372,77 @@ abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
 
   private val resultType = DoubleType
 
-  private val preCount = AttributeReference("preCount", resultType)()
-  private val currentCount = AttributeReference("currentCount", resultType)()
-  private val preAvg = AttributeReference("preAvg", resultType)()
-  private val currentAvg = AttributeReference("currentAvg", resultType)()
-  private val currentMk = AttributeReference("currentMk", resultType)()
+  private val count = AttributeReference("count", resultType)()
+  private val avg = AttributeReference("avg", resultType)()
+  private val mk = AttributeReference("mk", resultType)()
 
-  override val aggBufferAttributes = preCount :: currentCount :: preAvg ::
-                                  currentAvg :: currentMk :: Nil
+  override val aggBufferAttributes = count :: avg :: mk :: Nil
 
   override val initialValues = Seq(
-    /* preCount = */ Cast(Literal(0), resultType),
-    /* currentCount = */ Cast(Literal(0), resultType),
-    /* preAvg = */ Cast(Literal(0), resultType),
-    /* currentAvg = */ Cast(Literal(0), resultType),
-    /* currentMk = */ Cast(Literal(0), resultType)
+    /* count = */ Cast(Literal(0), resultType),
+    /* avg = */ Cast(Literal(0), resultType),
+    /* mk = */ Cast(Literal(0), resultType)
   )
 
   override val updateExpressions = {
+    val value = Cast(child, resultType)
+    val newCount = count + Cast(Literal(1), resultType)
 
     // update average
     // avg = avg + (value - avg)/count
-    def avgAdd: Expression = {
-      currentAvg + ((Cast(child, resultType) - currentAvg) / currentCount)
-    }
+    val newAvg = avg + (value - avg) / newCount
 
     // update sum of square of difference from mean
     // Mk = Mk + (value - preAvg) * (value - updatedAvg)
-    def mkAdd: Expression = {
-      val delta1 = Cast(child, resultType) - preAvg
-      val delta2 = Cast(child, resultType) - currentAvg
-      currentMk + (delta1 * delta2)
-    }
+    val newMk = mk + (value - avg) * (value - newAvg)
 
     Seq(
-      /* preCount = */ If(IsNull(child), preCount, currentCount),
-      /* currentCount = */ If(IsNull(child), currentCount,
-                           Add(currentCount, Cast(Literal(1), resultType))),
-      /* preAvg = */ If(IsNull(child), preAvg, currentAvg),
-      /* currentAvg = */ If(IsNull(child), currentAvg, avgAdd),
-      /* currentMk = */ If(IsNull(child), currentMk, mkAdd)
+      /* count = */ If(IsNull(child), count, newCount),
+      /* avg = */ If(IsNull(child), avg, newAvg),
+      /* mk = */ If(IsNull(child), mk, newMk)
     )
   }
 
   override val mergeExpressions = {
 
     // count merge
-    def countMerge: Expression = {
-      currentCount.left + currentCount.right
-    }
+    val newCount = count.left + count.right
 
     // average merge
-    def avgMerge: Expression = {
-      ((currentAvg.left * preCount) + (currentAvg.right * currentCount.right)) /
-      (preCount + currentCount.right)
-    }
+    val newAvg = ((avg.left * count.left) + (avg.right * count.right)) / newCount
 
     // update sum of square differences
-    def mkMerge: Expression = {
-      val avgDelta = currentAvg.right - preAvg
-      val mkDelta = (avgDelta * avgDelta) * (preCount * currentCount.right) /
-        (preCount + currentCount.right)
-
-      currentMk.left + currentMk.right + mkDelta
+    val newMk = {
+      val avgDelta = avg.right - avg.left
+      val mkDelta = (avgDelta * avgDelta) * (count.left * count.right) / newCount
+      mk.left + mk.right + mkDelta
     }
 
     Seq(
-      /* preCount = */ If(IsNull(currentCount.left),
-                         Cast(Literal(0), resultType), currentCount.left),
-      /* currentCount = */ If(IsNull(currentCount.left), currentCount.right,
-                             If(IsNull(currentCount.right), currentCount.left, countMerge)),
-      /* preAvg = */ If(IsNull(currentAvg.left), Cast(Literal(0), resultType), currentAvg.left),
-      /* currentAvg = */ If(IsNull(currentAvg.left), currentAvg.right,
-                           If(IsNull(currentAvg.right), currentAvg.left, avgMerge)),
-      /* currentMk = */ If(IsNull(currentMk.left), currentMk.right,
-                          If(IsNull(currentMk.right), currentMk.left, mkMerge))
+      /* count = */ If(IsNull(count.left), count.right,
+                       If(IsNull(count.right), count.left, newCount)),
+      /* avg = */ If(IsNull(avg.left), avg.right,
+                     If(IsNull(avg.right), avg.left, newAvg)),
+      /* mk = */ If(IsNull(mk.left), mk.right,
+                    If(IsNull(mk.right), mk.left, newMk))
     )
   }
 
   override val evaluateExpression = {
-    // when currentCount == 0, return null
-    // when currentCount == 1, return 0
-    // when currentCount >1
-    // stddev_samp = sqrt (currentMk/(currentCount -1))
-    // stddev_pop = sqrt (currentMk/currentCount)
-    val varCol = {
+    // when count == 0, return null
+    // when count == 1, return 0
+    // when count >1
+    // stddev_samp = sqrt (mk/(count -1))
+    // stddev_pop = sqrt (mk/count)
+    val varCol =
       if (isSample) {
-        currentMk / Cast((currentCount - Cast(Literal(1), resultType)), resultType)
-      }
-      else {
-        currentMk / currentCount
+        mk / Cast((count - Cast(Literal(1), resultType)), resultType)
+      } else {
+        mk / count
       }
-    }
 
-    If(EqualTo(currentCount, Cast(Literal(0), resultType)), Cast(Literal(null), resultType),
-      If(EqualTo(currentCount, Cast(Literal(1), resultType)), Cast(Literal(0), resultType),
+    If(EqualTo(count, Cast(Literal(0), resultType)), Cast(Literal(null), resultType),
+      If(EqualTo(count, Cast(Literal(1), resultType)), Cast(Literal(0), resultType),
         Cast(Sqrt(varCol), resultType)))
   }
 }
@@ -499,30 +475,30 @@ case class Sum(child: Expression) extends DeclarativeAggregate {
 
   private val sumDataType = resultType
 
-  private val currentSum = AttributeReference("currentSum", sumDataType)()
+  private val sum = AttributeReference("sum", sumDataType)()
 
   private val zero = Cast(Literal(0), sumDataType)
 
-  override val aggBufferAttributes = currentSum :: Nil
+  override val aggBufferAttributes = sum :: Nil
 
   override val initialValues = Seq(
-    /* currentSum = */ Literal.create(null, sumDataType)
+    /* sum = */ Literal.create(null, sumDataType)
   )
 
   override val updateExpressions = Seq(
-    /* currentSum = */
-    Coalesce(Seq(Add(Coalesce(Seq(currentSum, zero)), Cast(child, sumDataType)), currentSum))
+    /* sum = */
+    Coalesce(Seq(Add(Coalesce(Seq(sum, zero)), Cast(child, sumDataType)), sum))
   )
 
   override val mergeExpressions = {
-    val add = Add(Coalesce(Seq(currentSum.left, zero)), Cast(currentSum.right, sumDataType))
+    val add = Add(Coalesce(Seq(sum.left, zero)), Cast(sum.right, sumDataType))
     Seq(
-      /* currentSum = */
-      Coalesce(Seq(add, currentSum.left))
+      /* sum = */
+      Coalesce(Seq(add, sum.left))
     )
   }
 
-  override val evaluateExpression = Cast(currentSum, resultType)
+  override val evaluateExpression = Cast(sum, resultType)
 }
 
 /**
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala
index e8ee64756d5d0..4b66069b5f55a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala
@@ -44,28 +44,42 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], () => Mu
       case (NoOp, _) => ""
       case (e, i) =>
         val evaluationCode = e.gen(ctx)
+        val isNull = s"isNull_$i"
+        val value = s"value_$i"
+        ctx.addMutableState("boolean", isNull, s"this.$isNull = true;")
+        ctx.addMutableState(ctx.javaType(e.dataType), value,
+          s"this.$value = ${ctx.defaultValue(e.dataType)};")
+        s"""
+          ${evaluationCode.code}
+          this.$isNull = ${evaluationCode.isNull};
+          this.$value = ${evaluationCode.value};
+         """
+    }
+    val updates = expressions.zipWithIndex.map {
+      case (NoOp, _) => ""
+      case (e, i) =>
         if (e.dataType.isInstanceOf[DecimalType]) {
           // Can't call setNullAt on DecimalType, because we need to keep the offset
           s"""
-            ${evaluationCode.code}
-            if (${evaluationCode.isNull}) {
+            if (this.isNull_$i) {
               ${ctx.setColumn("mutableRow", e.dataType, i, null)};
             } else {
-              ${ctx.setColumn("mutableRow", e.dataType, i, evaluationCode.value)};
+              ${ctx.setColumn("mutableRow", e.dataType, i, s"this.value_$i")};
             }
           """
         } else {
           s"""
-            ${evaluationCode.code}
-            if (${evaluationCode.isNull}) {
+            if (this.isNull_$i) {
               mutableRow.setNullAt($i);
             } else {
-              ${ctx.setColumn("mutableRow", e.dataType, i, evaluationCode.value)};
+              ${ctx.setColumn("mutableRow", e.dataType, i, s"this.value_$i")};
             }
           """
         }
     }
+
     val allProjections = ctx.splitExpressions(ctx.INPUT_ROW, projectionCodes)
+    val allUpdates = ctx.splitExpressions(ctx.INPUT_ROW, updates)
 
     val code = s"""
       public Object generate($exprType[] expr) {
@@ -98,6 +112,8 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], () => Mu
         public Object apply(Object _i) {
           InternalRow ${ctx.INPUT_ROW} = (InternalRow) _i;
           $allProjections
+          // copy all the results into MutableRow
+          $allUpdates
           return mutableRow;
         }
       }

From 425ff03f5ac4f3ddda1ba06656e620d5426f4209 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Tue, 3 Nov 2015 12:47:39 +0100
Subject: [PATCH 022/518] [SPARK-11436] [SQL] rebind right encoder when join 2
 datasets

When we join 2 datasets, we will combine 2 encoders into a tupled one, and use it as the encoder for the jioned dataset. Assume both of the 2 encoders are flat, their `constructExpression`s both reference to the first element of input row. However, when we combine 2 encoders, the schema of input row changed,  now the right encoder should reference to second element of input row. So we should rebind right encoder to let it know the new schema of input row before combine it.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9391 from cloud-fan/join and squashes the following commits:

846d3ab [Wenchen Fan] rebind right encoder when join 2 datasets
---
 .../src/main/scala/org/apache/spark/sql/Dataset.scala     | 4 +++-
 .../test/scala/org/apache/spark/sql/DatasetSuite.scala    | 8 ++++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index e0ab5f593e933..ed98a2541598f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -390,7 +390,9 @@ class Dataset[T] private(
     val rightEncoder =
       if (other.encoder.flat) other.encoder else other.encoder.nested(rightData.toAttribute)
     implicit val tuple2Encoder: Encoder[(T, U)] =
-      ExpressionEncoder.tuple(leftEncoder, rightEncoder)
+      ExpressionEncoder.tuple(
+        leftEncoder,
+        rightEncoder.rebind(right.output, left.output ++ right.output))
 
     withPlan[(T, U)](other) { (left, right) =>
       Project(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 993e6d269ee03..95b8d05cf4414 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -214,4 +214,12 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       cogrouped,
       1 -> "a#", 2 -> "#q", 3 -> "abcfoo#w", 5 -> "hello#er")
   }
+
+  test("SPARK-11436: we should rebind right encoder when join 2 datasets") {
+    val ds1 = Seq("1", "2").toDS().as("a")
+    val ds2 = Seq(2, 3).toDS().as("b")
+
+    val joined = ds1.joinWith(ds2, $"a.value" === $"b.value")
+    checkAnswer(joined, ("2", 2))
+  }
 }

From b86f2cab67989f09ba1ba8604e52cd4b1e44e436 Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Tue, 3 Nov 2015 13:02:17 +0100
Subject: [PATCH 023/518] [SPARK-11404] [SQL] Support for groupBy using column
 expressions

This PR adds a new method `groupBy(cols: Column*)` to `Dataset` that allows users to group using column expressions instead of a lambda function.  Since the return type of these expressions is not known at compile time, we just set the key type as a generic `Row`.  If the user would like to work the key in a type-safe way, they can call `grouped.asKey[Type]`, which is also added in this PR.

```scala
val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
val grouped = ds.groupBy($"_1").asKey[String]
val agged = grouped.mapGroups { case (g, iter) =>
  Iterator((g, iter.map(_._2).sum))
}

agged.collect()

res0: Array(("a", 30), ("b", 3), ("c", 1))
```

Author: Michael Armbrust <michael@databricks.com>

Closes #9359 from marmbrus/columnGroupBy and squashes the following commits:

bbcb03b [Michael Armbrust] Update DatasetSuite.scala
8fd2908 [Michael Armbrust] Update DatasetSuite.scala
0b0e2f8 [Michael Armbrust] [SPARK-11404] [SQL] Support for groupBy using column expressions
---
 .../scala/org/apache/spark/sql/Dataset.scala  | 36 ++++++++++++--
 .../org/apache/spark/sql/GroupedDataset.scala | 28 +++++++++--
 .../org/apache/spark/sql/DatasetSuite.scala   | 48 +++++++++++++++++++
 3 files changed, 106 insertions(+), 6 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index ed98a2541598f..7b75aeec4cf3a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql
 
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAlias
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.Inner
@@ -78,9 +79,17 @@ class Dataset[T] private(
    * ************* */
 
   /**
-   * Returns a new `Dataset` where each record has been mapped on to the specified type.
-   * TODO: should bind here...
-   * TODO: document binding rules
+   * Returns a new `Dataset` where each record has been mapped on to the specified type.  The
+   * method used to map columns depend on the type of `U`:
+   *  - When `U` is a class, fields for the class will be mapped to columns of the same name
+   *    (case sensitivity is determined by `spark.sql.caseSensitive`)
+   *  - When `U` is a tuple, the columns will be be mapped by ordinal (i.e. the first column will
+   *    be assigned to `_1`).
+   *  - When `U` is a primitive type (i.e. String, Int, etc). then the first column of the
+   *    [[DataFrame]] will be used.
+   *
+   * If the schema of the [[DataFrame]] does not match the desired `U` type, you can use `select`
+   * along with `alias` or `as` to rearrange or rename as required.
    * @since 1.6.0
    */
   def as[U : Encoder]: Dataset[U] = {
@@ -225,6 +234,27 @@ class Dataset[T] private(
       withGroupingKey.newColumns)
   }
 
+  /**
+   * Returns a [[GroupedDataset]] where the data is grouped by the given [[Column]] expressions.
+   * @since 1.6.0
+   */
+  @scala.annotation.varargs
+  def groupBy(cols: Column*): GroupedDataset[Row, T] = {
+    val withKeyColumns = logicalPlan.output ++ cols.map(_.expr).map(UnresolvedAlias)
+    val withKey = Project(withKeyColumns, logicalPlan)
+    val executed = sqlContext.executePlan(withKey)
+
+    val dataAttributes = executed.analyzed.output.dropRight(cols.size)
+    val keyAttributes = executed.analyzed.output.takeRight(cols.size)
+
+    new GroupedDataset(
+      RowEncoder(keyAttributes.toStructType),
+      encoderFor[T],
+      executed,
+      dataAttributes,
+      keyAttributes)
+  }
+
   /* ****************** *
    *  Typed Relational  *
    * ****************** */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 612f2b60cd405..96d6e9dd548e5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql
 
-import org.apache.spark.sql.catalyst.encoders.Encoder
+import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor, Encoder}
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.QueryExecution
@@ -34,11 +34,33 @@ class GroupedDataset[K, T] private[sql](
     private val dataAttributes: Seq[Attribute],
     private val groupingAttributes: Seq[Attribute]) extends Serializable {
 
-  private implicit def kEnc = kEncoder
-  private implicit def tEnc = tEncoder
+  private implicit val kEnc = kEncoder match {
+    case e: ExpressionEncoder[K] => e.resolve(groupingAttributes)
+    case other =>
+      throw new UnsupportedOperationException("Only expression encoders are currently supported")
+  }
+
+  private implicit val tEnc = tEncoder match {
+    case e: ExpressionEncoder[T] => e.resolve(dataAttributes)
+    case other =>
+      throw new UnsupportedOperationException("Only expression encoders are currently supported")
+  }
+
   private def logicalPlan = queryExecution.analyzed
   private def sqlContext = queryExecution.sqlContext
 
+  /**
+   * Returns a new [[GroupedDataset]] where the type of the key has been mapped to the specified
+   * type. The mapping of key columns to the type follows the same rules as `as` on [[Dataset]].
+   */
+  def asKey[L : Encoder]: GroupedDataset[L, T] =
+    new GroupedDataset(
+      encoderFor[L],
+      tEncoder,
+      queryExecution,
+      dataAttributes,
+      groupingAttributes)
+
   /**
    * Returns a [[Dataset]] that contains each unique key.
    */
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 95b8d05cf4414..5973fa7f2a76b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -203,6 +203,54 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       ("a", 30), ("b", 3), ("c", 1))
   }
 
+  test("groupBy columns, mapGroups") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+    val grouped = ds.groupBy($"_1")
+    val agged = grouped.mapGroups { case (g, iter) =>
+      Iterator((g.getString(0), iter.map(_._2).sum))
+    }
+
+    checkAnswer(
+      agged,
+      ("a", 30), ("b", 3), ("c", 1))
+  }
+
+  test("groupBy columns asKey, mapGroups") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+    val grouped = ds.groupBy($"_1").asKey[String]
+    val agged = grouped.mapGroups { case (g, iter) =>
+      Iterator((g, iter.map(_._2).sum))
+    }
+
+    checkAnswer(
+      agged,
+      ("a", 30), ("b", 3), ("c", 1))
+  }
+
+  test("groupBy columns asKey tuple, mapGroups") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+    val grouped = ds.groupBy($"_1", lit(1)).asKey[(String, Int)]
+    val agged = grouped.mapGroups { case (g, iter) =>
+      Iterator((g, iter.map(_._2).sum))
+    }
+
+    checkAnswer(
+      agged,
+      (("a", 1), 30), (("b", 1), 3), (("c", 1), 1))
+  }
+
+  test("groupBy columns asKey class, mapGroups") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+    val grouped = ds.groupBy($"_1".as("a"), lit(1).as("b")).asKey[ClassData]
+    val agged = grouped.mapGroups { case (g, iter) =>
+      Iterator((g, iter.map(_._2).sum))
+    }
+
+    checkAnswer(
+      agged,
+      (ClassData("a", 1), 30), (ClassData("b", 1), 3), (ClassData("c", 1), 1))
+  }
+
   test("cogroup") {
     val ds1 = Seq(1 -> "a", 3 -> "abc", 5 -> "hello", 3 -> "foo").toDS()
     val ds2 = Seq(2 -> "q", 3 -> "w", 5 -> "e", 5 -> "r").toDS()

From 233e534ac43ea25ac1b0e6a985f6928d46c5d03a Mon Sep 17 00:00:00 2001
From: Jacek Lewandowski <lewandowski.jacek@gmail.com>
Date: Tue, 3 Nov 2015 12:46:11 +0000
Subject: [PATCH 024/518] [SPARK-11344] Made ApplicationDescription and
 DriverDescription case classes

DriverDescription refactored to case class because it included no mutable fields.

ApplicationDescription had one mutable field, which was appUiUrl. This field was set by the driver to point to the driver web UI. Master was modifying this field when the application was removed to redirect requests to history server. This was wrong because objects which are sent over the wire should be immutable. Now appUiUrl is immutable in ApplicationDescription and always points to the driver UI even if it is already shutdown. The UI url which master exposes to the user and modifies dynamically is now included into ApplicationInfo - a data object which describes the application state internally in master. That URL in ApplicationInfo is initialised with the value from ApplicationDescription.

ApplicationDescription also included value user, which is now a part of case class fields.

Author: Jacek Lewandowski <lewandowski.jacek@gmail.com>

Closes #9299 from jacek-lewandowski/SPARK-11344.
---
 .../spark/deploy/ApplicationDescription.scala | 33 ++++++-------------
 .../spark/deploy/DriverDescription.scala      | 21 ++++--------
 .../spark/deploy/master/ApplicationInfo.scala |  7 ++++
 .../apache/spark/deploy/master/Master.scala   | 12 ++++---
 .../deploy/master/ui/ApplicationPage.scala    |  2 +-
 .../spark/deploy/master/ui/MasterPage.scala   |  2 +-
 .../apache/spark/deploy/DeployTestUtils.scala |  3 +-
 7 files changed, 34 insertions(+), 46 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/ApplicationDescription.scala b/core/src/main/scala/org/apache/spark/deploy/ApplicationDescription.scala
index ae99432f5ce86..78bbd5c03f4a6 100644
--- a/core/src/main/scala/org/apache/spark/deploy/ApplicationDescription.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/ApplicationDescription.scala
@@ -19,30 +19,17 @@ package org.apache.spark.deploy
 
 import java.net.URI
 
-private[spark] class ApplicationDescription(
-    val name: String,
-    val maxCores: Option[Int],
-    val memoryPerExecutorMB: Int,
-    val command: Command,
-    var appUiUrl: String,
-    val eventLogDir: Option[URI] = None,
+private[spark] case class ApplicationDescription(
+    name: String,
+    maxCores: Option[Int],
+    memoryPerExecutorMB: Int,
+    command: Command,
+    appUiUrl: String,
+    eventLogDir: Option[URI] = None,
     // short name of compression codec used when writing event logs, if any (e.g. lzf)
-    val eventLogCodec: Option[String] = None,
-    val coresPerExecutor: Option[Int] = None)
-  extends Serializable {
-
-  val user = System.getProperty("user.name", "<unknown>")
-
-  def copy(
-      name: String = name,
-      maxCores: Option[Int] = maxCores,
-      memoryPerExecutorMB: Int = memoryPerExecutorMB,
-      command: Command = command,
-      appUiUrl: String = appUiUrl,
-      eventLogDir: Option[URI] = eventLogDir,
-      eventLogCodec: Option[String] = eventLogCodec): ApplicationDescription =
-    new ApplicationDescription(
-      name, maxCores, memoryPerExecutorMB, command, appUiUrl, eventLogDir, eventLogCodec)
+    eventLogCodec: Option[String] = None,
+    coresPerExecutor: Option[Int] = None,
+    user: String = System.getProperty("user.name", "<unknown>")) {
 
   override def toString: String = "ApplicationDescription(" + name + ")"
 }
diff --git a/core/src/main/scala/org/apache/spark/deploy/DriverDescription.scala b/core/src/main/scala/org/apache/spark/deploy/DriverDescription.scala
index 659fb434a80f5..1f5626ab5a896 100644
--- a/core/src/main/scala/org/apache/spark/deploy/DriverDescription.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/DriverDescription.scala
@@ -17,21 +17,12 @@
 
 package org.apache.spark.deploy
 
-private[deploy] class DriverDescription(
-    val jarUrl: String,
-    val mem: Int,
-    val cores: Int,
-    val supervise: Boolean,
-    val command: Command)
-  extends Serializable {
-
-  def copy(
-      jarUrl: String = jarUrl,
-      mem: Int = mem,
-      cores: Int = cores,
-      supervise: Boolean = supervise,
-      command: Command = command): DriverDescription =
-    new DriverDescription(jarUrl, mem, cores, supervise, command)
+private[deploy] case class DriverDescription(
+    jarUrl: String,
+    mem: Int,
+    cores: Int,
+    supervise: Boolean,
+    command: Command) {
 
   override def toString: String = s"DriverDescription (${command.mainClass})"
 }
diff --git a/core/src/main/scala/org/apache/spark/deploy/master/ApplicationInfo.scala b/core/src/main/scala/org/apache/spark/deploy/master/ApplicationInfo.scala
index b40d20f9f7868..ac553b71115df 100644
--- a/core/src/main/scala/org/apache/spark/deploy/master/ApplicationInfo.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/master/ApplicationInfo.scala
@@ -41,6 +41,7 @@ private[spark] class ApplicationInfo(
   @transient var coresGranted: Int = _
   @transient var endTime: Long = _
   @transient var appSource: ApplicationSource = _
+  @transient @volatile var appUIUrlAtHistoryServer: Option[String] = None
 
   // A cap on the number of executors this application can have at any given time.
   // By default, this is infinite. Only after the first allocation request is issued by the
@@ -135,4 +136,10 @@ private[spark] class ApplicationInfo(
     }
   }
 
+  /**
+   * Returns the original application UI url unless there is its address at history server
+   * is defined
+   */
+  def curAppUIUrl: String = appUIUrlAtHistoryServer.getOrElse(desc.appUiUrl)
+
 }
diff --git a/core/src/main/scala/org/apache/spark/deploy/master/Master.scala b/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
index 6715d6c70f497..b25a487806c7f 100644
--- a/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
@@ -768,7 +768,8 @@ private[deploy] class Master(
       ApplicationInfo = {
     val now = System.currentTimeMillis()
     val date = new Date(now)
-    new ApplicationInfo(now, newApplicationId(date), desc, date, driver, defaultCores)
+    val appId = newApplicationId(date)
+    new ApplicationInfo(now, appId, desc, date, driver, defaultCores)
   }
 
   private def registerApplication(app: ApplicationInfo): Unit = {
@@ -920,7 +921,7 @@ private[deploy] class Master(
       val eventLogDir = app.desc.eventLogDir
         .getOrElse {
           // Event logging is not enabled for this application
-          app.desc.appUiUrl = notFoundBasePath
+          app.appUIUrlAtHistoryServer = Some(notFoundBasePath)
           return None
         }
 
@@ -954,7 +955,7 @@ private[deploy] class Master(
       appIdToUI(app.id) = ui
       webUi.attachSparkUI(ui)
       // Application UI is successfully rebuilt, so link the Master UI to it
-      app.desc.appUiUrl = ui.basePath
+      app.appUIUrlAtHistoryServer = Some(ui.basePath)
       Some(ui)
     } catch {
       case fnf: FileNotFoundException =>
@@ -964,7 +965,7 @@ private[deploy] class Master(
         logWarning(msg)
         msg += " Did you specify the correct logging directory?"
         msg = URLEncoder.encode(msg, "UTF-8")
-        app.desc.appUiUrl = notFoundBasePath + s"?msg=$msg&title=$title"
+        app.appUIUrlAtHistoryServer = Some(notFoundBasePath + s"?msg=$msg&title=$title")
         None
       case e: Exception =>
         // Relay exception message to application UI page
@@ -973,7 +974,8 @@ private[deploy] class Master(
         var msg = s"Exception in replaying log for application $appName!"
         logError(msg, e)
         msg = URLEncoder.encode(msg, "UTF-8")
-        app.desc.appUiUrl = notFoundBasePath + s"?msg=$msg&exception=$exception&title=$title"
+        app.appUIUrlAtHistoryServer =
+            Some(notFoundBasePath + s"?msg=$msg&exception=$exception&title=$title")
         None
     }
   }
diff --git a/core/src/main/scala/org/apache/spark/deploy/master/ui/ApplicationPage.scala b/core/src/main/scala/org/apache/spark/deploy/master/ui/ApplicationPage.scala
index e28e7e379ac91..f405aa2bdc8b3 100644
--- a/core/src/main/scala/org/apache/spark/deploy/master/ui/ApplicationPage.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/master/ui/ApplicationPage.scala
@@ -76,7 +76,7 @@ private[ui] class ApplicationPage(parent: MasterWebUI) extends WebUIPage("app")
             </li>
             <li><strong>Submit Date:</strong> {app.submitDate}</li>
             <li><strong>State:</strong> {app.state}</li>
-            <li><strong><a href={app.desc.appUiUrl}>Application Detail UI</a></strong></li>
+            <li><strong><a href={app.curAppUIUrl}>Application Detail UI</a></strong></li>
           </ul>
         </div>
       </div>
diff --git a/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala b/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala
index c3e20ebf8d6eb..ee539dd1f5113 100644
--- a/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala
@@ -206,7 +206,7 @@ private[ui] class MasterPage(parent: MasterWebUI) extends WebUIPage("") {
         {killLink}
       </td>
       <td>
-        <a href={app.desc.appUiUrl}>{app.desc.name}</a>
+        <a href={app.curAppUIUrl}>{app.desc.name}</a>
       </td>
       <td>
         {app.coresGranted}
diff --git a/core/src/test/scala/org/apache/spark/deploy/DeployTestUtils.scala b/core/src/test/scala/org/apache/spark/deploy/DeployTestUtils.scala
index 967aa0976f0ce..3164760b08a71 100644
--- a/core/src/test/scala/org/apache/spark/deploy/DeployTestUtils.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/DeployTestUtils.scala
@@ -31,8 +31,9 @@ private[deploy] object DeployTestUtils {
   }
 
   def createAppInfo() : ApplicationInfo = {
+    val appDesc = createAppDesc()
     val appInfo = new ApplicationInfo(JsonConstants.appInfoStartTime,
-      "id", createAppDesc(), JsonConstants.submitDate, null, Int.MaxValue)
+      "id", appDesc, JsonConstants.submitDate, null, Int.MaxValue)
     appInfo.endTime = JsonConstants.currTimeInMillis
     appInfo
   }

From d188a67762dfc09929e30931509be5851e29dfa5 Mon Sep 17 00:00:00 2001
From: Daoyuan Wang <daoyuan.wang@intel.com>
Date: Tue, 3 Nov 2015 22:30:23 +0800
Subject: [PATCH 025/518] [SPARK-10533][SQL] handle scientific notation in
 sqlParser

https://issues.apache.org/jira/browse/SPARK-10533

val df = sqlContext.createDataFrame(Seq(("a",1.0),("b",2.0),("c",3.0)))
df.filter("_2 < 2.0e1").show

Scientific notation didn't work.

Author: Daoyuan Wang <daoyuan.wang@intel.com>

Closes #9085 from adrian-wang/scinotation.
---
 .../sql/catalyst/AbstractSparkSQLParser.scala     | 15 +++++++++++++--
 .../org/apache/spark/sql/catalyst/SqlParser.scala | 11 +++++++++++
 .../org/apache/spark/sql/DataFrameSuite.scala     | 11 ++++++++---
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/AbstractSparkSQLParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/AbstractSparkSQLParser.scala
index 2bac08eac4fe2..04ac4f20c66ec 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/AbstractSparkSQLParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/AbstractSparkSQLParser.scala
@@ -82,6 +82,10 @@ class SqlLexical extends StdLexical {
     override def toString: String = chars
   }
 
+  case class DecimalLit(chars: String) extends Token {
+    override def toString: String = chars
+  }
+
   /* This is a work around to support the lazy setting */
   def initialize(keywords: Seq[String]): Unit = {
     reserved.clear()
@@ -102,8 +106,12 @@ class SqlLexical extends StdLexical {
   }
 
   override lazy val token: Parser[Token] =
-    ( identChar ~ (identChar | digit).* ^^
-      { case first ~ rest => processIdent((first :: rest).mkString) }
+    ( rep1(digit) ~ ('.' ~> digit.*).? ~ (exp ~> sign.? ~ rep1(digit)) ^^ {
+        case i ~ None ~ (sig ~ rest) =>
+          DecimalLit(i.mkString + "e" + sig.mkString + rest.mkString)
+        case i ~ Some(d) ~ (sig ~ rest) =>
+          DecimalLit(i.mkString + "." + d.mkString + "e" + sig.mkString + rest.mkString)
+      }
     | digit.* ~ identChar ~ (identChar | digit).* ^^
       { case first ~ middle ~ rest => processIdent((first ++ (middle :: rest)).mkString) }
     | rep1(digit) ~ ('.' ~> digit.*).? ^^ {
@@ -125,6 +133,9 @@ class SqlLexical extends StdLexical {
 
   override def identChar: Parser[Elem] = letter | elem('_')
 
+  private lazy val sign: Parser[Elem] = elem("s", c => c == '+' || c == '-')
+  private lazy val exp: Parser[Elem] = elem("e", c => c == 'E' || c == 'e')
+
   override def whitespace: Parser[Any] =
     ( whitespaceChar
     | '/' ~ '*' ~ comment
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
index d7567e8613e3c..1ba559d9e3b18 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
@@ -337,6 +337,9 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
     | sign.? ~ unsignedFloat ^^ {
       case s ~ f => Literal(toDecimalOrDouble(s.getOrElse("") + f))
     }
+    | sign.? ~ unsignedDecimal ^^ {
+      case s ~ d => Literal(toDecimalOrDouble(s.getOrElse("") + d))
+    }
     )
 
   protected lazy val unsignedFloat: Parser[String] =
@@ -344,6 +347,14 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
     | elem("decimal", _.isInstanceOf[lexical.FloatLit]) ^^ (_.chars)
     )
 
+  protected lazy val unsignedDecimal: Parser[String] =
+    ( "." ~> decimalLit ^^ { u => "0." + u }
+    | elem("scientific_notation", _.isInstanceOf[lexical.DecimalLit]) ^^ (_.chars)
+    )
+
+  def decimalLit: Parser[String] =
+    elem("scientific_notation", _.isInstanceOf[lexical.DecimalLit]) ^^ (_.chars)
+
   protected lazy val sign: Parser[String] = ("+" | "-")
 
   protected lazy val integral: Parser[String] =
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index 6b86c5951b413..a883bcb7b1012 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -177,9 +177,14 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
   }
 
   test("filterExpr") {
-    checkAnswer(
-      testData.filter("key > 90"),
-      testData.collect().filter(_.getInt(0) > 90).toSeq)
+    val res = testData.collect().filter(_.getInt(0) > 90).toSeq
+    checkAnswer(testData.filter("key > 90"), res)
+    checkAnswer(testData.filter("key > 9.0e1"), res)
+    checkAnswer(testData.filter("key > .9e+2"), res)
+    checkAnswer(testData.filter("key > 0.9e+2"), res)
+    checkAnswer(testData.filter("key > 900e-1"), res)
+    checkAnswer(testData.filter("key > 900.0E-1"), res)
+    checkAnswer(testData.filter("key > 9.e+1"), res)
   }
 
   test("filterExpr using where") {

From 57446eb69ceb6b8856ab22b54abb22b47b80f841 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 3 Nov 2015 07:06:00 -0800
Subject: [PATCH 026/518] [SPARK-11256] Mark all
 Stage/ResultStage/ShuffleMapStage internal state as private.

Author: Reynold Xin <rxin@databricks.com>

Closes #9219 from rxin/stage-cleanup1.
---
 .../apache/spark/scheduler/DAGScheduler.scala | 33 +++++-----
 .../apache/spark/scheduler/ResultStage.scala  | 19 +++++-
 .../spark/scheduler/ShuffleMapStage.scala     | 61 +++++++++++++------
 .../org/apache/spark/scheduler/Stage.scala    |  5 +-
 4 files changed, 80 insertions(+), 38 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index 995862ece5944..5673fbf2c8fea 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -23,7 +23,7 @@ import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicInteger
 
 import scala.collection.Map
-import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet, Stack}
+import scala.collection.mutable.{HashMap, HashSet, Stack}
 import scala.concurrent.duration._
 import scala.language.existentials
 import scala.language.postfixOps
@@ -535,10 +535,8 @@ class DAGScheduler(
     jobIdToActiveJob -= job.jobId
     activeJobs -= job
     job.finalStage match {
-      case r: ResultStage =>
-        r.resultOfJob = None
-      case m: ShuffleMapStage =>
-        m.mapStageJobs = m.mapStageJobs.filter(_ != job)
+      case r: ResultStage => r.removeActiveJob()
+      case m: ShuffleMapStage => m.removeActiveJob(job)
     }
   }
 
@@ -848,7 +846,7 @@ class DAGScheduler(
     val jobSubmissionTime = clock.getTimeMillis()
     jobIdToActiveJob(jobId) = job
     activeJobs += job
-    finalStage.resultOfJob = Some(job)
+    finalStage.setActiveJob(job)
     val stageIds = jobIdToStageIds(jobId).toArray
     val stageInfos = stageIds.flatMap(id => stageIdToStage.get(id).map(_.latestInfo))
     listenerBus.post(
@@ -880,7 +878,7 @@ class DAGScheduler(
     val job = new ActiveJob(jobId, finalStage, callSite, listener, properties)
     clearCacheLocs()
     logInfo("Got map stage job %s (%s) with %d output partitions".format(
-      jobId, callSite.shortForm, dependency.rdd.partitions.size))
+      jobId, callSite.shortForm, dependency.rdd.partitions.length))
     logInfo("Final stage: " + finalStage + " (" + finalStage.name + ")")
     logInfo("Parents of final stage: " + finalStage.parents)
     logInfo("Missing parents: " + getMissingParentStages(finalStage))
@@ -888,7 +886,7 @@ class DAGScheduler(
     val jobSubmissionTime = clock.getTimeMillis()
     jobIdToActiveJob(jobId) = job
     activeJobs += job
-    finalStage.mapStageJobs = job :: finalStage.mapStageJobs
+    finalStage.addActiveJob(job)
     val stageIds = jobIdToStageIds(jobId).toArray
     val stageInfos = stageIds.flatMap(id => stageIdToStage.get(id).map(_.latestInfo))
     listenerBus.post(
@@ -950,12 +948,12 @@ class DAGScheduler(
     // will be posted, which should always come after a corresponding SparkListenerStageSubmitted
     // event.
     outputCommitCoordinator.stageStart(stage.id)
-    val taskIdToLocations = try {
+    val taskIdToLocations: Map[Int, Seq[TaskLocation]] = try {
       stage match {
         case s: ShuffleMapStage =>
           partitionsToCompute.map { id => (id, getPreferredLocs(stage.rdd, id))}.toMap
         case s: ResultStage =>
-          val job = s.resultOfJob.get
+          val job = s.activeJob.get
           partitionsToCompute.map { id =>
             val p = s.partitions(id)
             (id, getPreferredLocs(stage.rdd, p))
@@ -1016,7 +1014,7 @@ class DAGScheduler(
           }
 
         case stage: ResultStage =>
-          val job = stage.resultOfJob.get
+          val job = stage.activeJob.get
           partitionsToCompute.map { id =>
             val p: Int = stage.partitions(id)
             val part = stage.rdd.partitions(p)
@@ -1132,7 +1130,7 @@ class DAGScheduler(
             // Cast to ResultStage here because it's part of the ResultTask
             // TODO Refactor this out to a function that accepts a ResultStage
             val resultStage = stage.asInstanceOf[ResultStage]
-            resultStage.resultOfJob match {
+            resultStage.activeJob match {
               case Some(job) =>
                 if (!job.finished(rt.outputId)) {
                   updateAccumulators(event)
@@ -1187,7 +1185,7 @@ class DAGScheduler(
               //       we registered these map outputs.
               mapOutputTracker.registerMapOutputs(
                 shuffleStage.shuffleDep.shuffleId,
-                shuffleStage.outputLocs.map(_.headOption.orNull),
+                shuffleStage.outputLocInMapOutputTrackerFormat(),
                 changeEpoch = true)
 
               clearCacheLocs()
@@ -1197,8 +1195,7 @@ class DAGScheduler(
                 // TODO: Lower-level scheduler should also deal with this
                 logInfo("Resubmitting " + shuffleStage + " (" + shuffleStage.name +
                   ") because some of its tasks had failed: " +
-                  shuffleStage.outputLocs.zipWithIndex.filter(_._1.isEmpty)
-                    .map(_._2).mkString(", "))
+                  shuffleStage.findMissingPartitions().mkString(", "))
                 submitStage(shuffleStage)
               } else {
                 // Mark any map-stage jobs waiting on this stage as finished
@@ -1312,8 +1309,10 @@ class DAGScheduler(
         // TODO: This will be really slow if we keep accumulating shuffle map stages
         for ((shuffleId, stage) <- shuffleToMapStage) {
           stage.removeOutputsOnExecutor(execId)
-          val locs = stage.outputLocs.map(_.headOption.orNull)
-          mapOutputTracker.registerMapOutputs(shuffleId, locs, changeEpoch = true)
+          mapOutputTracker.registerMapOutputs(
+            shuffleId,
+            stage.outputLocInMapOutputTrackerFormat(),
+            changeEpoch = true)
         }
         if (shuffleToMapStage.isEmpty) {
           mapOutputTracker.incrementEpoch()
diff --git a/core/src/main/scala/org/apache/spark/scheduler/ResultStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ResultStage.scala
index c1d86af7e8fb5..d1687830ff7bf 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/ResultStage.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/ResultStage.scala
@@ -41,10 +41,25 @@ private[spark] class ResultStage(
    * The active job for this result stage. Will be empty if the job has already finished
    * (e.g., because the job was cancelled).
    */
-  var resultOfJob: Option[ActiveJob] = None
+  private[this] var _activeJob: Option[ActiveJob] = None
 
+  def activeJob: Option[ActiveJob] = _activeJob
+
+  def setActiveJob(job: ActiveJob): Unit = {
+    _activeJob = Option(job)
+  }
+
+  def removeActiveJob(): Unit = {
+    _activeJob = None
+  }
+
+  /**
+   * Returns the sequence of partition ids that are missing (i.e. needs to be computed).
+   *
+   * This can only be called when there is an active job.
+   */
   override def findMissingPartitions(): Seq[Int] = {
-    val job = resultOfJob.get
+    val job = activeJob.get
     (0 until job.numPartitions).filter(id => !job.finished(id))
   }
 
diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala
index 3832d99eddaef..51416e5ce97fc 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala
@@ -43,35 +43,53 @@ private[spark] class ShuffleMapStage(
     val shuffleDep: ShuffleDependency[_, _, _])
   extends Stage(id, rdd, numTasks, parents, firstJobId, callSite) {
 
+  private[this] var _mapStageJobs: List[ActiveJob] = Nil
+
+  private[this] var _numAvailableOutputs: Int = 0
+
+  /**
+   * List of [[MapStatus]] for each partition. The index of the array is the map partition id,
+   * and each value in the array is the list of possible [[MapStatus]] for a partition
+   * (a single task might run multiple times).
+   */
+  private[this] val outputLocs = Array.fill[List[MapStatus]](numPartitions)(Nil)
+
   override def toString: String = "ShuffleMapStage " + id
 
-  /** Running map-stage jobs that were submitted to execute this stage independently (if any) */
-  var mapStageJobs: List[ActiveJob] = Nil
+  /**
+   * Returns the list of active jobs,
+   * i.e. map-stage jobs that were submitted to execute this stage independently (if any).
+   */
+  def mapStageJobs: Seq[ActiveJob] = _mapStageJobs
+
+  /** Adds the job to the active job list. */
+  def addActiveJob(job: ActiveJob): Unit = {
+    _mapStageJobs = job :: _mapStageJobs
+  }
+
+  /** Removes the job from the active job list. */
+  def removeActiveJob(job: ActiveJob): Unit = {
+    _mapStageJobs = _mapStageJobs.filter(_ != job)
+  }
 
   /**
    * Number of partitions that have shuffle outputs.
    * When this reaches [[numPartitions]], this map stage is ready.
    * This should be kept consistent as `outputLocs.filter(!_.isEmpty).size`.
    */
-  var numAvailableOutputs: Int = 0
+  def numAvailableOutputs: Int = _numAvailableOutputs
 
   /**
    * Returns true if the map stage is ready, i.e. all partitions have shuffle outputs.
    * This should be the same as `outputLocs.contains(Nil)`.
    */
-  def isAvailable: Boolean = numAvailableOutputs == numPartitions
-
-  /**
-   * List of [[MapStatus]] for each partition. The index of the array is the map partition id,
-   * and each value in the array is the list of possible [[MapStatus]] for a partition
-   * (a single task might run multiple times).
-   */
-  val outputLocs = Array.fill[List[MapStatus]](numPartitions)(Nil)
+  def isAvailable: Boolean = _numAvailableOutputs == numPartitions
 
+  /** Returns the sequence of partition ids that are missing (i.e. needs to be computed). */
   override def findMissingPartitions(): Seq[Int] = {
     val missing = (0 until numPartitions).filter(id => outputLocs(id).isEmpty)
-    assert(missing.size == numPartitions - numAvailableOutputs,
-      s"${missing.size} missing, expected ${numPartitions - numAvailableOutputs}")
+    assert(missing.size == numPartitions - _numAvailableOutputs,
+      s"${missing.size} missing, expected ${numPartitions - _numAvailableOutputs}")
     missing
   }
 
@@ -79,7 +97,7 @@ private[spark] class ShuffleMapStage(
     val prevList = outputLocs(partition)
     outputLocs(partition) = status :: prevList
     if (prevList == Nil) {
-      numAvailableOutputs += 1
+      _numAvailableOutputs += 1
     }
   }
 
@@ -88,10 +106,19 @@ private[spark] class ShuffleMapStage(
     val newList = prevList.filterNot(_.location == bmAddress)
     outputLocs(partition) = newList
     if (prevList != Nil && newList == Nil) {
-      numAvailableOutputs -= 1
+      _numAvailableOutputs -= 1
     }
   }
 
+  /**
+   * Returns an array of [[MapStatus]] (index by partition id). For each partition, the returned
+   * value contains only one (i.e. the first) [[MapStatus]]. If there is no entry for the partition,
+   * that position is filled with null.
+   */
+  def outputLocInMapOutputTrackerFormat(): Array[MapStatus] = {
+    outputLocs.map(_.headOption.orNull)
+  }
+
   /**
    * Removes all shuffle outputs associated with this executor. Note that this will also remove
    * outputs which are served by an external shuffle server (if one exists), as they are still
@@ -105,12 +132,12 @@ private[spark] class ShuffleMapStage(
       outputLocs(partition) = newList
       if (prevList != Nil && newList == Nil) {
         becameUnavailable = true
-        numAvailableOutputs -= 1
+        _numAvailableOutputs -= 1
       }
     }
     if (becameUnavailable) {
       logInfo("%s is now unavailable on executor %s (%d/%d, %s)".format(
-        this, execId, numAvailableOutputs, numPartitions, isAvailable))
+        this, execId, _numAvailableOutputs, numPartitions, isAvailable))
     }
   }
 }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/Stage.scala b/core/src/main/scala/org/apache/spark/scheduler/Stage.scala
index 5ce4a484344f1..7ea24a217bd39 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/Stage.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/Stage.scala
@@ -71,8 +71,8 @@ private[scheduler] abstract class Stage(
   /** The ID to use for the next new attempt for this stage. */
   private var nextAttemptId: Int = 0
 
-  val name = callSite.shortForm
-  val details = callSite.longForm
+  val name: String = callSite.shortForm
+  val details: String = callSite.longForm
 
   private var _internalAccumulators: Seq[Accumulator[Long]] = Seq.empty
 
@@ -134,6 +134,7 @@ private[scheduler] abstract class Stage(
   def latestInfo: StageInfo = _latestInfo
 
   override final def hashCode(): Int = id
+
   override final def equals(other: Any): Boolean = other match {
     case stage: Stage => stage != null && stage.id == id
     case _ => false

From d6035d97c91fe78b1336ade48134252915263ea6 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Tue, 3 Nov 2015 07:41:50 -0800
Subject: [PATCH 027/518] [SPARK-10304] [SQL] Partition discovery should throw
 an exception if the dir structure is invalid

JIRA: https://issues.apache.org/jira/browse/SPARK-10304

This patch detects if the structure of partition directories is not valid.

The test cases are from #8547. Thanks zhzhan.

cc liancheng

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #8840 from viirya/detect_invalid_part_dir.
---
 .../datasources/PartitioningUtils.scala       | 36 +++++++++++++------
 .../ParquetPartitionDiscoverySuite.scala      | 36 +++++++++++++++++--
 2 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
index 0a2007e15843c..628c5e18936c5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
@@ -77,9 +77,11 @@ private[sql] object PartitioningUtils {
       defaultPartitionName: String,
       typeInference: Boolean): PartitionSpec = {
     // First, we need to parse every partition's path and see if we can find partition values.
-    val pathsWithPartitionValues = paths.flatMap { path =>
-      parsePartition(path, defaultPartitionName, typeInference).map(path -> _)
-    }
+    val (partitionValues, optBasePaths) = paths.map { path =>
+      parsePartition(path, defaultPartitionName, typeInference)
+    }.unzip
+
+    val pathsWithPartitionValues = paths.zip(partitionValues).flatMap(x => x._2.map(x._1 -> _))
 
     if (pathsWithPartitionValues.isEmpty) {
       // This dataset is not partitioned.
@@ -87,6 +89,12 @@ private[sql] object PartitioningUtils {
     } else {
       // This dataset is partitioned. We need to check whether all partitions have the same
       // partition columns and resolve potential type conflicts.
+      val basePaths = optBasePaths.flatMap(x => x)
+      assert(
+        basePaths.distinct.size == 1,
+        "Conflicting directory structures detected. Suspicious paths:\b" +
+          basePaths.mkString("\n\t", "\n\t", "\n\n"))
+
       val resolvedPartitionValues = resolvePartitions(pathsWithPartitionValues)
 
       // Creates the StructType which represents the partition columns.
@@ -110,12 +118,12 @@ private[sql] object PartitioningUtils {
   }
 
   /**
-   * Parses a single partition, returns column names and values of each partition column.  For
-   * example, given:
+   * Parses a single partition, returns column names and values of each partition column, also
+   * the base path.  For example, given:
    * {{{
    *   path = hdfs://<host>:<port>/path/to/partition/a=42/b=hello/c=3.14
    * }}}
-   * it returns:
+   * it returns the partition:
    * {{{
    *   PartitionValues(
    *     Seq("a", "b", "c"),
@@ -124,34 +132,40 @@ private[sql] object PartitioningUtils {
    *       Literal.create("hello", StringType),
    *       Literal.create(3.14, FloatType)))
    * }}}
+   * and the base path:
+   * {{{
+   *   /path/to/partition
+   * }}}
    */
   private[sql] def parsePartition(
       path: Path,
       defaultPartitionName: String,
-      typeInference: Boolean): Option[PartitionValues] = {
+      typeInference: Boolean): (Option[PartitionValues], Option[Path]) = {
     val columns = ArrayBuffer.empty[(String, Literal)]
     // Old Hadoop versions don't have `Path.isRoot`
     var finished = path.getParent == null
     var chopped = path
+    var basePath = path
 
     while (!finished) {
       // Sometimes (e.g., when speculative task is enabled), temporary directories may be left
       // uncleaned.  Here we simply ignore them.
       if (chopped.getName.toLowerCase == "_temporary") {
-        return None
+        return (None, None)
       }
 
       val maybeColumn = parsePartitionColumn(chopped.getName, defaultPartitionName, typeInference)
       maybeColumn.foreach(columns += _)
+      basePath = chopped
       chopped = chopped.getParent
-      finished = maybeColumn.isEmpty || chopped.getParent == null
+      finished = (maybeColumn.isEmpty && !columns.isEmpty) || chopped.getParent == null
     }
 
     if (columns.isEmpty) {
-      None
+      (None, Some(path))
     } else {
       val (columnNames, values) = columns.reverse.unzip
-      Some(PartitionValues(columnNames, values))
+      (Some(PartitionValues(columnNames, values)), Some(basePath))
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
index 3a23b8ed66808..67b6a37fa502e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
@@ -58,14 +58,46 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
     check(defaultPartitionName, Literal.create(null, NullType))
   }
 
+  test("parse invalid partitioned directories") {
+    // Invalid
+    var paths = Seq(
+      "hdfs://host:9000/invalidPath",
+      "hdfs://host:9000/path/a=10/b=20",
+      "hdfs://host:9000/path/a=10.5/b=hello")
+
+    var exception = intercept[AssertionError] {
+      parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
+    }
+    assert(exception.getMessage().contains("Conflicting directory structures detected"))
+
+    // Valid
+    paths = Seq(
+      "hdfs://host:9000/path/_temporary",
+      "hdfs://host:9000/path/a=10/b=20",
+      "hdfs://host:9000/path/_temporary/path")
+
+    parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
+
+    // Invalid
+    paths = Seq(
+      "hdfs://host:9000/path/_temporary",
+      "hdfs://host:9000/path/a=10/b=20",
+      "hdfs://host:9000/path/path1")
+
+    exception = intercept[AssertionError] {
+      parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
+    }
+    assert(exception.getMessage().contains("Conflicting directory structures detected"))
+  }
+
   test("parse partition") {
     def check(path: String, expected: Option[PartitionValues]): Unit = {
-      assert(expected === parsePartition(new Path(path), defaultPartitionName, true))
+      assert(expected === parsePartition(new Path(path), defaultPartitionName, true)._1)
     }
 
     def checkThrows[T <: Throwable: Manifest](path: String, expected: String): Unit = {
       val message = intercept[T] {
-        parsePartition(new Path(path), defaultPartitionName, true).get
+        parsePartition(new Path(path), defaultPartitionName, true)
       }.getMessage
 
       assert(message.contains(expected))

From d6f10aa7ea2806c0fbcfc31d7dee91d28319fab7 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Tue, 3 Nov 2015 08:29:07 -0800
Subject: [PATCH 028/518] [SPARK-9836][ML] Provide R-like summary statistics
 for OLS via normal equation solver

https://issues.apache.org/jira/browse/SPARK-9836

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9413 from yanboliang/spark-9836.
---
 .../spark/ml/optim/WeightedLeastSquares.scala |  15 +-
 .../ml/regression/LinearRegression.scala      |  90 +++++++++++-
 .../mllib/linalg/CholeskyDecomposition.scala  |  16 +++
 .../ml/regression/LinearRegressionSuite.scala | 129 ++++++++++++++++++
 4 files changed, 243 insertions(+), 7 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/optim/WeightedLeastSquares.scala b/mllib/src/main/scala/org/apache/spark/ml/optim/WeightedLeastSquares.scala
index 3d64f7f296137..e612a2122ed62 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/optim/WeightedLeastSquares.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/optim/WeightedLeastSquares.scala
@@ -26,10 +26,12 @@ import org.apache.spark.rdd.RDD
  * Model fitted by [[WeightedLeastSquares]].
  * @param coefficients model coefficients
  * @param intercept model intercept
+ * @param diagInvAtWA diagonal of matrix (A^T * W * A)^-1
  */
 private[ml] class WeightedLeastSquaresModel(
     val coefficients: DenseVector,
-    val intercept: Double) extends Serializable
+    val intercept: Double,
+    val diagInvAtWA: DenseVector) extends Serializable
 
 /**
  * Weighted least squares solver via normal equation.
@@ -73,7 +75,9 @@ private[ml] class WeightedLeastSquares(
     val summary = instances.treeAggregate(new Aggregator)(_.add(_), _.merge(_))
     summary.validate()
     logInfo(s"Number of instances: ${summary.count}.")
+    val k = summary.k
     val triK = summary.triK
+    val wSum = summary.wSum
     val bBar = summary.bBar
     val bStd = summary.bStd
     val aBar = summary.aBar
@@ -109,6 +113,11 @@ private[ml] class WeightedLeastSquares(
 
     val x = new DenseVector(CholeskyDecomposition.solve(aaBar.values, abBar.values))
 
+    val aaInv = CholeskyDecomposition.inverse(aaBar.values, k)
+    // aaInv is a packed upper triangular matrix, here we get all elements on diagonal
+    val diagInvAtWA = new DenseVector((1 to k).map { i =>
+      aaInv(i + (i - 1) * i / 2 - 1) / wSum }.toArray)
+
     // compute intercept
     val intercept = if (fitIntercept) {
       bBar - BLAS.dot(aBar, x)
@@ -116,7 +125,7 @@ private[ml] class WeightedLeastSquares(
       0.0
     }
 
-    new WeightedLeastSquaresModel(x, intercept)
+    new WeightedLeastSquaresModel(x, intercept, diagInvAtWA)
   }
 }
 
@@ -131,7 +140,7 @@ private[ml] object WeightedLeastSquares {
     var k: Int = _
     var count: Long = _
     var triK: Int = _
-    private var wSum: Double = _
+    var wSum: Double = _
     private var wwSum: Double = _
     private var bSum: Double = _
     private var bbSum: Double = _
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index 6e9c7442b8110..c51e30483ab3d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -21,6 +21,7 @@ import scala.collection.mutable
 
 import breeze.linalg.{DenseVector => BDV}
 import breeze.optimize.{CachedDiffFunction, DiffFunction, LBFGS => BreezeLBFGS, OWLQN => BreezeOWLQN}
+import breeze.stats.distributions.StudentsT
 
 import org.apache.spark.{Logging, SparkException}
 import org.apache.spark.annotation.Experimental
@@ -36,7 +37,7 @@ import org.apache.spark.mllib.linalg.BLAS._
 import org.apache.spark.mllib.stat.MultivariateOnlineSummarizer
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, Row}
-import org.apache.spark.sql.functions.{col, udf, lit}
+import org.apache.spark.sql.functions._
 import org.apache.spark.storage.StorageLevel
 
 /**
@@ -173,8 +174,11 @@ class LinearRegression(override val uid: String)
         summaryModel.transform(dataset),
         predictionColName,
         $(labelCol),
+        summaryModel,
+        model.diagInvAtWA.toArray,
         $(featuresCol),
         Array(0D))
+
       return lrModel.setSummary(trainingSummary)
     }
 
@@ -221,6 +225,8 @@ class LinearRegression(override val uid: String)
         summaryModel.transform(dataset),
         predictionColName,
         $(labelCol),
+        model,
+        Array(0D),
         $(featuresCol),
         Array(0D))
       return copyValues(model.setSummary(trainingSummary))
@@ -316,6 +322,8 @@ class LinearRegression(override val uid: String)
       summaryModel.transform(dataset),
       predictionColName,
       $(labelCol),
+      model,
+      Array(0D),
       $(featuresCol),
       objectiveHistory)
     model.setSummary(trainingSummary)
@@ -371,7 +379,8 @@ class LinearRegressionModel private[ml] (
   private[regression] def evaluate(dataset: DataFrame): LinearRegressionSummary = {
     // Handle possible missing or invalid prediction columns
     val (summaryModel, predictionColName) = findSummaryModelAndPredictionCol()
-    new LinearRegressionSummary(summaryModel.transform(dataset), predictionColName, $(labelCol))
+    new LinearRegressionSummary(summaryModel.transform(dataset), predictionColName,
+      $(labelCol), this, Array(0D))
   }
 
   /**
@@ -412,9 +421,11 @@ class LinearRegressionTrainingSummary private[regression] (
     predictions: DataFrame,
     predictionCol: String,
     labelCol: String,
+    model: LinearRegressionModel,
+    diagInvAtWA: Array[Double],
     val featuresCol: String,
     val objectiveHistory: Array[Double])
-  extends LinearRegressionSummary(predictions, predictionCol, labelCol) {
+  extends LinearRegressionSummary(predictions, predictionCol, labelCol, model, diagInvAtWA) {
 
   /** Number of training iterations until termination */
   val totalIterations = objectiveHistory.length
@@ -430,7 +441,9 @@ class LinearRegressionTrainingSummary private[regression] (
 class LinearRegressionSummary private[regression] (
     @transient val predictions: DataFrame,
     val predictionCol: String,
-    val labelCol: String) extends Serializable {
+    val labelCol: String,
+    val model: LinearRegressionModel,
+    val diagInvAtWA: Array[Double]) extends Serializable {
 
   @transient private val metrics = new RegressionMetrics(
     predictions
@@ -474,6 +487,75 @@ class LinearRegressionSummary private[regression] (
     predictions.select(t(col(predictionCol), col(labelCol)).as("residuals"))
   }
 
+  /** Number of instances in DataFrame predictions */
+  lazy val numInstances: Long = predictions.count()
+
+  /** Degrees of freedom */
+  private val degreesOfFreedom: Long = if (model.getFitIntercept) {
+    numInstances - model.coefficients.size - 1
+  } else {
+    numInstances - model.coefficients.size
+  }
+
+  /**
+   * The weighted residuals, the usual residuals rescaled by
+   * the square root of the instance weights.
+   */
+  lazy val devianceResiduals: Array[Double] = {
+    val weighted = if (model.getWeightCol.isEmpty) lit(1.0) else sqrt(col(model.getWeightCol))
+    val dr = predictions.select(col(model.getLabelCol).minus(col(model.getPredictionCol))
+      .multiply(weighted).as("weightedResiduals"))
+      .select(min(col("weightedResiduals")).as("min"), max(col("weightedResiduals")).as("max"))
+      .first()
+    Array(dr.getDouble(0), dr.getDouble(1))
+  }
+
+  /**
+   * Standard error of estimated coefficients.
+   * Note that standard error of estimated intercept is not supported currently.
+   */
+  lazy val coefficientStandardErrors: Array[Double] = {
+    if (diagInvAtWA.length == 1 && diagInvAtWA(0) == 0) {
+      throw new UnsupportedOperationException(
+        "No Std. Error of coefficients available for this LinearRegressionModel")
+    } else {
+      val rss = if (model.getWeightCol.isEmpty) {
+        meanSquaredError * numInstances
+      } else {
+        val t = udf { (pred: Double, label: Double, weight: Double) =>
+          math.pow(label - pred, 2.0) * weight }
+        predictions.select(t(col(model.getPredictionCol), col(model.getLabelCol),
+          col(model.getWeightCol)).as("wse")).agg(sum(col("wse"))).first().getDouble(0)
+      }
+      val sigma2 = rss / degreesOfFreedom
+      diagInvAtWA.map(_ * sigma2).map(math.sqrt(_))
+    }
+  }
+
+  /** T-statistic of estimated coefficients.
+    * Note that t-statistic of estimated intercept is not supported currently.
+    */
+  lazy val tValues: Array[Double] = {
+    if (diagInvAtWA.length == 1 && diagInvAtWA(0) == 0) {
+      throw new UnsupportedOperationException(
+        "No t-statistic available for this LinearRegressionModel")
+    } else {
+      model.coefficients.toArray.zip(coefficientStandardErrors).map { x => x._1 / x._2 }
+    }
+  }
+
+  /** Two-sided p-value of estimated coefficients.
+    * Note that p-value of estimated intercept is not supported currently.
+    */
+  lazy val pValues: Array[Double] = {
+    if (diagInvAtWA.length == 1 && diagInvAtWA(0) == 0) {
+      throw new UnsupportedOperationException(
+        "No p-value available for this LinearRegressionModel")
+    } else {
+      tValues.map { x => 2.0 * (1.0 - StudentsT(degreesOfFreedom.toDouble).cdf(math.abs(x))) }
+    }
+  }
+
 }
 
 /**
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/linalg/CholeskyDecomposition.scala b/mllib/src/main/scala/org/apache/spark/mllib/linalg/CholeskyDecomposition.scala
index 66eb40b6f4a69..0cd371e9cce34 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/linalg/CholeskyDecomposition.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/linalg/CholeskyDecomposition.scala
@@ -40,4 +40,20 @@ private[spark] object CholeskyDecomposition {
     assert(code == 0, s"lapack.dpotrs returned $code.")
     bx
   }
+
+  /**
+   * Computes the inverse of a real symmetric positive definite matrix A
+   * using the Cholesky factorization A = U**T*U.
+   * The input arguments are modified in-place to store the inverse matrix.
+   * @param UAi the upper triangular factor U from the Cholesky factorization A = U**T*U
+   * @param k the dimension of A
+   * @return the upper triangle of the (symmetric) inverse of A
+   */
+  def inverse(UAi: Array[Double], k: Int): Array[Double] = {
+    val info = new intW(0)
+    lapack.dpptri("U", k, UAi, info)
+    val code = info.`val`
+    assert(code == 0, s"lapack.dpptri returned $code.")
+    UAi
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
index 235c796d785a6..fbf83e8922861 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
@@ -35,6 +35,7 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
   @transient var datasetWithDenseFeature: DataFrame = _
   @transient var datasetWithDenseFeatureWithoutIntercept: DataFrame = _
   @transient var datasetWithSparseFeature: DataFrame = _
+  @transient var datasetWithWeight: DataFrame = _
 
   /*
      In `LinearRegressionSuite`, we will make sure that the model trained by SparkML
@@ -73,6 +74,22 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         xMean = Seq.fill(featureSize)(r.nextDouble).toArray,
         xVariance = Seq.fill(featureSize)(r.nextDouble).toArray, nPoints = 200,
         seed, eps = 0.1, sparsity = 0.7), 2))
+
+    /*
+       R code:
+
+       A <- matrix(c(0, 1, 2, 3, 5, 7, 11, 13), 4, 2)
+       b <- c(17, 19, 23, 29)
+       w <- c(1, 2, 3, 4)
+       df <- as.data.frame(cbind(A, b))
+     */
+    datasetWithWeight = sqlContext.createDataFrame(
+      sc.parallelize(Seq(
+        Instance(17.0, 1.0, Vectors.dense(0.0, 5.0).toSparse),
+        Instance(19.0, 2.0, Vectors.dense(1.0, 7.0)),
+        Instance(23.0, 3.0, Vectors.dense(2.0, 11.0)),
+        Instance(29.0, 4.0, Vectors.dense(3.0, 13.0))
+      ), 2))
   }
 
   test("params") {
@@ -603,6 +620,16 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         // To clalify that the normal solver is used here.
         assert(model.summary.objectiveHistory.length == 1)
         assert(model.summary.objectiveHistory(0) == 0.0)
+        val devianceResidualsR = Array(-0.35566, 0.34504)
+        val seCoefR = Array(0.0011756, 0.0009032)
+        val tValsR = Array(3998, 7971)
+        val pValsR = Array(0, 0)
+        model.summary.devianceResiduals.zip(devianceResidualsR).foreach { x =>
+          assert(x._1 ~== x._2 absTol 1E-3) }
+        model.summary.coefficientStandardErrors.zip(seCoefR).foreach{ x =>
+          assert(x._1 ~== x._2 absTol 1E-3) }
+        model.summary.tValues.map(_.round).zip(tValsR).foreach{ x => assert(x._1 === x._2) }
+        model.summary.pValues.map(_.round).zip(pValsR).foreach{ x => assert(x._1 === x._2) }
       }
     }
   }
@@ -725,4 +752,106 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         .sliding(2)
         .forall(x => x(0) >= x(1)))
   }
+
+  test("linear regression summary with weighted samples and intercept by normal solver") {
+    /*
+       R code:
+
+       model <- glm(formula = "b ~ .", data = df, weights = w)
+       summary(model)
+
+       Call:
+       glm(formula = "b ~ .", data = df, weights = w)
+
+       Deviance Residuals:
+            1       2       3       4
+        1.920  -1.358  -1.109   0.960
+
+       Coefficients:
+                   Estimate Std. Error t value Pr(>|t|)
+       (Intercept)   18.080      9.608   1.882    0.311
+       V1             6.080      5.556   1.094    0.471
+       V2            -0.600      1.960  -0.306    0.811
+
+       (Dispersion parameter for gaussian family taken to be 7.68)
+
+           Null deviance: 202.00  on 3  degrees of freedom
+       Residual deviance:   7.68  on 1  degrees of freedom
+       AIC: 18.783
+
+       Number of Fisher Scoring iterations: 2
+     */
+
+    val model = new LinearRegression()
+      .setWeightCol("weight")
+      .setSolver("normal")
+      .fit(datasetWithWeight)
+    val coefficientsR = Vectors.dense(Array(6.080, -0.600))
+    val interceptR = 18.080
+    val devianceResidualsR = Array(-1.358, 1.920)
+    val seCoefR = Array(5.556, 1.960)
+    val tValsR = Array(1.094, -0.306)
+    val pValsR = Array(0.471, 0.811)
+
+    assert(model.coefficients ~== coefficientsR absTol 1E-3)
+    assert(model.intercept ~== interceptR absTol 1E-3)
+    model.summary.devianceResiduals.zip(devianceResidualsR).foreach { x =>
+      assert(x._1 ~== x._2 absTol 1E-3) }
+    model.summary.coefficientStandardErrors.zip(seCoefR).foreach{ x =>
+      assert(x._1 ~== x._2 absTol 1E-3) }
+    model.summary.tValues.zip(tValsR).foreach{ x => assert(x._1 ~== x._2 absTol 1E-3) }
+    model.summary.pValues.zip(pValsR).foreach{ x => assert(x._1 ~== x._2 absTol 1E-3) }
+  }
+
+  test("linear regression summary with weighted samples and w/o intercept by normal solver") {
+    /*
+       R code:
+
+       model <- glm(formula = "b ~ . -1", data = df, weights = w)
+       summary(model)
+
+       Call:
+       glm(formula = "b ~ . -1", data = df, weights = w)
+
+       Deviance Residuals:
+            1       2       3       4
+        1.950   2.344  -4.600   2.103
+
+       Coefficients:
+          Estimate Std. Error t value Pr(>|t|)
+       V1  -3.7271     2.9032  -1.284   0.3279
+       V2   3.0100     0.6022   4.998   0.0378 *
+       ---
+       Signif. codes:  0 ‘***’ 0.001 ‘**’ 0.01 ‘*’ 0.05 ‘.’ 0.1 ‘ ’ 1
+
+       (Dispersion parameter for gaussian family taken to be 17.4376)
+
+           Null deviance: 5962.000  on 4  degrees of freedom
+       Residual deviance:   34.875  on 2  degrees of freedom
+       AIC: 22.835
+
+       Number of Fisher Scoring iterations: 2
+     */
+
+    val model = new LinearRegression()
+      .setWeightCol("weight")
+      .setSolver("normal")
+      .setFitIntercept(false)
+      .fit(datasetWithWeight)
+    val coefficientsR = Vectors.dense(Array(-3.7271, 3.0100))
+    val interceptR = 0.0
+    val devianceResidualsR = Array(-4.600, 2.344)
+    val seCoefR = Array(2.9032, 0.6022)
+    val tValsR = Array(-1.284, 4.998)
+    val pValsR = Array(0.3279, 0.0378)
+
+    assert(model.coefficients ~== coefficientsR absTol 1E-3)
+    assert(model.intercept === interceptR)
+    model.summary.devianceResiduals.zip(devianceResidualsR).foreach { x =>
+      assert(x._1 ~== x._2 absTol 1E-3) }
+    model.summary.coefficientStandardErrors.zip(seCoefR).foreach{ x =>
+      assert(x._1 ~== x._2 absTol 1E-3) }
+    model.summary.tValues.zip(tValsR).foreach{ x => assert(x._1 ~== x._2 absTol 1E-3) }
+    model.summary.pValues.zip(pValsR).foreach{ x => assert(x._1 ~== x._2 absTol 1E-3) }
+  }
 }

From 3434572b141075f00698d94e6ee80febd3093c3b Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Tue, 3 Nov 2015 08:31:16 -0800
Subject: [PATCH 029/518] [MINOR][ML] Fix naming conventions of
 AFTSurvivalRegression coefficients

Rename ```regressionCoefficients``` back to ```coefficients```, and name ```weights``` to ```parameters```.
See discussion [here](https://github.com/apache/spark/pull/9311/files#diff-e277fd0bc21f825d3196b4551c01fe5fR230). mengxr vectorijk dbtsai

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9431 from yanboliang/aft-coefficients.
---
 .../ml/regression/AFTSurvivalRegression.scala | 38 +++++++++----------
 .../AFTSurvivalRegressionSuite.scala          | 12 +++---
 2 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
index 4dbbc7d39931b..b7d095872ffa5 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
@@ -200,17 +200,17 @@ class AFTSurvivalRegression @Since("1.6.0") (@Since("1.6.0") override val uid: S
 
     val numFeatures = dataset.select($(featuresCol)).take(1)(0).getAs[Vector](0).size
     /*
-       The coefficients vector has three parts:
+       The parameters vector has three parts:
        the first element: Double, log(sigma), the log of scale parameter
        the second element: Double, intercept of the beta parameter
        the third to the end elements: Doubles, regression coefficients vector of the beta parameter
      */
-    val initialCoefficients = Vectors.zeros(numFeatures + 2)
+    val initialParameters = Vectors.zeros(numFeatures + 2)
 
     val states = optimizer.iterations(new CachedDiffFunction(costFun),
-      initialCoefficients.toBreeze.toDenseVector)
+      initialParameters.toBreeze.toDenseVector)
 
-    val coefficients = {
+    val parameters = {
       val arrayBuilder = mutable.ArrayBuilder.make[Double]
       var state: optimizer.State = null
       while (states.hasNext) {
@@ -227,10 +227,10 @@ class AFTSurvivalRegression @Since("1.6.0") (@Since("1.6.0") override val uid: S
 
     if (handlePersistence) instances.unpersist()
 
-    val regressionCoefficients = Vectors.dense(coefficients.slice(2, coefficients.length))
-    val intercept = coefficients(1)
-    val scale = math.exp(coefficients(0))
-    val model = new AFTSurvivalRegressionModel(uid, regressionCoefficients, intercept, scale)
+    val coefficients = Vectors.dense(parameters.slice(2, parameters.length))
+    val intercept = parameters(1)
+    val scale = math.exp(parameters(0))
+    val model = new AFTSurvivalRegressionModel(uid, coefficients, intercept, scale)
     copyValues(model.setParent(this))
   }
 
@@ -251,7 +251,7 @@ class AFTSurvivalRegression @Since("1.6.0") (@Since("1.6.0") override val uid: S
 @Since("1.6.0")
 class AFTSurvivalRegressionModel private[ml] (
     @Since("1.6.0") override val uid: String,
-    @Since("1.6.0") val regressionCoefficients: Vector,
+    @Since("1.6.0") val coefficients: Vector,
     @Since("1.6.0") val intercept: Double,
     @Since("1.6.0") val scale: Double)
   extends Model[AFTSurvivalRegressionModel] with AFTSurvivalRegressionParams {
@@ -275,7 +275,7 @@ class AFTSurvivalRegressionModel private[ml] (
   @Since("1.6.0")
   def predictQuantiles(features: Vector): Vector = {
     // scale parameter for the Weibull distribution of lifetime
-    val lambda = math.exp(BLAS.dot(regressionCoefficients, features) + intercept)
+    val lambda = math.exp(BLAS.dot(coefficients, features) + intercept)
     // shape parameter for the Weibull distribution of lifetime
     val k = 1 / scale
     val quantiles = $(quantileProbabilities).map {
@@ -286,7 +286,7 @@ class AFTSurvivalRegressionModel private[ml] (
 
   @Since("1.6.0")
   def predict(features: Vector): Double = {
-    math.exp(BLAS.dot(regressionCoefficients, features) + intercept)
+    math.exp(BLAS.dot(coefficients, features) + intercept)
   }
 
   @Since("1.6.0")
@@ -309,7 +309,7 @@ class AFTSurvivalRegressionModel private[ml] (
 
   @Since("1.6.0")
   override def copy(extra: ParamMap): AFTSurvivalRegressionModel = {
-    copyValues(new AFTSurvivalRegressionModel(uid, regressionCoefficients, intercept, scale), extra)
+    copyValues(new AFTSurvivalRegressionModel(uid, coefficients, intercept, scale), extra)
       .setParent(parent)
   }
 }
@@ -369,17 +369,17 @@ class AFTSurvivalRegressionModel private[ml] (
  *   \frac{\partial (-\iota)}{\partial (\log\sigma)}=
  *   \sum_{i=1}^{n}[\delta_{i}+(\delta_{i}-e^{\epsilon_{i}})\epsilon_{i}]
  * }}}
- * @param coefficients including three part: The log of scale parameter, the intercept and
+ * @param parameters including three part: The log of scale parameter, the intercept and
  *                regression coefficients corresponding to the features.
  * @param fitIntercept Whether to fit an intercept term.
  */
-private class AFTAggregator(coefficients: BDV[Double], fitIntercept: Boolean)
+private class AFTAggregator(parameters: BDV[Double], fitIntercept: Boolean)
   extends Serializable {
 
   // beta is the intercept and regression coefficients to the covariates
-  private val beta = coefficients.slice(1, coefficients.length)
+  private val beta = parameters.slice(1, parameters.length)
   // sigma is the scale parameter of the AFT model
-  private val sigma = math.exp(coefficients(0))
+  private val sigma = math.exp(parameters(0))
 
   private var totalCnt: Long = 0L
   private var lossSum = 0.0
@@ -449,15 +449,15 @@ private class AFTAggregator(coefficients: BDV[Double], fitIntercept: Boolean)
 
 /**
  * AFTCostFun implements Breeze's DiffFunction[T] for AFT cost.
- * It returns the loss and gradient at a particular point (coefficients).
+ * It returns the loss and gradient at a particular point (parameters).
  * It's used in Breeze's convex optimization routines.
  */
 private class AFTCostFun(data: RDD[AFTPoint], fitIntercept: Boolean)
   extends DiffFunction[BDV[Double]] {
 
-  override def calculate(coefficients: BDV[Double]): (Double, BDV[Double]) = {
+  override def calculate(parameters: BDV[Double]): (Double, BDV[Double]) = {
 
-    val aftAggregator = data.treeAggregate(new AFTAggregator(coefficients, fitIntercept))(
+    val aftAggregator = data.treeAggregate(new AFTAggregator(parameters, fitIntercept))(
       seqOp = (c, v) => (c, v) match {
         case (aggregator, instance) => aggregator.add(instance)
       },
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
index c0f791bce13d1..359f31027172b 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
@@ -141,12 +141,12 @@ class AFTSurvivalRegressionSuite extends SparkFunSuite with MLlibTestSparkContex
        Number of Newton-Raphson Iterations: 5
        n= 1000
      */
-    val regressionCoefficientsR = Vectors.dense(-0.039)
+    val coefficientsR = Vectors.dense(-0.039)
     val interceptR = 1.759
     val scaleR = 1.41
 
     assert(model.intercept ~== interceptR relTol 1E-3)
-    assert(model.regressionCoefficients ~== regressionCoefficientsR relTol 1E-3)
+    assert(model.coefficients ~== coefficientsR relTol 1E-3)
     assert(model.scale ~== scaleR relTol 1E-3)
 
     /*
@@ -212,12 +212,12 @@ class AFTSurvivalRegressionSuite extends SparkFunSuite with MLlibTestSparkContex
        Number of Newton-Raphson Iterations: 5
        n= 1000
      */
-    val regressionCoefficientsR = Vectors.dense(-0.0844, 0.0677)
+    val coefficientsR = Vectors.dense(-0.0844, 0.0677)
     val interceptR = 1.9206
     val scaleR = 0.977
 
     assert(model.intercept ~== interceptR relTol 1E-3)
-    assert(model.regressionCoefficients ~== regressionCoefficientsR relTol 1E-3)
+    assert(model.coefficients ~== coefficientsR relTol 1E-3)
     assert(model.scale ~== scaleR relTol 1E-3)
 
     /*
@@ -282,12 +282,12 @@ class AFTSurvivalRegressionSuite extends SparkFunSuite with MLlibTestSparkContex
        Number of Newton-Raphson Iterations: 6
        n= 1000
      */
-    val regressionCoefficientsR = Vectors.dense(0.896, -0.709)
+    val coefficientsR = Vectors.dense(0.896, -0.709)
     val interceptR = 0.0
     val scaleR = 1.52
 
     assert(model.intercept === interceptR)
-    assert(model.regressionCoefficients ~== regressionCoefficientsR relTol 1E-3)
+    assert(model.coefficients ~== coefficientsR relTol 1E-3)
     assert(model.scale ~== scaleR relTol 1E-3)
 
     /*

From f54ff19b1edd4903950cb334987a447445fa97ef Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Tue, 3 Nov 2015 08:32:37 -0800
Subject: [PATCH 030/518] [SPARK-11349][ML] Support transform string label for
 RFormula

Currently ```RFormula``` can only handle label with ```NumericType``` or ```BinaryType``` (cast it to ```DoubleType``` as the label of Linear Regression training), we should also support label of ```StringType``` which is needed for Logistic Regression (glm with family = "binomial").
For label of ```StringType```, we should use ```StringIndexer``` to transform it to 0-based index.

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9302 from yanboliang/spark-11349.
---
 .../apache/spark/ml/feature/RFormula.scala    | 10 +++++++++-
 .../spark/ml/feature/RFormulaSuite.scala      | 19 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
index f9b840097f3ed..5c43a41bee3b4 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
@@ -132,6 +132,14 @@ class RFormula(override val uid: String) extends Estimator[RFormulaModel] with R
       .setOutputCol($(featuresCol))
     encoderStages += new VectorAttributeRewriter($(featuresCol), prefixesToRewrite.toMap)
     encoderStages += new ColumnPruner(tempColumns.toSet)
+
+    if (dataset.schema.fieldNames.contains(resolvedFormula.label) &&
+      dataset.schema(resolvedFormula.label).dataType == StringType) {
+      encoderStages += new StringIndexer()
+        .setInputCol(resolvedFormula.label)
+        .setOutputCol($(labelCol))
+    }
+
     val pipelineModel = new Pipeline(uid).setStages(encoderStages.toArray).fit(dataset)
     copyValues(new RFormulaModel(uid, resolvedFormula, pipelineModel).setParent(this))
   }
@@ -172,7 +180,7 @@ class RFormulaModel private[feature](
   override def transformSchema(schema: StructType): StructType = {
     checkCanTransform(schema)
     val withFeatures = pipelineModel.transformSchema(schema)
-    if (hasLabelCol(schema)) {
+    if (hasLabelCol(withFeatures)) {
       withFeatures
     } else if (schema.exists(_.name == resolvedFormula.label)) {
       val nullable = schema(resolvedFormula.label).dataType match {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/RFormulaSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/RFormulaSuite.scala
index b56013008b116..dc20a5ec2152d 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/RFormulaSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/RFormulaSuite.scala
@@ -107,6 +107,25 @@ class RFormulaSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(result.collect() === expected.collect())
   }
 
+  test("index string label") {
+    val formula = new RFormula().setFormula("id ~ a + b")
+    val original = sqlContext.createDataFrame(
+      Seq(("male", "foo", 4), ("female", "bar", 4), ("female", "bar", 5), ("male", "baz", 5))
+    ).toDF("id", "a", "b")
+    val model = formula.fit(original)
+    val result = model.transform(original)
+    val resultSchema = model.transformSchema(original.schema)
+    val expected = sqlContext.createDataFrame(
+      Seq(
+        ("male", "foo", 4, Vectors.dense(0.0, 1.0, 4.0), 1.0),
+        ("female", "bar", 4, Vectors.dense(1.0, 0.0, 4.0), 0.0),
+        ("female", "bar", 5, Vectors.dense(1.0, 0.0, 5.0), 0.0),
+        ("male", "baz", 5, Vectors.dense(0.0, 0.0, 5.0), 1.0))
+    ).toDF("id", "a", "b", "features", "label")
+    // assert(result.schema.toString == resultSchema.toString)
+    assert(result.collect() === expected.collect())
+  }
+
   test("attribute generation") {
     val formula = new RFormula().setFormula("id ~ a + b")
     val original = sqlContext.createDataFrame(

From b2e4b314d989de8cad012bbddba703b31d8378a4 Mon Sep 17 00:00:00 2001
From: Mark Grover <grover.markgrover@gmail.com>
Date: Tue, 3 Nov 2015 08:51:40 -0800
Subject: [PATCH 031/518] [SPARK-9790][YARN] Expose in WebUI if NodeManager is
 the reason why executors were killed.

Author: Mark Grover <grover.markgrover@gmail.com>

Closes #8093 from markgrover/nm2.
---
 .../main/scala/org/apache/spark/TaskEndReason.scala   |  8 ++++++--
 .../scala/org/apache/spark/rpc/RpcEndpointRef.scala   |  4 ++--
 .../org/apache/spark/scheduler/TaskSetManager.scala   |  4 ++--
 .../cluster/CoarseGrainedSchedulerBackend.scala       |  5 +++--
 .../scheduler/cluster/YarnSchedulerBackend.scala      |  1 +
 .../scala/org/apache/spark/util/JsonProtocol.scala    | 11 ++++++++---
 .../spark/ui/jobs/JobProgressListenerSuite.scala      |  2 +-
 .../org/apache/spark/util/JsonProtocolSuite.scala     | 11 ++++++-----
 8 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/TaskEndReason.scala b/core/src/main/scala/org/apache/spark/TaskEndReason.scala
index 18278b292ff5a..13241b77bf97b 100644
--- a/core/src/main/scala/org/apache/spark/TaskEndReason.scala
+++ b/core/src/main/scala/org/apache/spark/TaskEndReason.scala
@@ -223,8 +223,10 @@ case class TaskCommitDenied(
  * the task crashed the JVM.
  */
 @DeveloperApi
-case class ExecutorLostFailure(execId: String, exitCausedByApp: Boolean = true)
-  extends TaskFailedReason {
+case class ExecutorLostFailure(
+    execId: String,
+    exitCausedByApp: Boolean = true,
+    reason: Option[String]) extends TaskFailedReason {
   override def toErrorString: String = {
     val exitBehavior = if (exitCausedByApp) {
       "caused by one of the running tasks"
@@ -232,6 +234,8 @@ case class ExecutorLostFailure(execId: String, exitCausedByApp: Boolean = true)
       "unrelated to the running tasks"
     }
     s"ExecutorLostFailure (executor ${execId} exited due to an issue ${exitBehavior})"
+    s"ExecutorLostFailure (executor ${execId} exited ${exitBehavior})" +
+      reason.map { r => s" Reason: $r" }.getOrElse("")
   }
 
   override def countTowardsTaskFailures: Boolean = exitCausedByApp
diff --git a/core/src/main/scala/org/apache/spark/rpc/RpcEndpointRef.scala b/core/src/main/scala/org/apache/spark/rpc/RpcEndpointRef.scala
index f25710bb5bd6e..623da3e9c11b8 100644
--- a/core/src/main/scala/org/apache/spark/rpc/RpcEndpointRef.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/RpcEndpointRef.scala
@@ -67,7 +67,7 @@ private[spark] abstract class RpcEndpointRef(conf: SparkConf)
    * The default `timeout` will be used in every trial of calling `sendWithReply`. Because this
    * method retries, the message handling in the receiver side should be idempotent.
    *
-   * Note: this is a blocking action which may cost a lot of time,  so don't call it in an message
+   * Note: this is a blocking action which may cost a lot of time,  so don't call it in a message
    * loop of [[RpcEndpoint]].
    *
    * @param message the message to send
@@ -82,7 +82,7 @@ private[spark] abstract class RpcEndpointRef(conf: SparkConf)
    * retries. `timeout` will be used in every trial of calling `sendWithReply`. Because this method
    * retries, the message handling in the receiver side should be idempotent.
    *
-   * Note: this is a blocking action which may cost a lot of time, so don't call it in an message
+   * Note: this is a blocking action which may cost a lot of time, so don't call it in a message
    * loop of [[RpcEndpoint]].
    *
    * @param message the message to send
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
index 9b3fad9012abc..114468c48c44c 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -802,8 +802,8 @@ private[spark] class TaskSetManager(
         case exited: ExecutorExited => exited.exitCausedByApp
         case _ => true
       }
-      handleFailedTask(
-        tid, TaskState.FAILED, ExecutorLostFailure(info.executorId, exitCausedByApp))
+      handleFailedTask(tid, TaskState.FAILED, ExecutorLostFailure(info.executorId, exitCausedByApp,
+        Some(reason.toString)))
     }
     // recalculate valid locality levels and waits when executor is lost
     recomputeLocality()
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
index 439a11927026b..ebce5021b19dc 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
@@ -125,7 +125,6 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
             // Ignoring the task kill since the executor is not registered.
             logWarning(s"Attempted to kill task $taskId for unknown executor $executorId.")
         }
-
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
@@ -195,7 +194,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     override def onDisconnected(remoteAddress: RpcAddress): Unit = {
       addressToExecutorId
         .get(remoteAddress)
-        .foreach(removeExecutor(_, SlaveLost("remote Rpc client disassociated")))
+        .foreach(removeExecutor(_, SlaveLost("Remote RPC client disassociated. Likely due to " +
+          "containers exceeding thresholds, or network issues. Check driver logs for WARN " +
+          "messages.")))
     }
 
     // Make fake resource offers on just one executor
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
index cb24072d7d941..d75d6f673e84e 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
@@ -175,6 +175,7 @@ private[spark] abstract class YarnSchedulerBackend(
         addWebUIFilter(filterName, filterParams, proxyBase)
 
       case RemoveExecutor(executorId, reason) =>
+        logWarning(reason.toString)
         removeExecutor(executorId, reason)
     }
 
diff --git a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
index ad6615c1124d0..ee2eb58cf5e2a 100644
--- a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
+++ b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
@@ -367,9 +367,10 @@ private[spark] object JsonProtocol {
         ("Job ID" -> taskCommitDenied.jobID) ~
         ("Partition ID" -> taskCommitDenied.partitionID) ~
         ("Attempt Number" -> taskCommitDenied.attemptNumber)
-      case ExecutorLostFailure(executorId, exitCausedByApp) =>
+      case ExecutorLostFailure(executorId, exitCausedByApp, reason) =>
         ("Executor ID" -> executorId) ~
-        ("Exit Caused By App" -> exitCausedByApp)
+        ("Exit Caused By App" -> exitCausedByApp) ~
+        ("Loss Reason" -> reason.map(_.toString))
       case _ => Utils.emptyJson
     }
     ("Reason" -> reason) ~ json
@@ -812,7 +813,11 @@ private[spark] object JsonProtocol {
       case `executorLostFailure` =>
         val exitCausedByApp = Utils.jsonOption(json \ "Exit Caused By App").map(_.extract[Boolean])
         val executorId = Utils.jsonOption(json \ "Executor ID").map(_.extract[String])
-        ExecutorLostFailure(executorId.getOrElse("Unknown"), exitCausedByApp.getOrElse(true))
+        val reason = Utils.jsonOption(json \ "Loss Reason").map(_.extract[String])
+        ExecutorLostFailure(
+          executorId.getOrElse("Unknown"),
+          exitCausedByApp.getOrElse(true),
+          reason)
       case `unknownReason` => UnknownReason
     }
   }
diff --git a/core/src/test/scala/org/apache/spark/ui/jobs/JobProgressListenerSuite.scala b/core/src/test/scala/org/apache/spark/ui/jobs/JobProgressListenerSuite.scala
index b140387d309f3..e02f5a1b20fe3 100644
--- a/core/src/test/scala/org/apache/spark/ui/jobs/JobProgressListenerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ui/jobs/JobProgressListenerSuite.scala
@@ -243,7 +243,7 @@ class JobProgressListenerSuite extends SparkFunSuite with LocalSparkContext with
       ExceptionFailure("Exception", "description", null, null, None, None),
       TaskResultLost,
       TaskKilled,
-      ExecutorLostFailure("0"),
+      ExecutorLostFailure("0", true, Some("Induced failure")),
       UnknownReason)
     var failCount = 0
     for (reason <- taskFailedReasons) {
diff --git a/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala b/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala
index 86137f259c13d..953456c2caa89 100644
--- a/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala
@@ -152,7 +152,7 @@ class JsonProtocolSuite extends SparkFunSuite {
     testTaskEndReason(TaskResultLost)
     testTaskEndReason(TaskKilled)
     testTaskEndReason(TaskCommitDenied(2, 3, 4))
-    testTaskEndReason(ExecutorLostFailure("100", true))
+    testTaskEndReason(ExecutorLostFailure("100", true, Some("Induced failure")))
     testTaskEndReason(UnknownReason)
 
     // BlockId
@@ -296,10 +296,10 @@ class JsonProtocolSuite extends SparkFunSuite {
 
   test("ExecutorLostFailure backward compatibility") {
     // ExecutorLostFailure in Spark 1.1.0 does not have an "Executor ID" property.
-    val executorLostFailure = ExecutorLostFailure("100", true)
+    val executorLostFailure = ExecutorLostFailure("100", true, Some("Induced failure"))
     val oldEvent = JsonProtocol.taskEndReasonToJson(executorLostFailure)
       .removeField({ _._1 == "Executor ID" })
-    val expectedExecutorLostFailure = ExecutorLostFailure("Unknown", true)
+    val expectedExecutorLostFailure = ExecutorLostFailure("Unknown", true, Some("Induced failure"))
     assert(expectedExecutorLostFailure === JsonProtocol.taskEndReasonFromJson(oldEvent))
   }
 
@@ -603,10 +603,11 @@ class JsonProtocolSuite extends SparkFunSuite {
         assert(jobId1 === jobId2)
         assert(partitionId1 === partitionId2)
         assert(attemptNumber1 === attemptNumber2)
-      case (ExecutorLostFailure(execId1, exit1CausedByApp),
-          ExecutorLostFailure(execId2, exit2CausedByApp)) =>
+      case (ExecutorLostFailure(execId1, exit1CausedByApp, reason1),
+          ExecutorLostFailure(execId2, exit2CausedByApp, reason2)) =>
         assert(execId1 === execId2)
         assert(exit1CausedByApp === exit2CausedByApp)
+        assert(reason1 === reason2)
       case (UnknownReason, UnknownReason) =>
       case _ => fail("Task end reasons don't match in types!")
     }

From ebf8b0b48deaad64f7ca27051caee763451e2623 Mon Sep 17 00:00:00 2001
From: Cheng Lian <lian@databricks.com>
Date: Tue, 3 Nov 2015 10:07:45 -0800
Subject: [PATCH 032/518] [SPARK-10978][SQL] Allow data sources to eliminate
 filters

This PR adds a new method `unhandledFilters` to `BaseRelation`. Data sources which implement this method properly may avoid the overhead of defensive filtering done by Spark SQL.

Author: Cheng Lian <lian@databricks.com>

Closes #9399 from liancheng/spark-10978.unhandled-filters.
---
 .../datasources/DataSourceStrategy.scala      | 131 ++++++++++++++----
 .../apache/spark/sql/sources/interfaces.scala |   9 ++
 .../parquet/ParquetFilterSuite.scala          |   2 +-
 .../spark/sql/sources/FilteredScanSuite.scala | 129 ++++++++++++-----
 .../SimpleTextHadoopFsRelationSuite.scala     |  47 ++++++-
 .../sql/sources/SimpleTextRelation.scala      |  65 ++++++++-
 6 files changed, 315 insertions(+), 68 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
index 65859865c8fbc..7265d6a4de2e6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
@@ -43,7 +43,8 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
         l,
         projects,
         filters,
-        (a, f) => toCatalystRDD(l, a, t.buildScan(a, f))) :: Nil
+        (requestedColumns, allPredicates, _) =>
+          toCatalystRDD(l, requestedColumns, t.buildScan(requestedColumns, allPredicates))) :: Nil
 
     case PhysicalOperation(projects, filters, l @ LogicalRelation(t: PrunedFilteredScan, _)) =>
       pruneFilterProject(
@@ -266,47 +267,81 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
       relation,
       projects,
       filterPredicates,
-      (requestedColumns, pushedFilters) => {
-        scanBuilder(requestedColumns, selectFilters(pushedFilters).toArray)
+      (requestedColumns, _, pushedFilters) => {
+        scanBuilder(requestedColumns, pushedFilters.toArray)
       })
   }
 
-  // Based on Catalyst expressions.
+  // Based on Catalyst expressions. The `scanBuilder` function accepts three arguments:
+  //
+  //  1. A `Seq[Attribute]`, containing all required column attributes. Used to handle relation
+  //     traits that support column pruning (e.g. `PrunedScan` and `PrunedFilteredScan`).
+  //
+  //  2. A `Seq[Expression]`, containing all gathered Catalyst filter expressions, only used for
+  //     `CatalystScan`.
+  //
+  //  3. A `Seq[Filter]`, containing all data source `Filter`s that are converted from (possibly a
+  //     subset of) Catalyst filter expressions and can be handled by `relation`.  Used to handle
+  //     relation traits (`CatalystScan` excluded) that support filter push-down (e.g.
+  //     `PrunedFilteredScan` and `HadoopFsRelation`).
+  //
+  // Note that 2 and 3 shouldn't be used together.
   protected def pruneFilterProjectRaw(
-      relation: LogicalRelation,
-      projects: Seq[NamedExpression],
-      filterPredicates: Seq[Expression],
-      scanBuilder: (Seq[Attribute], Seq[Expression]) => RDD[InternalRow]) = {
+    relation: LogicalRelation,
+    projects: Seq[NamedExpression],
+    filterPredicates: Seq[Expression],
+    scanBuilder: (Seq[Attribute], Seq[Expression], Seq[Filter]) => RDD[InternalRow]) = {
 
     val projectSet = AttributeSet(projects.flatMap(_.references))
     val filterSet = AttributeSet(filterPredicates.flatMap(_.references))
-    val filterCondition = filterPredicates.reduceLeftOption(expressions.And)
 
-    val pushedFilters = filterPredicates.map { _ transform {
+    val candidatePredicates = filterPredicates.map { _ transform {
       case a: AttributeReference => relation.attributeMap(a) // Match original case of attributes.
     }}
 
+    val (unhandledPredicates, pushedFilters) =
+      selectFilters(relation.relation, candidatePredicates)
+
+    // A set of column attributes that are only referenced by pushed down filters.  We can eliminate
+    // them from requested columns.
+    val handledSet = {
+      val handledPredicates = filterPredicates.filterNot(unhandledPredicates.contains)
+      val unhandledSet = AttributeSet(unhandledPredicates.flatMap(_.references))
+      AttributeSet(handledPredicates.flatMap(_.references)) --
+        (projectSet ++ unhandledSet).map(relation.attributeMap)
+    }
+
+    // Combines all Catalyst filter `Expression`s that are either not convertible to data source
+    // `Filter`s or cannot be handled by `relation`.
+    val filterCondition = unhandledPredicates.reduceLeftOption(expressions.And)
+
     if (projects.map(_.toAttribute) == projects &&
         projectSet.size == projects.size &&
         filterSet.subsetOf(projectSet)) {
       // When it is possible to just use column pruning to get the right projection and
       // when the columns of this projection are enough to evaluate all filter conditions,
       // just do a scan followed by a filter, with no extra project.
-      val requestedColumns =
-        projects.asInstanceOf[Seq[Attribute]] // Safe due to if above.
-          .map(relation.attributeMap)            // Match original case of attributes.
+      val requestedColumns = projects
+        // Safe due to if above.
+        .asInstanceOf[Seq[Attribute]]
+        // Match original case of attributes.
+        .map(relation.attributeMap)
+        // Don't request columns that are only referenced by pushed filters.
+        .filterNot(handledSet.contains)
 
       val scan = execution.PhysicalRDD.createFromDataSource(
         projects.map(_.toAttribute),
-        scanBuilder(requestedColumns, pushedFilters),
+        scanBuilder(requestedColumns, candidatePredicates, pushedFilters),
         relation.relation)
       filterCondition.map(execution.Filter(_, scan)).getOrElse(scan)
     } else {
-      val requestedColumns = (projectSet ++ filterSet).map(relation.attributeMap).toSeq
+      // Don't request columns that are only referenced by pushed filters.
+      val requestedColumns =
+        (projectSet ++ filterSet -- handledSet).map(relation.attributeMap).toSeq
 
       val scan = execution.PhysicalRDD.createFromDataSource(
         requestedColumns,
-        scanBuilder(requestedColumns, pushedFilters),
+        scanBuilder(requestedColumns, candidatePredicates, pushedFilters),
         relation.relation)
       execution.Project(projects, filterCondition.map(execution.Filter(_, scan)).getOrElse(scan))
     }
@@ -334,11 +369,12 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
   }
 
   /**
-   * Selects Catalyst predicate [[Expression]]s which are convertible into data source [[Filter]]s,
-   * and convert them.
+   * Tries to translate a Catalyst [[Expression]] into data source [[Filter]].
+   *
+   * @return a `Some[Filter]` if the input [[Expression]] is convertible, otherwise a `None`.
    */
-  protected[sql] def selectFilters(filters: Seq[Expression]) = {
-    def translate(predicate: Expression): Option[Filter] = predicate match {
+  protected[sql] def translateFilter(predicate: Expression): Option[Filter] = {
+    predicate match {
       case expressions.EqualTo(a: Attribute, Literal(v, t)) =>
         Some(sources.EqualTo(a.name, convertToScala(v, t)))
       case expressions.EqualTo(Literal(v, t), a: Attribute) =>
@@ -387,16 +423,16 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
         Some(sources.IsNotNull(a.name))
 
       case expressions.And(left, right) =>
-        (translate(left) ++ translate(right)).reduceOption(sources.And)
+        (translateFilter(left) ++ translateFilter(right)).reduceOption(sources.And)
 
       case expressions.Or(left, right) =>
         for {
-          leftFilter <- translate(left)
-          rightFilter <- translate(right)
+          leftFilter <- translateFilter(left)
+          rightFilter <- translateFilter(right)
         } yield sources.Or(leftFilter, rightFilter)
 
       case expressions.Not(child) =>
-        translate(child).map(sources.Not)
+        translateFilter(child).map(sources.Not)
 
       case expressions.StartsWith(a: Attribute, Literal(v: UTF8String, StringType)) =>
         Some(sources.StringStartsWith(a.name, v.toString))
@@ -409,7 +445,52 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
 
       case _ => None
     }
+  }
+
+  /**
+   * Selects Catalyst predicate [[Expression]]s which are convertible into data source [[Filter]]s
+   * and can be handled by `relation`.
+   *
+   * @return A pair of `Seq[Expression]` and `Seq[Filter]`. The first element contains all Catalyst
+   *         predicate [[Expression]]s that are either not convertible or cannot be handled by
+   *         `relation`. The second element contains all converted data source [[Filter]]s that can
+   *        be handled by `relation`.
+   */
+  protected[sql] def selectFilters(
+    relation: BaseRelation,
+    predicates: Seq[Expression]): (Seq[Expression], Seq[Filter]) = {
+
+    // For conciseness, all Catalyst filter expressions of type `expressions.Expression` below are
+    // called `predicate`s, while all data source filters of type `sources.Filter` are simply called
+    // `filter`s.
+
+    val translated: Seq[(Expression, Filter)] =
+      for {
+        predicate <- predicates
+        filter <- translateFilter(predicate)
+      } yield predicate -> filter
+
+    // A map from original Catalyst expressions to corresponding translated data source filters.
+    val translatedMap: Map[Expression, Filter] = translated.toMap
+
+    // Catalyst predicate expressions that cannot be translated to data source filters.
+    val unrecognizedPredicates = predicates.filterNot(translatedMap.contains)
+
+    // Data source filters that cannot be handled by `relation`
+    val unhandledFilters = relation.unhandledFilters(translatedMap.values.toArray).toSet
+
+    val (unhandled, handled) = translated.partition {
+      case (predicate, filter) =>
+        unhandledFilters.contains(filter)
+    }
+
+    // Catalyst predicate expressions that can be translated to data source filters, but cannot be
+    // handled by `relation`.
+    val (unhandledPredicates, _) = unhandled.unzip
+
+    // Translated data source filters that can be handled by `relation`
+    val (_, handledFilters) = handled.unzip
 
-    filters.flatMap(translate)
+    (unrecognizedPredicates ++ unhandledPredicates, handledFilters)
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index 7a553511483ff..e296d631f0f30 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -233,6 +233,15 @@ abstract class BaseRelation {
    * @since 1.4.0
    */
   def needConversion: Boolean = true
+
+  /**
+   * Given an array of [[Filter]]s, returns an array of [[Filter]]s that this data source relation
+   * cannot handle.  Spark SQL will apply all returned [[Filter]]s against rows returned by this
+   * data source relation.
+   *
+   * @since 1.6.0
+   */
+  def unhandledFilters(filters: Array[Filter]): Array[Filter] = filters
 }
 
 /**
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
index f88ddc77a6a4e..c24c9f025dad7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
@@ -59,7 +59,7 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
       }.flatten
       assert(analyzedPredicate.nonEmpty)
 
-      val selectedFilters = DataSourceStrategy.selectFilters(analyzedPredicate)
+      val selectedFilters = analyzedPredicate.flatMap(DataSourceStrategy.translateFilter)
       assert(selectedFilters.nonEmpty)
 
       selectedFilters.foreach { pred =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
index 68ce37c00077e..7541e723029bf 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.sources
 
+import org.apache.spark.sql.execution.datasources.LogicalRelation
+
 import scala.language.existentials
 
 import org.apache.spark.rdd.RDD
@@ -44,16 +46,39 @@ case class SimpleFilteredScan(from: Int, to: Int)(@transient val sqlContext: SQL
       StructField("b", IntegerType, nullable = false) ::
       StructField("c", StringType, nullable = false) :: Nil)
 
+  override def unhandledFilters(filters: Array[Filter]): Array[Filter] = {
+    def unhandled(filter: Filter): Boolean = {
+      filter match {
+        case EqualTo(col, v) => col == "b"
+        case EqualNullSafe(col, v) => col == "b"
+        case LessThan(col, v: Int) => col == "b"
+        case LessThanOrEqual(col, v: Int) => col == "b"
+        case GreaterThan(col, v: Int) => col == "b"
+        case GreaterThanOrEqual(col, v: Int) => col == "b"
+        case In(col, values) => col == "b"
+        case IsNull(col) => col == "b"
+        case IsNotNull(col) => col == "b"
+        case Not(pred) => unhandled(pred)
+        case And(left, right) => unhandled(left) || unhandled(right)
+        case Or(left, right) => unhandled(left) || unhandled(right)
+        case _ => false
+      }
+    }
+
+    filters.filter(unhandled)
+  }
+
   override def buildScan(requiredColumns: Array[String], filters: Array[Filter]): RDD[Row] = {
     val rowBuilders = requiredColumns.map {
       case "a" => (i: Int) => Seq(i)
       case "b" => (i: Int) => Seq(i * 2)
       case "c" => (i: Int) =>
         val c = (i - 1 + 'a').toChar.toString
-        Seq(c * 5 + c.toUpperCase() * 5)
+        Seq(c * 5 + c.toUpperCase * 5)
     }
 
     FiltersPushed.list = filters
+    ColumnsRequired.set = requiredColumns.toSet
 
     // Predicate test on integer column
     def translateFilterOnA(filter: Filter): Int => Boolean = filter match {
@@ -86,9 +111,8 @@ case class SimpleFilteredScan(from: Int, to: Int)(@transient val sqlContext: SQL
     }
 
     def eval(a: Int) = {
-      val c = (a - 1 + 'a').toChar.toString * 5 + (a - 1 + 'a').toChar.toString.toUpperCase() * 5
-      !filters.map(translateFilterOnA(_)(a)).contains(false) &&
-        !filters.map(translateFilterOnC(_)(c)).contains(false)
+      val c = (a - 1 + 'a').toChar.toString * 5 + (a - 1 + 'a').toChar.toString.toUpperCase * 5
+      filters.forall(translateFilterOnA(_)(a)) && filters.forall(translateFilterOnC(_)(c))
     }
 
     sqlContext.sparkContext.parallelize(from to to).filter(eval).map(i =>
@@ -101,6 +125,11 @@ object FiltersPushed {
   var list: Seq[Filter] = Nil
 }
 
+// Used together with `SimpleFilteredScan` to check pushed columns.
+object ColumnsRequired {
+  var set: Set[String] = Set.empty
+}
+
 class FilteredScanSuite extends DataSourceTest with SharedSQLContext {
   protected override lazy val sql = caseInsensitiveContext.sql _
 
@@ -115,12 +144,15 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext {
         |  to '10'
         |)
       """.stripMargin)
+
+    // UDF for testing filter push-down
+    caseInsensitiveContext.udf.register("udf_gt3", (_: Int) > 3)
   }
 
   sqlTest(
     "SELECT * FROM oneToTenFiltered",
     (1 to 10).map(i => Row(i, i * 2, (i - 1 + 'a').toChar.toString * 5
-      + (i - 1 + 'a').toChar.toString.toUpperCase() * 5)).toSeq)
+      + (i - 1 + 'a').toChar.toString.toUpperCase * 5)).toSeq)
 
   sqlTest(
     "SELECT a, b FROM oneToTenFiltered",
@@ -202,49 +234,64 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext {
     "SELECT a, b, c FROM oneToTenFiltered WHERE c like '%eE%'",
     Seq(Row(5, 5 * 2, "e" * 5 + "E" * 5)))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE A = 1", 1)
-  testPushDown("SELECT a FROM oneToTenFiltered WHERE A = 1", 1)
-  testPushDown("SELECT b FROM oneToTenFiltered WHERE A = 1", 1)
-  testPushDown("SELECT a, b FROM oneToTenFiltered WHERE A = 1", 1)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a = 1", 1)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE 1 = a", 1)
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE A = 1", 1, Set("a", "b", "c"))
+  testPushDown("SELECT a FROM oneToTenFiltered WHERE A = 1", 1, Set("a"))
+  testPushDown("SELECT b FROM oneToTenFiltered WHERE A = 1", 1, Set("b"))
+  testPushDown("SELECT a, b FROM oneToTenFiltered WHERE A = 1", 1, Set("a", "b"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a = 1", 1, Set("a", "b", "c"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE 1 = a", 1, Set("a", "b", "c"))
+
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a > 1", 9, Set("a", "b", "c"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a >= 2", 9, Set("a", "b", "c"))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a > 1", 9)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a >= 2", 9)
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE 1 < a", 9, Set("a", "b", "c"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE 2 <= a", 9, Set("a", "b", "c"))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE 1 < a", 9)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE 2 <= a", 9)
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE 1 > a", 0, Set("a", "b", "c"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE 2 >= a", 2, Set("a", "b", "c"))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE 1 > a", 0)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE 2 >= a", 2)
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a < 1", 0, Set("a", "b", "c"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a <= 2", 2, Set("a", "b", "c"))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a < 1", 0)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a <= 2", 2)
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a > 1 AND a < 10", 8, Set("a", "b", "c"))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a > 1 AND a < 10", 8)
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a IN (1,3,5)", 3, Set("a", "b", "c"))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a IN (1,3,5)", 3)
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a = 20", 0, Set("a", "b", "c"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE b = 1", 10, Set("a", "b", "c"))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a = 20", 0)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE b = 1", 10)
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a < 5 AND a > 1", 3, Set("a", "b", "c"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE a < 3 OR a > 8", 4, Set("a", "b", "c"))
+  testPushDown("SELECT * FROM oneToTenFiltered WHERE NOT (a < 6)", 5, Set("a", "b", "c"))
 
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a < 5 AND a > 1", 3)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE a < 3 OR a > 8", 4)
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE NOT (a < 6)", 5)
+  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like 'c%'", 1, Set("a", "b", "c"))
+  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like 'C%'", 0, Set("a", "b", "c"))
 
-  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like 'c%'", 1)
-  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like 'C%'", 0)
+  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like '%D'", 1, Set("a", "b", "c"))
+  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like '%d'", 0, Set("a", "b", "c"))
 
-  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like '%D'", 1)
-  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like '%d'", 0)
+  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like '%eE%'", 1, Set("a", "b", "c"))
+  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like '%Ee%'", 0, Set("a", "b", "c"))
 
-  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like '%eE%'", 1)
-  testPushDown("SELECT a, b, c FROM oneToTenFiltered WHERE c like '%Ee%'", 0)
+  testPushDown("SELECT c FROM oneToTenFiltered WHERE c = 'aaaaaAAAAA'", 1, Set("c"))
+  testPushDown("SELECT c FROM oneToTenFiltered WHERE c IN ('aaaaaAAAAA', 'foo')", 1, Set("c"))
 
-  testPushDown("SELECT c FROM oneToTenFiltered WHERE c = 'aaaaaAAAAA'", 1)
-  testPushDown("SELECT c FROM oneToTenFiltered WHERE c IN ('aaaaaAAAAA', 'foo')", 1)
+  // Columns only referenced by UDF filter must be required, as UDF filters can't be pushed down.
+  testPushDown("SELECT c FROM oneToTenFiltered WHERE udf_gt3(A)", 10, Set("a", "c"))
 
-  def testPushDown(sqlString: String, expectedCount: Int): Unit = {
+  // A query with an unconvertible filter, an unhandled filter, and a handled filter.
+  testPushDown(
+    """SELECT a
+      |  FROM oneToTenFiltered
+      | WHERE udf_gt3(b)
+      |   AND b < 16
+      |   AND c IN ('bbbbbBBBBB', 'cccccCCCCC', 'dddddDDDDD', 'foo')
+    """.stripMargin.split("\n").map(_.trim).mkString(" "), 3, Set("a", "b"))
+
+  def testPushDown(
+      sqlString: String,
+      expectedCount: Int,
+      requiredColumnNames: Set[String]): Unit = {
     test(s"PushDown Returns $expectedCount: $sqlString") {
       val queryExecution = sql(sqlString).queryExecution
       val rawPlan = queryExecution.executedPlan.collect {
@@ -254,6 +301,17 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext {
         case _ => fail(s"More than one PhysicalRDD found\n$queryExecution")
       }
       val rawCount = rawPlan.execute().count()
+      assert(ColumnsRequired.set === requiredColumnNames)
+
+      assert {
+        val table = caseInsensitiveContext.table("oneToTenFiltered")
+        val relation = table.queryExecution.logical.collectFirst {
+          case LogicalRelation(r, _) => r
+        }.get
+
+        // `relation` should be able to handle all pushed filters
+        relation.unhandledFilters(FiltersPushed.list.toArray).isEmpty
+      }
 
       if (rawCount != expectedCount) {
         fail(
@@ -264,4 +322,3 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext {
     }
   }
 }
-
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
index a3a124488d983..d945408341fc9 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
@@ -18,11 +18,16 @@
 package org.apache.spark.sql.sources
 
 import org.apache.hadoop.fs.Path
-
 import org.apache.spark.deploy.SparkHadoopUtil
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.execution.PhysicalRDD
+import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
 
 class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest {
+  import testImplicits._
+
   override val dataSourceName: String = classOf[SimpleTextSource].getCanonicalName
 
   // We have a very limited number of supported types at here since it is just for a
@@ -64,4 +69,44 @@ class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest {
           .load(file.getCanonicalPath))
     }
   }
+
+  private val writer = testDF.write.option("dataSchema", dataSchema.json).format(dataSourceName)
+  private val reader = sqlContext.read.option("dataSchema", dataSchema.json).format(dataSourceName)
+
+  test("unhandledFilters") {
+    withTempPath { dir =>
+
+      val path = dir.getCanonicalPath
+      writer.save(s"$path/p=0")
+      writer.save(s"$path/p=1")
+
+      val isOdd = udf((_: Int) % 2 == 1)
+      val df = reader.load(path)
+        .filter(
+          // This filter is inconvertible
+          isOdd('a) &&
+            // This filter is convertible but unhandled
+            'a > 1 &&
+            // This filter is convertible and handled
+            'b > "val_1" &&
+            // This filter references a partiiton column, won't be pushed down
+            'p === 1
+        ).select('a, 'p)
+      val rawScan = df.queryExecution.executedPlan collect {
+        case p: PhysicalRDD => p
+      } match {
+        case Seq(p) => p
+      }
+
+      val outputSchema = new StructType().add("a", IntegerType).add("p", IntegerType)
+
+      assertResult(Set((2, 1), (3, 1))) {
+        rawScan.execute().collect()
+          .map { CatalystTypeConverters.convertToScala(_, outputSchema) }
+          .map { case Row(a, p) => (a, p) }.toSet
+      }
+
+      checkAnswer(df, Row(3, 1))
+    }
+  }
 }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
index aeaaa3e1c5220..da09e1b00ae48 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
@@ -18,7 +18,6 @@
 package org.apache.spark.sql.sources
 
 import java.text.NumberFormat
-import java.util.UUID
 
 import com.google.common.base.Objects
 import org.apache.hadoop.fs.{FileStatus, Path}
@@ -26,12 +25,12 @@ import org.apache.hadoop.io.{NullWritable, Text}
 import org.apache.hadoop.mapreduce.lib.output.{FileOutputFormat, TextOutputFormat}
 import org.apache.hadoop.mapreduce.{Job, RecordWriter, TaskAttemptContext}
 
-import org.apache.spark.rdd.RDD
 import org.apache.spark.deploy.SparkHadoopUtil
-import org.apache.spark.sql.catalyst.CatalystTypeConverters
-import org.apache.spark.sql.catalyst.expressions.{Cast, Literal}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.{CatalystTypeConverters, expressions}
 import org.apache.spark.sql.types.{DataType, StructType}
-import org.apache.spark.sql.{Row, SQLContext}
+import org.apache.spark.sql.{Row, SQLContext, sources}
 
 /**
  * A simple example [[HadoopFsRelationProvider]].
@@ -124,6 +123,53 @@ class SimpleTextRelation(
     }
   }
 
+  override def buildScan(
+      requiredColumns: Array[String],
+      filters: Array[Filter],
+      inputFiles: Array[FileStatus]): RDD[Row] = {
+
+    val fields = this.dataSchema.map(_.dataType)
+    val inputAttributes = this.dataSchema.toAttributes
+    val outputAttributes = requiredColumns.flatMap(name => inputAttributes.find(_.name == name))
+    val dataSchema = this.dataSchema
+
+    val inputPaths = inputFiles.map(_.getPath).mkString(",")
+    sparkContext.textFile(inputPaths).mapPartitions { iterator =>
+      // Constructs a filter predicate to simulate filter push-down
+      val predicate = {
+        val filterCondition: Expression = filters.collect {
+          // According to `unhandledFilters`, `SimpleTextRelation` only handles `GreaterThan` filter
+          case sources.GreaterThan(column, value) =>
+            val dataType = dataSchema(column).dataType
+            val literal = Literal.create(value, dataType)
+            val attribute = inputAttributes.find(_.name == column).get
+            expressions.GreaterThan(attribute, literal)
+        }.reduceOption(expressions.And).getOrElse(Literal(true))
+        InterpretedPredicate.create(filterCondition, inputAttributes)
+      }
+
+      // Uses a simple projection to simulate column pruning
+      val projection = new InterpretedMutableProjection(outputAttributes, inputAttributes)
+      val toScala = {
+        val requiredSchema = StructType.fromAttributes(outputAttributes)
+        CatalystTypeConverters.createToScalaConverter(requiredSchema)
+      }
+
+      iterator.map { record =>
+        new GenericInternalRow(record.split(",", -1).zip(fields).map {
+          case (v, dataType) =>
+            val value = if (v == "") null else v
+            // `Cast`ed values are always of internal types (e.g. UTF8String instead of String)
+            Cast(Literal(value), dataType).eval()
+        })
+      }.filter { row =>
+        predicate(row)
+      }.map { row =>
+        toScala(projection(row)).asInstanceOf[Row]
+      }
+    }
+  }
+
   override def prepareJobForWrite(job: Job): OutputWriterFactory = new OutputWriterFactory {
     job.setOutputFormatClass(classOf[TextOutputFormat[_, _]])
 
@@ -134,6 +180,15 @@ class SimpleTextRelation(
       new SimpleTextOutputWriter(path, context)
     }
   }
+
+  // `SimpleTextRelation` only handles `GreaterThan` filter.  This is used to test filter push-down
+  // and `BaseRelation.unhandledFilters()`.
+  override def unhandledFilters(filters: Array[Filter]): Array[Filter] = {
+    filters.filter {
+      case _: GreaterThan => false
+      case _ => true
+    }
+  }
 }
 
 /**

From a9676cc7107c5df6c62a58668c4d95ced1238370 Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Tue, 3 Nov 2015 11:53:10 -0800
Subject: [PATCH 033/518] [SPARK-11407][SPARKR] Add doc for running from
 RStudio

![image](https://cloud.githubusercontent.com/assets/8969467/10871746/612ba44a-80a4-11e5-99a0-40b9931dee52.png)
(This is without css, but you get the idea)
shivaram

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9401 from felixcheung/rstudioprogrammingguide.
---
 docs/sparkr.md | 46 +++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 43 insertions(+), 3 deletions(-)

diff --git a/docs/sparkr.md b/docs/sparkr.md
index 497a276679f3b..437bd4756c276 100644
--- a/docs/sparkr.md
+++ b/docs/sparkr.md
@@ -30,14 +30,22 @@ The entry point into SparkR is the `SparkContext` which connects your R program
 You can create a `SparkContext` using `sparkR.init` and pass in options such as the application name
 , any spark packages depended on, etc. Further, to work with DataFrames we will need a `SQLContext`,
 which can be created from the  SparkContext. If you are working from the `sparkR` shell, the
-`SQLContext` and `SparkContext` should already be created for you.
+`SQLContext` and `SparkContext` should already be created for you, and you would not need to call
+`sparkR.init`.
 
+<div data-lang="r" markdown="1">
 {% highlight r %}
 sc <- sparkR.init()
 sqlContext <- sparkRSQL.init(sc)
 {% endhighlight %}
+</div>
+
+## Starting Up from RStudio
 
-In the event you are creating `SparkContext` instead of using `sparkR` shell or `spark-submit`, you 
+You can also start SparkR from RStudio. You can connect your R program to a Spark cluster from
+RStudio, R shell, Rscript or other R IDEs. To start, make sure SPARK_HOME is set in environment
+(you can check [Sys.getenv](https://stat.ethz.ch/R-manual/R-devel/library/base/html/Sys.getenv.html)),
+load the SparkR package, and call `sparkR.init` as below. In addition to calling `sparkR.init`, you
 could also specify certain Spark driver properties. Normally these
 [Application properties](configuration.html#application-properties) and
 [Runtime Environment](configuration.html#runtime-environment) cannot be set programmatically, as the
@@ -45,9 +53,41 @@ driver JVM process would have been started, in this case SparkR takes care of th
 them, pass them as you would other configuration properties in the `sparkEnvir` argument to
 `sparkR.init()`.
 
+<div data-lang="r" markdown="1">
 {% highlight r %}
-sc <- sparkR.init("local[*]", "SparkR", "/home/spark", list(spark.driver.memory="2g"))
+if (nchar(Sys.getenv("SPARK_HOME")) < 1) {
+  Sys.setenv(SPARK_HOME = "/home/spark")
+}
+library(SparkR, lib.loc = c(file.path(Sys.getenv("SPARK_HOME"), "R", "lib")))
+sc <- sparkR.init(master = "local[*]", sparkEnvir = list(spark.driver.memory="2g"))
 {% endhighlight %}
+</div>
+
+The following options can be set in `sparkEnvir` with `sparkR.init` from RStudio:
+
+<table class="table">
+  <tr><th>Property Name</th><th>Property group</th><th><code>spark-submit</code> equivalent</th></tr>
+  <tr>
+    <td><code>spark.driver.memory</code></td>
+    <td>Application Properties</td>
+    <td><code>--driver-memory</code></td>
+  </tr>
+  <tr>
+    <td><code>spark.driver.extraClassPath</code></td>
+    <td>Runtime Environment</td>
+    <td><code>--driver-class-path</code></td>
+  </tr>
+  <tr>
+    <td><code>spark.driver.extraJavaOptions</code></td>
+    <td>Runtime Environment</td>
+    <td><code>--driver-java-options</code></td>
+  </tr>
+  <tr>
+    <td><code>spark.driver.extraLibraryPath</code></td>
+    <td>Runtime Environment</td>
+    <td><code>--driver-library-path</code></td>
+  </tr>
+</table>
 
 </div>
 

From 1d04dc95c0d3caa485936e65b0493bcc9719f27e Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Tue, 3 Nov 2015 13:33:46 -0800
Subject: [PATCH 034/518] [SPARK-11467][SQL] add Python API for stddev/variance

Add Python API for stddev/stddev_pop/stddev_samp/variance/var_pop/var_samp/skewness/kurtosis

Author: Davies Liu <davies@databricks.com>

Closes #9424 from davies/py_var.
---
 python/pyspark/sql/functions.py               | 17 ++++
 python/pyspark/sql/group.py                   | 88 +++++++++++++++++++
 .../org/apache/spark/sql/functions.scala      | 67 --------------
 3 files changed, 105 insertions(+), 67 deletions(-)

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index fa04f4cd83b6f..2f7c2f4aacd47 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -122,6 +122,21 @@ def _():
     'bitwiseNOT': 'Computes bitwise not.',
 }
 
+_functions_1_6 = {
+    # unary math functions
+    "stddev": "Aggregate function: returns the unbiased sample standard deviation of" +
+              " the expression in a group.",
+    "stddev_samp": "Aggregate function: returns the unbiased sample standard deviation of" +
+              " the expression in a group.",
+    "stddev_pop": "Aggregate function: returns population standard deviation of" +
+              " the expression in a group.",
+    "variance": "Aggregate function: returns the population variance of the values in a group.",
+    "var_samp": "Aggregate function: returns the unbiased variance of the values in a group.",
+    "var_pop":  "Aggregate function: returns the population variance of the values in a group.",
+    "skewness": "Aggregate function: returns the skewness of the values in a group.",
+    "kurtosis": "Aggregate function: returns the kurtosis of the values in a group."
+}
+
 # math functions that take two arguments as input
 _binary_mathfunctions = {
     'atan2': 'Returns the angle theta from the conversion of rectangular coordinates (x, y) to' +
@@ -172,6 +187,8 @@ def _():
     globals()[_name] = since(1.4)(_create_binary_mathfunction(_name, _doc))
 for _name, _doc in _window_functions.items():
     globals()[_name] = since(1.4)(_create_window_function(_name, _doc))
+for _name, _doc in _functions_1_6.items():
+    globals()[_name] = since(1.6)(_create_function(_name, _doc))
 del _name, _doc
 
 
diff --git a/python/pyspark/sql/group.py b/python/pyspark/sql/group.py
index 71c0bccc5eeff..946b53e71c2c6 100644
--- a/python/pyspark/sql/group.py
+++ b/python/pyspark/sql/group.py
@@ -167,6 +167,94 @@ def sum(self, *cols):
         [Row(sum(age)=7, sum(height)=165)]
         """
 
+    @df_varargs_api
+    @since(1.6)
+    def stddev(self, *cols):
+        """Compute the sample standard deviation for each numeric columns for each group.
+
+        :param cols: list of column names (string). Non-numeric columns are ignored.
+
+        >>> df3.groupBy().stddev('age', 'height').collect()
+        [Row(STDDEV(age)=2.12..., STDDEV(height)=3.53...)]
+        """
+
+    @df_varargs_api
+    @since(1.6)
+    def stddev_samp(self, *cols):
+        """Compute the sample standard deviation for each numeric columns for each group.
+
+        :param cols: list of column names (string). Non-numeric columns are ignored.
+
+        >>> df3.groupBy().stddev_samp('age', 'height').collect()
+        [Row(STDDEV_SAMP(age)=2.12..., STDDEV_SAMP(height)=3.53...)]
+        """
+
+    @df_varargs_api
+    @since(1.6)
+    def stddev_pop(self, *cols):
+        """Compute the population standard deviation for each numeric columns for each group.
+
+        :param cols: list of column names (string). Non-numeric columns are ignored.
+
+        >>> df3.groupBy().stddev_pop('age', 'height').collect()
+        [Row(STDDEV_POP(age)=1.5, STDDEV_POP(height)=2.5)]
+        """
+
+    @df_varargs_api
+    @since(1.6)
+    def variance(self, *cols):
+        """Compute the sample variance for each numeric columns for each group.
+
+        :param cols: list of column names (string). Non-numeric columns are ignored.
+
+        >>> df3.groupBy().variance('age', 'height').collect()
+        [Row(VARIANCE(age)=2.25, VARIANCE(height)=6.25)]
+        """
+
+    @df_varargs_api
+    @since(1.6)
+    def var_pop(self, *cols):
+        """Compute the sample variance for each numeric columns for each group.
+
+        :param cols: list of column names (string). Non-numeric columns are ignored.
+
+        >>> df3.groupBy().var_pop('age', 'height').collect()
+        [Row(VAR_POP(age)=2.25, VAR_POP(height)=6.25)]
+        """
+
+    @df_varargs_api
+    @since(1.6)
+    def var_samp(self, *cols):
+        """Compute the sample variance for each numeric columns for each group.
+
+        :param cols: list of column names (string). Non-numeric columns are ignored.
+
+        >>> df3.groupBy().var_samp('age', 'height').collect()
+        [Row(VAR_SAMP(age)=4.5, VAR_SAMP(height)=12.5)]
+        """
+
+    @df_varargs_api
+    @since(1.6)
+    def skewness(self, *cols):
+        """Compute the skewness for each numeric columns for each group.
+
+        :param cols: list of column names (string). Non-numeric columns are ignored.
+
+        >>> df3.groupBy().skewness('age', 'height').collect()
+        [Row(SKEWNESS(age)=0.0, SKEWNESS(height)=0.0)]
+        """
+
+    @df_varargs_api
+    @since(1.6)
+    def kurtosis(self, *cols):
+        """Compute the kurtosis for each numeric columns for each group.
+
+        :param cols: list of column names (string). Non-numeric columns are ignored.
+
+        >>> df3.groupBy().kurtosis('age', 'height').collect()
+        [Row(KURTOSIS(age)=-2.0, KURTOSIS(height)=-2.0)]
+        """
+
 
 def _test():
     import doctest
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 5a5c695e6ab3b..c8c52831668cd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -254,14 +254,6 @@ object functions {
    */
   def kurtosis(e: Column): Column = Kurtosis(e.expr)
 
-  /**
-   * Aggregate function: returns the kurtosis of the values in a group.
-   *
-   * @group agg_funcs
-   * @since 1.6.0
-   */
-  def kurtosis(columnName: String): Column = kurtosis(Column(columnName))
-
   /**
    * Aggregate function: returns the last value in a group.
    *
@@ -336,14 +328,6 @@ object functions {
    */
   def skewness(e: Column): Column = Skewness(e.expr)
 
-  /**
-   * Aggregate function: returns the skewness of the values in a group.
-   *
-   * @group agg_funcs
-   * @since 1.6.0
-   */
-  def skewness(columnName: String): Column = skewness(Column(columnName))
-
   /**
    * Aggregate function: returns the unbiased sample standard deviation of
    * the expression in a group.
@@ -353,15 +337,6 @@ object functions {
    */
   def stddev(e: Column): Column = Stddev(e.expr)
 
-  /**
-   * Aggregate function: returns the unbiased sample standard deviation of
-   * the expression in a group.
-   *
-   * @group agg_funcs
-   * @since 1.6.0
-   */
-  def stddev(columnName: String): Column = stddev(Column(columnName))
-
   /**
    * Aggregate function: returns the unbiased sample standard deviation of
    * the expression in a group.
@@ -371,15 +346,6 @@ object functions {
    */
   def stddev_samp(e: Column): Column = StddevSamp(e.expr)
 
-  /**
-   * Aggregate function: returns the unbiased sample standard deviation of
-   * the expression in a group.
-   *
-   * @group agg_funcs
-   * @since 1.6.0
-   */
-  def stddev_samp(columnName: String): Column = stddev_samp(Column(columnName))
-
   /**
    * Aggregate function: returns the population standard deviation of
    * the expression in a group.
@@ -389,15 +355,6 @@ object functions {
    */
   def stddev_pop(e: Column): Column = StddevPop(e.expr)
 
-  /**
-   * Aggregate function: returns the population standard deviation of
-   * the expression in a group.
-   *
-   * @group agg_funcs
-   * @since 1.6.0
-   */
-  def stddev_pop(columnName: String): Column = stddev_pop(Column(columnName))
-
   /**
    * Aggregate function: returns the sum of all values in the expression.
    *
@@ -438,14 +395,6 @@ object functions {
    */
   def variance(e: Column): Column = Variance(e.expr)
 
-  /**
-   * Aggregate function: returns the population variance of the values in a group.
-   *
-   * @group agg_funcs
-   * @since 1.6.0
-   */
-  def variance(columnName: String): Column = variance(Column(columnName))
-
   /**
    * Aggregate function: returns the unbiased variance of the values in a group.
    *
@@ -454,14 +403,6 @@ object functions {
    */
   def var_samp(e: Column): Column = VarianceSamp(e.expr)
 
-  /**
-   * Aggregate function: returns the unbiased variance of the values in a group.
-   *
-   * @group agg_funcs
-   * @since 1.6.0
-   */
-  def var_samp(columnName: String): Column = var_samp(Column(columnName))
-
   /**
    * Aggregate function: returns the population variance of the values in a group.
    *
@@ -470,14 +411,6 @@ object functions {
    */
   def var_pop(e: Column): Column = VariancePop(e.expr)
 
-  /**
-   * Aggregate function: returns the population variance of the values in a group.
-   *
-   * @group agg_funcs
-   * @since 1.6.0
-   */
-  def var_pop(columnName: String): Column = var_pop(Column(columnName))
-
   //////////////////////////////////////////////////////////////////////////////////////////////
   // Window functions
   //////////////////////////////////////////////////////////////////////////////////////////////

From f6fcb4874ce20a1daa91b7434cf9c0254a89e979 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 4 Nov 2015 00:15:50 +0100
Subject: [PATCH 035/518] [SPARK-11477] [SQL] support create Dataset from RDD

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9434 from cloud-fan/rdd2ds and squashes the following commits:

0892d72 [Wenchen Fan] support create Dataset from RDD
---
 .../src/main/scala/org/apache/spark/sql/SQLContext.scala | 9 +++++++++
 .../main/scala/org/apache/spark/sql/SQLImplicits.scala   | 4 ++++
 .../test/scala/org/apache/spark/sql/DatasetSuite.scala   | 7 +++++++
 3 files changed, 20 insertions(+)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index 2cb94430e6178..5ad3871093fc8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -499,6 +499,15 @@ class SQLContext private[sql](
     new Dataset[T](this, plan)
   }
 
+  def createDataset[T : Encoder](data: RDD[T]): Dataset[T] = {
+    val enc = encoderFor[T]
+    val attributes = enc.schema.toAttributes
+    val encoded = data.map(d => enc.toRow(d))
+    val plan = LogicalRDD(attributes, encoded)(self)
+
+    new Dataset[T](this, plan)
+  }
+
   /**
    * Creates a DataFrame from an RDD[Row]. User can specify whether the input rows should be
    * converted to Catalyst rows.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
index f460a86414c41..f2904e270811e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
@@ -48,6 +48,10 @@ abstract class SQLImplicits {
   implicit def newBooleanEncoder: Encoder[Boolean] = ExpressionEncoder[Boolean](flat = true)
   implicit def newStringEncoder: Encoder[String] = ExpressionEncoder[String](flat = true)
 
+  implicit def rddToDatasetHolder[T : Encoder](rdd: RDD[T]): DatasetHolder[T] = {
+    DatasetHolder(_sqlContext.createDataset(rdd))
+  }
+
   implicit def localSeqToDatasetHolder[T : Encoder](s: Seq[T]): DatasetHolder[T] = {
     DatasetHolder(_sqlContext.createDataset(s))
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 5973fa7f2a76b..3e9b621cfd67f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -34,6 +34,13 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       data: _*)
   }
 
+  test("toDS with RDD") {
+    val ds = sparkContext.makeRDD(Seq("a", "b", "c"), 3).toDS()
+    checkAnswer(
+      ds.mapPartitions(_ => Iterator(1)),
+      1, 1, 1)
+  }
+
   test("as tuple") {
     val data = Seq(("a", 1), ("b", 2)).toDF("a", "b")
     checkAnswer(

From 680b4e7bca935dc1569f35fa319bdfb01a12f7e0 Mon Sep 17 00:00:00 2001
From: Jacek Laskowski <jacek.laskowski@deepsense.io>
Date: Tue, 3 Nov 2015 15:26:35 -0800
Subject: [PATCH 036/518] Fix typo in WebUI

Author: Jacek Laskowski <jacek.laskowski@deepsense.io>

Closes #9444 from jaceklaskowski/TImely-fix.
---
 core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala b/core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala
index 712782d27b3cf..51425e599e748 100644
--- a/core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala
+++ b/core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala
@@ -49,7 +49,7 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
             ("shuffle-read-time-proportion", "Shuffle Read Time"),
             ("executor-runtime-proportion", "Executor Computing Time"),
             ("shuffle-write-time-proportion", "Shuffle Write Time"),
-            ("serialization-time-proportion", "Result Serialization TIme"),
+            ("serialization-time-proportion", "Result Serialization Time"),
             ("getting-result-time-proportion", "Getting Result Time"))
 
           legendPairs.zipWithIndex.map {

From 53e9cee3e4e845d1f875c487215c0f22503347b1 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Tue, 3 Nov 2015 16:26:28 -0800
Subject: [PATCH 037/518] [SPARK-11466][CORE] Avoid mockito in multi-threaded
 FsHistoryProviderSuite test.

The test functionality should be the same, but without using mockito; logs don't
really say anything useful but I suspect it may be the cause of the flakiness,
since updating mocks when multiple threads may be using it doesn't work very
well. It also allows some other cleanup (= less test code in FsHistoryProvider).

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9425 from vanzin/SPARK-11466.
---
 .../deploy/history/FsHistoryProvider.scala    | 31 ++++++--------
 .../history/FsHistoryProviderSuite.scala      | 42 +++++++++----------
 2 files changed, 34 insertions(+), 39 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala b/core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala
index 24aa386c7212b..718efc4f3bd5e 100644
--- a/core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala
@@ -113,35 +113,30 @@ private[history] class FsHistoryProvider(conf: SparkConf, clock: Clock)
   }
 
   // Conf option used for testing the initialization code.
-  val initThread = if (!conf.getBoolean("spark.history.testing.skipInitialize", false)) {
-      initialize(None)
-    } else {
-      null
-    }
+  val initThread = initialize()
 
-  private[history] def initialize(errorHandler: Option[Thread.UncaughtExceptionHandler]): Thread = {
+  private[history] def initialize(): Thread = {
     if (!isFsInSafeMode()) {
       startPolling()
-      return null
+      null
+    } else {
+      startSafeModeCheckThread(None)
     }
+  }
 
+  private[history] def startSafeModeCheckThread(
+      errorHandler: Option[Thread.UncaughtExceptionHandler]): Thread = {
     // Cannot probe anything while the FS is in safe mode, so spawn a new thread that will wait
     // for the FS to leave safe mode before enabling polling. This allows the main history server
     // UI to be shown (so that the user can see the HDFS status).
-    //
-    // The synchronization in the run() method is needed because of the tests; mockito can
-    // misbehave if the test is modifying the mocked methods while the thread is calling
-    // them.
     val initThread = new Thread(new Runnable() {
       override def run(): Unit = {
         try {
-          clock.synchronized {
-            while (isFsInSafeMode()) {
-              logInfo("HDFS is still in safe mode. Waiting...")
-              val deadline = clock.getTimeMillis() +
-                TimeUnit.SECONDS.toMillis(SAFEMODE_CHECK_INTERVAL_S)
-              clock.waitTillTime(deadline)
-            }
+          while (isFsInSafeMode()) {
+            logInfo("HDFS is still in safe mode. Waiting...")
+            val deadline = clock.getTimeMillis() +
+              TimeUnit.SECONDS.toMillis(SAFEMODE_CHECK_INTERVAL_S)
+            clock.waitTillTime(deadline)
           }
           startPolling()
         } catch {
diff --git a/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala b/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala
index 833aab14ca2da..5cab17f8a38f5 100644
--- a/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala
@@ -41,7 +41,7 @@ import org.scalatest.concurrent.Eventually._
 import org.apache.spark.{Logging, SparkConf, SparkFunSuite}
 import org.apache.spark.io._
 import org.apache.spark.scheduler._
-import org.apache.spark.util.{JsonProtocol, ManualClock, Utils}
+import org.apache.spark.util.{Clock, JsonProtocol, ManualClock, Utils}
 
 class FsHistoryProviderSuite extends SparkFunSuite with BeforeAndAfter with Matchers with Logging {
 
@@ -423,22 +423,16 @@ class FsHistoryProviderSuite extends SparkFunSuite with BeforeAndAfter with Matc
 
   test("provider waits for safe mode to finish before initializing") {
     val clock = new ManualClock()
-    val conf = createTestConf().set("spark.history.testing.skipInitialize", "true")
-    val provider = spy(new FsHistoryProvider(conf, clock))
-    doReturn(true).when(provider).isFsInSafeMode()
-
-    val initThread = provider.initialize(None)
+    val provider = new SafeModeTestProvider(createTestConf(), clock)
+    val initThread = provider.initialize()
     try {
       provider.getConfig().keys should contain ("HDFS State")
 
       clock.setTime(5000)
       provider.getConfig().keys should contain ("HDFS State")
 
-      // Synchronization needed because of mockito.
-      clock.synchronized {
-        doReturn(false).when(provider).isFsInSafeMode()
-        clock.setTime(10000)
-      }
+      provider.inSafeMode = false
+      clock.setTime(10000)
 
       eventually(timeout(1 second), interval(10 millis)) {
         provider.getConfig().keys should not contain ("HDFS State")
@@ -451,18 +445,12 @@ class FsHistoryProviderSuite extends SparkFunSuite with BeforeAndAfter with Matc
   test("provider reports error after FS leaves safe mode") {
     testDir.delete()
     val clock = new ManualClock()
-    val conf = createTestConf().set("spark.history.testing.skipInitialize", "true")
-    val provider = spy(new FsHistoryProvider(conf, clock))
-    doReturn(true).when(provider).isFsInSafeMode()
-
+    val provider = new SafeModeTestProvider(createTestConf(), clock)
     val errorHandler = mock(classOf[Thread.UncaughtExceptionHandler])
-    val initThread = provider.initialize(Some(errorHandler))
+    val initThread = provider.startSafeModeCheckThread(Some(errorHandler))
     try {
-      // Synchronization needed because of mockito.
-      clock.synchronized {
-        doReturn(false).when(provider).isFsInSafeMode()
-        clock.setTime(10000)
-      }
+      provider.inSafeMode = false
+      clock.setTime(10000)
 
       eventually(timeout(1 second), interval(10 millis)) {
         verify(errorHandler).uncaughtException(any(), any())
@@ -530,4 +518,16 @@ class FsHistoryProviderSuite extends SparkFunSuite with BeforeAndAfter with Matc
     log
   }
 
+  private class SafeModeTestProvider(conf: SparkConf, clock: Clock)
+    extends FsHistoryProvider(conf, clock) {
+
+    @volatile var inSafeMode = true
+
+    // Skip initialization so that we can manually start the safe mode check thread.
+    private[history] override def initialize(): Thread = null
+
+    private[history] override def isFsInSafeMode(): Boolean = inSafeMode
+
+  }
+
 }

From 5051262d4ca6a2c529c9b1ba86d54cce60a7af17 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 3 Nov 2015 16:27:56 -0800
Subject: [PATCH 038/518] [SPARK-11489][SQL] Only include common first order
 statistics in GroupedData

We added a bunch of higher order statistics such as skewness and kurtosis to GroupedData. I don't think they are common enough to justify being listed, since users can always use the normal statistics aggregate functions.

That is to say, after this change, we won't support
```scala
df.groupBy("key").kurtosis("colA", "colB")
```

However, we will still support
```scala
df.groupBy("key").agg(kurtosis(col("colA")), kurtosis(col("colB")))
```

Author: Reynold Xin <rxin@databricks.com>

Closes #9446 from rxin/SPARK-11489.
---
 python/pyspark/sql/group.py                   |  88 -----------
 .../org/apache/spark/sql/GroupedData.scala    | 146 ++++--------------
 .../apache/spark/sql/JavaDataFrameSuite.java  |   1 -
 3 files changed, 28 insertions(+), 207 deletions(-)

diff --git a/python/pyspark/sql/group.py b/python/pyspark/sql/group.py
index 946b53e71c2c6..71c0bccc5eeff 100644
--- a/python/pyspark/sql/group.py
+++ b/python/pyspark/sql/group.py
@@ -167,94 +167,6 @@ def sum(self, *cols):
         [Row(sum(age)=7, sum(height)=165)]
         """
 
-    @df_varargs_api
-    @since(1.6)
-    def stddev(self, *cols):
-        """Compute the sample standard deviation for each numeric columns for each group.
-
-        :param cols: list of column names (string). Non-numeric columns are ignored.
-
-        >>> df3.groupBy().stddev('age', 'height').collect()
-        [Row(STDDEV(age)=2.12..., STDDEV(height)=3.53...)]
-        """
-
-    @df_varargs_api
-    @since(1.6)
-    def stddev_samp(self, *cols):
-        """Compute the sample standard deviation for each numeric columns for each group.
-
-        :param cols: list of column names (string). Non-numeric columns are ignored.
-
-        >>> df3.groupBy().stddev_samp('age', 'height').collect()
-        [Row(STDDEV_SAMP(age)=2.12..., STDDEV_SAMP(height)=3.53...)]
-        """
-
-    @df_varargs_api
-    @since(1.6)
-    def stddev_pop(self, *cols):
-        """Compute the population standard deviation for each numeric columns for each group.
-
-        :param cols: list of column names (string). Non-numeric columns are ignored.
-
-        >>> df3.groupBy().stddev_pop('age', 'height').collect()
-        [Row(STDDEV_POP(age)=1.5, STDDEV_POP(height)=2.5)]
-        """
-
-    @df_varargs_api
-    @since(1.6)
-    def variance(self, *cols):
-        """Compute the sample variance for each numeric columns for each group.
-
-        :param cols: list of column names (string). Non-numeric columns are ignored.
-
-        >>> df3.groupBy().variance('age', 'height').collect()
-        [Row(VARIANCE(age)=2.25, VARIANCE(height)=6.25)]
-        """
-
-    @df_varargs_api
-    @since(1.6)
-    def var_pop(self, *cols):
-        """Compute the sample variance for each numeric columns for each group.
-
-        :param cols: list of column names (string). Non-numeric columns are ignored.
-
-        >>> df3.groupBy().var_pop('age', 'height').collect()
-        [Row(VAR_POP(age)=2.25, VAR_POP(height)=6.25)]
-        """
-
-    @df_varargs_api
-    @since(1.6)
-    def var_samp(self, *cols):
-        """Compute the sample variance for each numeric columns for each group.
-
-        :param cols: list of column names (string). Non-numeric columns are ignored.
-
-        >>> df3.groupBy().var_samp('age', 'height').collect()
-        [Row(VAR_SAMP(age)=4.5, VAR_SAMP(height)=12.5)]
-        """
-
-    @df_varargs_api
-    @since(1.6)
-    def skewness(self, *cols):
-        """Compute the skewness for each numeric columns for each group.
-
-        :param cols: list of column names (string). Non-numeric columns are ignored.
-
-        >>> df3.groupBy().skewness('age', 'height').collect()
-        [Row(SKEWNESS(age)=0.0, SKEWNESS(height)=0.0)]
-        """
-
-    @df_varargs_api
-    @since(1.6)
-    def kurtosis(self, *cols):
-        """Compute the kurtosis for each numeric columns for each group.
-
-        :param cols: list of column names (string). Non-numeric columns are ignored.
-
-        >>> df3.groupBy().kurtosis('age', 'height').collect()
-        [Row(KURTOSIS(age)=-2.0, KURTOSIS(height)=-2.0)]
-        """
-
 
 def _test():
     import doctest
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
index dc96384a4d28d..c2b2a4013d510 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
@@ -26,42 +26,14 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Rollup, Cube, Aggregate}
 import org.apache.spark.sql.types.NumericType
 
-/**
- * Companion object for GroupedData
- */
-private[sql] object GroupedData {
-  def apply(
-      df: DataFrame,
-      groupingExprs: Seq[Expression],
-      groupType: GroupType): GroupedData = {
-    new GroupedData(df, groupingExprs, groupType: GroupType)
-  }
-
-  /**
-   * The Grouping Type
-   */
-  private[sql] trait GroupType
-
-  /**
-   * To indicate it's the GroupBy
-   */
-  private[sql] object GroupByType extends GroupType
-
-  /**
-   * To indicate it's the CUBE
-   */
-  private[sql] object CubeType extends GroupType
-
-  /**
-   * To indicate it's the ROLLUP
-   */
-  private[sql] object RollupType extends GroupType
-}
 
 /**
  * :: Experimental ::
  * A set of methods for aggregations on a [[DataFrame]], created by [[DataFrame.groupBy]].
  *
+ * The main method is the agg function, which has multiple variants. This class also contains
+ * convenience some first order statistics such as mean, sum for convenience.
+ *
  * @since 1.3.0
  */
 @Experimental
@@ -124,7 +96,7 @@ class GroupedData protected[sql](
       case "avg" | "average" | "mean" => Average
       case "max" => Max
       case "min" => Min
-      case "stddev" => Stddev
+      case "stddev" | "std" => Stddev
       case "stddev_pop" => StddevPop
       case "stddev_samp" => StddevSamp
       case "variance" => Variance
@@ -255,30 +227,6 @@ class GroupedData protected[sql](
     aggregateNumericColumns(colNames : _*)(Average)
   }
 
-  /**
-   * Compute the skewness for each numeric columns for each group.
-   * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the skewness values for them.
-   *
-   * @since 1.6.0
-   */
-  @scala.annotation.varargs
-  def skewness(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(Skewness)
-  }
-
-  /**
-   * Compute the kurtosis for each numeric columns for each group.
-   * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the kurtosis values for them.
-   *
-   * @since 1.6.0
-   */
-  @scala.annotation.varargs
-  def kurtosis(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(Kurtosis)
-  }
-
   /**
    * Compute the max value for each numeric columns for each group.
    * The resulting [[DataFrame]] will also contain the grouping columns.
@@ -316,86 +264,48 @@ class GroupedData protected[sql](
   }
 
   /**
-   * Compute the sample standard deviation for each numeric columns for each group.
+   * Compute the sum for each numeric columns for each group.
    * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the stddev for them.
+   * When specified columns are given, only compute the sum for them.
    *
-   * @since 1.6.0
+   * @since 1.3.0
    */
   @scala.annotation.varargs
-  def stddev(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(Stddev)
+  def sum(colNames: String*): DataFrame = {
+    aggregateNumericColumns(colNames : _*)(Sum)
   }
+}
 
-  /**
-   * Compute the population standard deviation for each numeric columns for each group.
-   * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the stddev for them.
-   *
-   * @since 1.6.0
-   */
-  @scala.annotation.varargs
-  def stddev_pop(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(StddevPop)
-  }
 
-  /**
-   * Compute the sample standard deviation for each numeric columns for each group.
-   * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the stddev for them.
-   *
-   * @since 1.6.0
-   */
-  @scala.annotation.varargs
-  def stddev_samp(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(StddevSamp)
+/**
+ * Companion object for GroupedData.
+ */
+private[sql] object GroupedData {
+
+  def apply(
+      df: DataFrame,
+      groupingExprs: Seq[Expression],
+      groupType: GroupType): GroupedData = {
+    new GroupedData(df, groupingExprs, groupType: GroupType)
   }
 
   /**
-   * Compute the sum for each numeric columns for each group.
-   * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the sum for them.
-   *
-   * @since 1.3.0
+   * The Grouping Type
    */
-  @scala.annotation.varargs
-  def sum(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(Sum)
-  }
+  private[sql] trait GroupType
 
   /**
-   * Compute the sample variance for each numeric columns for each group.
-   * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the variance for them.
-   *
-   * @since 1.6.0
+   * To indicate it's the GroupBy
    */
-  @scala.annotation.varargs
-  def variance(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(Variance)
-  }
+  private[sql] object GroupByType extends GroupType
 
   /**
-   * Compute the population variance for each numeric columns for each group.
-   * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the variance for them.
-   *
-   * @since 1.6.0
+   * To indicate it's the CUBE
    */
-  @scala.annotation.varargs
-  def var_pop(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(VariancePop)
-  }
+  private[sql] object CubeType extends GroupType
 
   /**
-   * Compute the sample variance for each numeric columns for each group.
-   * The resulting [[DataFrame]] will also contain the grouping columns.
-   * When specified columns are given, only compute the variance for them.
-   *
-   * @since 1.6.0
+   * To indicate it's the ROLLUP
    */
-  @scala.annotation.varargs
-  def var_samp(colNames: String*): DataFrame = {
-    aggregateNumericColumns(colNames : _*)(VarianceSamp)
-  }
+  private[sql] object RollupType extends GroupType
 }
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index a1a3fdbb486ea..49f516e86d754 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -91,7 +91,6 @@ public void testVarargMethods() {
     df.groupBy().mean("key");
     df.groupBy().max("key");
     df.groupBy().min("key");
-    df.groupBy().stddev("key");
     df.groupBy().sum("key");
 
     // Varargs in column expressions

From d648a4ad546eb05deab1005e92b815b2cbea621b Mon Sep 17 00:00:00 2001
From: lewuathe <lewuathe@me.com>
Date: Tue, 3 Nov 2015 16:38:22 -0800
Subject: [PATCH 039/518] [DOC] Missing link to R DataFrame API doc

Author: lewuathe <lewuathe@me.com>
Author: Lewuathe <lewuathe@me.com>

Closes #9394 from Lewuathe/missing-link-to-R-dataframe.
---
 R/pkg/R/DataFrame.R           | 105 +++++++++++++++++++++++++++++++---
 docs/sql-programming-guide.md |   2 +-
 2 files changed, 98 insertions(+), 9 deletions(-)

diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index 87a2c66ffd2a9..df5bc8137187b 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -23,15 +23,23 @@ NULL
 setOldClass("jobj")
 
 #' @title S4 class that represents a DataFrame
-#' @description DataFrames can be created using functions like
-#'              \code{jsonFile}, \code{table} etc.
+#' @description DataFrames can be created using functions like \link{createDataFrame},
+#'              \link{jsonFile}, \link{table} etc.
+#' @family dataframe_funcs
 #' @rdname DataFrame
-#' @seealso jsonFile, table
 #' @docType class
 #'
 #' @slot env An R environment that stores bookkeeping states of the DataFrame
 #' @slot sdf A Java object reference to the backing Scala DataFrame
+#' @seealso \link{createDataFrame}, \link{jsonFile}, \link{table}
+#' @seealso \url{https://spark.apache.org/docs/latest/sparkr.html#sparkr-dataframes}
 #' @export
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' sqlContext <- sparkRSQL.init(sc)
+#' df <- createDataFrame(sqlContext, faithful)
+#'}
 setClass("DataFrame",
          slots = list(env = "environment",
                       sdf = "jobj"))
@@ -46,7 +54,6 @@ setMethod("initialize", "DataFrame", function(.Object, sdf, isCached) {
 
 #' @rdname DataFrame
 #' @export
-#'
 #' @param sdf A Java object reference to the backing Scala DataFrame
 #' @param isCached TRUE if the dataFrame is cached
 dataFrame <- function(sdf, isCached = FALSE) {
@@ -61,6 +68,7 @@ dataFrame <- function(sdf, isCached = FALSE) {
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname printSchema
 #' @name printSchema
 #' @export
@@ -85,6 +93,7 @@ setMethod("printSchema",
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname schema
 #' @name schema
 #' @export
@@ -108,6 +117,7 @@ setMethod("schema",
 #'
 #' @param x A SparkSQL DataFrame
 #' @param extended Logical. If extended is False, explain() only prints the physical plan.
+#' @family dataframe_funcs
 #' @rdname explain
 #' @name explain
 #' @export
@@ -138,6 +148,7 @@ setMethod("explain",
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname isLocal
 #' @name isLocal
 #' @export
@@ -162,6 +173,7 @@ setMethod("isLocal",
 #' @param x A SparkSQL DataFrame
 #' @param numRows The number of rows to print. Defaults to 20.
 #'
+#' @family dataframe_funcs
 #' @rdname showDF
 #' @name showDF
 #' @export
@@ -186,6 +198,7 @@ setMethod("showDF",
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname show
 #' @name show
 #' @export
@@ -212,6 +225,7 @@ setMethod("show", "DataFrame",
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname dtypes
 #' @name dtypes
 #' @export
@@ -237,6 +251,7 @@ setMethod("dtypes",
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname columns
 #' @name columns
 #' @aliases names
@@ -257,6 +272,7 @@ setMethod("columns",
             })
           })
 
+#' @family dataframe_funcs
 #' @rdname columns
 #' @name names
 setMethod("names",
@@ -265,6 +281,7 @@ setMethod("names",
             columns(x)
           })
 
+#' @family dataframe_funcs
 #' @rdname columns
 #' @name names<-
 setMethod("names<-",
@@ -283,6 +300,7 @@ setMethod("names<-",
 #' @param x A SparkSQL DataFrame
 #' @param tableName A character vector containing the name of the table
 #'
+#' @family dataframe_funcs
 #' @rdname registerTempTable
 #' @name registerTempTable
 #' @export
@@ -310,6 +328,7 @@ setMethod("registerTempTable",
 #' @param overwrite A logical argument indicating whether or not to overwrite
 #' the existing rows in the table.
 #'
+#' @family dataframe_funcs
 #' @rdname insertInto
 #' @name insertInto
 #' @export
@@ -334,6 +353,7 @@ setMethod("insertInto",
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname cache
 #' @name cache
 #' @export
@@ -360,6 +380,8 @@ setMethod("cache",
 #' \url{http://spark.apache.org/docs/latest/programming-guide.html#rdd-persistence}.
 #'
 #' @param x The DataFrame to persist
+#'
+#' @family dataframe_funcs
 #' @rdname persist
 #' @name persist
 #' @export
@@ -386,6 +408,8 @@ setMethod("persist",
 #'
 #' @param x The DataFrame to unpersist
 #' @param blocking Whether to block until all blocks are deleted
+#'
+#' @family dataframe_funcs
 #' @rdname unpersist-methods
 #' @name unpersist
 #' @export
@@ -412,6 +436,8 @@ setMethod("unpersist",
 #'
 #' @param x A SparkSQL DataFrame
 #' @param numPartitions The number of partitions to use.
+#'
+#' @family dataframe_funcs
 #' @rdname repartition
 #' @name repartition
 #' @export
@@ -435,8 +461,10 @@ setMethod("repartition",
 # Convert the rows of a DataFrame into JSON objects and return an RDD where
 # each element contains a JSON string.
 #
-#@param x A SparkSQL DataFrame
+# @param x A SparkSQL DataFrame
 # @return A StringRRDD of JSON objects
+#
+# @family dataframe_funcs
 # @rdname tojson
 # @export
 # @examples
@@ -462,6 +490,8 @@ setMethod("toJSON",
 #'
 #' @param x A SparkSQL DataFrame
 #' @param path The directory where the file is saved
+#'
+#' @family dataframe_funcs
 #' @rdname saveAsParquetFile
 #' @name saveAsParquetFile
 #' @export
@@ -484,6 +514,8 @@ setMethod("saveAsParquetFile",
 #' Return a new DataFrame containing the distinct rows in this DataFrame.
 #'
 #' @param x A SparkSQL DataFrame
+#'
+#' @family dataframe_funcs
 #' @rdname distinct
 #' @name distinct
 #' @export
@@ -506,6 +538,7 @@ setMethod("distinct",
 #
 #' @description Returns a new DataFrame containing distinct rows in this DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname unique
 #' @name unique
 #' @aliases distinct
@@ -522,6 +555,8 @@ setMethod("unique",
 #' @param x A SparkSQL DataFrame
 #' @param withReplacement Sampling with replacement or not
 #' @param fraction The (rough) sample target fraction
+#'
+#' @family dataframe_funcs
 #' @rdname sample
 #' @aliases sample_frac
 #' @export
@@ -545,6 +580,7 @@ setMethod("sample",
             dataFrame(sdf)
           })
 
+#' @family dataframe_funcs
 #' @rdname sample
 #' @name sample_frac
 setMethod("sample_frac",
@@ -560,6 +596,7 @@ setMethod("sample_frac",
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname count
 #' @name count
 #' @aliases nrow
@@ -583,6 +620,7 @@ setMethod("count",
 #'
 #' @name nrow
 #'
+#' @family dataframe_funcs
 #' @rdname nrow
 #' @aliases count
 setMethod("nrow",
@@ -595,6 +633,7 @@ setMethod("nrow",
 #'
 #' @param x a SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname ncol
 #' @name ncol
 #' @export
@@ -615,6 +654,7 @@ setMethod("ncol",
 #' Returns the dimentions (number of rows and columns) of a DataFrame
 #' @param x a SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname dim
 #' @name dim
 #' @export
@@ -637,6 +677,8 @@ setMethod("dim",
 #' @param x A SparkSQL DataFrame
 #' @param stringsAsFactors (Optional) A logical indicating whether or not string columns
 #' should be converted to factors. FALSE by default.
+#'
+#' @family dataframe_funcs
 #' @rdname collect
 #' @name collect
 #' @export
@@ -704,6 +746,7 @@ setMethod("collect",
 #' @param num The number of rows to return
 #' @return A new DataFrame containing the number of rows specified.
 #'
+#' @family dataframe_funcs
 #' @rdname limit
 #' @name limit
 #' @export
@@ -724,6 +767,7 @@ setMethod("limit",
 
 #' Take the first NUM rows of a DataFrame and return a the results as a data.frame
 #'
+#' @family dataframe_funcs
 #' @rdname take
 #' @name take
 #' @export
@@ -752,6 +796,7 @@ setMethod("take",
 #' @param num The number of rows to return. Default is 6.
 #' @return A data.frame
 #'
+#' @family dataframe_funcs
 #' @rdname head
 #' @name head
 #' @export
@@ -774,6 +819,7 @@ setMethod("head",
 #'
 #' @param x A SparkSQL DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname first
 #' @name first
 #' @export
@@ -797,6 +843,7 @@ setMethod("first",
 #
 # @param x A Spark DataFrame
 #
+# @family dataframe_funcs
 # @rdname DataFrame
 # @export
 # @examples
@@ -827,6 +874,7 @@ setMethod("toRDD",
 #' @return a GroupedData
 #' @seealso GroupedData
 #' @aliases group_by
+#' @family dataframe_funcs
 #' @rdname groupBy
 #' @name groupBy
 #' @export
@@ -851,6 +899,7 @@ setMethod("groupBy",
              groupedData(sgd)
            })
 
+#' @family dataframe_funcs
 #' @rdname groupBy
 #' @name group_by
 setMethod("group_by",
@@ -864,6 +913,7 @@ setMethod("group_by",
 #' Compute aggregates by specifying a list of columns
 #'
 #' @param x a DataFrame
+#' @family dataframe_funcs
 #' @rdname agg
 #' @name agg
 #' @aliases summarize
@@ -874,6 +924,7 @@ setMethod("agg",
             agg(groupBy(x), ...)
           })
 
+#' @family dataframe_funcs
 #' @rdname agg
 #' @name summarize
 setMethod("summarize",
@@ -889,6 +940,7 @@ setMethod("summarize",
 # the requested map function.                                                     #
 ###################################################################################
 
+# @family dataframe_funcs
 # @rdname lapply
 setMethod("lapply",
           signature(X = "DataFrame", FUN = "function"),
@@ -897,6 +949,7 @@ setMethod("lapply",
             lapply(rdd, FUN)
           })
 
+# @family dataframe_funcs
 # @rdname lapply
 setMethod("map",
           signature(X = "DataFrame", FUN = "function"),
@@ -904,6 +957,7 @@ setMethod("map",
             lapply(X, FUN)
           })
 
+# @family dataframe_funcs
 # @rdname flatMap
 setMethod("flatMap",
           signature(X = "DataFrame", FUN = "function"),
@@ -911,7 +965,7 @@ setMethod("flatMap",
             rdd <- toRDD(X)
             flatMap(rdd, FUN)
           })
-
+# @family dataframe_funcs
 # @rdname lapplyPartition
 setMethod("lapplyPartition",
           signature(X = "DataFrame", FUN = "function"),
@@ -920,6 +974,7 @@ setMethod("lapplyPartition",
             lapplyPartition(rdd, FUN)
           })
 
+# @family dataframe_funcs
 # @rdname lapplyPartition
 setMethod("mapPartitions",
           signature(X = "DataFrame", FUN = "function"),
@@ -927,6 +982,7 @@ setMethod("mapPartitions",
             lapplyPartition(X, FUN)
           })
 
+# @family dataframe_funcs
 # @rdname foreach
 setMethod("foreach",
           signature(x = "DataFrame", func = "function"),
@@ -935,6 +991,7 @@ setMethod("foreach",
             foreach(rdd, func)
           })
 
+# @family dataframe_funcs
 # @rdname foreach
 setMethod("foreachPartition",
           signature(x = "DataFrame", func = "function"),
@@ -1034,6 +1091,7 @@ setMethod("[", signature(x = "DataFrame", i = "Column"),
 #' @param select expression for the single Column or a list of columns to select from the DataFrame
 #' @return A new DataFrame containing only the rows that meet the condition with selected columns
 #' @export
+#' @family dataframe_funcs
 #' @rdname subset
 #' @name subset
 #' @aliases [
@@ -1064,6 +1122,7 @@ setMethod("subset", signature(x = "DataFrame"),
 #' @param col A list of columns or single Column or name
 #' @return A new DataFrame with selected columns
 #' @export
+#' @family dataframe_funcs
 #' @rdname select
 #' @name select
 #' @family subsetting functions
@@ -1091,6 +1150,7 @@ setMethod("select", signature(x = "DataFrame", col = "character"),
             }
           })
 
+#' @family dataframe_funcs
 #' @rdname select
 #' @export
 setMethod("select", signature(x = "DataFrame", col = "Column"),
@@ -1102,6 +1162,7 @@ setMethod("select", signature(x = "DataFrame", col = "Column"),
             dataFrame(sdf)
           })
 
+#' @family dataframe_funcs
 #' @rdname select
 #' @export
 setMethod("select",
@@ -1126,6 +1187,7 @@ setMethod("select",
 #' @param expr A string containing a SQL expression
 #' @param ... Additional expressions
 #' @return A DataFrame
+#' @family dataframe_funcs
 #' @rdname selectExpr
 #' @name selectExpr
 #' @export
@@ -1153,6 +1215,7 @@ setMethod("selectExpr",
 #' @param colName A string containing the name of the new column.
 #' @param col A Column expression.
 #' @return A DataFrame with the new column added.
+#' @family dataframe_funcs
 #' @rdname withColumn
 #' @name withColumn
 #' @aliases mutate transform
@@ -1178,6 +1241,7 @@ setMethod("withColumn",
 #' @param .data A DataFrame
 #' @param col a named argument of the form name = col
 #' @return A new DataFrame with the new columns added.
+#' @family dataframe_funcs
 #' @rdname withColumn
 #' @name mutate
 #' @aliases withColumn transform
@@ -1211,6 +1275,7 @@ setMethod("mutate",
           })
 
 #' @export
+#' @family dataframe_funcs
 #' @rdname withColumn
 #' @name transform
 #' @aliases withColumn mutate
@@ -1228,6 +1293,7 @@ setMethod("transform",
 #' @param existingCol The name of the column you want to change.
 #' @param newCol The new column name.
 #' @return A DataFrame with the column name changed.
+#' @family dataframe_funcs
 #' @rdname withColumnRenamed
 #' @name withColumnRenamed
 #' @export
@@ -1259,6 +1325,7 @@ setMethod("withColumnRenamed",
 #' @param x A DataFrame
 #' @param newCol A named pair of the form new_column_name = existing_column
 #' @return A DataFrame with the column name changed.
+#' @family dataframe_funcs
 #' @rdname withColumnRenamed
 #' @name rename
 #' @aliases withColumnRenamed
@@ -1303,6 +1370,7 @@ setClassUnion("characterOrColumn", c("character", "Column"))
 #' @param decreasing A logical argument indicating sorting order for columns when
 #'                   a character vector is specified for col
 #' @return A DataFrame where all elements are sorted.
+#' @family dataframe_funcs
 #' @rdname arrange
 #' @name arrange
 #' @aliases orderby
@@ -1329,6 +1397,7 @@ setMethod("arrange",
             dataFrame(sdf)
           })
 
+#' @family dataframe_funcs
 #' @rdname arrange
 #' @export
 setMethod("arrange",
@@ -1360,6 +1429,7 @@ setMethod("arrange",
             do.call("arrange", c(x, jcols))
           })
 
+#' @family dataframe_funcs
 #' @rdname arrange
 #' @name orderby
 setMethod("orderBy",
@@ -1376,6 +1446,7 @@ setMethod("orderBy",
 #' @param condition The condition to filter on. This may either be a Column expression
 #' or a string containing a SQL statement
 #' @return A DataFrame containing only the rows that meet the condition.
+#' @family dataframe_funcs
 #' @rdname filter
 #' @name filter
 #' @family subsetting functions
@@ -1399,6 +1470,7 @@ setMethod("filter",
             dataFrame(sdf)
           })
 
+#' @family dataframe_funcs
 #' @rdname filter
 #' @name where
 setMethod("where",
@@ -1419,6 +1491,7 @@ setMethod("where",
 #' 'inner', 'outer', 'full', 'fullouter', leftouter', 'left_outer', 'left',
 #' 'right_outer', 'rightouter', 'right', and 'leftsemi'. The default joinType is "inner".
 #' @return A DataFrame containing the result of the join operation.
+#' @family dataframe_funcs
 #' @rdname join
 #' @name join
 #' @export
@@ -1477,6 +1550,7 @@ setMethod("join",
 #'   be returned. If all.x is set to FALSE and all.y is set to TRUE, a right
 #'   outer join will be returned. If all.x and all.y are set to TRUE, a full
 #'   outer join will be returned.
+#' @family dataframe_funcs
 #' @rdname merge
 #' @export
 #' @examples
@@ -1608,6 +1682,7 @@ generateAliasesForIntersectedCols <- function (x, intersectedColNames, suffix) {
 #' @param x A Spark DataFrame
 #' @param y A Spark DataFrame
 #' @return A DataFrame containing the result of the union.
+#' @family dataframe_funcs
 #' @rdname unionAll
 #' @name unionAll
 #' @export
@@ -1627,9 +1702,10 @@ setMethod("unionAll",
           })
 
 #' @title Union two or more DataFrames
-#
+#'
 #' @description Returns a new DataFrame containing rows of all parameters.
-#
+#'
+#' @family dataframe_funcs
 #' @rdname rbind
 #' @name rbind
 #' @aliases unionAll
@@ -1651,6 +1727,7 @@ setMethod("rbind",
 #' @param x A Spark DataFrame
 #' @param y A Spark DataFrame
 #' @return A DataFrame containing the result of the intersect.
+#' @family dataframe_funcs
 #' @rdname intersect
 #' @name intersect
 #' @export
@@ -1677,6 +1754,7 @@ setMethod("intersect",
 #' @param x A Spark DataFrame
 #' @param y A Spark DataFrame
 #' @return A DataFrame containing the result of the except operation.
+#' @family dataframe_funcs
 #' @rdname except
 #' @name except
 #' @export
@@ -1716,6 +1794,7 @@ setMethod("except",
 #' @param source A name for external data source
 #' @param mode One of 'append', 'overwrite', 'error', 'ignore' save mode
 #'
+#' @family dataframe_funcs
 #' @rdname write.df
 #' @name write.df
 #' @aliases saveDF
@@ -1751,6 +1830,7 @@ setMethod("write.df",
             callJMethod(df@sdf, "save", source, jmode, options)
           })
 
+#' @family dataframe_funcs
 #' @rdname write.df
 #' @name saveDF
 #' @export
@@ -1781,6 +1861,7 @@ setMethod("saveDF",
 #' @param source A name for external data source
 #' @param mode One of 'append', 'overwrite', 'error', 'ignore' save mode
 #'
+#' @family dataframe_funcs
 #' @rdname saveAsTable
 #' @name saveAsTable
 #' @export
@@ -1821,6 +1902,7 @@ setMethod("saveAsTable",
 #' @param col A string of name
 #' @param ... Additional expressions
 #' @return A DataFrame
+#' @family dataframe_funcs
 #' @rdname describe
 #' @name describe
 #' @aliases summary
@@ -1843,6 +1925,7 @@ setMethod("describe",
             dataFrame(sdf)
           })
 
+#' @family dataframe_funcs
 #' @rdname describe
 #' @name describe
 setMethod("describe",
@@ -1857,6 +1940,7 @@ setMethod("describe",
 #'
 #' @description Computes statistics for numeric columns of the DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname summary
 #' @name summary
 setMethod("summary",
@@ -1881,6 +1965,7 @@ setMethod("summary",
 #' @param cols Optional list of column names to consider.
 #' @return A DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname nafunctions
 #' @name dropna
 #' @aliases na.omit
@@ -1910,6 +1995,7 @@ setMethod("dropna",
             dataFrame(sdf)
           })
 
+#' @family dataframe_funcs
 #' @rdname nafunctions
 #' @name na.omit
 #' @export
@@ -1937,6 +2023,7 @@ setMethod("na.omit",
 #'             column is simply ignored.
 #' @return A DataFrame
 #'
+#' @family dataframe_funcs
 #' @rdname nafunctions
 #' @name fillna
 #' @export
@@ -2000,6 +2087,7 @@ setMethod("fillna",
 #' @title Download data from a DataFrame into a data.frame
 #' @param x a DataFrame
 #' @return a data.frame
+#' @family dataframe_funcs
 #' @rdname as.data.frame
 #' @examples \dontrun{
 #'
@@ -2020,6 +2108,7 @@ setMethod("as.data.frame",
 #' the DataFrame is searched by R when evaluating a variable, so columns in
 #' the DataFrame can be accessed by simply giving their names.
 #'
+#' @family dataframe_funcs
 #' @rdname attach
 #' @title Attach DataFrame to R search path
 #' @param what (DataFrame) The DataFrame to attach
diff --git a/docs/sql-programming-guide.md b/docs/sql-programming-guide.md
index f07c9573696ed..510b3599721a3 100644
--- a/docs/sql-programming-guide.md
+++ b/docs/sql-programming-guide.md
@@ -160,7 +160,7 @@ showDF(df)
 
 ## DataFrame Operations
 
-DataFrames provide a domain-specific language for structured data manipulation in [Scala](api/scala/index.html#org.apache.spark.sql.DataFrame), [Java](api/java/index.html?org/apache/spark/sql/DataFrame.html), and [Python](api/python/pyspark.sql.html#pyspark.sql.DataFrame).
+DataFrames provide a domain-specific language for structured data manipulation in [Scala](api/scala/index.html#org.apache.spark.sql.DataFrame), [Java](api/java/index.html?org/apache/spark/sql/DataFrame.html), [Python](api/python/pyspark.sql.html#pyspark.sql.DataFrame) and [R](api/R/DataFrame.html).
 
 Here we include some basic examples of structured data processing using DataFrames:
 

From e352de0db2789919e1e0385b79f29b508a6b2b77 Mon Sep 17 00:00:00 2001
From: Nong <nong@cloudera.com>
Date: Tue, 3 Nov 2015 16:44:37 -0800
Subject: [PATCH 040/518] [SPARK-11329] [SQL] Cleanup from spark-11329 fix.

Author: Nong <nong@cloudera.com>

Closes #9442 from nongli/spark-11483.
---
 .../apache/spark/sql/catalyst/SqlParser.scala |  4 +-
 .../sql/catalyst/analysis/unresolved.scala    | 18 +----
 .../scala/org/apache/spark/sql/Column.scala   |  6 +-
 .../org/apache/spark/sql/SQLQuerySuite.scala  | 79 +++++++++++--------
 4 files changed, 55 insertions(+), 52 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
index 1ba559d9e3b18..440e9e28fa783 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
@@ -477,8 +477,8 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
 
   protected lazy val baseExpression: Parser[Expression] =
     ( "*" ^^^ UnresolvedStar(None)
-    | (ident <~ "."). + <~ "*" ^^ { case target => { UnresolvedStar(Option(target)) }
-    } | primary
+    | (ident <~ "."). + <~ "*" ^^ { case target => UnresolvedStar(Option(target))}
+    | primary
    )
 
   protected lazy val signedPrimary: Parser[Expression] =
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
index 6975662e2b738..eae17c86ddc7a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
@@ -183,28 +183,16 @@ case class UnresolvedStar(target: Option[Seq[String]]) extends Star with Unevalu
       case None => input.output
       // If there is a table, pick out attributes that are part of this table.
       case Some(t) => if (t.size == 1) {
-        input.output.filter(_.qualifiers.filter(resolver(_, t.head)).nonEmpty)
+        input.output.filter(_.qualifiers.exists(resolver(_, t.head)))
       } else {
         List()
       }
     }
-    if (!expandedAttributes.isEmpty) {
-      if (expandedAttributes.forall(_.isInstanceOf[NamedExpression])) {
-        return expandedAttributes
-      } else {
-        require(expandedAttributes.size == input.output.size)
-        expandedAttributes.zip(input.output).map {
-          case (e, originalAttribute) =>
-            Alias(e, originalAttribute.name)(qualifiers = originalAttribute.qualifiers)
-        }
-      }
-      return expandedAttributes
-    }
-
-    require(target.isDefined)
+    if (expandedAttributes.nonEmpty) return expandedAttributes
 
     // Try to resolve it as a struct expansion. If there is a conflict and both are possible,
     // (i.e. [name].* is both a table and a struct), the struct path can always be qualified.
+    require(target.isDefined)
     val attribute = input.resolve(target.get, resolver)
     if (attribute.isDefined) {
       // This target resolved to an attribute in child. It must be a struct. Expand it.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index 3cde9d6cb4708..c73f696962de5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -60,8 +60,10 @@ class Column(protected[sql] val expr: Expression) extends Logging {
 
   def this(name: String) = this(name match {
     case "*" => UnresolvedStar(None)
-    case _ if name.endsWith(".*") => UnresolvedStar(Some(UnresolvedAttribute.parseAttributeName(
-      name.substring(0, name.length - 2))))
+    case _ if name.endsWith(".*") => {
+      val parts = UnresolvedAttribute.parseAttributeName(name.substring(0, name.length - 2))
+      UnresolvedStar(Some(parts))
+    }
     case _ => UnresolvedAttribute.quotedString(name)
   })
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index ee54bff24b196..6388a8b9c3720 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.DefaultParserDialect
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry
 import org.apache.spark.sql.catalyst.errors.DialectException
 import org.apache.spark.sql.execution.aggregate
-import org.apache.spark.sql.execution.joins.{SortMergeJoin, CartesianProduct}
+import org.apache.spark.sql.execution.joins.{CartesianProduct, SortMergeJoin}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SQLTestData._
 import org.apache.spark.sql.test.{SharedSQLContext, TestSQLContext}
@@ -1956,7 +1956,8 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
 
     // Try with a registered table.
     sql("select struct(a, b) as record from testData2").registerTempTable("structTable")
-    checkAnswer(sql("SELECT record.* FROM structTable"),
+    checkAnswer(
+      sql("SELECT record.* FROM structTable"),
       Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
 
     checkAnswer(sql(
@@ -2019,50 +2020,62 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
       Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
 
     // Try with a registered table
-    nestedStructData.registerTempTable("nestedStructTable")
-    checkAnswer(sql("SELECT record.* FROM nestedStructTable"),
-      nestedStructData.select($"record.*"))
-    checkAnswer(sql("SELECT record.r1 FROM nestedStructTable"),
-      nestedStructData.select($"record.r1"))
-    checkAnswer(sql("SELECT record.r1.* FROM nestedStructTable"),
-      nestedStructData.select($"record.r1.*"))
-
-    // Create paths with unusual characters.
+    withTempTable("nestedStructTable") {
+      nestedStructData.registerTempTable("nestedStructTable")
+      checkAnswer(
+        sql("SELECT record.* FROM nestedStructTable"),
+        nestedStructData.select($"record.*"))
+      checkAnswer(
+        sql("SELECT record.r1 FROM nestedStructTable"),
+        nestedStructData.select($"record.r1"))
+      checkAnswer(
+        sql("SELECT record.r1.* FROM nestedStructTable"),
+        nestedStructData.select($"record.r1.*"))
+
+      // Try resolving something not there.
+      assert(intercept[AnalysisException](sql("SELECT abc.* FROM nestedStructTable"))
+        .getMessage.contains("cannot resolve"))
+    }
+
+    // Create paths with unusual characters
     val specialCharacterPath = sql(
       """
         | SELECT struct(`col$.a_`, `a.b.c.`) as `r&&b.c` FROM
         |   (SELECT struct(a, b) as `col$.a_`, struct(b, a) as `a.b.c.` FROM testData2) tmp
       """.stripMargin)
-    specialCharacterPath.registerTempTable("specialCharacterTable")
-    checkAnswer(specialCharacterPath.select($"`r&&b.c`.*"),
-      nestedStructData.select($"record.*"))
-    checkAnswer(sql("SELECT `r&&b.c`.`col$.a_` FROM specialCharacterTable"),
-      nestedStructData.select($"record.r1"))
-    checkAnswer(sql("SELECT `r&&b.c`.`a.b.c.` FROM specialCharacterTable"),
-      nestedStructData.select($"record.r2"))
-    checkAnswer(sql("SELECT `r&&b.c`.`col$.a_`.* FROM specialCharacterTable"),
-      nestedStructData.select($"record.r1.*"))
+    withTempTable("specialCharacterTable") {
+      specialCharacterPath.registerTempTable("specialCharacterTable")
+      checkAnswer(
+        specialCharacterPath.select($"`r&&b.c`.*"),
+        nestedStructData.select($"record.*"))
+      checkAnswer(
+        sql("SELECT `r&&b.c`.`col$.a_` FROM specialCharacterTable"),
+        nestedStructData.select($"record.r1"))
+      checkAnswer(
+        sql("SELECT `r&&b.c`.`a.b.c.` FROM specialCharacterTable"),
+        nestedStructData.select($"record.r2"))
+      checkAnswer(
+        sql("SELECT `r&&b.c`.`col$.a_`.* FROM specialCharacterTable"),
+        nestedStructData.select($"record.r1.*"))
+    }
 
     // Try star expanding a scalar. This should fail.
     assert(intercept[AnalysisException](sql("select a.* from testData2")).getMessage.contains(
       "Can only star expand struct data types."))
-
-    // Try resolving something not there.
-    assert(intercept[AnalysisException](sql("SELECT abc.* FROM nestedStructTable"))
-      .getMessage.contains("cannot resolve"))
   }
 
-
   test("Struct Star Expansion - Name conflict") {
     // Create a data set that contains a naming conflict
     val nameConflict = sql("SELECT struct(a, b) as nameConflict, a as a FROM testData2")
-    nameConflict.registerTempTable("nameConflict")
-    // Unqualified should resolve to table.
-    checkAnswer(sql("SELECT nameConflict.* FROM nameConflict"),
-      Row(Row(1, 1), 1) :: Row(Row(1, 2), 1) :: Row(Row(2, 1), 2) :: Row(Row(2, 2), 2) ::
-        Row(Row(3, 1), 3) :: Row(Row(3, 2), 3) :: Nil)
-    // Qualify the struct type with the table name.
-    checkAnswer(sql("SELECT nameConflict.nameConflict.* FROM nameConflict"),
-      Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
+    withTempTable("nameConflict") {
+      nameConflict.registerTempTable("nameConflict")
+      // Unqualified should resolve to table.
+      checkAnswer(sql("SELECT nameConflict.* FROM nameConflict"),
+        Row(Row(1, 1), 1) :: Row(Row(1, 2), 1) :: Row(Row(2, 1), 2) :: Row(Row(2, 2), 2) ::
+          Row(Row(3, 1), 3) :: Row(Row(3, 2), 3) :: Nil)
+      // Qualify the struct type with the table name.
+      checkAnswer(sql("SELECT nameConflict.nameConflict.* FROM nameConflict"),
+        Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
+    }
   }
 }

From 2692bdb7dbf36d6247f595d5fd0cb9cda89e1fdd Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Tue, 3 Nov 2015 20:25:58 -0800
Subject: [PATCH 041/518] [SPARK-11455][SQL] fix case sensitivity of partition
 by

depend on `caseSensitive` to do column name equality check, instead of just `==`

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9410 from cloud-fan/partition.
---
 .../datasources/PartitioningUtils.scala       |  7 ++---
 .../datasources/ResolvedDataSource.scala      | 27 ++++++++++++++-----
 .../sql/execution/datasources/rules.scala     |  6 +++--
 .../org/apache/spark/sql/DataFrameSuite.scala | 10 +++++++
 4 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
index 628c5e18936c5..16dc23661c070 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
@@ -287,10 +287,11 @@ private[sql] object PartitioningUtils {
 
   def validatePartitionColumnDataTypes(
       schema: StructType,
-      partitionColumns: Array[String]): Unit = {
+      partitionColumns: Array[String],
+      caseSensitive: Boolean): Unit = {
 
-    ResolvedDataSource.partitionColumnsSchema(schema, partitionColumns).foreach { field =>
-      field.dataType match {
+    ResolvedDataSource.partitionColumnsSchema(schema, partitionColumns, caseSensitive).foreach {
+      field => field.dataType match {
         case _: AtomicType => // OK
         case _ => throw new AnalysisException(s"Cannot use ${field.dataType} for partition column")
       }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ResolvedDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ResolvedDataSource.scala
index 54beabbf63b5f..86a306b8f941d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ResolvedDataSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ResolvedDataSource.scala
@@ -99,7 +99,8 @@ object ResolvedDataSource extends Logging {
           val maybePartitionsSchema = if (partitionColumns.isEmpty) {
             None
           } else {
-            Some(partitionColumnsSchema(schema, partitionColumns))
+            Some(partitionColumnsSchema(
+              schema, partitionColumns, sqlContext.conf.caseSensitiveAnalysis))
           }
 
           val caseInsensitiveOptions = new CaseInsensitiveMap(options)
@@ -172,14 +173,24 @@ object ResolvedDataSource extends Logging {
 
   def partitionColumnsSchema(
       schema: StructType,
-      partitionColumns: Array[String]): StructType = {
+      partitionColumns: Array[String],
+      caseSensitive: Boolean): StructType = {
+    val equality = columnNameEquality(caseSensitive)
     StructType(partitionColumns.map { col =>
-      schema.find(_.name == col).getOrElse {
+      schema.find(f => equality(f.name, col)).getOrElse {
         throw new RuntimeException(s"Partition column $col not found in schema $schema")
       }
     }).asNullable
   }
 
+  private def columnNameEquality(caseSensitive: Boolean): (String, String) => Boolean = {
+    if (caseSensitive) {
+      org.apache.spark.sql.catalyst.analysis.caseSensitiveResolution
+    } else {
+      org.apache.spark.sql.catalyst.analysis.caseInsensitiveResolution
+    }
+  }
+
   /** Create a [[ResolvedDataSource]] for saving the content of the given DataFrame. */
   def apply(
       sqlContext: SQLContext,
@@ -207,14 +218,18 @@ object ResolvedDataSource extends Logging {
           path.makeQualified(fs.getUri, fs.getWorkingDirectory)
         }
 
-        PartitioningUtils.validatePartitionColumnDataTypes(data.schema, partitionColumns)
+        val caseSensitive = sqlContext.conf.caseSensitiveAnalysis
+        PartitioningUtils.validatePartitionColumnDataTypes(
+          data.schema, partitionColumns, caseSensitive)
 
-        val dataSchema = StructType(data.schema.filterNot(f => partitionColumns.contains(f.name)))
+        val equality = columnNameEquality(caseSensitive)
+        val dataSchema = StructType(
+          data.schema.filterNot(f => partitionColumns.exists(equality(_, f.name))))
         val r = dataSource.createRelation(
           sqlContext,
           Array(outputPath.toString),
           Some(dataSchema.asNullable),
-          Some(partitionColumnsSchema(data.schema, partitionColumns)),
+          Some(partitionColumnsSchema(data.schema, partitionColumns, caseSensitive)),
           caseInsensitiveOptions)
 
         // For partitioned relation r, r.schema's column ordering can be different from the column
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/rules.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/rules.scala
index abc016bf020d9..1a8e7ab202dc2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/rules.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/rules.scala
@@ -140,7 +140,8 @@ private[sql] case class PreWriteCheck(catalog: Catalog) extends (LogicalPlan =>
           // OK
         }
 
-        PartitioningUtils.validatePartitionColumnDataTypes(r.schema, part.keySet.toArray)
+        PartitioningUtils.validatePartitionColumnDataTypes(
+          r.schema, part.keySet.toArray, catalog.conf.caseSensitiveAnalysis)
 
         // Get all input data source relations of the query.
         val srcRelations = query.collect {
@@ -190,7 +191,8 @@ private[sql] case class PreWriteCheck(catalog: Catalog) extends (LogicalPlan =>
           // OK
         }
 
-        PartitioningUtils.validatePartitionColumnDataTypes(query.schema, partitionColumns)
+        PartitioningUtils.validatePartitionColumnDataTypes(
+          query.schema, partitionColumns, catalog.conf.caseSensitiveAnalysis)
 
       case _ => // OK
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index a883bcb7b1012..a9e6413423118 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -1118,4 +1118,14 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
       if (!allSequential) throw new SparkException("Partition should contain all sequential values")
     })
   }
+
+  test("fix case sensitivity of partition by") {
+    withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") {
+      withTempPath { path =>
+        val p = path.getAbsolutePath
+        Seq(2012 -> "a").toDF("year", "val").write.partitionBy("yEAr").parquet(p)
+        checkAnswer(sqlContext.read.parquet(p).select("YeaR"), Row(2012))
+      }
+    }
+  }
 }

From 8aff36e91de0fee2f3f56c6d240bb203b5bb48ba Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Wed, 4 Nov 2015 10:49:34 +0000
Subject: [PATCH 042/518] [SPARK-2960][DEPLOY] Support executing Spark from
 symlinks (reopen)

This PR is based on the work of roji to support running Spark scripts from symlinks. Thanks for the great work roji . Would you mind taking a look at this PR, thanks a lot.

For releases like HDP and others, normally it will expose the Spark executables as symlinks and put in `PATH`, but current Spark's scripts do not support finding real path from symlink recursively, this will make spark fail to execute from symlink. This PR try to solve this issue by finding the absolute path from symlink.

Instead of using `readlink -f` like what this PR (https://github.com/apache/spark/pull/2386) implemented is that `-f` is not support for Mac, so here manually seeking the path through loop.

I've tested with Mac and Linux (Cent OS), looks fine.

This PR did not fix the scripts under `sbin` folder, not sure if it needs to be fixed also?

Please help to review, any comment is greatly appreciated.

Author: jerryshao <sshao@hortonworks.com>
Author: Shay Rojansky <roji@roji.org>

Closes #8669 from jerryshao/SPARK-2960.
---
 bin/beeline                         |  8 +++++---
 bin/load-spark-env.sh               | 32 ++++++++++++++++-------------
 bin/pyspark                         | 14 +++++++------
 bin/run-example                     | 18 ++++++++--------
 bin/spark-class                     | 15 +++++++-------
 bin/spark-shell                     |  9 +++++---
 bin/spark-sql                       |  7 +++++--
 bin/spark-submit                    |  6 ++++--
 bin/sparkR                          |  9 +++++---
 sbin/slaves.sh                      |  9 ++++----
 sbin/spark-config.sh                | 23 +++++++--------------
 sbin/spark-daemon.sh                | 23 +++++++++++----------
 sbin/spark-daemons.sh               |  9 ++++----
 sbin/start-all.sh                   | 11 +++++-----
 sbin/start-history-server.sh        | 11 +++++-----
 sbin/start-master.sh                | 17 +++++++--------
 sbin/start-mesos-dispatcher.sh      | 11 +++++-----
 sbin/start-mesos-shuffle-service.sh | 11 +++++-----
 sbin/start-shuffle-service.sh       | 11 +++++-----
 sbin/start-slave.sh                 | 18 ++++++++--------
 sbin/start-slaves.sh                | 19 ++++++++---------
 sbin/start-thriftserver.sh          | 11 +++++-----
 sbin/stop-all.sh                    | 14 ++++++-------
 sbin/stop-history-server.sh         |  7 ++++---
 sbin/stop-master.sh                 | 13 ++++++------
 sbin/stop-mesos-dispatcher.sh       |  9 ++++----
 sbin/stop-mesos-shuffle-service.sh  |  7 ++++---
 sbin/stop-shuffle-service.sh        |  7 ++++---
 sbin/stop-slave.sh                  | 15 +++++++-------
 sbin/stop-slaves.sh                 | 15 +++++++-------
 sbin/stop-thriftserver.sh           |  7 ++++---
 31 files changed, 213 insertions(+), 183 deletions(-)

diff --git a/bin/beeline b/bin/beeline
index 3fcb6df34339d..1627626941a73 100755
--- a/bin/beeline
+++ b/bin/beeline
@@ -23,8 +23,10 @@
 # Enter posix mode for bash
 set -o posix
 
-# Figure out where Spark is installed
-FWDIR="$(cd "`dirname "$0"`"/..; pwd)"
+# Figure out if SPARK_HOME is set
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
 CLASS="org.apache.hive.beeline.BeeLine"
-exec "$FWDIR/bin/spark-class" $CLASS "$@"
+exec "${SPARK_HOME}/bin/spark-class" $CLASS "$@"
diff --git a/bin/load-spark-env.sh b/bin/load-spark-env.sh
index 95779e9ddbb18..eaea964ed5b3d 100644
--- a/bin/load-spark-env.sh
+++ b/bin/load-spark-env.sh
@@ -20,13 +20,17 @@
 # This script loads spark-env.sh if it exists, and ensures it is only loaded once.
 # spark-env.sh is loaded from SPARK_CONF_DIR if set, or within the current directory's
 # conf/ subdirectory.
-FWDIR="$(cd "`dirname "$0"`"/..; pwd)"
+
+# Figure out where Spark is installed
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
 if [ -z "$SPARK_ENV_LOADED" ]; then
   export SPARK_ENV_LOADED=1
 
   # Returns the parent of the directory this script lives in.
-  parent_dir="$(cd "`dirname "$0"`"/..; pwd)"
+  parent_dir="${SPARK_HOME}"
 
   user_conf_dir="${SPARK_CONF_DIR:-"$parent_dir"/conf}"
 
@@ -42,18 +46,18 @@ fi
 
 if [ -z "$SPARK_SCALA_VERSION" ]; then
 
-    ASSEMBLY_DIR2="$FWDIR/assembly/target/scala-2.11"
-    ASSEMBLY_DIR1="$FWDIR/assembly/target/scala-2.10"
+  ASSEMBLY_DIR2="${SPARK_HOME}/assembly/target/scala-2.11"
+  ASSEMBLY_DIR1="${SPARK_HOME}/assembly/target/scala-2.10"
 
-    if [[ -d "$ASSEMBLY_DIR2" && -d "$ASSEMBLY_DIR1" ]]; then
-        echo -e "Presence of build for both scala versions(SCALA 2.10 and SCALA 2.11) detected." 1>&2
-        echo -e 'Either clean one of them or, export SPARK_SCALA_VERSION=2.11 in spark-env.sh.' 1>&2
-        exit 1
-    fi
+  if [[ -d "$ASSEMBLY_DIR2" && -d "$ASSEMBLY_DIR1" ]]; then
+    echo -e "Presence of build for both scala versions(SCALA 2.10 and SCALA 2.11) detected." 1>&2
+    echo -e 'Either clean one of them or, export SPARK_SCALA_VERSION=2.11 in spark-env.sh.' 1>&2
+    exit 1
+  fi
 
-    if [ -d "$ASSEMBLY_DIR2" ]; then
-        export SPARK_SCALA_VERSION="2.11"
-    else
-        export SPARK_SCALA_VERSION="2.10"
-    fi
+  if [ -d "$ASSEMBLY_DIR2" ]; then
+    export SPARK_SCALA_VERSION="2.11"
+  else
+    export SPARK_SCALA_VERSION="2.10"
+  fi
 fi
diff --git a/bin/pyspark b/bin/pyspark
index 18012ee4a0b4f..5eaa17d3c2016 100755
--- a/bin/pyspark
+++ b/bin/pyspark
@@ -17,9 +17,11 @@
 # limitations under the License.
 #
 
-export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-source "$SPARK_HOME"/bin/load-spark-env.sh
+source "${SPARK_HOME}"/bin/load-spark-env.sh
 export _SPARK_CMD_USAGE="Usage: ./bin/pyspark [options]"
 
 # In Spark <= 1.1, setting IPYTHON=1 would cause the driver to be launched using the `ipython`
@@ -64,12 +66,12 @@ fi
 export PYSPARK_PYTHON
 
 # Add the PySpark classes to the Python path:
-export PYTHONPATH="$SPARK_HOME/python/:$PYTHONPATH"
-export PYTHONPATH="$SPARK_HOME/python/lib/py4j-0.9-src.zip:$PYTHONPATH"
+export PYTHONPATH="${SPARK_HOME}/python/:$PYTHONPATH"
+export PYTHONPATH="${SPARK_HOME}/python/lib/py4j-0.9-src.zip:$PYTHONPATH"
 
 # Load the PySpark shell.py script when ./pyspark is used interactively:
 export OLD_PYTHONSTARTUP="$PYTHONSTARTUP"
-export PYTHONSTARTUP="$SPARK_HOME/python/pyspark/shell.py"
+export PYTHONSTARTUP="${SPARK_HOME}/python/pyspark/shell.py"
 
 # For pyspark tests
 if [[ -n "$SPARK_TESTING" ]]; then
@@ -82,4 +84,4 @@ fi
 
 export PYSPARK_DRIVER_PYTHON
 export PYSPARK_DRIVER_PYTHON_OPTS
-exec "$SPARK_HOME"/bin/spark-submit pyspark-shell-main --name "PySparkShell" "$@"
+exec "${SPARK_HOME}"/bin/spark-submit pyspark-shell-main --name "PySparkShell" "$@"
diff --git a/bin/run-example b/bin/run-example
index 798e2caeb88ce..e1b0d5789bed6 100755
--- a/bin/run-example
+++ b/bin/run-example
@@ -17,11 +17,13 @@
 # limitations under the License.
 #
 
-FWDIR="$(cd "`dirname "$0"`"/..; pwd)"
-export SPARK_HOME="$FWDIR"
-EXAMPLES_DIR="$FWDIR"/examples
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
+
+EXAMPLES_DIR="${SPARK_HOME}"/examples
 
-. "$FWDIR"/bin/load-spark-env.sh
+. "${SPARK_HOME}"/bin/load-spark-env.sh
 
 if [ -n "$1" ]; then
   EXAMPLE_CLASS="$1"
@@ -34,8 +36,8 @@ else
   exit 1
 fi
 
-if [ -f "$FWDIR/RELEASE" ]; then
-  JAR_PATH="${FWDIR}/lib"
+if [ -f "${SPARK_HOME}/RELEASE" ]; then
+  JAR_PATH="${SPARK_HOME}/lib"
 else
   JAR_PATH="${EXAMPLES_DIR}/target/scala-${SPARK_SCALA_VERSION}"
 fi
@@ -44,7 +46,7 @@ JAR_COUNT=0
 
 for f in "${JAR_PATH}"/spark-examples-*hadoop*.jar; do
   if [[ ! -e "$f" ]]; then
-    echo "Failed to find Spark examples assembly in $FWDIR/lib or $FWDIR/examples/target" 1>&2
+    echo "Failed to find Spark examples assembly in ${SPARK_HOME}/lib or ${SPARK_HOME}/examples/target" 1>&2
     echo "You need to build Spark before running this program" 1>&2
     exit 1
   fi
@@ -67,7 +69,7 @@ if [[ ! $EXAMPLE_CLASS == org.apache.spark.examples* ]]; then
   EXAMPLE_CLASS="org.apache.spark.examples.$EXAMPLE_CLASS"
 fi
 
-exec "$FWDIR"/bin/spark-submit \
+exec "${SPARK_HOME}"/bin/spark-submit \
   --master $EXAMPLE_MASTER \
   --class $EXAMPLE_CLASS \
   "$SPARK_EXAMPLES_JAR" \
diff --git a/bin/spark-class b/bin/spark-class
index 8cae6ccbabe7c..87d06693af4fe 100755
--- a/bin/spark-class
+++ b/bin/spark-class
@@ -17,10 +17,11 @@
 # limitations under the License.
 #
 
-# Figure out where Spark is installed
-export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$SPARK_HOME"/bin/load-spark-env.sh
+. "${SPARK_HOME}"/bin/load-spark-env.sh
 
 # Find the java binary
 if [ -n "${JAVA_HOME}" ]; then
@@ -36,10 +37,10 @@ fi
 
 # Find assembly jar
 SPARK_ASSEMBLY_JAR=
-if [ -f "$SPARK_HOME/RELEASE" ]; then
-  ASSEMBLY_DIR="$SPARK_HOME/lib"
+if [ -f "${SPARK_HOME}/RELEASE" ]; then
+  ASSEMBLY_DIR="${SPARK_HOME}/lib"
 else
-  ASSEMBLY_DIR="$SPARK_HOME/assembly/target/scala-$SPARK_SCALA_VERSION"
+  ASSEMBLY_DIR="${SPARK_HOME}/assembly/target/scala-$SPARK_SCALA_VERSION"
 fi
 
 GREP_OPTIONS=
@@ -65,7 +66,7 @@ LAUNCH_CLASSPATH="$SPARK_ASSEMBLY_JAR"
 
 # Add the launcher build dir to the classpath if requested.
 if [ -n "$SPARK_PREPEND_CLASSES" ]; then
-  LAUNCH_CLASSPATH="$SPARK_HOME/launcher/target/scala-$SPARK_SCALA_VERSION/classes:$LAUNCH_CLASSPATH"
+  LAUNCH_CLASSPATH="${SPARK_HOME}/launcher/target/scala-$SPARK_SCALA_VERSION/classes:$LAUNCH_CLASSPATH"
 fi
 
 export _SPARK_ASSEMBLY="$SPARK_ASSEMBLY_JAR"
diff --git a/bin/spark-shell b/bin/spark-shell
index 00ab7afd118b5..6583b5bd880ee 100755
--- a/bin/spark-shell
+++ b/bin/spark-shell
@@ -28,7 +28,10 @@ esac
 # Enter posix mode for bash
 set -o posix
 
-export FWDIR="$(cd "`dirname "$0"`"/..; pwd)"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
+
 export _SPARK_CMD_USAGE="Usage: ./bin/spark-shell [options]"
 
 # SPARK-4161: scala does not assume use of the java classpath,
@@ -47,11 +50,11 @@ function main() {
     # (see https://github.com/sbt/sbt/issues/562).
     stty -icanon min 1 -echo > /dev/null 2>&1
     export SPARK_SUBMIT_OPTS="$SPARK_SUBMIT_OPTS -Djline.terminal=unix"
-    "$FWDIR"/bin/spark-submit --class org.apache.spark.repl.Main --name "Spark shell" "$@"
+    "${SPARK_HOME}"/bin/spark-submit --class org.apache.spark.repl.Main --name "Spark shell" "$@"
     stty icanon echo > /dev/null 2>&1
   else
     export SPARK_SUBMIT_OPTS
-    "$FWDIR"/bin/spark-submit --class org.apache.spark.repl.Main --name "Spark shell" "$@"
+    "${SPARK_HOME}"/bin/spark-submit --class org.apache.spark.repl.Main --name "Spark shell" "$@"
   fi
 }
 
diff --git a/bin/spark-sql b/bin/spark-sql
index 4ea7bc6e39c07..970d12cbf51dd 100755
--- a/bin/spark-sql
+++ b/bin/spark-sql
@@ -17,6 +17,9 @@
 # limitations under the License.
 #
 
-export FWDIR="$(cd "`dirname "$0"`"/..; pwd)"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
+
 export _SPARK_CMD_USAGE="Usage: ./bin/spark-sql [options] [cli option]"
-exec "$FWDIR"/bin/spark-submit --class org.apache.spark.sql.hive.thriftserver.SparkSQLCLIDriver "$@"
+exec "${SPARK_HOME}"/bin/spark-submit --class org.apache.spark.sql.hive.thriftserver.SparkSQLCLIDriver "$@"
diff --git a/bin/spark-submit b/bin/spark-submit
index 255378b0f077c..023f9c162f4b8 100755
--- a/bin/spark-submit
+++ b/bin/spark-submit
@@ -17,9 +17,11 @@
 # limitations under the License.
 #
 
-SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
 # disable randomized hash for string in Python 3.3+
 export PYTHONHASHSEED=0
 
-exec "$SPARK_HOME"/bin/spark-class org.apache.spark.deploy.SparkSubmit "$@"
+exec "${SPARK_HOME}"/bin/spark-class org.apache.spark.deploy.SparkSubmit "$@"
diff --git a/bin/sparkR b/bin/sparkR
index 464c29f369424..2c07a82e2173b 100755
--- a/bin/sparkR
+++ b/bin/sparkR
@@ -17,7 +17,10 @@
 # limitations under the License.
 #
 
-export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
-source "$SPARK_HOME"/bin/load-spark-env.sh
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
+
+source "${SPARK_HOME}"/bin/load-spark-env.sh
 export _SPARK_CMD_USAGE="Usage: ./bin/sparkR [options]"
-exec "$SPARK_HOME"/bin/spark-submit sparkr-shell-main "$@"
+exec "${SPARK_HOME}"/bin/spark-submit sparkr-shell-main "$@"
diff --git a/sbin/slaves.sh b/sbin/slaves.sh
index cdad47ee2e594..c971aa3296b09 100755
--- a/sbin/slaves.sh
+++ b/sbin/slaves.sh
@@ -36,10 +36,11 @@ if [ $# -le 0 ]; then
   exit 1
 fi
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
 # If the slaves file is specified in the command line,
 # then it takes precedence over the definition in
@@ -65,7 +66,7 @@ then
   shift
 fi
 
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
 if [ "$HOSTLIST" = "" ]; then
   if [ "$SPARK_SLAVES" = "" ]; then
diff --git a/sbin/spark-config.sh b/sbin/spark-config.sh
index e6bf544c14799..d8d9d00d64ebc 100755
--- a/sbin/spark-config.sh
+++ b/sbin/spark-config.sh
@@ -19,21 +19,12 @@
 # should not be executable directly
 # also should not be passed any arguments, since we need original $*
 
-# resolve links - $0 may be a softlink
-this="${BASH_SOURCE:-$0}"
-common_bin="$(cd -P -- "$(dirname -- "$this")" && pwd -P)"
-script="$(basename -- "$this")"
-this="$common_bin/$script"
+# symlink and absolute path should rely on SPARK_HOME to resolve
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-# convert relative path to absolute path
-config_bin="`dirname "$this"`"
-script="`basename "$this"`"
-config_bin="`cd "$config_bin"; pwd`"
-this="$config_bin/$script"
-
-export SPARK_PREFIX="`dirname "$this"`"/..
-export SPARK_HOME="${SPARK_PREFIX}"
-export SPARK_CONF_DIR="${SPARK_CONF_DIR:-"$SPARK_HOME/conf"}"
+export SPARK_CONF_DIR="${SPARK_CONF_DIR:-"${SPARK_HOME}/conf"}"
 # Add the PySpark classes to the PYTHONPATH:
-export PYTHONPATH="$SPARK_HOME/python:$PYTHONPATH"
-export PYTHONPATH="$SPARK_HOME/python/lib/py4j-0.9-src.zip:$PYTHONPATH"
+export PYTHONPATH="${SPARK_HOME}/python:${PYTHONPATH}"
+export PYTHONPATH="${SPARK_HOME}/python/lib/py4j-0.9-src.zip:${PYTHONPATH}"
diff --git a/sbin/spark-daemon.sh b/sbin/spark-daemon.sh
index 0fbe795822fbf..6ab57df409529 100755
--- a/sbin/spark-daemon.sh
+++ b/sbin/spark-daemon.sh
@@ -37,10 +37,11 @@ if [ $# -le 1 ]; then
   exit 1
 fi
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
 # get arguments
 
@@ -86,7 +87,7 @@ spark_rotate_log ()
     fi
 }
 
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
 if [ "$SPARK_IDENT_STRING" = "" ]; then
   export SPARK_IDENT_STRING="$USER"
@@ -97,7 +98,7 @@ export SPARK_PRINT_LAUNCH_COMMAND="1"
 
 # get log directory
 if [ "$SPARK_LOG_DIR" = "" ]; then
-  export SPARK_LOG_DIR="$SPARK_HOME/logs"
+  export SPARK_LOG_DIR="${SPARK_HOME}/logs"
 fi
 mkdir -p "$SPARK_LOG_DIR"
 touch "$SPARK_LOG_DIR"/.spark_test > /dev/null 2>&1
@@ -137,7 +138,7 @@ run_command() {
 
   if [ "$SPARK_MASTER" != "" ]; then
     echo rsync from "$SPARK_MASTER"
-    rsync -a -e ssh --delete --exclude=.svn --exclude='logs/*' --exclude='contrib/hod/logs/*' "$SPARK_MASTER/" "$SPARK_HOME"
+    rsync -a -e ssh --delete --exclude=.svn --exclude='logs/*' --exclude='contrib/hod/logs/*' "$SPARK_MASTER/" "${SPARK_HOME}"
   fi
 
   spark_rotate_log "$log"
@@ -145,12 +146,12 @@ run_command() {
 
   case "$mode" in
     (class)
-      nohup nice -n "$SPARK_NICENESS" "$SPARK_PREFIX"/bin/spark-class $command "$@" >> "$log" 2>&1 < /dev/null &
+      nohup nice -n "$SPARK_NICENESS" "${SPARK_HOME}"/bin/spark-class $command "$@" >> "$log" 2>&1 < /dev/null &
       newpid="$!"
       ;;
 
     (submit)
-      nohup nice -n "$SPARK_NICENESS" "$SPARK_PREFIX"/bin/spark-submit --class $command "$@" >> "$log" 2>&1 < /dev/null &
+      nohup nice -n "$SPARK_NICENESS" "${SPARK_HOME}"/bin/spark-submit --class $command "$@" >> "$log" 2>&1 < /dev/null &
       newpid="$!"
       ;;
 
@@ -205,13 +206,13 @@ case $option in
       else
         echo $pid file is present but $command not running
         exit 1
-      fi  
+      fi
     else
       echo $command not running.
       exit 2
-    fi  
+    fi
     ;;
-  
+
   (*)
     echo $usage
     exit 1
diff --git a/sbin/spark-daemons.sh b/sbin/spark-daemons.sh
index 5d9f2bb51cae0..dec2f4432df39 100755
--- a/sbin/spark-daemons.sh
+++ b/sbin/spark-daemons.sh
@@ -27,9 +27,10 @@ if [ $# -le 1 ]; then
   exit 1
 fi
 
-sbin=`dirname "$0"`
-sbin=`cd "$sbin"; pwd`
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
-exec "$sbin/slaves.sh" cd "$SPARK_HOME" \; "$sbin/spark-daemon.sh" "$@"
+exec "${SPARK_HOME}/sbin/slaves.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/sbin/spark-daemon.sh" "$@"
diff --git a/sbin/start-all.sh b/sbin/start-all.sh
index 1baf57cea09ee..6217f9bf28e3d 100755
--- a/sbin/start-all.sh
+++ b/sbin/start-all.sh
@@ -21,8 +21,9 @@
 # Starts the master on this node.
 # Starts a worker on each node specified in conf/slaves
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
 TACHYON_STR=""
 
@@ -36,10 +37,10 @@ shift
 done
 
 # Load the Spark configuration
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
 # Start Master
-"$sbin"/start-master.sh $TACHYON_STR
+"${SPARK_HOME}/sbin"/start-master.sh $TACHYON_STR
 
 # Start Workers
-"$sbin"/start-slaves.sh $TACHYON_STR
+"${SPARK_HOME}/sbin"/start-slaves.sh $TACHYON_STR
diff --git a/sbin/start-history-server.sh b/sbin/start-history-server.sh
index 9034e5715cc85..6851d99b7e8f4 100755
--- a/sbin/start-history-server.sh
+++ b/sbin/start-history-server.sh
@@ -24,10 +24,11 @@
 # Use the SPARK_HISTORY_OPTS environment variable to set history server configuration.
 #
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
-exec "$sbin"/spark-daemon.sh start org.apache.spark.deploy.history.HistoryServer 1 $@
+exec "${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.history.HistoryServer 1 $@
diff --git a/sbin/start-master.sh b/sbin/start-master.sh
index a7f5d5702fd80..c20e19a8412df 100755
--- a/sbin/start-master.sh
+++ b/sbin/start-master.sh
@@ -19,8 +19,9 @@
 
 # Starts the master on the machine this script is executed on.
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
 ORIGINAL_ARGS="$@"
 
@@ -39,9 +40,9 @@ case $1 in
 shift
 done
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
 if [ "$SPARK_MASTER_PORT" = "" ]; then
   SPARK_MASTER_PORT=7077
@@ -55,12 +56,12 @@ if [ "$SPARK_MASTER_WEBUI_PORT" = "" ]; then
   SPARK_MASTER_WEBUI_PORT=8080
 fi
 
-"$sbin"/spark-daemon.sh start org.apache.spark.deploy.master.Master 1 \
+"${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.master.Master 1 \
   --ip $SPARK_MASTER_IP --port $SPARK_MASTER_PORT --webui-port $SPARK_MASTER_WEBUI_PORT \
   $ORIGINAL_ARGS
 
 if [ "$START_TACHYON" == "true" ]; then
-  "$sbin"/../tachyon/bin/tachyon bootstrap-conf $SPARK_MASTER_IP
-  "$sbin"/../tachyon/bin/tachyon format -s
-  "$sbin"/../tachyon/bin/tachyon-start.sh master
+  "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon bootstrap-conf $SPARK_MASTER_IP
+  "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon format -s
+  "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon-start.sh master
 fi
diff --git a/sbin/start-mesos-dispatcher.sh b/sbin/start-mesos-dispatcher.sh
index ef1fc573d5c65..4777e1668c703 100755
--- a/sbin/start-mesos-dispatcher.sh
+++ b/sbin/start-mesos-dispatcher.sh
@@ -21,12 +21,13 @@
 # Rest server to handle driver requests for Mesos cluster mode.
 # Only one cluster dispatcher is needed per Mesos cluster.
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
 if [ "$SPARK_MESOS_DISPATCHER_PORT" = "" ]; then
   SPARK_MESOS_DISPATCHER_PORT=7077
@@ -37,4 +38,4 @@ if [ "$SPARK_MESOS_DISPATCHER_HOST" = "" ]; then
 fi
 
 
-"$sbin"/spark-daemon.sh start org.apache.spark.deploy.mesos.MesosClusterDispatcher 1 --host $SPARK_MESOS_DISPATCHER_HOST --port $SPARK_MESOS_DISPATCHER_PORT "$@"
+"${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.mesos.MesosClusterDispatcher 1 --host $SPARK_MESOS_DISPATCHER_HOST --port $SPARK_MESOS_DISPATCHER_PORT "$@"
diff --git a/sbin/start-mesos-shuffle-service.sh b/sbin/start-mesos-shuffle-service.sh
index 64580762c5dc4..1845845676029 100755
--- a/sbin/start-mesos-shuffle-service.sh
+++ b/sbin/start-mesos-shuffle-service.sh
@@ -26,10 +26,11 @@
 # Use the SPARK_SHUFFLE_OPTS environment variable to set shuffle service configuration.
 #
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
-exec "$sbin"/spark-daemon.sh start org.apache.spark.deploy.mesos.MesosExternalShuffleService 1
+exec "${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.mesos.MesosExternalShuffleService 1
diff --git a/sbin/start-shuffle-service.sh b/sbin/start-shuffle-service.sh
index 4fddcf7f95d40..793e165be6c78 100755
--- a/sbin/start-shuffle-service.sh
+++ b/sbin/start-shuffle-service.sh
@@ -24,10 +24,11 @@
 # Use the SPARK_SHUFFLE_OPTS environment variable to set shuffle server configuration.
 #
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
-exec "$sbin"/spark-daemon.sh start org.apache.spark.deploy.ExternalShuffleService 1
+exec "${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.ExternalShuffleService 1
diff --git a/sbin/start-slave.sh b/sbin/start-slave.sh
index 4c919ff76a8f5..21455648d1c6d 100755
--- a/sbin/start-slave.sh
+++ b/sbin/start-slave.sh
@@ -21,14 +21,14 @@
 #
 # Environment Variables
 #
-#   SPARK_WORKER_INSTANCES  The number of worker instances to run on this 
+#   SPARK_WORKER_INSTANCES  The number of worker instances to run on this
 #                           slave.  Default is 1.
-#   SPARK_WORKER_PORT       The base port number for the first worker. If set, 
+#   SPARK_WORKER_PORT       The base port number for the first worker. If set,
 #                           subsequent workers will increment this number.  If
 #                           unset, Spark will find a valid port number, but
 #                           with no guarantee of a predictable pattern.
 #   SPARK_WORKER_WEBUI_PORT The base port for the web interface of the first
-#                           worker.  Subsequent workers will increment this 
+#                           worker.  Subsequent workers will increment this
 #                           number.  Default is 8081.
 
 usage="Usage: start-slave.sh <spark-master-URL> where <spark-master-URL> is like spark://localhost:7077"
@@ -39,12 +39,13 @@ if [ $# -lt 1 ]; then
   exit 1
 fi
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
 # First argument should be the master; we need to store it aside because we may
 # need to insert arguments between it and the other arguments
@@ -71,7 +72,7 @@ function start_instance {
   fi
   WEBUI_PORT=$(( $SPARK_WORKER_WEBUI_PORT + $WORKER_NUM - 1 ))
 
-  "$sbin"/spark-daemon.sh start org.apache.spark.deploy.worker.Worker $WORKER_NUM \
+  "${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.worker.Worker $WORKER_NUM \
      --webui-port "$WEBUI_PORT" $PORT_FLAG $PORT_NUM $MASTER "$@"
 }
 
@@ -82,4 +83,3 @@ else
     start_instance $(( 1 + $i )) "$@"
   done
 fi
-
diff --git a/sbin/start-slaves.sh b/sbin/start-slaves.sh
index 24d6268815ed3..51ca81e053b70 100755
--- a/sbin/start-slaves.sh
+++ b/sbin/start-slaves.sh
@@ -19,16 +19,16 @@
 
 # Starts a slave instance on each machine specified in the conf/slaves file.
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
-
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
 START_TACHYON=false
 
 while (( "$#" )); do
 case $1 in
     --with-tachyon)
-      if [ ! -e "$sbin"/../tachyon/bin/tachyon ]; then
+      if [ ! -e "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon ]; then
         echo "Error: --with-tachyon specified, but tachyon not found."
         exit -1
       fi
@@ -38,9 +38,8 @@ case $1 in
 shift
 done
 
-. "$sbin/spark-config.sh"
-
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
 # Find the port number for the master
 if [ "$SPARK_MASTER_PORT" = "" ]; then
@@ -52,11 +51,11 @@ if [ "$SPARK_MASTER_IP" = "" ]; then
 fi
 
 if [ "$START_TACHYON" == "true" ]; then
-  "$sbin/slaves.sh" cd "$SPARK_HOME" \; "$sbin"/../tachyon/bin/tachyon bootstrap-conf "$SPARK_MASTER_IP"
+  "${SPARK_HOME}/sbin/slaves.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon bootstrap-conf "$SPARK_MASTER_IP"
 
   # set -t so we can call sudo
-  SPARK_SSH_OPTS="-o StrictHostKeyChecking=no -t" "$sbin/slaves.sh" cd "$SPARK_HOME" \; "$sbin/../tachyon/bin/tachyon-start.sh" worker SudoMount \; sleep 1
+  SPARK_SSH_OPTS="-o StrictHostKeyChecking=no -t" "${SPARK_HOME}/sbin/slaves.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/tachyon/bin/tachyon-start.sh" worker SudoMount \; sleep 1
 fi
 
 # Launch the slaves
-"$sbin/slaves.sh" cd "$SPARK_HOME" \; "$sbin/start-slave.sh" "spark://$SPARK_MASTER_IP:$SPARK_MASTER_PORT"
+"${SPARK_HOME}/sbin/slaves.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/sbin/start-slave.sh" "spark://$SPARK_MASTER_IP:$SPARK_MASTER_PORT"
diff --git a/sbin/start-thriftserver.sh b/sbin/start-thriftserver.sh
index 5b0aeb177fff3..ad7e7c5277eb1 100755
--- a/sbin/start-thriftserver.sh
+++ b/sbin/start-thriftserver.sh
@@ -23,8 +23,9 @@
 # Enter posix mode for bash
 set -o posix
 
-# Figure out where Spark is installed
-FWDIR="$(cd "`dirname "$0"`"/..; pwd)"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
 # NOTE: This exact class name is matched downstream by SparkSubmit.
 # Any changes need to be reflected there.
@@ -39,10 +40,10 @@ function usage {
   pattern+="\|======="
   pattern+="\|--help"
 
-  "$FWDIR"/bin/spark-submit --help 2>&1 | grep -v Usage 1>&2
+  "${SPARK_HOME}"/bin/spark-submit --help 2>&1 | grep -v Usage 1>&2
   echo
   echo "Thrift server options:"
-  "$FWDIR"/bin/spark-class $CLASS --help 2>&1 | grep -v "$pattern" 1>&2
+  "${SPARK_HOME}"/bin/spark-class $CLASS --help 2>&1 | grep -v "$pattern" 1>&2
 }
 
 if [[ "$@" = *--help ]] || [[ "$@" = *-h ]]; then
@@ -52,4 +53,4 @@ fi
 
 export SUBMIT_USAGE_FUNCTION=usage
 
-exec "$FWDIR"/sbin/spark-daemon.sh submit $CLASS 1 "$@"
+exec "${SPARK_HOME}"/sbin/spark-daemon.sh submit $CLASS 1 "$@"
diff --git a/sbin/stop-all.sh b/sbin/stop-all.sh
index 1a9abe07db844..4e476ca05cb05 100755
--- a/sbin/stop-all.sh
+++ b/sbin/stop-all.sh
@@ -20,23 +20,23 @@
 # Stop all spark daemons.
 # Run this on the master node.
 
-
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
 # Load the Spark configuration
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
 # Stop the slaves, then the master
-"$sbin"/stop-slaves.sh
-"$sbin"/stop-master.sh
+"${SPARK_HOME}/sbin"/stop-slaves.sh
+"${SPARK_HOME}/sbin"/stop-master.sh
 
 if [ "$1" == "--wait" ]
 then
   printf "Waiting for workers to shut down..."
   while true
   do
-    running=`$sbin/slaves.sh ps -ef | grep -v grep | grep deploy.worker.Worker`
+    running=`${SPARK_HOME}/sbin/slaves.sh ps -ef | grep -v grep | grep deploy.worker.Worker`
     if [ -z "$running" ]
     then
       printf "\nAll workers successfully shut down.\n"
diff --git a/sbin/stop-history-server.sh b/sbin/stop-history-server.sh
index 6e6056359510f..14e3af4be910a 100755
--- a/sbin/stop-history-server.sh
+++ b/sbin/stop-history-server.sh
@@ -19,7 +19,8 @@
 
 # Stops the history server on the machine this script is executed on.
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-"$sbin"/spark-daemon.sh stop org.apache.spark.deploy.history.HistoryServer 1
+"${SPARK_HOME}/sbin/spark-daemon.sh" stop org.apache.spark.deploy.history.HistoryServer 1
diff --git a/sbin/stop-master.sh b/sbin/stop-master.sh
index 729702d92191e..e57962bb354d9 100755
--- a/sbin/stop-master.sh
+++ b/sbin/stop-master.sh
@@ -19,13 +19,14 @@
 
 # Stops the master on the machine this script is executed on.
 
-sbin=`dirname "$0"`
-sbin=`cd "$sbin"; pwd`
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
-"$sbin"/spark-daemon.sh stop org.apache.spark.deploy.master.Master 1
+"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.master.Master 1
 
-if [ -e "$sbin"/../tachyon/bin/tachyon ]; then
-  "$sbin"/../tachyon/bin/tachyon killAll tachyon.master.Master
+if [ -e "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon ]; then
+  "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon killAll tachyon.master.Master
 fi
diff --git a/sbin/stop-mesos-dispatcher.sh b/sbin/stop-mesos-dispatcher.sh
index cb65d95b5e524..5c0b4e051db38 100755
--- a/sbin/stop-mesos-dispatcher.sh
+++ b/sbin/stop-mesos-dispatcher.sh
@@ -18,10 +18,11 @@
 #
 # Stop the Mesos Cluster dispatcher on the machine this script is executed on.
 
-sbin=`dirname "$0"`
-sbin=`cd "$sbin"; pwd`
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
-"$sbin"/spark-daemon.sh stop org.apache.spark.deploy.mesos.MesosClusterDispatcher 1
+"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.mesos.MesosClusterDispatcher 1
 
diff --git a/sbin/stop-mesos-shuffle-service.sh b/sbin/stop-mesos-shuffle-service.sh
index 0e965d5ec5886..d23cad375e1bd 100755
--- a/sbin/stop-mesos-shuffle-service.sh
+++ b/sbin/stop-mesos-shuffle-service.sh
@@ -19,7 +19,8 @@
 
 # Stops the Mesos external shuffle service on the machine this script is executed on.
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-"$sbin"/spark-daemon.sh stop org.apache.spark.deploy.mesos.MesosExternalShuffleService 1
+"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.mesos.MesosExternalShuffleService 1
diff --git a/sbin/stop-shuffle-service.sh b/sbin/stop-shuffle-service.sh
index 4cb6891ae27fa..50d69cf34e0a5 100755
--- a/sbin/stop-shuffle-service.sh
+++ b/sbin/stop-shuffle-service.sh
@@ -19,7 +19,8 @@
 
 # Stops the external shuffle service on the machine this script is executed on.
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-"$sbin"/spark-daemon.sh stop org.apache.spark.deploy.ExternalShuffleService 1
+"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.ExternalShuffleService 1
diff --git a/sbin/stop-slave.sh b/sbin/stop-slave.sh
index 3d1da5b254f2a..685bcf59b33aa 100755
--- a/sbin/stop-slave.sh
+++ b/sbin/stop-slave.sh
@@ -21,23 +21,24 @@
 #
 # Environment variables
 #
-#   SPARK_WORKER_INSTANCES The number of worker instances that should be 
+#   SPARK_WORKER_INSTANCES The number of worker instances that should be
 #                          running on this slave.  Default is 1.
 
 # Usage: stop-slave.sh
 #   Stops all slaves on this worker machine
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
 if [ "$SPARK_WORKER_INSTANCES" = "" ]; then
-  "$sbin"/spark-daemon.sh stop org.apache.spark.deploy.worker.Worker 1
+  "${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.worker.Worker 1
 else
   for ((i=0; i<$SPARK_WORKER_INSTANCES; i++)); do
-    "$sbin"/spark-daemon.sh stop org.apache.spark.deploy.worker.Worker $(( $i + 1 ))
+    "${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.deploy.worker.Worker $(( $i + 1 ))
   done
 fi
diff --git a/sbin/stop-slaves.sh b/sbin/stop-slaves.sh
index 54c9bd46803a9..63956377629d6 100755
--- a/sbin/stop-slaves.sh
+++ b/sbin/stop-slaves.sh
@@ -17,16 +17,17 @@
 # limitations under the License.
 #
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-. "$sbin/spark-config.sh"
+. "${SPARK_HOME}/sbin/spark-config.sh"
 
-. "$SPARK_PREFIX/bin/load-spark-env.sh"
+. "${SPARK_HOME}/bin/load-spark-env.sh"
 
 # do before the below calls as they exec
-if [ -e "$sbin"/../tachyon/bin/tachyon ]; then
-  "$sbin/slaves.sh" cd "$SPARK_HOME" \; "$sbin"/../tachyon/bin/tachyon killAll tachyon.worker.Worker
+if [ -e "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon ]; then
+  "${SPARK_HOME}/sbin/slaves.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon killAll tachyon.worker.Worker
 fi
 
-"$sbin/slaves.sh" cd "$SPARK_HOME" \; "$sbin"/stop-slave.sh
+"${SPARK_HOME}/sbin/slaves.sh" cd "${SPARK_HOME}" \; "${SPARK_HOME}/sbin"/stop-slave.sh
diff --git a/sbin/stop-thriftserver.sh b/sbin/stop-thriftserver.sh
index 4031a00d4a689..cf45058f882a0 100755
--- a/sbin/stop-thriftserver.sh
+++ b/sbin/stop-thriftserver.sh
@@ -19,7 +19,8 @@
 
 # Stops the thrift server on the machine this script is executed on.
 
-sbin="`dirname "$0"`"
-sbin="`cd "$sbin"; pwd`"
+if [ -z "${SPARK_HOME}" ]; then
+  export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
+fi
 
-"$sbin"/spark-daemon.sh stop org.apache.spark.sql.hive.thriftserver.HiveThriftServer2 1
+"${SPARK_HOME}/sbin"/spark-daemon.sh stop org.apache.spark.sql.hive.thriftserver.HiveThriftServer2 1

From c09e5139874fb3626e005c8240cca5308b902ef3 Mon Sep 17 00:00:00 2001
From: tedyu <yuzhihong@gmail.com>
Date: Wed, 4 Nov 2015 10:51:40 +0000
Subject: [PATCH 043/518] [SPARK-11442] Reduce numSlices for local metrics test
 of SparkListenerSuite

In the thread, http://search-hadoop.com/m/q3RTtcQiFSlTxeP/test+failed+due+to+OOME&subj=test+failed+due+to+OOME, it was discussed that memory consumption for SparkListenerSuite should be brought down.

This is an attempt in that direction by reducing numSlices for local metrics test.

Author: tedyu <yuzhihong@gmail.com>

Closes #9384 from tedyu/master.
---
 .../org/apache/spark/scheduler/SparkListenerSuite.scala  | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
index a9652d7e7d0b0..53102b9f1c936 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
@@ -212,14 +212,15 @@ class SparkListenerSuite extends SparkFunSuite with LocalSparkContext with Match
       i
     }
 
-    val d = sc.parallelize(0 to 1e4.toInt, 64).map(w)
+    val numSlices = 16
+    val d = sc.parallelize(0 to 1e3.toInt, numSlices).map(w)
     d.count()
     sc.listenerBus.waitUntilEmpty(WAIT_TIMEOUT_MILLIS)
     listener.stageInfos.size should be (1)
 
     val d2 = d.map { i => w(i) -> i * 2 }.setName("shuffle input 1")
     val d3 = d.map { i => w(i) -> (0 to (i % 5)) }.setName("shuffle input 2")
-    val d4 = d2.cogroup(d3, 64).map { case (k, (v1, v2)) =>
+    val d4 = d2.cogroup(d3, numSlices).map { case (k, (v1, v2)) =>
       w(k) -> (v1.size, v2.size)
     }
     d4.setName("A Cogroup")
@@ -258,8 +259,8 @@ class SparkListenerSuite extends SparkFunSuite with LocalSparkContext with Match
         if (stageInfo.rddInfos.exists(_.name == d4.name)) {
           taskMetrics.shuffleReadMetrics should be ('defined)
           val sm = taskMetrics.shuffleReadMetrics.get
-          sm.totalBlocksFetched should be (128)
-          sm.localBlocksFetched should be (128)
+          sm.totalBlocksFetched should be (2*numSlices)
+          sm.localBlocksFetched should be (2*numSlices)
           sm.remoteBlocksFetched should be (0)
           sm.remoteBytesRead should be (0L)
         }

From e328b69c31821e4b27673d7ef6182ab3b7a05ca8 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Wed, 4 Nov 2015 08:28:33 -0800
Subject: [PATCH 044/518] [SPARK-9492][ML][R] LogisticRegression in R should
 provide model statistics

Like ml ```LinearRegression```, ```LogisticRegression``` should provide a training summary including feature names and their coefficients.

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9303 from yanboliang/spark-9492.
---
 R/pkg/inst/tests/test_mllib.R                   | 17 +++++++++++++++++
 .../ml/classification/LogisticRegression.scala  | 17 +++++++++++++----
 .../org/apache/spark/ml/r/SparkRWrappers.scala  |  7 ++++---
 project/MimaExcludes.scala                      |  4 +++-
 4 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/R/pkg/inst/tests/test_mllib.R b/R/pkg/inst/tests/test_mllib.R
index 3331ce738358c..032cfef061fd3 100644
--- a/R/pkg/inst/tests/test_mllib.R
+++ b/R/pkg/inst/tests/test_mllib.R
@@ -67,3 +67,20 @@ test_that("summary coefficients match with native glm", {
     as.character(stats$features) ==
     c("(Intercept)", "Sepal_Length", "Species_versicolor", "Species_virginica")))
 })
+
+test_that("summary coefficients match with native glm of family 'binomial'", {
+  df <- createDataFrame(sqlContext, iris)
+  training <- filter(df, df$Species != "setosa")
+  stats <- summary(glm(Species ~ Sepal_Length + Sepal_Width, data = training,
+    family = "binomial"))
+  coefs <- as.vector(stats$coefficients)
+
+  rTraining <- iris[iris$Species %in% c("versicolor","virginica"),]
+  rCoefs <- as.vector(coef(glm(Species ~ Sepal.Length + Sepal.Width, data = rTraining,
+    family = binomial(link = "logit"))))
+
+  expect_true(all(abs(rCoefs - coefs) < 1e-4))
+  expect_true(all(
+    as.character(stats$features) ==
+    c("(Intercept)", "Sepal_Length", "Sepal_Width")))
+})
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
index a1335e7a1bde8..f5fca686df144 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
@@ -378,6 +378,7 @@ class LogisticRegression(override val uid: String)
       model.transform(dataset),
       $(probabilityCol),
       $(labelCol),
+      $(featuresCol),
       objectiveHistory)
     model.setSummary(logRegSummary)
   }
@@ -452,7 +453,8 @@ class LogisticRegressionModel private[ml] (
    */
   // TODO: decide on a good name before exposing to public API
   private[classification] def evaluate(dataset: DataFrame): LogisticRegressionSummary = {
-    new BinaryLogisticRegressionSummary(this.transform(dataset), $(probabilityCol), $(labelCol))
+    new BinaryLogisticRegressionSummary(
+      this.transform(dataset), $(probabilityCol), $(labelCol), $(featuresCol))
   }
 
   /**
@@ -614,9 +616,12 @@ sealed trait LogisticRegressionSummary extends Serializable {
   /** Field in "predictions" which gives the calibrated probability of each instance as a vector. */
   def probabilityCol: String
 
-  /** Field in "predictions" which gives the the true label of each instance. */
+  /** Field in "predictions" which gives the true label of each instance. */
   def labelCol: String
 
+  /** Field in "predictions" which gives the features of each instance as a vector. */
+  def featuresCol: String
+
 }
 
 /**
@@ -626,6 +631,7 @@ sealed trait LogisticRegressionSummary extends Serializable {
  * @param probabilityCol field in "predictions" which gives the calibrated probability of
  *                       each instance as a vector.
  * @param labelCol field in "predictions" which gives the true label of each instance.
+ * @param featuresCol field in "predictions" which gives the features of each instance as a vector.
  * @param objectiveHistory objective function (scaled loss + regularization) at each iteration.
  */
 @Experimental
@@ -633,8 +639,9 @@ class BinaryLogisticRegressionTrainingSummary private[classification] (
     predictions: DataFrame,
     probabilityCol: String,
     labelCol: String,
+    featuresCol: String,
     val objectiveHistory: Array[Double])
-  extends BinaryLogisticRegressionSummary(predictions, probabilityCol, labelCol)
+  extends BinaryLogisticRegressionSummary(predictions, probabilityCol, labelCol, featuresCol)
   with LogisticRegressionTrainingSummary {
 
 }
@@ -646,12 +653,14 @@ class BinaryLogisticRegressionTrainingSummary private[classification] (
  * @param probabilityCol field in "predictions" which gives the calibrated probability of
  *                       each instance.
  * @param labelCol field in "predictions" which gives the true label of each instance.
+ * @param featuresCol field in "predictions" which gives the features of each instance as a vector.
  */
 @Experimental
 class BinaryLogisticRegressionSummary private[classification] (
     @transient override val predictions: DataFrame,
     override val probabilityCol: String,
-    override val labelCol: String) extends LogisticRegressionSummary {
+    override val labelCol: String,
+    override val featuresCol: String) extends LogisticRegressionSummary {
 
   private val sqlContext = predictions.sqlContext
   import sqlContext.implicits._
diff --git a/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala b/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
index 24f76de806d8f..5be2f86936211 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
@@ -66,9 +66,10 @@ private[r] object SparkRWrappers {
         val attrs = AttributeGroup.fromStructField(
           m.summary.predictions.schema(m.summary.featuresCol))
         Array("(Intercept)") ++ attrs.attributes.get.map(_.name.get)
-      case _: LogisticRegressionModel =>
-        throw new UnsupportedOperationException(
-          "No features names available for LogisticRegressionModel")  // SPARK-9492
+      case m: LogisticRegressionModel =>
+        val attrs = AttributeGroup.fromStructField(
+          m.summary.predictions.schema(m.summary.featuresCol))
+        Array("(Intercept)") ++ attrs.attributes.get.map(_.name.get)
     }
   }
 }
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index ec0e44b7f2d66..eeef96c378bdb 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -59,7 +59,9 @@ object MimaExcludes {
         ProblemFilters.exclude[MissingMethodProblem](
           "org.apache.spark.ml.classification.LogisticAggregator.add"),
         ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.ml.classification.LogisticAggregator.count")
+          "org.apache.spark.ml.classification.LogisticAggregator.count"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.ml.classification.LogisticRegressionSummary.featuresCol")
       ) ++ Seq(
         // SPARK-10381 Fix types / units in private AskPermissionToCommitOutput RPC message.
         // This class is marked as `private` but MiMa still seems to be confused by the change.

From 820064e613609bbf7edd726d982da1de60bf417a Mon Sep 17 00:00:00 2001
From: Pravin Gadakh <pravingadakh177@gmail.com>
Date: Wed, 4 Nov 2015 08:32:08 -0800
Subject: [PATCH 045/518] [SPARK-11380][DOCS] Replace example code in
 mllib-frequent-pattern-mining.md using include_example

Author: Pravin Gadakh <pravingadakh177@gmail.com>
Author: Pravin Gadakh <prgadakh@in.ibm.com>

Closes #9340 from pravingadakh/SPARK-11380.
---
 docs/mllib-frequent-pattern-mining.md         | 168 +-----------------
 .../mllib/JavaAssociationRulesExample.java    |  56 ++++++
 .../examples/mllib/JavaPrefixSpanExample.java |  55 ++++++
 .../examples/mllib/JavaSimpleFPGrowth.java    |  71 ++++++++
 .../src/main/python/mllib/fpgrowth_example.py |  33 ++++
 .../mllib/AssociationRulesExample.scala       |  54 ++++++
 .../examples/mllib/PrefixSpanExample.scala    |  52 ++++++
 .../spark/examples/mllib/SimpleFPGrowth.scala |  59 ++++++
 8 files changed, 387 insertions(+), 161 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaAssociationRulesExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaPrefixSpanExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaSimpleFPGrowth.java
 create mode 100644 examples/src/main/python/mllib/fpgrowth_example.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/AssociationRulesExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/PrefixSpanExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/SimpleFPGrowth.scala

diff --git a/docs/mllib-frequent-pattern-mining.md b/docs/mllib-frequent-pattern-mining.md
index f749eb4f2ff4f..fe42896a05d8e 100644
--- a/docs/mllib-frequent-pattern-mining.md
+++ b/docs/mllib-frequent-pattern-mining.md
@@ -52,31 +52,7 @@ details) from `transactions`.
 
 Refer to the [`FPGrowth` Scala docs](api/scala/index.html#org.apache.spark.mllib.fpm.FPGrowth) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.rdd.RDD
-import org.apache.spark.mllib.fpm.FPGrowth
-
-val data = sc.textFile("data/mllib/sample_fpgrowth.txt")
-
-val transactions: RDD[Array[String]] = data.map(s => s.trim.split(' '))
-
-val fpg = new FPGrowth()
-  .setMinSupport(0.2)
-  .setNumPartitions(10)
-val model = fpg.run(transactions)
-
-model.freqItemsets.collect().foreach { itemset =>
-  println(itemset.items.mkString("[", ",", "]") + ", " + itemset.freq)
-}
-
-val minConfidence = 0.8
-model.generateAssociationRules(minConfidence).collect().foreach { rule =>
-  println(
-    rule.antecedent.mkString("[", ",", "]")
-      + " => " + rule.consequent .mkString("[", ",", "]")
-      + ", " + rule.confidence)
-}
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/SimpleFPGrowth.scala %}
 
 </div>
 
@@ -95,46 +71,7 @@ details) from `transactions`.
 
 Refer to the [`FPGrowth` Java docs](api/java/org/apache/spark/mllib/fpm/FPGrowth.html) for details on the API.
 
-{% highlight java %}
-import java.util.Arrays;
-import java.util.List;
-
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.mllib.fpm.AssociationRules;
-import org.apache.spark.mllib.fpm.FPGrowth;
-import org.apache.spark.mllib.fpm.FPGrowthModel;
-
-SparkConf conf = new SparkConf().setAppName("FP-growth Example");
-JavaSparkContext sc = new JavaSparkContext(conf);
-
-JavaRDD<String> data = sc.textFile("data/mllib/sample_fpgrowth.txt");
-
-JavaRDD<List<String>> transactions = data.map(
-  new Function<String, List<String>>() {
-    public List<String> call(String line) {
-      String[] parts = line.split(" ");
-      return Arrays.asList(parts);
-    }
-  }
-);
-
-FPGrowth fpg = new FPGrowth()
-  .setMinSupport(0.2)
-  .setNumPartitions(10);
-FPGrowthModel<String> model = fpg.run(transactions);
-
-for (FPGrowth.FreqItemset<String> itemset: model.freqItemsets().toJavaRDD().collect()) {
-  System.out.println("[" + itemset.javaItems() + "], " + itemset.freq());
-}
-
-double minConfidence = 0.8;
-for (AssociationRules.Rule<String> rule
-    : model.generateAssociationRules(minConfidence).toJavaRDD().collect()) {
-  System.out.println(
-    rule.javaAntecedent() + " => " + rule.javaConsequent() + ", " + rule.confidence());
-}
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaSimpleFPGrowth.java %}
 
 </div>
 
@@ -149,19 +86,7 @@ that stores the frequent itemsets with their frequencies.
 
 Refer to the [`FPGrowth` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.fpm.FPGrowth) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.fpm import FPGrowth
-
-data = sc.textFile("data/mllib/sample_fpgrowth.txt")
-
-transactions = data.map(lambda line: line.strip().split(' '))
-
-model = FPGrowth.train(transactions, minSupport=0.2, numPartitions=10)
-
-result = model.freqItemsets().collect()
-for fi in result:
-    print(fi)
-{% endhighlight %}
+{% include_example python/mllib/fpgrowth_example.py %}
 
 </div>
 
@@ -177,27 +102,7 @@ that have a single item as the consequent.
 
 Refer to the [`AssociationRules` Scala docs](api/java/org/apache/spark/mllib/fpm/AssociationRules.html) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.rdd.RDD
-import org.apache.spark.mllib.fpm.AssociationRules
-import org.apache.spark.mllib.fpm.FPGrowth.FreqItemset
-
-val freqItemsets = sc.parallelize(Seq(
-  new FreqItemset(Array("a"), 15L),
-  new FreqItemset(Array("b"), 35L),
-  new FreqItemset(Array("a", "b"), 12L)
-));
-
-val ar = new AssociationRules()
-  .setMinConfidence(0.8)
-val results = ar.run(freqItemsets)
-
-results.collect().foreach { rule =>
-  println("[" + rule.antecedent.mkString(",")
-    + "=>"
-    + rule.consequent.mkString(",") + "]," + rule.confidence)
-}
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/AssociationRulesExample.scala %}
 
 </div>
 
@@ -208,29 +113,7 @@ that have a single item as the consequent.
 
 Refer to the [`AssociationRules` Java docs](api/java/org/apache/spark/mllib/fpm/AssociationRules.html) for details on the API.
 
-{% highlight java %}
-import java.util.Arrays;
-
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.mllib.fpm.AssociationRules;
-import org.apache.spark.mllib.fpm.FPGrowth.FreqItemset;
-
-JavaRDD<FPGrowth.FreqItemset<String>> freqItemsets = sc.parallelize(Arrays.asList(
-  new FreqItemset<String>(new String[] {"a"}, 15L),
-  new FreqItemset<String>(new String[] {"b"}, 35L),
-  new FreqItemset<String>(new String[] {"a", "b"}, 12L)
-));
-
-AssociationRules arules = new AssociationRules()
-  .setMinConfidence(0.8);
-JavaRDD<AssociationRules.Rule<String>> results = arules.run(freqItemsets);
-
-for (AssociationRules.Rule<String> rule: results.collect()) {
-  System.out.println(
-    rule.javaAntecedent() + " => " + rule.javaConsequent() + ", " + rule.confidence());
-}
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaAssociationRulesExample.java %}
 
 </div>
 </div>
@@ -278,24 +161,7 @@ that stores the frequent sequences with their frequencies.
 
 Refer to the [`PrefixSpan` Scala docs](api/scala/index.html#org.apache.spark.mllib.fpm.PrefixSpan) and [`PrefixSpanModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.fpm.PrefixSpanModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.fpm.PrefixSpan
-
-val sequences = sc.parallelize(Seq(
-    Array(Array(1, 2), Array(3)),
-    Array(Array(1), Array(3, 2), Array(1, 2)),
-    Array(Array(1, 2), Array(5)),
-    Array(Array(6))
-  ), 2).cache()
-val prefixSpan = new PrefixSpan()
-  .setMinSupport(0.5)
-  .setMaxPatternLength(5)
-val model = prefixSpan.run(sequences)
-model.freqSequences.collect().foreach { freqSequence =>
-println(
-  freqSequence.sequence.map(_.mkString("[", ", ", "]")).mkString("[", ", ", "]") + ", " + freqSequence.freq)
-}
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/PrefixSpanExample.scala %}
 
 </div>
 
@@ -309,27 +175,7 @@ that stores the frequent sequences with their frequencies.
 
 Refer to the [`PrefixSpan` Java docs](api/java/org/apache/spark/mllib/fpm/PrefixSpan.html) and [`PrefixSpanModel` Java docs](api/java/org/apache/spark/mllib/fpm/PrefixSpanModel.html) for details on the API.
 
-{% highlight java %}
-import java.util.Arrays;
-import java.util.List;
-
-import org.apache.spark.mllib.fpm.PrefixSpan;
-import org.apache.spark.mllib.fpm.PrefixSpanModel;
-
-JavaRDD<List<List<Integer>>> sequences = sc.parallelize(Arrays.asList(
-  Arrays.asList(Arrays.asList(1, 2), Arrays.asList(3)),
-  Arrays.asList(Arrays.asList(1), Arrays.asList(3, 2), Arrays.asList(1, 2)),
-  Arrays.asList(Arrays.asList(1, 2), Arrays.asList(5)),
-  Arrays.asList(Arrays.asList(6))
-), 2);
-PrefixSpan prefixSpan = new PrefixSpan()
-  .setMinSupport(0.5)
-  .setMaxPatternLength(5);
-PrefixSpanModel<Integer> model = prefixSpan.run(sequences);
-for (PrefixSpan.FreqSequence<Integer> freqSeq: model.freqSequences().toJavaRDD().collect()) {
-  System.out.println(freqSeq.javaSequence() + ", " + freqSeq.freq());
-}
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaPrefixSpanExample.java %}
 
 </div>
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaAssociationRulesExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaAssociationRulesExample.java
new file mode 100644
index 0000000000000..4d0f989819ace
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaAssociationRulesExample.java
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.Arrays;
+
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.mllib.fpm.AssociationRules;
+import org.apache.spark.mllib.fpm.FPGrowth;
+import org.apache.spark.mllib.fpm.FPGrowth.FreqItemset;
+// $example off$
+
+import org.apache.spark.SparkConf;
+
+public class JavaAssociationRulesExample {
+
+  public static void main(String[] args) {
+
+    SparkConf sparkConf = new SparkConf().setAppName("JavaAssociationRulesExample");
+    JavaSparkContext sc = new JavaSparkContext(sparkConf);
+
+    // $example on$
+    JavaRDD<FPGrowth.FreqItemset<String>> freqItemsets = sc.parallelize(Arrays.asList(
+      new FreqItemset<String>(new String[] {"a"}, 15L),
+      new FreqItemset<String>(new String[] {"b"}, 35L),
+      new FreqItemset<String>(new String[] {"a", "b"}, 12L)
+    ));
+
+    AssociationRules arules = new AssociationRules()
+      .setMinConfidence(0.8);
+    JavaRDD<AssociationRules.Rule<String>> results = arules.run(freqItemsets);
+
+    for (AssociationRules.Rule<String> rule : results.collect()) {
+      System.out.println(
+        rule.javaAntecedent() + " => " + rule.javaConsequent() + ", " + rule.confidence());
+    }
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaPrefixSpanExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaPrefixSpanExample.java
new file mode 100644
index 0000000000000..68ec7c1e6ebe0
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaPrefixSpanExample.java
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.Arrays;
+import java.util.List;
+// $example off$
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
+import org.apache.spark.mllib.fpm.PrefixSpan;
+import org.apache.spark.mllib.fpm.PrefixSpanModel;
+// $example off$
+import org.apache.spark.SparkConf;
+
+public class JavaPrefixSpanExample {
+
+  public static void main(String[] args) {
+
+    SparkConf sparkConf = new SparkConf().setAppName("JavaPrefixSpanExample");
+    JavaSparkContext sc = new JavaSparkContext(sparkConf);
+
+    // $example on$
+    JavaRDD<List<List<Integer>>> sequences = sc.parallelize(Arrays.asList(
+      Arrays.asList(Arrays.asList(1, 2), Arrays.asList(3)),
+      Arrays.asList(Arrays.asList(1), Arrays.asList(3, 2), Arrays.asList(1, 2)),
+      Arrays.asList(Arrays.asList(1, 2), Arrays.asList(5)),
+      Arrays.asList(Arrays.asList(6))
+    ), 2);
+    PrefixSpan prefixSpan = new PrefixSpan()
+      .setMinSupport(0.5)
+      .setMaxPatternLength(5);
+    PrefixSpanModel<Integer> model = prefixSpan.run(sequences);
+    for (PrefixSpan.FreqSequence<Integer> freqSeq: model.freqSequences().toJavaRDD().collect()) {
+      System.out.println(freqSeq.javaSequence() + ", " + freqSeq.freq());
+    }
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaSimpleFPGrowth.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaSimpleFPGrowth.java
new file mode 100644
index 0000000000000..72edaca5e95b1
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaSimpleFPGrowth.java
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.Arrays;
+import java.util.List;
+
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example off$
+import org.apache.spark.api.java.function.Function;
+// $example on$
+import org.apache.spark.mllib.fpm.AssociationRules;
+import org.apache.spark.mllib.fpm.FPGrowth;
+import org.apache.spark.mllib.fpm.FPGrowthModel;
+// $example off$
+
+import org.apache.spark.SparkConf;
+
+public class JavaSimpleFPGrowth {
+
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("FP-growth Example");
+    JavaSparkContext sc = new JavaSparkContext(conf);
+
+    // $example on$
+    JavaRDD<String> data = sc.textFile("data/mllib/sample_fpgrowth.txt");
+
+    JavaRDD<List<String>> transactions = data.map(
+      new Function<String, List<String>>() {
+        public List<String> call(String line) {
+          String[] parts = line.split(" ");
+          return Arrays.asList(parts);
+        }
+      }
+    );
+
+    FPGrowth fpg = new FPGrowth()
+      .setMinSupport(0.2)
+      .setNumPartitions(10);
+    FPGrowthModel<String> model = fpg.run(transactions);
+
+    for (FPGrowth.FreqItemset<String> itemset: model.freqItemsets().toJavaRDD().collect()) {
+      System.out.println("[" + itemset.javaItems() + "], " + itemset.freq());
+    }
+
+    double minConfidence = 0.8;
+    for (AssociationRules.Rule<String> rule
+      : model.generateAssociationRules(minConfidence).toJavaRDD().collect()) {
+      System.out.println(
+        rule.javaAntecedent() + " => " + rule.javaConsequent() + ", " + rule.confidence());
+    }
+    // $example off$
+  }
+}
diff --git a/examples/src/main/python/mllib/fpgrowth_example.py b/examples/src/main/python/mllib/fpgrowth_example.py
new file mode 100644
index 0000000000000..715f5268206cb
--- /dev/null
+++ b/examples/src/main/python/mllib/fpgrowth_example.py
@@ -0,0 +1,33 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# $example on$
+from pyspark.mllib.fpm import FPGrowth
+# $example off$
+from pyspark import SparkContext
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="FPGrowth")
+
+    # $example on$
+    data = sc.textFile("data/mllib/sample_fpgrowth.txt")
+    transactions = data.map(lambda line: line.strip().split(' '))
+    model = FPGrowth.train(transactions, minSupport=0.2, numPartitions=10)
+    result = model.freqItemsets().collect()
+    for fi in result:
+        print(fi)
+    # $example off$
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/AssociationRulesExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/AssociationRulesExample.scala
new file mode 100644
index 0000000000000..ca22ddafc3c48
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/AssociationRulesExample.scala
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.fpm.AssociationRules
+import org.apache.spark.mllib.fpm.FPGrowth.FreqItemset
+// $example off$
+
+import org.apache.spark.{SparkConf, SparkContext}
+
+object AssociationRulesExample {
+
+  def main(args: Array[String]) {
+    val conf = new SparkConf().setAppName("AssociationRulesExample")
+    val sc = new SparkContext(conf)
+
+    // $example on$
+    val freqItemsets = sc.parallelize(Seq(
+      new FreqItemset(Array("a"), 15L),
+      new FreqItemset(Array("b"), 35L),
+      new FreqItemset(Array("a", "b"), 12L)
+    ))
+
+    val ar = new AssociationRules()
+      .setMinConfidence(0.8)
+    val results = ar.run(freqItemsets)
+
+    results.collect().foreach { rule =>
+      println("[" + rule.antecedent.mkString(",")
+        + "=>"
+        + rule.consequent.mkString(",") + "]," + rule.confidence)
+    }
+    // $example off$
+  }
+
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/PrefixSpanExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/PrefixSpanExample.scala
new file mode 100644
index 0000000000000..d237232c430ca
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/PrefixSpanExample.scala
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.fpm.PrefixSpan
+// $example off$
+
+import org.apache.spark.{SparkConf, SparkContext}
+
+object PrefixSpanExample {
+
+  def main(args: Array[String]) {
+    val conf = new SparkConf().setAppName("PrefixSpanExample")
+    val sc = new SparkContext(conf)
+
+    // $example on$
+    val sequences = sc.parallelize(Seq(
+      Array(Array(1, 2), Array(3)),
+      Array(Array(1), Array(3, 2), Array(1, 2)),
+      Array(Array(1, 2), Array(5)),
+      Array(Array(6))
+    ), 2).cache()
+    val prefixSpan = new PrefixSpan()
+      .setMinSupport(0.5)
+      .setMaxPatternLength(5)
+    val model = prefixSpan.run(sequences)
+    model.freqSequences.collect().foreach { freqSequence =>
+      println(
+        freqSequence.sequence.map(_.mkString("[", ", ", "]")).mkString("[", ", ", "]") +
+          ", " + freqSequence.freq)
+    }
+    // $example off$
+  }
+}
+// scalastyle:off println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/SimpleFPGrowth.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/SimpleFPGrowth.scala
new file mode 100644
index 0000000000000..b4e06afa7410f
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/SimpleFPGrowth.scala
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.fpm.FPGrowth
+import org.apache.spark.rdd.RDD
+// $example off$
+
+import org.apache.spark.{SparkContext, SparkConf}
+
+object SimpleFPGrowth {
+
+  def main(args: Array[String]) {
+
+    val conf = new SparkConf().setAppName("SimpleFPGrowth")
+    val sc = new SparkContext(conf)
+
+    // $example on$
+    val data = sc.textFile("data/mllib/sample_fpgrowth.txt")
+
+    val transactions: RDD[Array[String]] = data.map(s => s.trim.split(' '))
+
+    val fpg = new FPGrowth()
+      .setMinSupport(0.2)
+      .setNumPartitions(10)
+    val model = fpg.run(transactions)
+
+    model.freqItemsets.collect().foreach { itemset =>
+      println(itemset.items.mkString("[", ",", "]") + ", " + itemset.freq)
+    }
+
+    val minConfidence = 0.8
+    model.generateAssociationRules(minConfidence).collect().foreach { rule =>
+      println(
+        rule.antecedent.mkString("[", ",", "]")
+          + " => " + rule.consequent .mkString("[", ",", "]")
+          + ", " + rule.confidence)
+    }
+    // $example off$
+  }
+}
+// scalastyle:on println

From 9b214cea896056e7d0a69ae9d3c282e1f027d5b9 Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Wed, 4 Nov 2015 08:36:55 -0800
Subject: [PATCH 046/518] [SPARK-11443] Reserve space lines

The trim_codeblock(lines) function in include_example.rb removes some blank lines in the code.

Author: Xusen Yin <yinxusen@gmail.com>

Closes #9400 from yinxusen/SPARK-11443.
---
 docs/_plugins/include_example.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/_plugins/include_example.rb b/docs/_plugins/include_example.rb
index 0f4184c7462be..6ee63a5ac69df 100644
--- a/docs/_plugins/include_example.rb
+++ b/docs/_plugins/include_example.rb
@@ -50,7 +50,7 @@ def trim_codeblock(lines)
         .map { |l| l[/\A */].size }
         .min
 
-      lines.map { |l| l[min_start_spaces .. -1] }
+      lines.map { |l| l.strip.size == 0 ? l : l[min_start_spaces .. -1] }
     end
 
     # Select lines according to labels in code. Currently we use "$example on$" and "$example off$"

From 8790ee6d69e50ca84eb849742be48f2476743b5b Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Wed, 4 Nov 2015 09:07:22 -0800
Subject: [PATCH 047/518] [SPARK-10622][CORE][YARN] Differentiate dead from
 "mostly dead" executors.

In YARN mode, when preemption is enabled, we may leave executors in a
zombie state while we wait to retrieve the reason for which the executor
exited. This is so that we don't account for failed tasks that were
running on a preempted executor.

The issue is that while we wait for this information, the scheduler
might decide to schedule tasks on the executor, which will never be
able to run them. Other side effects include the block manager still
considering the executor available to cache blocks, for example.

So, when we know that an executor went down but we don't know why,
stop everything related to the executor, except its running tasks.
Only when we know the reason for the exit (or give up waiting for
it) we do update the running tasks.

This is achieved by a new `disableExecutor()` method in the
`Schedulable` interface. For managers that do not behave like this
(i.e. every one but YARN), the existing `executorLost()` method
will behave the same way it did before.

On top of that change, a few minor changes that made debugging easier,
and fixed some other minor issues:
- The cluster-mode AM was printing a misleading log message every
  time an executor disconnected from the driver (because the akka
  actor system was shared between driver and AM).
- Avoid sending unnecessary requests for an executor's exit reason
  when we already know it was explicitly disabled / killed. This
  avoids both multiple requests, and unnecessary requests that would
  just cause warning messages on the AM (in the explicit kill case).
- Tone down a log message about the executor being lost when it
  exited normally (e.g. preemption)
- Wake up the AM monitor thread when requests for executor loss
  reasons arrive too, so that we can more quickly remove executors
  from this zombie state.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #8887 from vanzin/SPARK-10622.
---
 .../spark/scheduler/ExecutorLossReason.scala  |  9 +++
 .../spark/scheduler/TaskSchedulerImpl.scala   | 32 ++++++--
 .../CoarseGrainedSchedulerBackend.scala       | 37 ++++++++-
 .../cluster/YarnSchedulerBackend.scala        |  9 +--
 .../scheduler/TaskSchedulerImplSuite.scala    | 36 +++++++++
 .../spark/deploy/yarn/ApplicationMaster.scala | 79 +++++++++++--------
 .../spark/deploy/yarn/YarnAllocator.scala     |  5 ++
 7 files changed, 157 insertions(+), 50 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
index 33edf25043850..47a5cbff4930b 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
@@ -40,6 +40,15 @@ private[spark] object ExecutorExited {
   }
 }
 
+/**
+ * A loss reason that means we don't yet know why the executor exited.
+ *
+ * This is used by the task scheduler to remove state associated with the executor, but
+ * not yet fail any tasks that were running in the executor before the real loss reason
+ * is known.
+ */
+private [spark] object LossReasonPending extends ExecutorLossReason("Pending loss reason.")
+
 private[spark]
 case class SlaveLost(_message: String = "Slave lost")
   extends ExecutorLossReason(_message)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index 1c7bfe89c02ac..43d7d80b7aae1 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -468,11 +468,20 @@ private[spark] class TaskSchedulerImpl(
         removeExecutor(executorId, reason)
         failedExecutor = Some(executorId)
       } else {
-         // We may get multiple executorLost() calls with different loss reasons. For example, one
-         // may be triggered by a dropped connection from the slave while another may be a report
-         // of executor termination from Mesos. We produce log messages for both so we eventually
-         // report the termination reason.
-         logError("Lost an executor " + executorId + " (already removed): " + reason)
+         executorIdToHost.get(executorId) match {
+           case Some(_) =>
+             // If the host mapping still exists, it means we don't know the loss reason for the
+             // executor. So call removeExecutor() to update tasks running on that executor when
+             // the real loss reason is finally known.
+             removeExecutor(executorId, reason)
+
+           case None =>
+             // We may get multiple executorLost() calls with different loss reasons. For example,
+             // one may be triggered by a dropped connection from the slave while another may be a
+             // report of executor termination from Mesos. We produce log messages for both so we
+             // eventually report the termination reason.
+             logError("Lost an executor " + executorId + " (already removed): " + reason)
+         }
       }
     }
     // Call dagScheduler.executorLost without holding the lock on this to prevent deadlock
@@ -482,7 +491,11 @@ private[spark] class TaskSchedulerImpl(
     }
   }
 
-  /** Remove an executor from all our data structures and mark it as lost */
+  /**
+   * Remove an executor from all our data structures and mark it as lost. If the executor's loss
+   * reason is not yet known, do not yet remove its association with its host nor update the status
+   * of any running tasks, since the loss reason defines whether we'll fail those tasks.
+   */
   private def removeExecutor(executorId: String, reason: ExecutorLossReason) {
     activeExecutorIds -= executorId
     val host = executorIdToHost(executorId)
@@ -497,8 +510,11 @@ private[spark] class TaskSchedulerImpl(
         }
       }
     }
-    executorIdToHost -= executorId
-    rootPool.executorLost(executorId, host, reason)
+
+    if (reason != LossReasonPending) {
+      executorIdToHost -= executorId
+      rootPool.executorLost(executorId, host, reason)
+    }
   }
 
   def executorAdded(execId: String, host: String) {
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
index ebce5021b19dc..f71d98feac050 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
@@ -73,6 +73,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
   // The number of pending tasks which is locality required
   protected var localityAwareTasks = 0
 
+  // Executors that have been lost, but for which we don't yet know the real exit reason.
+  protected val executorsPendingLossReason = new HashSet[String]
+
   class DriverEndpoint(override val rpcEnv: RpcEnv, sparkProperties: Seq[(String, String)])
     extends ThreadSafeRpcEndpoint with Logging {
 
@@ -184,7 +187,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     // Make fake resource offers on all executors
     private def makeOffers() {
       // Filter out executors under killing
-      val activeExecutors = executorDataMap.filterKeys(!executorsPendingToRemove.contains(_))
+      val activeExecutors = executorDataMap.filterKeys(executorIsAlive)
       val workOffers = activeExecutors.map { case (id, executorData) =>
         new WorkerOffer(id, executorData.executorHost, executorData.freeCores)
       }.toSeq
@@ -202,7 +205,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     // Make fake resource offers on just one executor
     private def makeOffers(executorId: String) {
       // Filter out executors under killing
-      if (!executorsPendingToRemove.contains(executorId)) {
+      if (executorIsAlive(executorId)) {
         val executorData = executorDataMap(executorId)
         val workOffers = Seq(
           new WorkerOffer(executorId, executorData.executorHost, executorData.freeCores))
@@ -210,6 +213,11 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
       }
     }
 
+    private def executorIsAlive(executorId: String): Boolean = synchronized {
+      !executorsPendingToRemove.contains(executorId) &&
+        !executorsPendingLossReason.contains(executorId)
+    }
+
     // Launch tasks returned by a set of resource offers
     private def launchTasks(tasks: Seq[Seq[TaskDescription]]) {
       for (task <- tasks.flatten) {
@@ -246,6 +254,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
             addressToExecutorId -= executorInfo.executorAddress
             executorDataMap -= executorId
             executorsPendingToRemove -= executorId
+            executorsPendingLossReason -= executorId
           }
           totalCoreCount.addAndGet(-executorInfo.totalCores)
           totalRegisteredExecutors.addAndGet(-1)
@@ -256,6 +265,30 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
       }
     }
 
+    /**
+     * Stop making resource offers for the given executor. The executor is marked as lost with
+     * the loss reason still pending.
+     *
+     * @return Whether executor was alive.
+     */
+    protected def disableExecutor(executorId: String): Boolean = {
+      val shouldDisable = CoarseGrainedSchedulerBackend.this.synchronized {
+        if (executorIsAlive(executorId)) {
+          executorsPendingLossReason += executorId
+          true
+        } else {
+          false
+        }
+      }
+
+      if (shouldDisable) {
+        logInfo(s"Disabling executor $executorId.")
+        scheduler.executorLost(executorId, LossReasonPending)
+      }
+
+      shouldDisable
+    }
+
     override def onStop() {
       reviveThread.shutdownNow()
     }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
index d75d6f673e84e..80da37b09b590 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala
@@ -115,15 +115,12 @@ private[spark] abstract class YarnSchedulerBackend(
      * (e.g., preemption), according to the application master, then we pass that information down
      * to the TaskSetManager to inform the TaskSetManager that tasks on that lost executor should
      * not count towards a job failure.
-     *
-     * TODO there's a race condition where while we are querying the ApplicationMaster for
-     * the executor loss reason, there is the potential that tasks will be scheduled on
-     * the executor that failed. We should fix this by having this onDisconnected event
-     * also "blacklist" executors so that tasks are not assigned to them.
      */
     override def onDisconnected(rpcAddress: RpcAddress): Unit = {
       addressToExecutorId.get(rpcAddress).foreach { executorId =>
-        yarnSchedulerEndpoint.handleExecutorDisconnectedFromDriver(executorId, rpcAddress)
+        if (disableExecutor(executorId)) {
+          yarnSchedulerEndpoint.handleExecutorDisconnectedFromDriver(executorId, rpcAddress)
+        }
       }
     }
   }
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
index c2edd4c317d6e..2afb595e6f10d 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
@@ -237,4 +237,40 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext with L
     }
   }
 
+  test("tasks are not re-scheduled while executor loss reason is pending") {
+    sc = new SparkContext("local", "TaskSchedulerImplSuite")
+    val taskScheduler = new TaskSchedulerImpl(sc)
+    taskScheduler.initialize(new FakeSchedulerBackend)
+    // Need to initialize a DAGScheduler for the taskScheduler to use for callbacks.
+    new DAGScheduler(sc, taskScheduler) {
+      override def taskStarted(task: Task[_], taskInfo: TaskInfo) {}
+      override def executorAdded(execId: String, host: String) {}
+    }
+
+    val e0Offers = Seq(new WorkerOffer("executor0", "host0", 1))
+    val e1Offers = Seq(new WorkerOffer("executor1", "host0", 1))
+    val attempt1 = FakeTask.createTaskSet(1)
+
+    // submit attempt 1, offer resources, task gets scheduled
+    taskScheduler.submitTasks(attempt1)
+    val taskDescriptions = taskScheduler.resourceOffers(e0Offers).flatten
+    assert(1 === taskDescriptions.length)
+
+    // mark executor0 as dead but pending fail reason
+    taskScheduler.executorLost("executor0", LossReasonPending)
+
+    // offer some more resources on a different executor, nothing should change
+    val taskDescriptions2 = taskScheduler.resourceOffers(e1Offers).flatten
+    assert(0 === taskDescriptions2.length)
+
+    // provide the actual loss reason for executor0
+    taskScheduler.executorLost("executor0", SlaveLost("oops"))
+
+    // executor0's tasks should have failed now that the loss reason is known, so offering more
+    // resources should make them be scheduled on the new executor.
+    val taskDescriptions3 = taskScheduler.resourceOffers(e1Offers).flatten
+    assert(1 === taskDescriptions3.length)
+    assert("executor1" === taskDescriptions3(0).executorId)
+  }
+
 }
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
index 12ae350e4cef6..50ae7ffeec4c5 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
@@ -87,8 +87,27 @@ private[spark] class ApplicationMaster(
 
   @volatile private var reporterThread: Thread = _
   @volatile private var allocator: YarnAllocator = _
+
+  // Lock for controlling the allocator (heartbeat) thread.
   private val allocatorLock = new Object()
 
+  // Steady state heartbeat interval. We want to be reasonably responsive without causing too many
+  // requests to RM.
+  private val heartbeatInterval = {
+    // Ensure that progress is sent before YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS elapses.
+    val expiryInterval = yarnConf.getInt(YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS, 120000)
+    math.max(0, math.min(expiryInterval / 2,
+      sparkConf.getTimeAsMs("spark.yarn.scheduler.heartbeat.interval-ms", "3s")))
+  }
+
+  // Initial wait interval before allocator poll, to allow for quicker ramp up when executors are
+  // being requested.
+  private val initialAllocationInterval = math.min(heartbeatInterval,
+    sparkConf.getTimeAsMs("spark.yarn.scheduler.initial-allocation.interval", "200ms"))
+
+  // Next wait interval before allocator poll.
+  private var nextAllocationInterval = initialAllocationInterval
+
   // Fields used in client mode.
   private var rpcEnv: RpcEnv = null
   private var amEndpoint: RpcEndpointRef = _
@@ -332,19 +351,6 @@ private[spark] class ApplicationMaster(
   }
 
   private def launchReporterThread(): Thread = {
-    // Ensure that progress is sent before YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS elapses.
-    val expiryInterval = yarnConf.getInt(YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS, 120000)
-
-    // we want to be reasonably responsive without causing too many requests to RM.
-    val heartbeatInterval = math.max(0, math.min(expiryInterval / 2,
-      sparkConf.getTimeAsMs("spark.yarn.scheduler.heartbeat.interval-ms", "3s")))
-
-    // we want to check more frequently for pending containers
-    val initialAllocationInterval = math.min(heartbeatInterval,
-      sparkConf.getTimeAsMs("spark.yarn.scheduler.initial-allocation.interval", "200ms"))
-
-    var nextAllocationInterval = initialAllocationInterval
-
     // The number of failures in a row until Reporter thread give up
     val reporterMaxFailures = sparkConf.getInt("spark.yarn.scheduler.reporterThread.maxFailures", 5)
 
@@ -377,19 +383,19 @@ private[spark] class ApplicationMaster(
           }
           try {
             val numPendingAllocate = allocator.getPendingAllocate.size
-            val sleepInterval =
-              if (numPendingAllocate > 0) {
-                val currentAllocationInterval =
-                  math.min(heartbeatInterval, nextAllocationInterval)
-                nextAllocationInterval = currentAllocationInterval * 2 // avoid overflow
-                currentAllocationInterval
-              } else {
-                nextAllocationInterval = initialAllocationInterval
-                heartbeatInterval
-              }
-            logDebug(s"Number of pending allocations is $numPendingAllocate. " +
-                     s"Sleeping for $sleepInterval.")
             allocatorLock.synchronized {
+              val sleepInterval =
+                if (numPendingAllocate > 0 || allocator.getNumPendingLossReasonRequests > 0) {
+                  val currentAllocationInterval =
+                    math.min(heartbeatInterval, nextAllocationInterval)
+                  nextAllocationInterval = currentAllocationInterval * 2 // avoid overflow
+                  currentAllocationInterval
+                } else {
+                  nextAllocationInterval = initialAllocationInterval
+                  heartbeatInterval
+                }
+              logDebug(s"Number of pending allocations is $numPendingAllocate. " +
+                       s"Sleeping for $sleepInterval.")
               allocatorLock.wait(sleepInterval)
             }
           } catch {
@@ -560,6 +566,11 @@ private[spark] class ApplicationMaster(
     userThread
   }
 
+  private def resetAllocatorInterval(): Unit = allocatorLock.synchronized {
+    nextAllocationInterval = initialAllocationInterval
+    allocatorLock.notifyAll()
+  }
+
   /**
    * An [[RpcEndpoint]] that communicates with the driver's scheduler backend.
    */
@@ -581,11 +592,9 @@ private[spark] class ApplicationMaster(
       case RequestExecutors(requestedTotal, localityAwareTasks, hostToLocalTaskCount) =>
         Option(allocator) match {
           case Some(a) =>
-            allocatorLock.synchronized {
-              if (a.requestTotalExecutorsWithPreferredLocalities(requestedTotal,
-                localityAwareTasks, hostToLocalTaskCount)) {
-                allocatorLock.notifyAll()
-              }
+            if (a.requestTotalExecutorsWithPreferredLocalities(requestedTotal,
+              localityAwareTasks, hostToLocalTaskCount)) {
+              resetAllocatorInterval()
             }
 
           case None =>
@@ -603,17 +612,19 @@ private[spark] class ApplicationMaster(
 
       case GetExecutorLossReason(eid) =>
         Option(allocator) match {
-          case Some(a) => a.enqueueGetLossReasonRequest(eid, context)
-          case None => logWarning(s"Container allocator is not ready to find" +
-            s" executor loss reasons yet.")
+          case Some(a) =>
+            a.enqueueGetLossReasonRequest(eid, context)
+            resetAllocatorInterval()
+          case None =>
+            logWarning("Container allocator is not ready to find executor loss reasons yet.")
         }
     }
 
     override def onDisconnected(remoteAddress: RpcAddress): Unit = {
-      logInfo(s"Driver terminated or disconnected! Shutting down. $remoteAddress")
       // In cluster mode, do not rely on the disassociated event to exit
       // This avoids potentially reporting incorrect exit codes if the driver fails
       if (!isClusterMode) {
+        logInfo(s"Driver terminated or disconnected! Shutting down. $remoteAddress")
         finish(FinalApplicationStatus.SUCCEEDED, ApplicationMaster.EXIT_SUCCESS)
       }
     }
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index a0cf1b4aa469b..4d9e777cb4134 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -550,6 +550,10 @@ private[yarn] class YarnAllocator(
 
   private[yarn] def getNumUnexpectedContainerRelease = numUnexpectedContainerRelease
 
+  private[yarn] def getNumPendingLossReasonRequests: Int = synchronized {
+    pendingLossReasonRequests.size
+  }
+
   /**
    * Split the pending container requests into 3 groups based on current localities of pending
    * tasks.
@@ -582,6 +586,7 @@ private[yarn] class YarnAllocator(
 
     (localityMatched.toSeq, localityUnMatched.toSeq, localityFree.toSeq)
   }
+
 }
 
 private object YarnAllocator {

From 27feafccbd6945b000ca51b14c57912acbad9031 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Wed, 4 Nov 2015 09:11:54 -0800
Subject: [PATCH 048/518] [SPARK-11235][NETWORK] Add ability to stream data
 using network lib.

The current interface used to fetch shuffle data is not very efficient for
large buffers; it requires the receiver to buffer the entirety of the
contents being downloaded in memory before processing the data.

To use the network library to transfer large files (such as those that
can be added using SparkContext addJar / addFile), this change adds a
more efficient way of downloding data, by streaming the data and feeding
it to a callback as data arrives.

This is achieved by a custom frame decoder that replaces the current netty
one; this decoder allows entering a mode where framing is skipped and data
is instead provided directly to a callback. The existing netty classes
(ByteToMessageDecoder and LengthFieldBasedFrameDecoder) could not be reused
since their semantics do not allow for the interception approach the new
decoder uses.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9206 from vanzin/SPARK-11235.
---
 .../spark/network/TransportContext.java       |   3 +-
 .../spark/network/client/StreamCallback.java  |  40 +++
 .../network/client/StreamInterceptor.java     |  76 ++++
 .../spark/network/client/TransportClient.java |  41 +++
 .../client/TransportResponseHandler.java      |  47 ++-
 .../network/protocol/ChunkFetchSuccess.java   |  16 +-
 .../spark/network/protocol/Message.java       |   6 +-
 .../network/protocol/MessageDecoder.java      |   9 +
 .../network/protocol/MessageEncoder.java      |  27 +-
 .../network/protocol/ResponseWithBody.java    |  40 +++
 .../spark/network/protocol/StreamFailure.java |  80 +++++
 .../spark/network/protocol/StreamRequest.java |  78 +++++
 .../network/protocol/StreamResponse.java      |  91 +++++
 .../spark/network/server/StreamManager.java   |  13 +
 .../server/TransportRequestHandler.java       |  20 ++
 .../apache/spark/network/util/NettyUtils.java |   9 +-
 .../network/util/TransportFrameDecoder.java   | 154 +++++++++
 .../apache/spark/network/ProtocolSuite.java   |   8 +
 .../org/apache/spark/network/StreamSuite.java | 325 ++++++++++++++++++
 .../util/TransportFrameDecoderSuite.java      | 142 ++++++++
 20 files changed, 1196 insertions(+), 29 deletions(-)
 create mode 100644 network/common/src/main/java/org/apache/spark/network/client/StreamCallback.java
 create mode 100644 network/common/src/main/java/org/apache/spark/network/client/StreamInterceptor.java
 create mode 100644 network/common/src/main/java/org/apache/spark/network/protocol/ResponseWithBody.java
 create mode 100644 network/common/src/main/java/org/apache/spark/network/protocol/StreamFailure.java
 create mode 100644 network/common/src/main/java/org/apache/spark/network/protocol/StreamRequest.java
 create mode 100644 network/common/src/main/java/org/apache/spark/network/protocol/StreamResponse.java
 create mode 100644 network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
 create mode 100644 network/common/src/test/java/org/apache/spark/network/StreamSuite.java
 create mode 100644 network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java

diff --git a/network/common/src/main/java/org/apache/spark/network/TransportContext.java b/network/common/src/main/java/org/apache/spark/network/TransportContext.java
index b8d073fa16b4b..43900e6f2c972 100644
--- a/network/common/src/main/java/org/apache/spark/network/TransportContext.java
+++ b/network/common/src/main/java/org/apache/spark/network/TransportContext.java
@@ -39,6 +39,7 @@
 import org.apache.spark.network.server.TransportServerBootstrap;
 import org.apache.spark.network.util.NettyUtils;
 import org.apache.spark.network.util.TransportConf;
+import org.apache.spark.network.util.TransportFrameDecoder;
 
 /**
  * Contains the context to create a {@link TransportServer}, {@link TransportClientFactory}, and to
@@ -119,7 +120,7 @@ public TransportChannelHandler initializePipeline(
       TransportChannelHandler channelHandler = createChannelHandler(channel, channelRpcHandler);
       channel.pipeline()
         .addLast("encoder", encoder)
-        .addLast("frameDecoder", NettyUtils.createFrameDecoder())
+        .addLast(TransportFrameDecoder.HANDLER_NAME, NettyUtils.createFrameDecoder())
         .addLast("decoder", decoder)
         .addLast("idleStateHandler", new IdleStateHandler(0, 0, conf.connectionTimeoutMs() / 1000))
         // NOTE: Chunks are currently guaranteed to be returned in the order of request, but this
diff --git a/network/common/src/main/java/org/apache/spark/network/client/StreamCallback.java b/network/common/src/main/java/org/apache/spark/network/client/StreamCallback.java
new file mode 100644
index 0000000000000..093fada320cc3
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/client/StreamCallback.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.client;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+/**
+ * Callback for streaming data. Stream data will be offered to the {@link onData(ByteBuffer)}
+ * method as it arrives. Once all the stream data is received, {@link onComplete()} will be
+ * called.
+ * <p>
+ * The network library guarantees that a single thread will call these methods at a time, but
+ * different call may be made by different threads.
+ */
+public interface StreamCallback {
+  /** Called upon receipt of stream data. */
+  void onData(String streamId, ByteBuffer buf) throws IOException;
+
+  /** Called when all data from the stream has been received. */
+  void onComplete(String streamId) throws IOException;
+
+  /** Called if there's an error reading data from the stream. */
+  void onFailure(String streamId, Throwable cause) throws IOException;
+}
diff --git a/network/common/src/main/java/org/apache/spark/network/client/StreamInterceptor.java b/network/common/src/main/java/org/apache/spark/network/client/StreamInterceptor.java
new file mode 100644
index 0000000000000..02230a00e69fc
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/client/StreamInterceptor.java
@@ -0,0 +1,76 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.client;
+
+import java.nio.ByteBuffer;
+import java.nio.channels.ClosedChannelException;
+
+import io.netty.buffer.ByteBuf;
+
+import org.apache.spark.network.util.TransportFrameDecoder;
+
+/**
+ * An interceptor that is registered with the frame decoder to feed stream data to a
+ * callback.
+ */
+class StreamInterceptor implements TransportFrameDecoder.Interceptor {
+
+  private final String streamId;
+  private final long byteCount;
+  private final StreamCallback callback;
+
+  private volatile long bytesRead;
+
+  StreamInterceptor(String streamId, long byteCount, StreamCallback callback) {
+    this.streamId = streamId;
+    this.byteCount = byteCount;
+    this.callback = callback;
+    this.bytesRead = 0;
+  }
+
+  @Override
+  public void exceptionCaught(Throwable cause) throws Exception {
+    callback.onFailure(streamId, cause);
+  }
+
+  @Override
+  public void channelInactive() throws Exception {
+    callback.onFailure(streamId, new ClosedChannelException());
+  }
+
+  @Override
+  public boolean handle(ByteBuf buf) throws Exception {
+    int toRead = (int) Math.min(buf.readableBytes(), byteCount - bytesRead);
+    ByteBuffer nioBuffer = buf.readSlice(toRead).nioBuffer();
+
+    int available = nioBuffer.remaining();
+    callback.onData(streamId, nioBuffer);
+    bytesRead += available;
+    if (bytesRead > byteCount) {
+      RuntimeException re = new IllegalStateException(String.format(
+        "Read too many bytes? Expected %d, but read %d.", byteCount, bytesRead));
+      callback.onFailure(streamId, re);
+      throw re;
+    } else if (bytesRead == byteCount) {
+      callback.onComplete(streamId);
+    }
+
+    return bytesRead != byteCount;
+  }
+
+}
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java b/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
index fbb8bb6b2f6c3..a0ba223e340a2 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
@@ -38,6 +38,7 @@
 import org.apache.spark.network.protocol.ChunkFetchRequest;
 import org.apache.spark.network.protocol.RpcRequest;
 import org.apache.spark.network.protocol.StreamChunkId;
+import org.apache.spark.network.protocol.StreamRequest;
 import org.apache.spark.network.util.NettyUtils;
 
 /**
@@ -159,6 +160,46 @@ public void operationComplete(ChannelFuture future) throws Exception {
       });
   }
 
+  /**
+   * Request to stream the data with the given stream ID from the remote end.
+   *
+   * @param streamId The stream to fetch.
+   * @param callback Object to call with the stream data.
+   */
+  public void stream(final String streamId, final StreamCallback callback) {
+    final String serverAddr = NettyUtils.getRemoteAddress(channel);
+    final long startTime = System.currentTimeMillis();
+    logger.debug("Sending stream request for {} to {}", streamId, serverAddr);
+
+    // Need to synchronize here so that the callback is added to the queue and the RPC is
+    // written to the socket atomically, so that callbacks are called in the right order
+    // when responses arrive.
+    synchronized (this) {
+      handler.addStreamCallback(callback);
+      channel.writeAndFlush(new StreamRequest(streamId)).addListener(
+        new ChannelFutureListener() {
+          @Override
+          public void operationComplete(ChannelFuture future) throws Exception {
+            if (future.isSuccess()) {
+              long timeTaken = System.currentTimeMillis() - startTime;
+              logger.trace("Sending request for {} to {} took {} ms", streamId, serverAddr,
+                timeTaken);
+            } else {
+              String errorMsg = String.format("Failed to send request for %s to %s: %s", streamId,
+                serverAddr, future.cause());
+              logger.error(errorMsg, future.cause());
+              channel.close();
+              try {
+                callback.onFailure(streamId, new IOException(errorMsg, future.cause()));
+              } catch (Exception e) {
+                logger.error("Uncaught exception in RPC response callback handler!", e);
+              }
+            }
+          }
+        });
+    }
+  }
+
   /**
    * Sends an opaque message to the RpcHandler on the server-side. The callback will be invoked
    * with the server's response or upon any failure.
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
index 94fc21af5e606..ed3f36af58048 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
@@ -19,7 +19,9 @@
 
 import java.io.IOException;
 import java.util.Map;
+import java.util.Queue;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicLong;
 
 import io.netty.channel.Channel;
@@ -32,8 +34,11 @@
 import org.apache.spark.network.protocol.RpcFailure;
 import org.apache.spark.network.protocol.RpcResponse;
 import org.apache.spark.network.protocol.StreamChunkId;
+import org.apache.spark.network.protocol.StreamFailure;
+import org.apache.spark.network.protocol.StreamResponse;
 import org.apache.spark.network.server.MessageHandler;
 import org.apache.spark.network.util.NettyUtils;
+import org.apache.spark.network.util.TransportFrameDecoder;
 
 /**
  * Handler that processes server responses, in response to requests issued from a
@@ -50,6 +55,8 @@ public class TransportResponseHandler extends MessageHandler<ResponseMessage> {
 
   private final Map<Long, RpcResponseCallback> outstandingRpcs;
 
+  private final Queue<StreamCallback> streamCallbacks;
+
   /** Records the time (in system nanoseconds) that the last fetch or RPC request was sent. */
   private final AtomicLong timeOfLastRequestNs;
 
@@ -57,6 +64,7 @@ public TransportResponseHandler(Channel channel) {
     this.channel = channel;
     this.outstandingFetches = new ConcurrentHashMap<StreamChunkId, ChunkReceivedCallback>();
     this.outstandingRpcs = new ConcurrentHashMap<Long, RpcResponseCallback>();
+    this.streamCallbacks = new ConcurrentLinkedQueue<StreamCallback>();
     this.timeOfLastRequestNs = new AtomicLong(0);
   }
 
@@ -78,6 +86,10 @@ public void removeRpcRequest(long requestId) {
     outstandingRpcs.remove(requestId);
   }
 
+  public void addStreamCallback(StreamCallback callback) {
+    streamCallbacks.offer(callback);
+  }
+
   /**
    * Fire the failure callback for all outstanding requests. This is called when we have an
    * uncaught exception or pre-mature connection termination.
@@ -124,11 +136,11 @@ public void handle(ResponseMessage message) {
       if (listener == null) {
         logger.warn("Ignoring response for block {} from {} since it is not outstanding",
           resp.streamChunkId, remoteAddress);
-        resp.buffer.release();
+        resp.body.release();
       } else {
         outstandingFetches.remove(resp.streamChunkId);
-        listener.onSuccess(resp.streamChunkId.chunkIndex, resp.buffer);
-        resp.buffer.release();
+        listener.onSuccess(resp.streamChunkId.chunkIndex, resp.body);
+        resp.body.release();
       }
     } else if (message instanceof ChunkFetchFailure) {
       ChunkFetchFailure resp = (ChunkFetchFailure) message;
@@ -161,6 +173,34 @@ public void handle(ResponseMessage message) {
         outstandingRpcs.remove(resp.requestId);
         listener.onFailure(new RuntimeException(resp.errorString));
       }
+    } else if (message instanceof StreamResponse) {
+      StreamResponse resp = (StreamResponse) message;
+      StreamCallback callback = streamCallbacks.poll();
+      if (callback != null) {
+        StreamInterceptor interceptor = new StreamInterceptor(resp.streamId, resp.byteCount,
+          callback);
+        try {
+          TransportFrameDecoder frameDecoder = (TransportFrameDecoder)
+            channel.pipeline().get(TransportFrameDecoder.HANDLER_NAME);
+          frameDecoder.setInterceptor(interceptor);
+        } catch (Exception e) {
+          logger.error("Error installing stream handler.", e);
+        }
+      } else {
+        logger.error("Could not find callback for StreamResponse.");
+      }
+    } else if (message instanceof StreamFailure) {
+      StreamFailure resp = (StreamFailure) message;
+      StreamCallback callback = streamCallbacks.poll();
+      if (callback != null) {
+        try {
+          callback.onFailure(resp.streamId, new RuntimeException(resp.error));
+        } catch (IOException ioe) {
+          logger.warn("Error in stream failure handler.", ioe);
+        }
+      } else {
+        logger.warn("Stream failure with unknown callback: {}", resp.error);
+      }
     } else {
       throw new IllegalStateException("Unknown response type: " + message.type());
     }
@@ -175,4 +215,5 @@ public int numOutstandingRequests() {
   public long getTimeOfLastRequestNs() {
     return timeOfLastRequestNs.get();
   }
+
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchSuccess.java b/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchSuccess.java
index c962fb7ecf76d..e6a7e9a8b4145 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchSuccess.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchSuccess.java
@@ -30,13 +30,12 @@
  * may be written by Netty in a more efficient manner (i.e., zero-copy write).
  * Similarly, the client-side decoding will reuse the Netty ByteBuf as the buffer.
  */
-public final class ChunkFetchSuccess implements ResponseMessage {
+public final class ChunkFetchSuccess extends ResponseWithBody {
   public final StreamChunkId streamChunkId;
-  public final ManagedBuffer buffer;
 
   public ChunkFetchSuccess(StreamChunkId streamChunkId, ManagedBuffer buffer) {
+    super(buffer, true);
     this.streamChunkId = streamChunkId;
-    this.buffer = buffer;
   }
 
   @Override
@@ -53,6 +52,11 @@ public void encode(ByteBuf buf) {
     streamChunkId.encode(buf);
   }
 
+  @Override
+  public ResponseMessage createFailureResponse(String error) {
+    return new ChunkFetchFailure(streamChunkId, error);
+  }
+
   /** Decoding uses the given ByteBuf as our data, and will retain() it. */
   public static ChunkFetchSuccess decode(ByteBuf buf) {
     StreamChunkId streamChunkId = StreamChunkId.decode(buf);
@@ -63,14 +67,14 @@ public static ChunkFetchSuccess decode(ByteBuf buf) {
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(streamChunkId, buffer);
+    return Objects.hashCode(streamChunkId, body);
   }
 
   @Override
   public boolean equals(Object other) {
     if (other instanceof ChunkFetchSuccess) {
       ChunkFetchSuccess o = (ChunkFetchSuccess) other;
-      return streamChunkId.equals(o.streamChunkId) && buffer.equals(o.buffer);
+      return streamChunkId.equals(o.streamChunkId) && body.equals(o.body);
     }
     return false;
   }
@@ -79,7 +83,7 @@ public boolean equals(Object other) {
   public String toString() {
     return Objects.toStringHelper(this)
       .add("streamChunkId", streamChunkId)
-      .add("buffer", buffer)
+      .add("buffer", body)
       .toString();
   }
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/Message.java b/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
index d568370125fd4..d01598c20f16f 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
@@ -27,7 +27,8 @@ public interface Message extends Encodable {
   /** Preceding every serialized Message is its type, which allows us to deserialize it. */
   public static enum Type implements Encodable {
     ChunkFetchRequest(0), ChunkFetchSuccess(1), ChunkFetchFailure(2),
-    RpcRequest(3), RpcResponse(4), RpcFailure(5);
+    RpcRequest(3), RpcResponse(4), RpcFailure(5),
+    StreamRequest(6), StreamResponse(7), StreamFailure(8);
 
     private final byte id;
 
@@ -51,6 +52,9 @@ public static Type decode(ByteBuf buf) {
         case 3: return RpcRequest;
         case 4: return RpcResponse;
         case 5: return RpcFailure;
+        case 6: return StreamRequest;
+        case 7: return StreamResponse;
+        case 8: return StreamFailure;
         default: throw new IllegalArgumentException("Unknown message type: " + id);
       }
     }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/MessageDecoder.java b/network/common/src/main/java/org/apache/spark/network/protocol/MessageDecoder.java
index 81f8d7f96350f..3c04048f3821a 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/MessageDecoder.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/MessageDecoder.java
@@ -63,6 +63,15 @@ private Message decode(Message.Type msgType, ByteBuf in) {
       case RpcFailure:
         return RpcFailure.decode(in);
 
+      case StreamRequest:
+        return StreamRequest.decode(in);
+
+      case StreamResponse:
+        return StreamResponse.decode(in);
+
+      case StreamFailure:
+        return StreamFailure.decode(in);
+
       default:
         throw new IllegalArgumentException("Unexpected message type: " + msgType);
     }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java b/network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java
index 0f999f5dfe8d8..6cce97c807dc0 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java
@@ -45,27 +45,32 @@ public final class MessageEncoder extends MessageToMessageEncoder<Message> {
   public void encode(ChannelHandlerContext ctx, Message in, List<Object> out) {
     Object body = null;
     long bodyLength = 0;
+    boolean isBodyInFrame = false;
 
-    // Only ChunkFetchSuccesses have data besides the header.
+    // Detect ResponseWithBody messages and get the data buffer out of them.
     // The body is used in order to enable zero-copy transfer for the payload.
-    if (in instanceof ChunkFetchSuccess) {
-      ChunkFetchSuccess resp = (ChunkFetchSuccess) in;
+    if (in instanceof ResponseWithBody) {
+      ResponseWithBody resp = (ResponseWithBody) in;
       try {
-        bodyLength = resp.buffer.size();
-        body = resp.buffer.convertToNetty();
+        bodyLength = resp.body.size();
+        body = resp.body.convertToNetty();
+        isBodyInFrame = resp.isBodyInFrame;
       } catch (Exception e) {
-        // Re-encode this message as BlockFetchFailure.
-        logger.error(String.format("Error opening block %s for client %s",
-          resp.streamChunkId, ctx.channel().remoteAddress()), e);
-        encode(ctx, new ChunkFetchFailure(resp.streamChunkId, e.getMessage()), out);
+        // Re-encode this message as a failure response.
+        String error = e.getMessage() != null ? e.getMessage() : "null";
+        logger.error(String.format("Error processing %s for client %s",
+          resp, ctx.channel().remoteAddress()), e);
+        encode(ctx, resp.createFailureResponse(error), out);
         return;
       }
     }
 
     Message.Type msgType = in.type();
-    // All messages have the frame length, message type, and message itself.
+    // All messages have the frame length, message type, and message itself. The frame length
+    // may optionally include the length of the body data, depending on what message is being
+    // sent.
     int headerLength = 8 + msgType.encodedLength() + in.encodedLength();
-    long frameLength = headerLength + bodyLength;
+    long frameLength = headerLength + (isBodyInFrame ? bodyLength : 0);
     ByteBuf header = ctx.alloc().heapBuffer(headerLength);
     header.writeLong(frameLength);
     msgType.encode(header);
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/ResponseWithBody.java b/network/common/src/main/java/org/apache/spark/network/protocol/ResponseWithBody.java
new file mode 100644
index 0000000000000..67be77e39f711
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/ResponseWithBody.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.protocol;
+
+import com.google.common.base.Objects;
+import io.netty.buffer.ByteBuf;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
+
+/**
+ * Abstract class for response messages that contain a large data portion kept in a separate
+ * buffer. These messages are treated especially by MessageEncoder.
+ */
+public abstract class ResponseWithBody implements ResponseMessage {
+  public final ManagedBuffer body;
+  public final boolean isBodyInFrame;
+
+  protected ResponseWithBody(ManagedBuffer body, boolean isBodyInFrame) {
+    this.body = body;
+    this.isBodyInFrame = isBodyInFrame;
+  }
+
+  public abstract ResponseMessage createFailureResponse(String error);
+}
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/StreamFailure.java b/network/common/src/main/java/org/apache/spark/network/protocol/StreamFailure.java
new file mode 100644
index 0000000000000..e3dade2ebf905
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/StreamFailure.java
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.protocol;
+
+import com.google.common.base.Objects;
+import io.netty.buffer.ByteBuf;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
+
+/**
+ * Message indicating an error when transferring a stream.
+ */
+public final class StreamFailure implements ResponseMessage {
+  public final String streamId;
+  public final String error;
+
+  public StreamFailure(String streamId, String error) {
+    this.streamId = streamId;
+    this.error = error;
+  }
+
+  @Override
+  public Type type() { return Type.StreamFailure; }
+
+  @Override
+  public int encodedLength() {
+    return Encoders.Strings.encodedLength(streamId) + Encoders.Strings.encodedLength(error);
+  }
+
+  @Override
+  public void encode(ByteBuf buf) {
+    Encoders.Strings.encode(buf, streamId);
+    Encoders.Strings.encode(buf, error);
+  }
+
+  public static StreamFailure decode(ByteBuf buf) {
+    String streamId = Encoders.Strings.decode(buf);
+    String error = Encoders.Strings.decode(buf);
+    return new StreamFailure(streamId, error);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(streamId, error);
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (other instanceof StreamFailure) {
+      StreamFailure o = (StreamFailure) other;
+      return streamId.equals(o.streamId) && error.equals(o.error);
+    }
+    return false;
+  }
+
+  @Override
+  public String toString() {
+    return Objects.toStringHelper(this)
+      .add("streamId", streamId)
+      .add("error", error)
+      .toString();
+  }
+
+}
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/StreamRequest.java b/network/common/src/main/java/org/apache/spark/network/protocol/StreamRequest.java
new file mode 100644
index 0000000000000..821e8f53884d7
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/StreamRequest.java
@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.protocol;
+
+import com.google.common.base.Objects;
+import io.netty.buffer.ByteBuf;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
+
+/**
+ * Request to stream data from the remote end.
+ * <p>
+ * The stream ID is an arbitrary string that needs to be negotiated between the two endpoints before
+ * the data can be streamed.
+ */
+public final class StreamRequest implements RequestMessage {
+   public final String streamId;
+
+   public StreamRequest(String streamId) {
+     this.streamId = streamId;
+   }
+
+  @Override
+  public Type type() { return Type.StreamRequest; }
+
+  @Override
+  public int encodedLength() {
+    return Encoders.Strings.encodedLength(streamId);
+  }
+
+  @Override
+  public void encode(ByteBuf buf) {
+    Encoders.Strings.encode(buf, streamId);
+  }
+
+  public static StreamRequest decode(ByteBuf buf) {
+    String streamId = Encoders.Strings.decode(buf);
+    return new StreamRequest(streamId);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(streamId);
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (other instanceof StreamRequest) {
+      StreamRequest o = (StreamRequest) other;
+      return streamId.equals(o.streamId);
+    }
+    return false;
+  }
+
+  @Override
+  public String toString() {
+    return Objects.toStringHelper(this)
+      .add("streamId", streamId)
+      .toString();
+  }
+
+}
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/StreamResponse.java b/network/common/src/main/java/org/apache/spark/network/protocol/StreamResponse.java
new file mode 100644
index 0000000000000..ac5ab9a323a11
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/StreamResponse.java
@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.protocol;
+
+import com.google.common.base.Objects;
+import io.netty.buffer.ByteBuf;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
+
+/**
+ * Response to {@link StreamRequest} when the stream has been successfully opened.
+ * <p>
+ * Note the message itself does not contain the stream data. That is written separately by the
+ * sender. The receiver is expected to set a temporary channel handler that will consume the
+ * number of bytes this message says the stream has.
+ */
+public final class StreamResponse extends ResponseWithBody {
+   public final String streamId;
+   public final long byteCount;
+
+   public StreamResponse(String streamId, long byteCount, ManagedBuffer buffer) {
+     super(buffer, false);
+     this.streamId = streamId;
+     this.byteCount = byteCount;
+   }
+
+  @Override
+  public Type type() { return Type.StreamResponse; }
+
+  @Override
+  public int encodedLength() {
+    return 8 + Encoders.Strings.encodedLength(streamId);
+  }
+
+  /** Encoding does NOT include 'buffer' itself. See {@link MessageEncoder}. */
+  @Override
+  public void encode(ByteBuf buf) {
+    Encoders.Strings.encode(buf, streamId);
+    buf.writeLong(byteCount);
+  }
+
+  @Override
+  public ResponseMessage createFailureResponse(String error) {
+    return new StreamFailure(streamId, error);
+  }
+
+  public static StreamResponse decode(ByteBuf buf) {
+    String streamId = Encoders.Strings.decode(buf);
+    long byteCount = buf.readLong();
+    return new StreamResponse(streamId, byteCount, null);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(byteCount, streamId);
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (other instanceof StreamResponse) {
+      StreamResponse o = (StreamResponse) other;
+      return byteCount == o.byteCount && streamId.equals(o.streamId);
+    }
+    return false;
+  }
+
+  @Override
+  public String toString() {
+    return Objects.toStringHelper(this)
+      .add("streamId", streamId)
+      .add("byteCount", byteCount)
+      .toString();
+  }
+
+}
diff --git a/network/common/src/main/java/org/apache/spark/network/server/StreamManager.java b/network/common/src/main/java/org/apache/spark/network/server/StreamManager.java
index aaa677c965640..3f0155957a140 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/StreamManager.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/StreamManager.java
@@ -46,6 +46,19 @@ public abstract class StreamManager {
    */
   public abstract ManagedBuffer getChunk(long streamId, int chunkIndex);
 
+  /**
+   * Called in response to a stream() request. The returned data is streamed to the client
+   * through a single TCP connection.
+   *
+   * Note the <code>streamId</code> argument is not related to the similarly named argument in the
+   * {@link #getChunk(long, int)} method.
+   *
+   * @param streamId id of a stream that has been previously registered with the StreamManager.
+   */
+  public ManagedBuffer openStream(String streamId) {
+    throw new UnsupportedOperationException();
+  }
+
   /**
    * Associates a stream with a single client connection, which is guaranteed to be the only reader
    * of the stream. The getChunk() method will be called serially on this connection and once the
diff --git a/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java b/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
index 9b8b047b49a86..4f67bd573be21 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
@@ -35,6 +35,9 @@
 import org.apache.spark.network.protocol.ChunkFetchSuccess;
 import org.apache.spark.network.protocol.RpcFailure;
 import org.apache.spark.network.protocol.RpcResponse;
+import org.apache.spark.network.protocol.StreamFailure;
+import org.apache.spark.network.protocol.StreamRequest;
+import org.apache.spark.network.protocol.StreamResponse;
 import org.apache.spark.network.util.NettyUtils;
 
 /**
@@ -92,6 +95,8 @@ public void handle(RequestMessage request) {
       processFetchRequest((ChunkFetchRequest) request);
     } else if (request instanceof RpcRequest) {
       processRpcRequest((RpcRequest) request);
+    } else if (request instanceof StreamRequest) {
+      processStreamRequest((StreamRequest) request);
     } else {
       throw new IllegalArgumentException("Unknown request type: " + request);
     }
@@ -117,6 +122,21 @@ private void processFetchRequest(final ChunkFetchRequest req) {
     respond(new ChunkFetchSuccess(req.streamChunkId, buf));
   }
 
+  private void processStreamRequest(final StreamRequest req) {
+    final String client = NettyUtils.getRemoteAddress(channel);
+    ManagedBuffer buf;
+    try {
+      buf = streamManager.openStream(req.streamId);
+    } catch (Exception e) {
+      logger.error(String.format(
+        "Error opening stream %s for request from %s", req.streamId, client), e);
+      respond(new StreamFailure(req.streamId, Throwables.getStackTraceAsString(e)));
+      return;
+    }
+
+    respond(new StreamResponse(req.streamId, buf.size(), buf));
+  }
+
   private void processRpcRequest(final RpcRequest req) {
     try {
       rpcHandler.receive(reverseClient, req.message, new RpcResponseCallback() {
diff --git a/network/common/src/main/java/org/apache/spark/network/util/NettyUtils.java b/network/common/src/main/java/org/apache/spark/network/util/NettyUtils.java
index 26c6399ce7dbc..caa7260bc8281 100644
--- a/network/common/src/main/java/org/apache/spark/network/util/NettyUtils.java
+++ b/network/common/src/main/java/org/apache/spark/network/util/NettyUtils.java
@@ -89,13 +89,8 @@ public static Class<? extends ServerChannel> getServerChannelClass(IOMode mode)
    * Creates a LengthFieldBasedFrameDecoder where the first 8 bytes are the length of the frame.
    * This is used before all decoders.
    */
-  public static ByteToMessageDecoder createFrameDecoder() {
-    // maxFrameLength = 2G
-    // lengthFieldOffset = 0
-    // lengthFieldLength = 8
-    // lengthAdjustment = -8, i.e. exclude the 8 byte length itself
-    // initialBytesToStrip = 8, i.e. strip out the length field itself
-    return new LengthFieldBasedFrameDecoder(Integer.MAX_VALUE, 0, 8, -8, 8);
+  public static TransportFrameDecoder createFrameDecoder() {
+    return new TransportFrameDecoder();
   }
 
   /** Returns the remote address on the channel or "&lt;unknown remote&gt;" if none exists. */
diff --git a/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java b/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
new file mode 100644
index 0000000000000..272ea84e6180d
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
@@ -0,0 +1,154 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.util;
+
+import com.google.common.base.Preconditions;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelInboundHandlerAdapter;
+
+/**
+ * A customized frame decoder that allows intercepting raw data.
+ * <p>
+ * This behaves like Netty's frame decoder (with harcoded parameters that match this library's
+ * needs), except it allows an interceptor to be installed to read data directly before it's
+ * framed.
+ * <p>
+ * Unlike Netty's frame decoder, each frame is dispatched to child handlers as soon as it's
+ * decoded, instead of building as many frames as the current buffer allows and dispatching
+ * all of them. This allows a child handler to install an interceptor if needed.
+ * <p>
+ * If an interceptor is installed, framing stops, and data is instead fed directly to the
+ * interceptor. When the interceptor indicates that it doesn't need to read any more data,
+ * framing resumes. Interceptors should not hold references to the data buffers provided
+ * to their handle() method.
+ */
+public class TransportFrameDecoder extends ChannelInboundHandlerAdapter {
+
+  public static final String HANDLER_NAME = "frameDecoder";
+  private static final int LENGTH_SIZE = 8;
+  private static final int MAX_FRAME_SIZE = Integer.MAX_VALUE;
+
+  private CompositeByteBuf buffer;
+  private volatile Interceptor interceptor;
+
+  @Override
+  public void channelRead(ChannelHandlerContext ctx, Object data) throws Exception {
+    ByteBuf in = (ByteBuf) data;
+
+    if (buffer == null) {
+      buffer = in.alloc().compositeBuffer();
+    }
+
+    buffer.writeBytes(in);
+
+    while (buffer.isReadable()) {
+      feedInterceptor();
+      if (interceptor != null) {
+        continue;
+      }
+
+      ByteBuf frame = decodeNext();
+      if (frame != null) {
+        ctx.fireChannelRead(frame);
+      } else {
+        break;
+      }
+    }
+
+    // We can't discard read sub-buffers if there are other references to the buffer (e.g.
+    // through slices used for framing). This assumes that code that retains references
+    // will call retain() from the thread that called "fireChannelRead()" above, otherwise
+    // ref counting will go awry.
+    if (buffer != null && buffer.refCnt() == 1) {
+      buffer.discardReadComponents();
+    }
+  }
+
+  protected ByteBuf decodeNext() throws Exception {
+    if (buffer.readableBytes() < LENGTH_SIZE) {
+      return null;
+    }
+
+    int frameLen = (int) buffer.readLong() - LENGTH_SIZE;
+    if (buffer.readableBytes() < frameLen) {
+      buffer.readerIndex(buffer.readerIndex() - LENGTH_SIZE);
+      return null;
+    }
+
+    Preconditions.checkArgument(frameLen < MAX_FRAME_SIZE, "Too large frame: %s", frameLen);
+    Preconditions.checkArgument(frameLen > 0, "Frame length should be positive: %s", frameLen);
+
+    ByteBuf frame = buffer.readSlice(frameLen);
+    frame.retain();
+    return frame;
+  }
+
+  @Override
+  public void channelInactive(ChannelHandlerContext ctx) throws Exception {
+    if (buffer != null) {
+      if (buffer.isReadable()) {
+        feedInterceptor();
+      }
+      buffer.release();
+    }
+    if (interceptor != null) {
+      interceptor.channelInactive();
+    }
+    super.channelInactive(ctx);
+  }
+
+  @Override
+  public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
+    if (interceptor != null) {
+      interceptor.exceptionCaught(cause);
+    }
+    super.exceptionCaught(ctx, cause);
+  }
+
+  public void setInterceptor(Interceptor interceptor) {
+    Preconditions.checkState(this.interceptor == null, "Already have an interceptor.");
+    this.interceptor = interceptor;
+  }
+
+  private void feedInterceptor() throws Exception {
+    if (interceptor != null && !interceptor.handle(buffer)) {
+      interceptor = null;
+    }
+  }
+
+  public static interface Interceptor {
+
+    /**
+     * Handles data received from the remote end.
+     *
+     * @param data Buffer containing data.
+     * @return "true" if the interceptor expects more data, "false" to uninstall the interceptor.
+     */
+    boolean handle(ByteBuf data) throws Exception;
+
+    /** Called if an exception is thrown in the channel pipeline. */
+    void exceptionCaught(Throwable cause) throws Exception;
+
+    /** Called if the channel is closed and the interceptor is still installed. */
+    void channelInactive() throws Exception;
+
+  }
+
+}
diff --git a/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java b/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
index d500bc3c98a78..22b451fc0e60e 100644
--- a/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
@@ -39,6 +39,9 @@
 import org.apache.spark.network.protocol.RpcRequest;
 import org.apache.spark.network.protocol.RpcResponse;
 import org.apache.spark.network.protocol.StreamChunkId;
+import org.apache.spark.network.protocol.StreamFailure;
+import org.apache.spark.network.protocol.StreamRequest;
+import org.apache.spark.network.protocol.StreamResponse;
 import org.apache.spark.network.util.ByteArrayWritableChannel;
 import org.apache.spark.network.util.NettyUtils;
 
@@ -80,6 +83,7 @@ public void requests() {
     testClientToServer(new ChunkFetchRequest(new StreamChunkId(1, 2)));
     testClientToServer(new RpcRequest(12345, new byte[0]));
     testClientToServer(new RpcRequest(12345, new byte[100]));
+    testClientToServer(new StreamRequest("abcde"));
   }
 
   @Test
@@ -92,6 +96,10 @@ public void responses() {
     testServerToClient(new RpcResponse(12345, new byte[1000]));
     testServerToClient(new RpcFailure(0, "this is an error"));
     testServerToClient(new RpcFailure(0, ""));
+    // Note: buffer size must be "0" since StreamResponse's buffer is written differently to the
+    // channel and cannot be tested like this.
+    testServerToClient(new StreamResponse("anId", 12345L, new TestManagedBuffer(0)));
+    testServerToClient(new StreamFailure("anId", "this is an error"));
   }
 
   /**
diff --git a/network/common/src/test/java/org/apache/spark/network/StreamSuite.java b/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
new file mode 100644
index 0000000000000..6dcec831dec71
--- /dev/null
+++ b/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
@@ -0,0 +1,325 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network;
+
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import com.google.common.io.Files;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import static org.junit.Assert.*;
+
+import org.apache.spark.network.buffer.FileSegmentManagedBuffer;
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NioManagedBuffer;
+import org.apache.spark.network.client.RpcResponseCallback;
+import org.apache.spark.network.client.StreamCallback;
+import org.apache.spark.network.client.TransportClient;
+import org.apache.spark.network.client.TransportClientFactory;
+import org.apache.spark.network.server.RpcHandler;
+import org.apache.spark.network.server.StreamManager;
+import org.apache.spark.network.server.TransportServer;
+import org.apache.spark.network.util.SystemPropertyConfigProvider;
+import org.apache.spark.network.util.TransportConf;
+
+public class StreamSuite {
+  private static final String[] STREAMS = { "largeBuffer", "smallBuffer", "file" };
+
+  private static TransportServer server;
+  private static TransportClientFactory clientFactory;
+  private static File testFile;
+  private static File tempDir;
+
+  private static ByteBuffer smallBuffer;
+  private static ByteBuffer largeBuffer;
+
+  private static ByteBuffer createBuffer(int bufSize) {
+    ByteBuffer buf = ByteBuffer.allocate(bufSize);
+    for (int i = 0; i < bufSize; i ++) {
+      buf.put((byte) i);
+    }
+    buf.flip();
+    return buf;
+  }
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    tempDir = Files.createTempDir();
+    smallBuffer = createBuffer(100);
+    largeBuffer = createBuffer(100000);
+
+    testFile = File.createTempFile("stream-test-file", "txt", tempDir);
+    FileOutputStream fp = new FileOutputStream(testFile);
+    try {
+      Random rnd = new Random();
+      for (int i = 0; i < 512; i++) {
+        byte[] fileContent = new byte[1024];
+        rnd.nextBytes(fileContent);
+        fp.write(fileContent);
+      }
+    } finally {
+      fp.close();
+    }
+
+    final TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+    final StreamManager streamManager = new StreamManager() {
+      @Override
+      public ManagedBuffer getChunk(long streamId, int chunkIndex) {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public ManagedBuffer openStream(String streamId) {
+        switch (streamId) {
+          case "largeBuffer":
+            return new NioManagedBuffer(largeBuffer);
+          case "smallBuffer":
+            return new NioManagedBuffer(smallBuffer);
+          case "file":
+            return new FileSegmentManagedBuffer(conf, testFile, 0, testFile.length());
+          default:
+            throw new IllegalArgumentException("Invalid stream: " + streamId);
+        }
+      }
+    };
+    RpcHandler handler = new RpcHandler() {
+      @Override
+      public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+        throw new UnsupportedOperationException();
+      }
+
+      @Override
+      public StreamManager getStreamManager() {
+        return streamManager;
+      }
+    };
+    TransportContext context = new TransportContext(conf, handler);
+    server = context.createServer();
+    clientFactory = context.createClientFactory();
+  }
+
+  @AfterClass
+  public static void tearDown() {
+    server.close();
+    clientFactory.close();
+    if (tempDir != null) {
+      for (File f : tempDir.listFiles()) {
+        f.delete();
+      }
+      tempDir.delete();
+    }
+  }
+
+  @Test
+  public void testSingleStream() throws Throwable {
+    TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
+    try {
+      StreamTask task = new StreamTask(client, "largeBuffer", TimeUnit.SECONDS.toMillis(5));
+      task.run();
+      task.check();
+    } finally {
+      client.close();
+    }
+  }
+
+  @Test
+  public void testMultipleStreams() throws Throwable {
+    TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
+    try {
+      for (int i = 0; i < 20; i++) {
+        StreamTask task = new StreamTask(client, STREAMS[i % STREAMS.length],
+          TimeUnit.SECONDS.toMillis(5));
+        task.run();
+        task.check();
+      }
+    } finally {
+      client.close();
+    }
+  }
+
+  @Test
+  public void testConcurrentStreams() throws Throwable {
+    ExecutorService executor = Executors.newFixedThreadPool(20);
+    TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
+
+    try {
+      List<StreamTask> tasks = new ArrayList<>();
+      for (int i = 0; i < 20; i++) {
+        StreamTask task = new StreamTask(client, STREAMS[i % STREAMS.length],
+          TimeUnit.SECONDS.toMillis(20));
+        tasks.add(task);
+        executor.submit(task);
+      }
+
+      executor.shutdown();
+      assertTrue("Timed out waiting for tasks.", executor.awaitTermination(30, TimeUnit.SECONDS));
+      for (StreamTask task : tasks) {
+        task.check();
+      }
+    } finally {
+      executor.shutdownNow();
+      client.close();
+    }
+  }
+
+  private static class StreamTask implements Runnable {
+
+    private final TransportClient client;
+    private final String streamId;
+    private final long timeoutMs;
+    private Throwable error;
+
+    StreamTask(TransportClient client, String streamId, long timeoutMs) {
+      this.client = client;
+      this.streamId = streamId;
+      this.timeoutMs = timeoutMs;
+    }
+
+    @Override
+    public void run() {
+      ByteBuffer srcBuffer = null;
+      OutputStream out = null;
+      File outFile = null;
+      try {
+        ByteArrayOutputStream baos = null;
+
+        switch (streamId) {
+          case "largeBuffer":
+            baos = new ByteArrayOutputStream();
+            out = baos;
+            srcBuffer = largeBuffer;
+            break;
+          case "smallBuffer":
+            baos = new ByteArrayOutputStream();
+            out = baos;
+            srcBuffer = smallBuffer;
+            break;
+          case "file":
+            outFile = File.createTempFile("data", ".tmp", tempDir);
+            out = new FileOutputStream(outFile);
+            break;
+          default:
+            throw new IllegalArgumentException(streamId);
+        }
+
+        TestCallback callback = new TestCallback(out);
+        client.stream(streamId, callback);
+        waitForCompletion(callback);
+
+        if (srcBuffer == null) {
+          assertTrue("File stream did not match.", Files.equal(testFile, outFile));
+        } else {
+          ByteBuffer base;
+          synchronized (srcBuffer) {
+            base = srcBuffer.duplicate();
+          }
+          byte[] result = baos.toByteArray();
+          byte[] expected = new byte[base.remaining()];
+          base.get(expected);
+          assertEquals(expected.length, result.length);
+          assertTrue("buffers don't match", Arrays.equals(expected, result));
+        }
+      } catch (Throwable t) {
+        error = t;
+      } finally {
+        if (out != null) {
+          try {
+            out.close();
+          } catch (Exception e) {
+            // ignore.
+          }
+        }
+        if (outFile != null) {
+          outFile.delete();
+        }
+      }
+    }
+
+    public void check() throws Throwable {
+      if (error != null) {
+        throw error;
+      }
+    }
+
+    private void waitForCompletion(TestCallback callback) throws Exception {
+      long now = System.currentTimeMillis();
+      long deadline = now + timeoutMs;
+      synchronized (callback) {
+        while (!callback.completed && now < deadline) {
+          callback.wait(deadline - now);
+          now = System.currentTimeMillis();
+        }
+      }
+      assertTrue("Timed out waiting for stream.", callback.completed);
+      assertNull(callback.error);
+    }
+
+  }
+
+  private static class TestCallback implements StreamCallback {
+
+    private final OutputStream out;
+    public volatile boolean completed;
+    public volatile Throwable error;
+
+    TestCallback(OutputStream out) {
+      this.out = out;
+      this.completed = false;
+    }
+
+    @Override
+    public void onData(String streamId, ByteBuffer buf) throws IOException {
+      byte[] tmp = new byte[buf.remaining()];
+      buf.get(tmp);
+      out.write(tmp);
+    }
+
+    @Override
+    public void onComplete(String streamId) throws IOException {
+      out.close();
+      synchronized (this) {
+        completed = true;
+        notifyAll();
+      }
+    }
+
+    @Override
+    public void onFailure(String streamId, Throwable cause) {
+      error = cause;
+      synchronized (this) {
+        completed = true;
+        notifyAll();
+      }
+    }
+
+  }
+
+}
diff --git a/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java b/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java
new file mode 100644
index 0000000000000..ca74f0a00cf9d
--- /dev/null
+++ b/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.util;
+
+import java.nio.ByteBuffer;
+import java.util.Random;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import io.netty.channel.ChannelHandlerContext;
+import org.junit.Test;
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+
+public class TransportFrameDecoderSuite {
+
+  @Test
+  public void testFrameDecoding() throws Exception {
+    Random rnd = new Random();
+    TransportFrameDecoder decoder = new TransportFrameDecoder();
+    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+
+    final int frameCount = 100;
+    ByteBuf data = Unpooled.buffer();
+    try {
+      for (int i = 0; i < frameCount; i++) {
+        byte[] frame = new byte[1024 * (rnd.nextInt(31) + 1)];
+        data.writeLong(frame.length + 8);
+        data.writeBytes(frame);
+      }
+
+      while (data.isReadable()) {
+        int size = rnd.nextInt(16 * 1024) + 256;
+        decoder.channelRead(ctx, data.readSlice(Math.min(data.readableBytes(), size)));
+      }
+
+      verify(ctx, times(frameCount)).fireChannelRead(any(ByteBuf.class));
+    } finally {
+      data.release();
+    }
+  }
+
+  @Test
+  public void testInterception() throws Exception {
+    final int interceptedReads = 3;
+    TransportFrameDecoder decoder = new TransportFrameDecoder();
+    TransportFrameDecoder.Interceptor interceptor = spy(new MockInterceptor(interceptedReads));
+    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+
+    byte[] data = new byte[8];
+    ByteBuf len = Unpooled.copyLong(8 + data.length);
+    ByteBuf dataBuf = Unpooled.wrappedBuffer(data);
+
+    try {
+      decoder.setInterceptor(interceptor);
+      for (int i = 0; i < interceptedReads; i++) {
+        decoder.channelRead(ctx, dataBuf);
+        dataBuf.release();
+        dataBuf = Unpooled.wrappedBuffer(data);
+      }
+      decoder.channelRead(ctx, len);
+      decoder.channelRead(ctx, dataBuf);
+      verify(interceptor, times(interceptedReads)).handle(any(ByteBuf.class));
+      verify(ctx).fireChannelRead(any(ByteBuffer.class));
+    } finally {
+      len.release();
+      dataBuf.release();
+    }
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testNegativeFrameSize() throws Exception {
+    testInvalidFrame(-1);
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testEmptyFrame() throws Exception {
+    // 8 because frame size includes the frame length.
+    testInvalidFrame(8);
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testLargeFrame() throws Exception {
+    // Frame length includes the frame size field, so need to add a few more bytes.
+    testInvalidFrame(Integer.MAX_VALUE + 9);
+  }
+
+  private void testInvalidFrame(long size) throws Exception {
+    TransportFrameDecoder decoder = new TransportFrameDecoder();
+    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+    ByteBuf frame = Unpooled.copyLong(size);
+    try {
+      decoder.channelRead(ctx, frame);
+    } finally {
+      frame.release();
+    }
+  }
+
+  private static class MockInterceptor implements TransportFrameDecoder.Interceptor {
+
+    private int remainingReads;
+
+    MockInterceptor(int readCount) {
+      this.remainingReads = readCount;
+    }
+
+    @Override
+    public boolean handle(ByteBuf data) throws Exception {
+      data.readerIndex(data.readerIndex() + data.readableBytes());
+      assertFalse(data.isReadable());
+      remainingReads -= 1;
+      return remainingReads != 0;
+    }
+
+    @Override
+    public void exceptionCaught(Throwable cause) throws Exception {
+
+    }
+
+    @Override
+    public void channelInactive() throws Exception {
+
+    }
+
+  }
+
+}

From cd1df662386c599a9d0968b9fc14f27b0883d285 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 4 Nov 2015 09:32:30 -0800
Subject: [PATCH 049/518] [SPARK-11485][SQL] Make DataFrameHolder and
 DatasetHolder public.

These two classes should be public, since they are used in public code.

Author: Reynold Xin <rxin@databricks.com>

Closes #9445 from rxin/SPARK-11485.
---
 project/MimaExcludes.scala                            |  3 +++
 .../scala/org/apache/spark/sql/DataFrameHolder.scala  |  7 ++++++-
 .../scala/org/apache/spark/sql/DatasetHolder.scala    | 11 ++++++++---
 .../scala/org/apache/spark/sql/SQLImplicits.scala     |  4 ++++
 4 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index eeef96c378bdb..90dc947d4e588 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -161,6 +161,9 @@ object MimaExcludes {
           "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$23"),
         ProblemFilters.exclude[MissingMethodProblem](
           "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$24")
+      ) ++ Seq(
+        // SPARK-11485
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.DataFrameHolder.df")
       )
     case v if v.startsWith("1.5") =>
       Seq(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameHolder.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameHolder.scala
index 2f19ec0403017..3b30337f1f877 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameHolder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameHolder.scala
@@ -20,9 +20,14 @@ package org.apache.spark.sql
 /**
  * A container for a [[DataFrame]], used for implicit conversions.
  *
+ * To use this, import implicit conversions in SQL:
+ * {{{
+ *   import sqlContext.implicits._
+ * }}}
+ *
  * @since 1.3.0
  */
-private[sql] case class DataFrameHolder(df: DataFrame) {
+case class DataFrameHolder private[sql](private val df: DataFrame) {
 
   // This is declared with parentheses to prevent the Scala compiler from treating
   // `rdd.toDF("1")` as invoking this toDF and then apply on the returned DataFrame.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala b/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala
index 17817cbcc5e05..45f0098b92887 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala
@@ -18,11 +18,16 @@
 package org.apache.spark.sql
 
 /**
- * A container for a [[DataFrame]], used for implicit conversions.
+ * A container for a [[Dataset]], used for implicit conversions.
  *
- * @since 1.3.0
+ * To use this, import implicit conversions in SQL:
+ * {{{
+ *   import sqlContext.implicits._
+ * }}}
+ *
+ * @since 1.6.0
  */
-private[sql] case class DatasetHolder[T](df: Dataset[T]) {
+case class DatasetHolder[T] private[sql](private val df: Dataset[T]) {
 
   // This is declared with parentheses to prevent the Scala compiler from treating
   // `rdd.toDF("1")` as invoking this toDF and then apply on the returned DataFrame.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
index f2904e270811e..6da46a5f7ef9a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
@@ -52,6 +52,10 @@ abstract class SQLImplicits {
     DatasetHolder(_sqlContext.createDataset(rdd))
   }
 
+  /**
+   * Creates a [[Dataset]] from a local Seq.
+   * @since 1.6.0
+   */
   implicit def localSeqToDatasetHolder[T : Encoder](s: Seq[T]): DatasetHolder[T] = {
     DatasetHolder(_sqlContext.createDataset(s))
   }

From e0fc9c7e59848cb78f8d598898bfca004a3710d8 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 4 Nov 2015 09:33:30 -0800
Subject: [PATCH 050/518] [SPARK-11197][SQL] add doc for run SQL on files
 directly

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9467 from cloud-fan/doc.
---
 docs/sql-programming-guide.md | 38 +++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/docs/sql-programming-guide.md b/docs/sql-programming-guide.md
index 510b3599721a3..2fe5c36338899 100644
--- a/docs/sql-programming-guide.md
+++ b/docs/sql-programming-guide.md
@@ -882,6 +882,44 @@ saveDF(select(df, "name", "age"), "namesAndAges.parquet", "parquet")
 </div>
 </div>
 
+### Run SQL on files directly
+
+Instead of using read API to load a file into DataFrame and query it, you can also query that
+file directly with SQL.
+
+<div class="codetabs">
+<div data-lang="scala"  markdown="1">
+
+{% highlight scala %}
+val df = sqlContext.sql("SELECT * FROM parquet.`examples/src/main/resources/users.parquet`")
+{% endhighlight %}
+
+</div>
+
+<div data-lang="java"  markdown="1">
+
+{% highlight java %}
+DataFrame df = sqlContext.sql("SELECT * FROM parquet.`examples/src/main/resources/users.parquet`");
+{% endhighlight %}
+</div>
+
+<div data-lang="python"  markdown="1">
+
+{% highlight python %}
+df = sqlContext.sql("SELECT * FROM parquet.`examples/src/main/resources/users.parquet`")
+{% endhighlight %}
+
+</div>
+
+<div data-lang="r"  markdown="1">
+
+{% highlight r %}
+df <- sql(sqlContext, "SELECT * FROM parquet.`examples/src/main/resources/users.parquet`")
+{% endhighlight %}
+
+</div>
+</div>
+
 ### Save Modes
 
 Save operations can optionally take a `SaveMode`, that specifies how to handle existing data if

From 3bd6f5d2ae503468de0e218d51c331e249a862bb Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 4 Nov 2015 09:34:52 -0800
Subject: [PATCH 051/518] [SPARK-11490][SQL] variance should alias var_samp
 instead of var_pop.

stddev is an alias for stddev_samp. variance should be consistent with stddev.

Also took the chance to remove internal Stddev and Variance, and only kept StddevSamp/StddevPop and VarianceSamp/VariancePop.

Author: Reynold Xin <rxin@databricks.com>

Closes #9449 from rxin/SPARK-11490.
---
 .../catalyst/analysis/FunctionRegistry.scala  |  4 +-
 .../catalyst/analysis/HiveTypeCoercion.scala  |  2 -
 .../spark/sql/catalyst/dsl/package.scala      |  8 ----
 .../expressions/aggregate/functions.scala     | 29 ------------
 .../expressions/aggregate/utils.scala         | 12 -----
 .../sql/catalyst/expressions/aggregates.scala | 45 +++++--------------
 .../org/apache/spark/sql/DataFrame.scala      |  2 +-
 .../org/apache/spark/sql/GroupedData.scala    |  4 +-
 .../org/apache/spark/sql/functions.scala      |  9 ++--
 .../spark/sql/DataFrameAggregateSuite.scala   | 17 +++----
 .../org/apache/spark/sql/SQLQuerySuite.scala  | 14 +++---
 11 files changed, 32 insertions(+), 114 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index 24c1a7b7ac5af..d4334d16289a5 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -187,11 +187,11 @@ object FunctionRegistry {
     expression[Max]("max"),
     expression[Average]("mean"),
     expression[Min]("min"),
-    expression[Stddev]("stddev"),
+    expression[StddevSamp]("stddev"),
     expression[StddevPop]("stddev_pop"),
     expression[StddevSamp]("stddev_samp"),
     expression[Sum]("sum"),
-    expression[Variance]("variance"),
+    expression[VarianceSamp]("variance"),
     expression[VariancePop]("var_pop"),
     expression[VarianceSamp]("var_samp"),
     expression[Skewness]("skewness"),
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
index 3c675672dab85..84e2b1366f626 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
@@ -297,10 +297,8 @@ object HiveTypeCoercion {
       case Sum(e @ StringType()) => Sum(Cast(e, DoubleType))
       case SumDistinct(e @ StringType()) => Sum(Cast(e, DoubleType))
       case Average(e @ StringType()) => Average(Cast(e, DoubleType))
-      case Stddev(e @ StringType()) => Stddev(Cast(e, DoubleType))
       case StddevPop(e @ StringType()) => StddevPop(Cast(e, DoubleType))
       case StddevSamp(e @ StringType()) => StddevSamp(Cast(e, DoubleType))
-      case Variance(e @ StringType()) => Variance(Cast(e, DoubleType))
       case VariancePop(e @ StringType()) => VariancePop(Cast(e, DoubleType))
       case VarianceSamp(e @ StringType()) => VarianceSamp(Cast(e, DoubleType))
       case Skewness(e @ StringType()) => Skewness(Cast(e, DoubleType))
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
index 787f67a297e33..d8df66430a695 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
@@ -159,14 +159,6 @@ package object dsl {
     def lower(e: Expression): Expression = Lower(e)
     def sqrt(e: Expression): Expression = Sqrt(e)
     def abs(e: Expression): Expression = Abs(e)
-    def stddev(e: Expression): Expression = Stddev(e)
-    def stddev_pop(e: Expression): Expression = StddevPop(e)
-    def stddev_samp(e: Expression): Expression = StddevSamp(e)
-    def variance(e: Expression): Expression = Variance(e)
-    def var_pop(e: Expression): Expression = VariancePop(e)
-    def var_samp(e: Expression): Expression = VarianceSamp(e)
-    def skewness(e: Expression): Expression = Skewness(e)
-    def kurtosis(e: Expression): Expression = Kurtosis(e)
 
     implicit class DslSymbol(sym: Symbol) extends ImplicitAttribute { def s: String = sym.name }
     // TODO more implicit class for literal?
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala
index f2c3eca095115..10dc5e64b7ec9 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala
@@ -328,13 +328,6 @@ case class Min(child: Expression) extends DeclarativeAggregate {
   override val evaluateExpression = min
 }
 
-// Compute the sample standard deviation of a column
-case class Stddev(child: Expression) extends StddevAgg(child) {
-
-  override def isSample: Boolean = true
-  override def prettyName: String = "stddev"
-}
-
 // Compute the population standard deviation of a column
 case class StddevPop(child: Expression) extends StddevAgg(child) {
 
@@ -1274,28 +1267,6 @@ abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate w
   }
 }
 
-case class Variance(child: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0) extends CentralMomentAgg(child) {
-
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
-    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
-
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
-    copy(inputAggBufferOffset = newInputAggBufferOffset)
-
-  override def prettyName: String = "variance"
-
-  override protected val momentOrder = 2
-
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
-    require(moments.length == momentOrder + 1,
-      s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
-
-    if (n == 0.0) Double.NaN else moments(2) / n
-  }
-}
-
 case class VarianceSamp(child: Expression,
     mutableAggBufferOffset: Int = 0,
     inputAggBufferOffset: Int = 0) extends CentralMomentAgg(child) {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/utils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/utils.scala
index 564174f9b64e4..644c6211d5f31 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/utils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/utils.scala
@@ -97,12 +97,6 @@ object Utils {
             mode = aggregate.Complete,
             isDistinct = false)
 
-        case expressions.Stddev(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Stddev(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
         case expressions.StddevPop(child) =>
           aggregate.AggregateExpression2(
             aggregateFunction = aggregate.StddevPop(child),
@@ -139,12 +133,6 @@ object Utils {
             mode = aggregate.Complete,
             isDistinct = false)
 
-        case expressions.Variance(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Variance(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
         case expressions.VariancePop(child) =>
           aggregate.AggregateExpression2(
             aggregateFunction = aggregate.VariancePop(child),
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
index bf59660c385ed..89d63abd9f272 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
@@ -785,13 +785,6 @@ abstract class StddevAgg1(child: Expression) extends UnaryExpression with Partia
 
 }
 
-// Compute the sample standard deviation of a column
-case class Stddev(child: Expression) extends StddevAgg1(child) {
-
-  override def toString: String = s"STDDEV($child)"
-  override def isSample: Boolean = true
-}
-
 // Compute the population standard deviation of a column
 case class StddevPop(child: Expression) extends StddevAgg1(child) {
 
@@ -807,20 +800,21 @@ case class StddevSamp(child: Expression) extends StddevAgg1(child) {
 }
 
 case class ComputePartialStd(child: Expression) extends UnaryExpression with AggregateExpression1 {
-    def this() = this(null)
-
-    override def children: Seq[Expression] = child :: Nil
-    override def nullable: Boolean = false
-    override def dataType: DataType = ArrayType(DoubleType)
-    override def toString: String = s"computePartialStddev($child)"
-    override def newInstance(): ComputePartialStdFunction =
-      new ComputePartialStdFunction(child, this)
+  def this() = this(null)
+
+  override def children: Seq[Expression] = child :: Nil
+  override def nullable: Boolean = false
+  override def dataType: DataType = ArrayType(DoubleType)
+  override def toString: String = s"computePartialStddev($child)"
+  override def newInstance(): ComputePartialStdFunction =
+    new ComputePartialStdFunction(child, this)
 }
 
 case class ComputePartialStdFunction (
     expr: Expression,
     base: AggregateExpression1
-) extends AggregateFunction1 {
+  ) extends AggregateFunction1 {
+
   def this() = this(null, null)  // Required for serialization
 
   private val computeType = DoubleType
@@ -1048,25 +1042,6 @@ case class Skewness(child: Expression) extends UnaryExpression with AggregateExp
   override def toString: String = s"SKEWNESS($child)"
 }
 
-// placeholder
-case class Variance(child: Expression) extends UnaryExpression with AggregateExpression1 {
-
-  override def newInstance(): AggregateFunction1 = {
-    throw new UnsupportedOperationException("AggregateExpression1 is no longer supported, " +
-      "please set spark.sql.useAggregate2 = true")
-  }
-
-  override def nullable: Boolean = false
-
-  override def dataType: DoubleType.type = DoubleType
-
-  override def foldable: Boolean = false
-
-  override def prettyName: String = "variance"
-
-  override def toString: String = s"VARIANCE($child)"
-}
-
 // placeholder
 case class VariancePop(child: Expression) extends UnaryExpression with AggregateExpression1 {
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index fc0ab632f9930..5e9c7efbbf160 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -1383,7 +1383,7 @@ class DataFrame private[sql](
     val statistics = List[(String, Expression => Expression)](
       "count" -> Count,
       "mean" -> Average,
-      "stddev" -> Stddev,
+      "stddev" -> StddevSamp,
       "min" -> Min,
       "max" -> Max)
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
index c2b2a4013d510..7cf66b65c8722 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
@@ -96,10 +96,10 @@ class GroupedData protected[sql](
       case "avg" | "average" | "mean" => Average
       case "max" => Max
       case "min" => Min
-      case "stddev" | "std" => Stddev
+      case "stddev" | "std" => StddevSamp
       case "stddev_pop" => StddevPop
       case "stddev_samp" => StddevSamp
-      case "variance" => Variance
+      case "variance" => VarianceSamp
       case "var_pop" => VariancePop
       case "var_samp" => VarianceSamp
       case "sum" => Sum
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index c8c52831668cd..c70c965a9b04c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -329,13 +329,12 @@ object functions {
   def skewness(e: Column): Column = Skewness(e.expr)
 
   /**
-   * Aggregate function: returns the unbiased sample standard deviation of
-   * the expression in a group.
+   * Aggregate function: alias for [[stddev_samp]].
    *
    * @group agg_funcs
    * @since 1.6.0
    */
-  def stddev(e: Column): Column = Stddev(e.expr)
+  def stddev(e: Column): Column = StddevSamp(e.expr)
 
   /**
    * Aggregate function: returns the unbiased sample standard deviation of
@@ -388,12 +387,12 @@ object functions {
   def sumDistinct(columnName: String): Column = sumDistinct(Column(columnName))
 
   /**
-   * Aggregate function: returns the population variance of the values in a group.
+   * Aggregate function: alias for [[var_samp]].
    *
    * @group agg_funcs
    * @since 1.6.0
    */
-  def variance(e: Column): Column = Variance(e.expr)
+  def variance(e: Column): Column = VarianceSamp(e.expr)
 
   /**
    * Aggregate function: returns the unbiased variance of the values in a group.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index 9b23977c765dc..b0e2ffaa60687 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -226,23 +226,18 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
     val absTol = 1e-8
 
     val sparkVariance = testData2.agg(variance('a))
-    val expectedVariance = Row(4.0 / 6.0)
-    checkAggregatesWithTol(sparkVariance, expectedVariance, absTol)
+    checkAggregatesWithTol(sparkVariance, Row(4.0 / 5.0), absTol)
     val sparkVariancePop = testData2.agg(var_pop('a))
-    checkAggregatesWithTol(sparkVariancePop, expectedVariance, absTol)
+    checkAggregatesWithTol(sparkVariancePop, Row(4.0 / 6.0), absTol)
 
     val sparkVarianceSamp = testData2.agg(var_samp('a))
-    val expectedVarianceSamp = Row(4.0 / 5.0)
-    checkAggregatesWithTol(sparkVarianceSamp, expectedVarianceSamp, absTol)
+    checkAggregatesWithTol(sparkVarianceSamp, Row(4.0 / 5.0), absTol)
 
     val sparkSkewness = testData2.agg(skewness('a))
-    val expectedSkewness = Row(0.0)
-    checkAggregatesWithTol(sparkSkewness, expectedSkewness, absTol)
+    checkAggregatesWithTol(sparkSkewness, Row(0.0), absTol)
 
     val sparkKurtosis = testData2.agg(kurtosis('a))
-    val expectedKurtosis = Row(-1.5)
-    checkAggregatesWithTol(sparkKurtosis, expectedKurtosis, absTol)
-
+    checkAggregatesWithTol(sparkKurtosis, Row(-1.5), absTol)
   }
 
   test("zero moments") {
@@ -251,7 +246,7 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
 
     checkAnswer(
       emptyTableData.agg(variance('a)),
-      Row(0.0))
+      Row(Double.NaN))
 
     checkAnswer(
       emptyTableData.agg(var_samp('a)),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index 6388a8b9c3720..5731a356243e5 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -536,7 +536,7 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
     checkAnswer(
       sql("SELECT SKEWNESS(a), KURTOSIS(a), MIN(a), MAX(a)," +
         "AVG(a), VARIANCE(a), STDDEV(a), SUM(a), COUNT(a) FROM nullInts"),
-      Row(0, -1.5, 1, 3, 2, 2.0 / 3.0, 1, 6, 3)
+      Row(0, -1.5, 1, 3, 2, 1.0, 1, 6, 3)
     )
   }
 
@@ -757,7 +757,7 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
   test("variance") {
     val absTol = 1e-8
     val sparkAnswer = sql("SELECT VARIANCE(a) FROM testData2")
-    val expectedAnswer = Row(4.0 / 6.0)
+    val expectedAnswer = Row(0.8)
     checkAggregatesWithTol(sparkAnswer, expectedAnswer, absTol)
   }
 
@@ -784,16 +784,16 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
 
   test("stddev agg") {
     checkAnswer(
-        sql("SELECT a, stddev(b), stddev_pop(b), stddev_samp(b) FROM testData2 GROUP BY a"),
+      sql("SELECT a, stddev(b), stddev_pop(b), stddev_samp(b) FROM testData2 GROUP BY a"),
       (1 to 3).map(i => Row(i, math.sqrt(1.0 / 2.0), math.sqrt(1.0 / 4.0), math.sqrt(1.0 / 2.0))))
   }
 
   test("variance agg") {
     val absTol = 1e-8
-    val sparkAnswer = sql("SELECT a, variance(b), var_samp(b), var_pop(b)" +
-      "FROM testData2 GROUP BY a")
-    val expectedAnswer = (1 to 3).map(i => Row(i, 1.0 / 4.0, 1.0 / 2.0, 1.0 / 4.0))
-    checkAggregatesWithTol(sparkAnswer, expectedAnswer, absTol)
+    checkAggregatesWithTol(
+      sql("SELECT a, variance(b), var_samp(b), var_pop(b) FROM testData2 GROUP BY a"),
+      (1 to 3).map(i => Row(i, 1.0 / 2.0, 1.0 / 2.0, 1.0 / 4.0)),
+      absTol)
   }
 
   test("skewness and kurtosis agg") {

From 987df4bfcafeca3633453c2d2f8e14d221fcef33 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 4 Nov 2015 10:04:51 -0800
Subject: [PATCH 052/518] Closes #9464


From de289bf279e14e47859b5fbcd70e97b9d0759f14 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Wed, 4 Nov 2015 10:56:32 -0800
Subject: [PATCH 053/518] [SPARK-10304][SQL] Following up checking valid dir
 structure for partition discovery

This patch follows up #8840.

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #9459 from viirya/detect_invalid_part_dir_following.
---
 .../datasources/PartitioningUtils.scala          | 14 +++++++++++++-
 .../parquet/ParquetPartitionDiscoverySuite.scala | 16 ++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
index 16dc23661c070..86bc3a1b6dab2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
@@ -81,6 +81,8 @@ private[sql] object PartitioningUtils {
       parsePartition(path, defaultPartitionName, typeInference)
     }.unzip
 
+    // We create pairs of (path -> path's partition value) here
+    // If the corresponding partition value is None, the pair will be skiped
     val pathsWithPartitionValues = paths.zip(partitionValues).flatMap(x => x._2.map(x._1 -> _))
 
     if (pathsWithPartitionValues.isEmpty) {
@@ -89,11 +91,21 @@ private[sql] object PartitioningUtils {
     } else {
       // This dataset is partitioned. We need to check whether all partitions have the same
       // partition columns and resolve potential type conflicts.
+
+      // Check if there is conflicting directory structure.
+      // For the paths such as:
+      // var paths = Seq(
+      //   "hdfs://host:9000/invalidPath",
+      //   "hdfs://host:9000/path/a=10/b=20",
+      //   "hdfs://host:9000/path/a=10.5/b=hello")
+      // It will be recognised as conflicting directory structure:
+      //   "hdfs://host:9000/invalidPath"
+      //   "hdfs://host:9000/path"
       val basePaths = optBasePaths.flatMap(x => x)
       assert(
         basePaths.distinct.size == 1,
         "Conflicting directory structures detected. Suspicious paths:\b" +
-          basePaths.mkString("\n\t", "\n\t", "\n\n"))
+          basePaths.distinct.mkString("\n\t", "\n\t", "\n\n"))
 
       val resolvedPartitionValues = resolvePartitions(pathsWithPartitionValues)
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
index 67b6a37fa502e..61cc0da50865c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
@@ -88,6 +88,22 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
       parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
     }
     assert(exception.getMessage().contains("Conflicting directory structures detected"))
+
+    // Invalid
+    // Conflicting directory structure:
+    // "hdfs://host:9000/tmp/tables/partitionedTable"
+    // "hdfs://host:9000/tmp/tables/nonPartitionedTable1"
+    // "hdfs://host:9000/tmp/tables/nonPartitionedTable2"
+    paths = Seq(
+      "hdfs://host:9000/tmp/tables/partitionedTable",
+      "hdfs://host:9000/tmp/tables/partitionedTable/p=1/",
+      "hdfs://host:9000/tmp/tables/nonPartitionedTable1",
+      "hdfs://host:9000/tmp/tables/nonPartitionedTable2")
+
+    exception = intercept[AssertionError] {
+      parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
+    }
+    assert(exception.getMessage().contains("Conflicting directory structures detected"))
   }
 
   test("parse partition") {

From abf5e4285d97b148a32cf22f5287511198175cb6 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 4 Nov 2015 12:33:47 -0800
Subject: [PATCH 054/518] [SPARK-11504][SQL] API audit for distributeBy and
 localSort

1. Renamed localSort -> sortWithinPartitions to avoid ambiguity in "local"
2. distributeBy -> repartition to match the existing repartition.

Author: Reynold Xin <rxin@databricks.com>

Closes #9470 from rxin/SPARK-11504.
---
 .../org/apache/spark/sql/DataFrame.scala      | 132 ++++++++++--------
 .../apache/spark/sql/CachedTableSuite.scala   |  20 ++-
 .../org/apache/spark/sql/DataFrameSuite.scala |  44 +++---
 3 files changed, 113 insertions(+), 83 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 5e9c7efbbf160..d3a2249d7006c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -241,18 +241,6 @@ class DataFrame private[sql](
     sb.toString()
   }
 
-  private[sql] def sortInternal(global: Boolean, sortExprs: Seq[Column]): DataFrame = {
-    val sortOrder: Seq[SortOrder] = sortExprs.map { col =>
-      col.expr match {
-        case expr: SortOrder =>
-          expr
-        case expr: Expression =>
-          SortOrder(expr, Ascending)
-      }
-    }
-    Sort(sortOrder, global = global, logicalPlan)
-  }
-
   override def toString: String = {
     try {
       schema.map(f => s"${f.name}: ${f.dataType.simpleString}").mkString("[", ", ", "]")
@@ -619,6 +607,32 @@ class DataFrame private[sql](
     plan.copy(condition = cond)
   }
 
+  /**
+   * Returns a new [[DataFrame]] with each partition sorted by the given expressions.
+   *
+   * This is the same operation as "SORT BY" in SQL (Hive QL).
+   *
+   * @group dfops
+   * @since 1.6.0
+   */
+  @scala.annotation.varargs
+  def sortWithinPartitions(sortCol: String, sortCols: String*): DataFrame = {
+    sortWithinPartitions(sortCol, sortCols : _*)
+  }
+
+  /**
+   * Returns a new [[DataFrame]] with each partition sorted by the given expressions.
+   *
+   * This is the same operation as "SORT BY" in SQL (Hive QL).
+   *
+   * @group dfops
+   * @since 1.6.0
+   */
+  @scala.annotation.varargs
+  def sortWithinPartitions(sortExprs: Column*): DataFrame = {
+    sortInternal(global = false, sortExprs)
+  }
+
   /**
    * Returns a new [[DataFrame]] sorted by the specified column, all in ascending order.
    * {{{
@@ -645,7 +659,7 @@ class DataFrame private[sql](
    */
   @scala.annotation.varargs
   def sort(sortExprs: Column*): DataFrame = {
-    sortInternal(true, sortExprs)
+    sortInternal(global = true, sortExprs)
   }
 
   /**
@@ -666,44 +680,6 @@ class DataFrame private[sql](
   @scala.annotation.varargs
   def orderBy(sortExprs: Column*): DataFrame = sort(sortExprs : _*)
 
-  /**
-   * Returns a new [[DataFrame]] partitioned by the given partitioning expressions into
-   * `numPartitions`. The resulting DataFrame is hash partitioned.
-   * @group dfops
-   * @since 1.6.0
-   */
-  def distributeBy(partitionExprs: Seq[Column], numPartitions: Int): DataFrame = {
-    RepartitionByExpression(partitionExprs.map { _.expr }, logicalPlan, Some(numPartitions))
-  }
-
-  /**
-   * Returns a new [[DataFrame]] partitioned by the given partitioning expressions preserving
-   * the existing number of partitions. The resulting DataFrame is hash partitioned.
-   * @group dfops
-   * @since 1.6.0
-   */
-  def distributeBy(partitionExprs: Seq[Column]): DataFrame = {
-    RepartitionByExpression(partitionExprs.map { _.expr }, logicalPlan, None)
-  }
-
-  /**
-   * Returns a new [[DataFrame]] with each partition sorted by the given expressions.
-   * @group dfops
-   * @since 1.6.0
-   */
-  @scala.annotation.varargs
-  def localSort(sortCol: String, sortCols: String*): DataFrame = localSort(sortCol, sortCols : _*)
-
-  /**
-   * Returns a new [[DataFrame]] with each partition sorted by the given expressions.
-   * @group dfops
-   * @since 1.6.0
-   */
-  @scala.annotation.varargs
-  def localSort(sortExprs: Column*): DataFrame = {
-    sortInternal(false, sortExprs)
-  }
-
   /**
    * Selects column based on the column name and return it as a [[Column]].
    * Note that the column name can also reference to a nested column like `a.b`.
@@ -798,7 +774,9 @@ class DataFrame private[sql](
    * SQL expressions.
    *
    * {{{
+   *   // The following are equivalent:
    *   df.selectExpr("colA", "colB as newName", "abs(colC)")
+   *   df.select(expr("colA"), expr("colB as newName"), expr("abs(colC)"))
    * }}}
    * @group dfops
    * @since 1.3.0
@@ -1524,13 +1502,41 @@ class DataFrame private[sql](
 
   /**
    * Returns a new [[DataFrame]] that has exactly `numPartitions` partitions.
-   * @group rdd
+   * @group dfops
    * @since 1.3.0
    */
   def repartition(numPartitions: Int): DataFrame = {
     Repartition(numPartitions, shuffle = true, logicalPlan)
   }
 
+  /**
+   * Returns a new [[DataFrame]] partitioned by the given partitioning expressions into
+   * `numPartitions`. The resulting DataFrame is hash partitioned.
+   *
+   * This is the same operation as "DISTRIBUTE BY" in SQL (Hive QL).
+   *
+   * @group dfops
+   * @since 1.6.0
+   */
+  @scala.annotation.varargs
+  def repartition(numPartitions: Int, partitionExprs: Column*): DataFrame = {
+    RepartitionByExpression(partitionExprs.map(_.expr), logicalPlan, Some(numPartitions))
+  }
+
+  /**
+   * Returns a new [[DataFrame]] partitioned by the given partitioning expressions preserving
+   * the existing number of partitions. The resulting DataFrame is hash partitioned.
+   *
+   * This is the same operation as "DISTRIBUTE BY" in SQL (Hive QL).
+   *
+   * @group dfops
+   * @since 1.6.0
+   */
+  @scala.annotation.varargs
+  def repartition(partitionExprs: Column*): DataFrame = {
+    RepartitionByExpression(partitionExprs.map(_.expr), logicalPlan, numPartitions = None)
+  }
+
   /**
    * Returns a new [[DataFrame]] that has exactly `numPartitions` partitions.
    * Similar to coalesce defined on an [[RDD]], this operation results in a narrow dependency, e.g.
@@ -2016,6 +2022,12 @@ class DataFrame private[sql](
     write.mode(SaveMode.Append).insertInto(tableName)
   }
 
+  ////////////////////////////////////////////////////////////////////////////
+  ////////////////////////////////////////////////////////////////////////////
+  // End of deprecated methods
+  ////////////////////////////////////////////////////////////////////////////
+  ////////////////////////////////////////////////////////////////////////////
+
   /**
    * Wrap a DataFrame action to track all Spark jobs in the body so that we can connect them with
    * an execution.
@@ -2045,10 +2057,16 @@ class DataFrame private[sql](
     }
   }
 
-  ////////////////////////////////////////////////////////////////////////////
-  ////////////////////////////////////////////////////////////////////////////
-  // End of deprecated methods
-  ////////////////////////////////////////////////////////////////////////////
-  ////////////////////////////////////////////////////////////////////////////
+  private def sortInternal(global: Boolean, sortExprs: Seq[Column]): DataFrame = {
+    val sortOrder: Seq[SortOrder] = sortExprs.map { col =>
+      col.expr match {
+        case expr: SortOrder =>
+          expr
+        case expr: Expression =>
+          SortOrder(expr, Ascending)
+      }
+    }
+    Sort(sortOrder, global = global, logicalPlan)
+  }
 
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
index 605954b105d1e..dbcb011f603f7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
@@ -379,8 +379,8 @@ class CachedTableSuite extends QueryTest with SharedSQLContext {
     // Set up two tables distributed in the same way. Try this with the data distributed into
     // different number of partitions.
     for (numPartitions <- 1 until 10 by 4) {
-      testData.distributeBy(Column("key") :: Nil, numPartitions).registerTempTable("t1")
-      testData2.distributeBy(Column("a") :: Nil, numPartitions).registerTempTable("t2")
+      testData.repartition(numPartitions, $"key").registerTempTable("t1")
+      testData2.repartition(numPartitions, $"a").registerTempTable("t2")
       sqlContext.cacheTable("t1")
       sqlContext.cacheTable("t2")
 
@@ -401,8 +401,20 @@ class CachedTableSuite extends QueryTest with SharedSQLContext {
     }
 
     // Distribute the tables into non-matching number of partitions. Need to shuffle.
-    testData.distributeBy(Column("key") :: Nil, 6).registerTempTable("t1")
-    testData2.distributeBy(Column("a") :: Nil, 3).registerTempTable("t2")
+    testData.repartition(6, $"key").registerTempTable("t1")
+    testData2.repartition(3, $"a").registerTempTable("t2")
+    sqlContext.cacheTable("t1")
+    sqlContext.cacheTable("t2")
+
+    verifyNumExchanges(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"), 2)
+    sqlContext.uncacheTable("t1")
+    sqlContext.uncacheTable("t2")
+    sqlContext.dropTempTable("t1")
+    sqlContext.dropTempTable("t2")
+
+    // One side of join is not partitioned in the desired way. Need to shuffle.
+    testData.repartition(6, $"value").registerTempTable("t1")
+    testData2.repartition(6, $"a").registerTempTable("t2")
     sqlContext.cacheTable("t1")
     sqlContext.cacheTable("t2")
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index a9e6413423118..84a616d0b9081 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -1044,79 +1044,79 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
   test("distributeBy and localSort") {
     val original = testData.repartition(1)
     assert(original.rdd.partitions.length == 1)
-    val df = original.distributeBy(Column("key") :: Nil, 5)
-    assert(df.rdd.partitions.length  == 5)
+    val df = original.repartition(5, $"key")
+    assert(df.rdd.partitions.length == 5)
     checkAnswer(original.select(), df.select())
 
-    val df2 = original.distributeBy(Column("key") :: Nil, 10)
-    assert(df2.rdd.partitions.length  == 10)
+    val df2 = original.repartition(10, $"key")
+    assert(df2.rdd.partitions.length == 10)
     checkAnswer(original.select(), df2.select())
 
     // Group by the column we are distributed by. This should generate a plan with no exchange
     // between the aggregates
-    val df3 = testData.distributeBy(Column("key") :: Nil).groupBy("key").count()
+    val df3 = testData.repartition($"key").groupBy("key").count()
     verifyNonExchangingAgg(df3)
-    verifyNonExchangingAgg(testData.distributeBy(Column("key") :: Column("value") :: Nil)
+    verifyNonExchangingAgg(testData.repartition($"key", $"value")
       .groupBy("key", "value").count())
 
     // Grouping by just the first distributeBy expr, need to exchange.
-    verifyExchangingAgg(testData.distributeBy(Column("key") :: Column("value") :: Nil)
+    verifyExchangingAgg(testData.repartition($"key", $"value")
       .groupBy("key").count())
 
     val data = sqlContext.sparkContext.parallelize(
       (1 to 100).map(i => TestData2(i % 10, i))).toDF()
 
     // Distribute and order by.
-    val df4 = data.distributeBy(Column("a") :: Nil).localSort($"b".desc)
+    val df4 = data.repartition($"a").sortWithinPartitions($"b".desc)
     // Walk each partition and verify that it is sorted descending and does not contain all
     // the values.
-    df4.rdd.foreachPartition(p => {
+    df4.rdd.foreachPartition { p =>
       var previousValue: Int = -1
       var allSequential: Boolean = true
-      p.foreach(r => {
+      p.foreach { r =>
         val v: Int = r.getInt(1)
         if (previousValue != -1) {
           if (previousValue < v) throw new SparkException("Partition is not ordered.")
           if (v + 1 != previousValue) allSequential = false
         }
         previousValue = v
-      })
+      }
       if (allSequential) throw new SparkException("Partition should not be globally ordered")
-    })
+    }
 
     // Distribute and order by with multiple order bys
-    val df5 = data.distributeBy(Column("a") :: Nil, 2).localSort($"b".asc, $"a".asc)
+    val df5 = data.repartition(2, $"a").sortWithinPartitions($"b".asc, $"a".asc)
     // Walk each partition and verify that it is sorted ascending
-    df5.rdd.foreachPartition(p => {
+    df5.rdd.foreachPartition { p =>
       var previousValue: Int = -1
       var allSequential: Boolean = true
-      p.foreach(r => {
+      p.foreach { r =>
         val v: Int = r.getInt(1)
         if (previousValue != -1) {
           if (previousValue > v) throw new SparkException("Partition is not ordered.")
           if (v - 1 != previousValue) allSequential = false
         }
         previousValue = v
-      })
+      }
       if (allSequential) throw new SparkException("Partition should not be all sequential")
-    })
+    }
 
     // Distribute into one partition and order by. This partition should contain all the values.
-    val df6 = data.distributeBy(Column("a") :: Nil, 1).localSort($"b".asc)
+    val df6 = data.repartition(1, $"a").sortWithinPartitions($"b".asc)
     // Walk each partition and verify that it is sorted descending and not globally sorted.
-    df6.rdd.foreachPartition(p => {
+    df6.rdd.foreachPartition { p =>
       var previousValue: Int = -1
       var allSequential: Boolean = true
-      p.foreach(r => {
+      p.foreach { r =>
         val v: Int = r.getInt(1)
         if (previousValue != -1) {
           if (previousValue > v) throw new SparkException("Partition is not ordered.")
           if (v - 1 != previousValue) allSequential = false
         }
         previousValue = v
-      })
+      }
       if (!allSequential) throw new SparkException("Partition should contain all sequential values")
-    })
+    }
   }
 
   test("fix case sensitivity of partition by") {

From d19f4fda63d0800a85b3e1c8379160bbbf17b6a3 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 4 Nov 2015 13:44:07 -0800
Subject: [PATCH 055/518] [SPARK-11505][SQL] Break aggregate functions into
 multiple files

functions.scala was getting pretty long. I broke it into multiple files.

I also added explicit data types for some public vals, and renamed aggregate function pretty names to lower case, which is more consistent with rest of the functions.

Author: Reynold Xin <rxin@databricks.com>

Closes #9471 from rxin/SPARK-11505.
---
 .../unsafe/sort/UnsafeExternalSorter.java     |   5 +-
 .../expressions/aggregate/Average.scala       |  86 ++
 .../aggregate/CentralMomentAgg.scala          | 230 +++++
 .../catalyst/expressions/aggregate/Corr.scala | 179 ++++
 .../expressions/aggregate/Count.scala         |  52 +
 .../expressions/aggregate/First.scala         |  92 ++
 ...ctions.scala => HyperLogLogPlusPlus.scala} | 933 ------------------
 .../expressions/aggregate/Kurtosis.scala      |  49 +
 .../catalyst/expressions/aggregate/Last.scala |  89 ++
 .../catalyst/expressions/aggregate/Max.scala  |  55 ++
 .../catalyst/expressions/aggregate/Min.scala  |  56 ++
 .../expressions/aggregate/Skewness.scala      |  48 +
 .../expressions/aggregate/Stddev.scala        | 134 +++
 .../catalyst/expressions/aggregate/Sum.scala  |  75 ++
 .../aggregate/{utils.scala => Utils.scala}    |   0
 .../expressions/aggregate/Variance.scala      |  66 ++
 .../sql/catalyst/expressions/aggregates.scala |  24 +-
 17 files changed, 1223 insertions(+), 950 deletions(-)
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala
 rename sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/{functions.scala => HyperLogLogPlusPlus.scala} (72%)
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Last.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
 rename sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/{utils.scala => Utils.scala} (100%)
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala

diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 49a5a4b13b70d..509fb0a044c0c 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -157,11 +157,14 @@ public void closeCurrentPage() {
    */
   @Override
   public long spill(long size, MemoryConsumer trigger) throws IOException {
+    assert(inMemSorter != null);
     if (trigger != this) {
       if (readingIterator != null) {
         return readingIterator.spill();
+      } else {
+
       }
-      return 0L;
+      return 0L; // this should throw exception
     }
 
     if (inMemSorter == null || inMemSorter.numRecords() <= 0) {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
new file mode 100644
index 0000000000000..c8c20ada5fbc7
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+case class Average(child: Expression) extends DeclarativeAggregate {
+
+  override def prettyName: String = "avg"
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  // Return data type.
+  override def dataType: DataType = resultType
+
+  // Expected input data type.
+  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
+  // new version at planning time (after analysis phase). For now, NullType is added at here
+  // to make it resolved when we have cases like `select avg(null)`.
+  // We can use our analyzer to cast NullType to the default data type of the NumericType once
+  // we remove the old aggregate functions. Then, we will not need NullType at here.
+  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
+
+  private val resultType = child.dataType match {
+    case DecimalType.Fixed(p, s) =>
+      DecimalType.bounded(p + 4, s + 4)
+    case _ => DoubleType
+  }
+
+  private val sumDataType = child.dataType match {
+    case _ @ DecimalType.Fixed(p, s) => DecimalType.bounded(p + 10, s)
+    case _ => DoubleType
+  }
+
+  private val sum = AttributeReference("sum", sumDataType)()
+  private val count = AttributeReference("count", LongType)()
+
+  override val aggBufferAttributes = sum :: count :: Nil
+
+  override val initialValues = Seq(
+    /* sum = */ Cast(Literal(0), sumDataType),
+    /* count = */ Literal(0L)
+  )
+
+  override val updateExpressions = Seq(
+    /* sum = */
+    Add(
+      sum,
+      Coalesce(Cast(child, sumDataType) :: Cast(Literal(0), sumDataType) :: Nil)),
+    /* count = */ If(IsNull(child), count, count + 1L)
+  )
+
+  override val mergeExpressions = Seq(
+    /* sum = */ sum.left + sum.right,
+    /* count = */ count.left + count.right
+  )
+
+  // If all input are nulls, count will be 0 and we will get null after the division.
+  override val evaluateExpression = child.dataType match {
+    case DecimalType.Fixed(p, s) =>
+      // increase the precision and scale to prevent precision loss
+      val dt = DecimalType.bounded(p + 14, s + 4)
+      Cast(Cast(sum, dt) / Cast(count, dt), resultType)
+    case _ =>
+      Cast(sum, resultType) / Cast(count, resultType)
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
new file mode 100644
index 0000000000000..ef08b025ff556
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
@@ -0,0 +1,230 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+/**
+ * A central moment is the expected value of a specified power of the deviation of a random
+ * variable from the mean. Central moments are often used to characterize the properties of about
+ * the shape of a distribution.
+ *
+ * This class implements online, one-pass algorithms for computing the central moments of a set of
+ * points.
+ *
+ * Behavior:
+ *  - null values are ignored
+ *  - returns `Double.NaN` when the column contains `Double.NaN` values
+ *
+ * References:
+ *  - Xiangrui Meng.  "Simpler Online Updates for Arbitrary-Order Central Moments."
+ *      2015. http://arxiv.org/abs/1510.04923
+ *
+ * @see [[https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
+ *     Algorithms for calculating variance (Wikipedia)]]
+ *
+ * @param child to compute central moments of.
+ */
+abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate with Serializable {
+
+  /**
+   * The central moment order to be computed.
+   */
+  protected def momentOrder: Int
+
+  override def children: Seq[Expression] = Seq(child)
+
+  override def nullable: Boolean = false
+
+  override def dataType: DataType = DoubleType
+
+  // Expected input data type.
+  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
+  // new version at planning time (after analysis phase). For now, NullType is added at here
+  // to make it resolved when we have cases like `select avg(null)`.
+  // We can use our analyzer to cast NullType to the default data type of the NumericType once
+  // we remove the old aggregate functions. Then, we will not need NullType at here.
+  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
+
+  override def aggBufferSchema: StructType = StructType.fromAttributes(aggBufferAttributes)
+
+  /**
+   * Size of aggregation buffer.
+   */
+  private[this] val bufferSize = 5
+
+  override val aggBufferAttributes: Seq[AttributeReference] = Seq.tabulate(bufferSize) { i =>
+    AttributeReference(s"M$i", DoubleType)()
+  }
+
+  // Note: although this simply copies aggBufferAttributes, this common code can not be placed
+  // in the superclass because that will lead to initialization ordering issues.
+  override val inputAggBufferAttributes: Seq[AttributeReference] =
+    aggBufferAttributes.map(_.newInstance())
+
+  // buffer offsets
+  private[this] val nOffset = mutableAggBufferOffset
+  private[this] val meanOffset = mutableAggBufferOffset + 1
+  private[this] val secondMomentOffset = mutableAggBufferOffset + 2
+  private[this] val thirdMomentOffset = mutableAggBufferOffset + 3
+  private[this] val fourthMomentOffset = mutableAggBufferOffset + 4
+
+  // frequently used values for online updates
+  private[this] var delta = 0.0
+  private[this] var deltaN = 0.0
+  private[this] var delta2 = 0.0
+  private[this] var deltaN2 = 0.0
+  private[this] var n = 0.0
+  private[this] var mean = 0.0
+  private[this] var m2 = 0.0
+  private[this] var m3 = 0.0
+  private[this] var m4 = 0.0
+
+  /**
+   * Initialize all moments to zero.
+   */
+  override def initialize(buffer: MutableRow): Unit = {
+    for (aggIndex <- 0 until bufferSize) {
+      buffer.setDouble(mutableAggBufferOffset + aggIndex, 0.0)
+    }
+  }
+
+  /**
+   * Update the central moments buffer.
+   */
+  override def update(buffer: MutableRow, input: InternalRow): Unit = {
+    val v = Cast(child, DoubleType).eval(input)
+    if (v != null) {
+      val updateValue = v match {
+        case d: Double => d
+      }
+
+      n = buffer.getDouble(nOffset)
+      mean = buffer.getDouble(meanOffset)
+
+      n += 1.0
+      buffer.setDouble(nOffset, n)
+      delta = updateValue - mean
+      deltaN = delta / n
+      mean += deltaN
+      buffer.setDouble(meanOffset, mean)
+
+      if (momentOrder >= 2) {
+        m2 = buffer.getDouble(secondMomentOffset)
+        m2 += delta * (delta - deltaN)
+        buffer.setDouble(secondMomentOffset, m2)
+      }
+
+      if (momentOrder >= 3) {
+        delta2 = delta * delta
+        deltaN2 = deltaN * deltaN
+        m3 = buffer.getDouble(thirdMomentOffset)
+        m3 += -3.0 * deltaN * m2 + delta * (delta2 - deltaN2)
+        buffer.setDouble(thirdMomentOffset, m3)
+      }
+
+      if (momentOrder >= 4) {
+        m4 = buffer.getDouble(fourthMomentOffset)
+        m4 += -4.0 * deltaN * m3 - 6.0 * deltaN2 * m2 +
+          delta * (delta * delta2 - deltaN * deltaN2)
+        buffer.setDouble(fourthMomentOffset, m4)
+      }
+    }
+  }
+
+  /**
+   * Merge two central moment buffers.
+   */
+  override def merge(buffer1: MutableRow, buffer2: InternalRow): Unit = {
+    val n1 = buffer1.getDouble(nOffset)
+    val n2 = buffer2.getDouble(inputAggBufferOffset)
+    val mean1 = buffer1.getDouble(meanOffset)
+    val mean2 = buffer2.getDouble(inputAggBufferOffset + 1)
+
+    var secondMoment1 = 0.0
+    var secondMoment2 = 0.0
+
+    var thirdMoment1 = 0.0
+    var thirdMoment2 = 0.0
+
+    var fourthMoment1 = 0.0
+    var fourthMoment2 = 0.0
+
+    n = n1 + n2
+    buffer1.setDouble(nOffset, n)
+    delta = mean2 - mean1
+    deltaN = if (n == 0.0) 0.0 else delta / n
+    mean = mean1 + deltaN * n2
+    buffer1.setDouble(mutableAggBufferOffset + 1, mean)
+
+    // higher order moments computed according to:
+    // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Higher-order_statistics
+    if (momentOrder >= 2) {
+      secondMoment1 = buffer1.getDouble(secondMomentOffset)
+      secondMoment2 = buffer2.getDouble(inputAggBufferOffset + 2)
+      m2 = secondMoment1 + secondMoment2 + delta * deltaN * n1 * n2
+      buffer1.setDouble(secondMomentOffset, m2)
+    }
+
+    if (momentOrder >= 3) {
+      thirdMoment1 = buffer1.getDouble(thirdMomentOffset)
+      thirdMoment2 = buffer2.getDouble(inputAggBufferOffset + 3)
+      m3 = thirdMoment1 + thirdMoment2 + deltaN * deltaN * delta * n1 * n2 *
+        (n1 - n2) + 3.0 * deltaN * (n1 * secondMoment2 - n2 * secondMoment1)
+      buffer1.setDouble(thirdMomentOffset, m3)
+    }
+
+    if (momentOrder >= 4) {
+      fourthMoment1 = buffer1.getDouble(fourthMomentOffset)
+      fourthMoment2 = buffer2.getDouble(inputAggBufferOffset + 4)
+      m4 = fourthMoment1 + fourthMoment2 + deltaN * deltaN * deltaN * delta * n1 *
+        n2 * (n1 * n1 - n1 * n2 + n2 * n2) + deltaN * deltaN * 6.0 *
+        (n1 * n1 * secondMoment2 + n2 * n2 * secondMoment1) +
+        4.0 * deltaN * (n1 * thirdMoment2 - n2 * thirdMoment1)
+      buffer1.setDouble(fourthMomentOffset, m4)
+    }
+  }
+
+  /**
+   * Compute aggregate statistic from sufficient moments.
+   * @param centralMoments Length `momentOrder + 1` array of central moments (un-normalized)
+   *                       needed to compute the aggregate stat.
+   */
+  def getStatistic(n: Double, mean: Double, centralMoments: Array[Double]): Double
+
+  override final def eval(buffer: InternalRow): Any = {
+    val n = buffer.getDouble(nOffset)
+    val mean = buffer.getDouble(meanOffset)
+    val moments = Array.ofDim[Double](momentOrder + 1)
+    moments(0) = 1.0
+    moments(1) = 0.0
+    if (momentOrder >= 2) {
+      moments(2) = buffer.getDouble(secondMomentOffset)
+    }
+    if (momentOrder >= 3) {
+      moments(3) = buffer.getDouble(thirdMomentOffset)
+    }
+    if (momentOrder >= 4) {
+      moments(4) = buffer.getDouble(fourthMomentOffset)
+    }
+
+    getStatistic(n, mean, moments)
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
new file mode 100644
index 0000000000000..832338378fb38
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
@@ -0,0 +1,179 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+/**
+ * Compute Pearson correlation between two expressions.
+ * When applied on empty data (i.e., count is zero), it returns NULL.
+ *
+ * Definition of Pearson correlation can be found at
+ * http://en.wikipedia.org/wiki/Pearson_product-moment_correlation_coefficient
+ */
+case class Corr(
+    left: Expression,
+    right: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0)
+  extends ImperativeAggregate {
+
+  override def children: Seq[Expression] = Seq(left, right)
+
+  override def nullable: Boolean = false
+
+  override def dataType: DataType = DoubleType
+
+  override def inputTypes: Seq[AbstractDataType] = Seq(DoubleType, DoubleType)
+
+  override def aggBufferSchema: StructType = StructType.fromAttributes(aggBufferAttributes)
+
+  override def inputAggBufferAttributes: Seq[AttributeReference] = {
+    aggBufferAttributes.map(_.newInstance())
+  }
+
+  override val aggBufferAttributes: Seq[AttributeReference] = Seq(
+    AttributeReference("xAvg", DoubleType)(),
+    AttributeReference("yAvg", DoubleType)(),
+    AttributeReference("Ck", DoubleType)(),
+    AttributeReference("MkX", DoubleType)(),
+    AttributeReference("MkY", DoubleType)(),
+    AttributeReference("count", LongType)())
+
+  // Local cache of mutableAggBufferOffset(s) that will be used in update and merge
+  private[this] val mutableAggBufferOffsetPlus1 = mutableAggBufferOffset + 1
+  private[this] val mutableAggBufferOffsetPlus2 = mutableAggBufferOffset + 2
+  private[this] val mutableAggBufferOffsetPlus3 = mutableAggBufferOffset + 3
+  private[this] val mutableAggBufferOffsetPlus4 = mutableAggBufferOffset + 4
+  private[this] val mutableAggBufferOffsetPlus5 = mutableAggBufferOffset + 5
+
+  // Local cache of inputAggBufferOffset(s) that will be used in update and merge
+  private[this] val inputAggBufferOffsetPlus1 = inputAggBufferOffset + 1
+  private[this] val inputAggBufferOffsetPlus2 = inputAggBufferOffset + 2
+  private[this] val inputAggBufferOffsetPlus3 = inputAggBufferOffset + 3
+  private[this] val inputAggBufferOffsetPlus4 = inputAggBufferOffset + 4
+  private[this] val inputAggBufferOffsetPlus5 = inputAggBufferOffset + 5
+
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+  override def initialize(buffer: MutableRow): Unit = {
+    buffer.setDouble(mutableAggBufferOffset, 0.0)
+    buffer.setDouble(mutableAggBufferOffsetPlus1, 0.0)
+    buffer.setDouble(mutableAggBufferOffsetPlus2, 0.0)
+    buffer.setDouble(mutableAggBufferOffsetPlus3, 0.0)
+    buffer.setDouble(mutableAggBufferOffsetPlus4, 0.0)
+    buffer.setLong(mutableAggBufferOffsetPlus5, 0L)
+  }
+
+  override def update(buffer: MutableRow, input: InternalRow): Unit = {
+    val leftEval = left.eval(input)
+    val rightEval = right.eval(input)
+
+    if (leftEval != null && rightEval != null) {
+      val x = leftEval.asInstanceOf[Double]
+      val y = rightEval.asInstanceOf[Double]
+
+      var xAvg = buffer.getDouble(mutableAggBufferOffset)
+      var yAvg = buffer.getDouble(mutableAggBufferOffsetPlus1)
+      var Ck = buffer.getDouble(mutableAggBufferOffsetPlus2)
+      var MkX = buffer.getDouble(mutableAggBufferOffsetPlus3)
+      var MkY = buffer.getDouble(mutableAggBufferOffsetPlus4)
+      var count = buffer.getLong(mutableAggBufferOffsetPlus5)
+
+      val deltaX = x - xAvg
+      val deltaY = y - yAvg
+      count += 1
+      xAvg += deltaX / count
+      yAvg += deltaY / count
+      Ck += deltaX * (y - yAvg)
+      MkX += deltaX * (x - xAvg)
+      MkY += deltaY * (y - yAvg)
+
+      buffer.setDouble(mutableAggBufferOffset, xAvg)
+      buffer.setDouble(mutableAggBufferOffsetPlus1, yAvg)
+      buffer.setDouble(mutableAggBufferOffsetPlus2, Ck)
+      buffer.setDouble(mutableAggBufferOffsetPlus3, MkX)
+      buffer.setDouble(mutableAggBufferOffsetPlus4, MkY)
+      buffer.setLong(mutableAggBufferOffsetPlus5, count)
+    }
+  }
+
+  // Merge counters from other partitions. Formula can be found at:
+  // http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
+  override def merge(buffer1: MutableRow, buffer2: InternalRow): Unit = {
+    val count2 = buffer2.getLong(inputAggBufferOffsetPlus5)
+
+    // We only go to merge two buffers if there is at least one record aggregated in buffer2.
+    // We don't need to check count in buffer1 because if count2 is more than zero, totalCount
+    // is more than zero too, then we won't get a divide by zero exception.
+    if (count2 > 0) {
+      var xAvg = buffer1.getDouble(mutableAggBufferOffset)
+      var yAvg = buffer1.getDouble(mutableAggBufferOffsetPlus1)
+      var Ck = buffer1.getDouble(mutableAggBufferOffsetPlus2)
+      var MkX = buffer1.getDouble(mutableAggBufferOffsetPlus3)
+      var MkY = buffer1.getDouble(mutableAggBufferOffsetPlus4)
+      var count = buffer1.getLong(mutableAggBufferOffsetPlus5)
+
+      val xAvg2 = buffer2.getDouble(inputAggBufferOffset)
+      val yAvg2 = buffer2.getDouble(inputAggBufferOffsetPlus1)
+      val Ck2 = buffer2.getDouble(inputAggBufferOffsetPlus2)
+      val MkX2 = buffer2.getDouble(inputAggBufferOffsetPlus3)
+      val MkY2 = buffer2.getDouble(inputAggBufferOffsetPlus4)
+
+      val totalCount = count + count2
+      val deltaX = xAvg - xAvg2
+      val deltaY = yAvg - yAvg2
+      Ck += Ck2 + deltaX * deltaY * count / totalCount * count2
+      xAvg = (xAvg * count + xAvg2 * count2) / totalCount
+      yAvg = (yAvg * count + yAvg2 * count2) / totalCount
+      MkX += MkX2 + deltaX * deltaX * count / totalCount * count2
+      MkY += MkY2 + deltaY * deltaY * count / totalCount * count2
+      count = totalCount
+
+      buffer1.setDouble(mutableAggBufferOffset, xAvg)
+      buffer1.setDouble(mutableAggBufferOffsetPlus1, yAvg)
+      buffer1.setDouble(mutableAggBufferOffsetPlus2, Ck)
+      buffer1.setDouble(mutableAggBufferOffsetPlus3, MkX)
+      buffer1.setDouble(mutableAggBufferOffsetPlus4, MkY)
+      buffer1.setLong(mutableAggBufferOffsetPlus5, count)
+    }
+  }
+
+  override def eval(buffer: InternalRow): Any = {
+    val count = buffer.getLong(mutableAggBufferOffsetPlus5)
+    if (count > 0) {
+      val Ck = buffer.getDouble(mutableAggBufferOffsetPlus2)
+      val MkX = buffer.getDouble(mutableAggBufferOffsetPlus3)
+      val MkY = buffer.getDouble(mutableAggBufferOffsetPlus4)
+      val corr = Ck / math.sqrt(MkX * MkY)
+      if (corr.isNaN) {
+        null
+      } else {
+        corr
+      }
+    } else {
+      null
+    }
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
new file mode 100644
index 0000000000000..54df96cd2446a
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
@@ -0,0 +1,52 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+case class Count(child: Expression) extends DeclarativeAggregate {
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = false
+
+  // Return data type.
+  override def dataType: DataType = LongType
+
+  // Expected input data type.
+  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
+
+  private val count = AttributeReference("count", LongType)()
+
+  override val aggBufferAttributes = count :: Nil
+
+  override val initialValues = Seq(
+    /* count = */ Literal(0L)
+  )
+
+  override val updateExpressions = Seq(
+    /* count = */ If(IsNull(child), count, count + 1L)
+  )
+
+  override val mergeExpressions = Seq(
+    /* count = */ count.left + count.right
+  )
+
+  override val evaluateExpression = Cast(count, LongType)
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala
new file mode 100644
index 0000000000000..9028143015853
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala
@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+/**
+ * Returns the first value of `child` for a group of rows. If the first value of `child`
+ * is `null`, it returns `null` (respecting nulls). Even if [[First]] is used on a already
+ * sorted column, if we do partial aggregation and final aggregation (when mergeExpression
+ * is used) its result will not be deterministic (unless the input table is sorted and has
+ * a single partition, and we use a single reducer to do the aggregation.).
+ */
+case class First(child: Expression, ignoreNullsExpr: Expression) extends DeclarativeAggregate {
+
+  def this(child: Expression) = this(child, Literal.create(false, BooleanType))
+
+  private val ignoreNulls: Boolean = ignoreNullsExpr match {
+    case Literal(b: Boolean, BooleanType) => b
+    case _ =>
+      throw new AnalysisException("The second argument of First should be a boolean literal.")
+  }
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  // First is not a deterministic function.
+  override def deterministic: Boolean = false
+
+  // Return data type.
+  override def dataType: DataType = child.dataType
+
+  // Expected input data type.
+  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
+
+  private val first = AttributeReference("first", child.dataType)()
+
+  private val valueSet = AttributeReference("valueSet", BooleanType)()
+
+  override val aggBufferAttributes: Seq[AttributeReference] = first :: valueSet :: Nil
+
+  override val initialValues: Seq[Literal] = Seq(
+    /* first = */ Literal.create(null, child.dataType),
+    /* valueSet = */ Literal.create(false, BooleanType)
+  )
+
+  override val updateExpressions: Seq[Expression] = {
+    if (ignoreNulls) {
+      Seq(
+        /* first = */ If(Or(valueSet, IsNull(child)), first, child),
+        /* valueSet = */ Or(valueSet, IsNotNull(child))
+      )
+    } else {
+      Seq(
+        /* first = */ If(valueSet, first, child),
+        /* valueSet = */ Literal.create(true, BooleanType)
+      )
+    }
+  }
+
+  override val mergeExpressions: Seq[Expression] = {
+    // For first, we can just check if valueSet.left is set to true. If it is set
+    // to true, we use first.right. If not, we use first.right (even if valueSet.right is
+    // false, we are safe to do so because first.right will be null in this case).
+    Seq(
+      /* first = */ If(valueSet.left, first.left, first.right),
+      /* valueSet = */ Or(valueSet.left, valueSet.right)
+    )
+  }
+
+  override val evaluateExpression: AttributeReference = first
+
+  override def toString: String = s"first($child)${if (ignoreNulls) " ignore nulls"}"
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
similarity index 72%
rename from sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala
rename to sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
index 10dc5e64b7ec9..8d341ee630bdb 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/functions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
@@ -22,636 +22,10 @@ import java.util
 
 import com.clearspring.analytics.hash.MurmurHash
 
-import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst._
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.types._
 
-case class Average(child: Expression) extends DeclarativeAggregate {
-
-  override def children: Seq[Expression] = child :: Nil
-
-  override def nullable: Boolean = true
-
-  // Return data type.
-  override def dataType: DataType = resultType
-
-  // Expected input data type.
-  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
-  // new version at planning time (after analysis phase). For now, NullType is added at here
-  // to make it resolved when we have cases like `select avg(null)`.
-  // We can use our analyzer to cast NullType to the default data type of the NumericType once
-  // we remove the old aggregate functions. Then, we will not need NullType at here.
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
-
-  private val resultType = child.dataType match {
-    case DecimalType.Fixed(p, s) =>
-      DecimalType.bounded(p + 4, s + 4)
-    case _ => DoubleType
-  }
-
-  private val sumDataType = child.dataType match {
-    case _ @ DecimalType.Fixed(p, s) => DecimalType.bounded(p + 10, s)
-    case _ => DoubleType
-  }
-
-  private val sum = AttributeReference("sum", sumDataType)()
-  private val count = AttributeReference("count", LongType)()
-
-  override val aggBufferAttributes = sum :: count :: Nil
-
-  override val initialValues = Seq(
-    /* sum = */ Cast(Literal(0), sumDataType),
-    /* count = */ Literal(0L)
-  )
-
-  override val updateExpressions = Seq(
-    /* sum = */
-    Add(
-      sum,
-      Coalesce(Cast(child, sumDataType) :: Cast(Literal(0), sumDataType) :: Nil)),
-    /* count = */ If(IsNull(child), count, count + 1L)
-  )
-
-  override val mergeExpressions = Seq(
-    /* sum = */ sum.left + sum.right,
-    /* count = */ count.left + count.right
-  )
-
-  // If all input are nulls, count will be 0 and we will get null after the division.
-  override val evaluateExpression = child.dataType match {
-    case DecimalType.Fixed(p, s) =>
-      // increase the precision and scale to prevent precision loss
-      val dt = DecimalType.bounded(p + 14, s + 4)
-      Cast(Cast(sum, dt) / Cast(count, dt), resultType)
-    case _ =>
-      Cast(sum, resultType) / Cast(count, resultType)
-  }
-}
-
-case class Count(child: Expression) extends DeclarativeAggregate {
-  override def children: Seq[Expression] = child :: Nil
-
-  override def nullable: Boolean = false
-
-  // Return data type.
-  override def dataType: DataType = LongType
-
-  // Expected input data type.
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
-
-  private val count = AttributeReference("count", LongType)()
-
-  override val aggBufferAttributes = count :: Nil
-
-  override val initialValues = Seq(
-    /* count = */ Literal(0L)
-  )
-
-  override val updateExpressions = Seq(
-    /* count = */ If(IsNull(child), count, count + 1L)
-  )
-
-  override val mergeExpressions = Seq(
-    /* count = */ count.left + count.right
-  )
-
-  override val evaluateExpression = Cast(count, LongType)
-}
-
-/**
- * Returns the first value of `child` for a group of rows. If the first value of `child`
- * is `null`, it returns `null` (respecting nulls). Even if [[First]] is used on a already
- * sorted column, if we do partial aggregation and final aggregation (when mergeExpression
- * is used) its result will not be deterministic (unless the input table is sorted and has
- * a single partition, and we use a single reducer to do the aggregation.).
- * @param child
- */
-case class First(child: Expression, ignoreNullsExpr: Expression) extends DeclarativeAggregate {
-
-  def this(child: Expression) = this(child, Literal.create(false, BooleanType))
-
-  private val ignoreNulls: Boolean = ignoreNullsExpr match {
-    case Literal(b: Boolean, BooleanType) => b
-    case _ =>
-      throw new AnalysisException("The second argument of First should be a boolean literal.")
-  }
-
-  override def children: Seq[Expression] = child :: Nil
-
-  override def nullable: Boolean = true
-
-  // First is not a deterministic function.
-  override def deterministic: Boolean = false
-
-  // Return data type.
-  override def dataType: DataType = child.dataType
-
-  // Expected input data type.
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
-
-  private val first = AttributeReference("first", child.dataType)()
-
-  private val valueSet = AttributeReference("valueSet", BooleanType)()
-
-  override val aggBufferAttributes = first :: valueSet :: Nil
-
-  override val initialValues = Seq(
-    /* first = */ Literal.create(null, child.dataType),
-    /* valueSet = */ Literal.create(false, BooleanType)
-  )
-
-  override val updateExpressions = {
-    if (ignoreNulls) {
-      Seq(
-        /* first = */ If(Or(valueSet, IsNull(child)), first, child),
-        /* valueSet = */ Or(valueSet, IsNotNull(child))
-      )
-    } else {
-      Seq(
-        /* first = */ If(valueSet, first, child),
-        /* valueSet = */ Literal.create(true, BooleanType)
-      )
-    }
-  }
-
-  override val mergeExpressions = {
-    // For first, we can just check if valueSet.left is set to true. If it is set
-    // to true, we use first.right. If not, we use first.right (even if valueSet.right is
-    // false, we are safe to do so because first.right will be null in this case).
-    Seq(
-      /* first = */ If(valueSet.left, first.left, first.right),
-      /* valueSet = */ Or(valueSet.left, valueSet.right)
-    )
-  }
-
-  override val evaluateExpression = first
-
-  override def toString: String = s"FIRST($child)${if (ignoreNulls) " IGNORE NULLS"}"
-}
-
-/**
- * Returns the last value of `child` for a group of rows. If the last value of `child`
- * is `null`, it returns `null` (respecting nulls). Even if [[Last]] is used on a already
- * sorted column, if we do partial aggregation and final aggregation (when mergeExpression
- * is used) its result will not be deterministic (unless the input table is sorted and has
- * a single partition, and we use a single reducer to do the aggregation.).
- * @param child
- */
-case class Last(child: Expression, ignoreNullsExpr: Expression) extends DeclarativeAggregate {
-
-  def this(child: Expression) = this(child, Literal.create(false, BooleanType))
-
-  private val ignoreNulls: Boolean = ignoreNullsExpr match {
-    case Literal(b: Boolean, BooleanType) => b
-    case _ =>
-      throw new AnalysisException("The second argument of First should be a boolean literal.")
-  }
-
-  override def children: Seq[Expression] = child :: Nil
-
-  override def nullable: Boolean = true
-
-  // Last is not a deterministic function.
-  override def deterministic: Boolean = false
-
-  // Return data type.
-  override def dataType: DataType = child.dataType
-
-  // Expected input data type.
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
-
-  private val last = AttributeReference("last", child.dataType)()
-
-  override val aggBufferAttributes = last :: Nil
-
-  override val initialValues = Seq(
-    /* last = */ Literal.create(null, child.dataType)
-  )
-
-  override val updateExpressions = {
-    if (ignoreNulls) {
-      Seq(
-        /* last = */ If(IsNull(child), last, child)
-      )
-    } else {
-      Seq(
-        /* last = */ child
-      )
-    }
-  }
-
-  override val mergeExpressions = {
-    if (ignoreNulls) {
-      Seq(
-        /* last = */ If(IsNull(last.right), last.left, last.right)
-      )
-    } else {
-      Seq(
-        /* last = */ last.right
-      )
-    }
-  }
-
-  override val evaluateExpression = last
-
-  override def toString: String = s"LAST($child)${if (ignoreNulls) " IGNORE NULLS"}"
-}
-
-case class Max(child: Expression) extends DeclarativeAggregate {
-
-  override def children: Seq[Expression] = child :: Nil
-
-  override def nullable: Boolean = true
-
-  // Return data type.
-  override def dataType: DataType = child.dataType
-
-  // Expected input data type.
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
-
-  private val max = AttributeReference("max", child.dataType)()
-
-  override val aggBufferAttributes = max :: Nil
-
-  override val initialValues = Seq(
-    /* max = */ Literal.create(null, child.dataType)
-  )
-
-  override val updateExpressions = Seq(
-    /* max = */ If(IsNull(child), max, If(IsNull(max), child, Greatest(Seq(max, child))))
-  )
-
-  override val mergeExpressions = {
-    val greatest = Greatest(Seq(max.left, max.right))
-    Seq(
-      /* max = */ If(IsNull(max.right), max.left, If(IsNull(max.left), max.right, greatest))
-    )
-  }
-
-  override val evaluateExpression = max
-}
-
-case class Min(child: Expression) extends DeclarativeAggregate {
-
-  override def children: Seq[Expression] = child :: Nil
-
-  override def nullable: Boolean = true
-
-  // Return data type.
-  override def dataType: DataType = child.dataType
-
-  // Expected input data type.
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
-
-  private val min = AttributeReference("min", child.dataType)()
-
-  override val aggBufferAttributes = min :: Nil
-
-  override val initialValues = Seq(
-    /* min = */ Literal.create(null, child.dataType)
-  )
-
-  override val updateExpressions = Seq(
-    /* min = */ If(IsNull(child), min, If(IsNull(min), child, Least(Seq(min, child))))
-  )
-
-  override val mergeExpressions = {
-    val least = Least(Seq(min.left, min.right))
-    Seq(
-      /* min = */ If(IsNull(min.right), min.left, If(IsNull(min.left), min.right, least))
-    )
-  }
-
-  override val evaluateExpression = min
-}
-
-// Compute the population standard deviation of a column
-case class StddevPop(child: Expression) extends StddevAgg(child) {
-
-  override def isSample: Boolean = false
-  override def prettyName: String = "stddev_pop"
-}
-
-// Compute the sample standard deviation of a column
-case class StddevSamp(child: Expression) extends StddevAgg(child) {
-
-  override def isSample: Boolean = true
-  override def prettyName: String = "stddev_samp"
-}
-
-// Compute standard deviation based on online algorithm specified here:
-// http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
-abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
-
-  override def children: Seq[Expression] = child :: Nil
-
-  override def nullable: Boolean = true
-
-  def isSample: Boolean
-
-  // Return data type.
-  override def dataType: DataType = resultType
-
-  // Expected input data type.
-  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
-  // new version at planning time (after analysis phase). For now, NullType is added at here
-  // to make it resolved when we have cases like `select stddev(null)`.
-  // We can use our analyzer to cast NullType to the default data type of the NumericType once
-  // we remove the old aggregate functions. Then, we will not need NullType at here.
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
-
-  private val resultType = DoubleType
-
-  private val count = AttributeReference("count", resultType)()
-  private val avg = AttributeReference("avg", resultType)()
-  private val mk = AttributeReference("mk", resultType)()
-
-  override val aggBufferAttributes = count :: avg :: mk :: Nil
-
-  override val initialValues = Seq(
-    /* count = */ Cast(Literal(0), resultType),
-    /* avg = */ Cast(Literal(0), resultType),
-    /* mk = */ Cast(Literal(0), resultType)
-  )
-
-  override val updateExpressions = {
-    val value = Cast(child, resultType)
-    val newCount = count + Cast(Literal(1), resultType)
-
-    // update average
-    // avg = avg + (value - avg)/count
-    val newAvg = avg + (value - avg) / newCount
-
-    // update sum of square of difference from mean
-    // Mk = Mk + (value - preAvg) * (value - updatedAvg)
-    val newMk = mk + (value - avg) * (value - newAvg)
-
-    Seq(
-      /* count = */ If(IsNull(child), count, newCount),
-      /* avg = */ If(IsNull(child), avg, newAvg),
-      /* mk = */ If(IsNull(child), mk, newMk)
-    )
-  }
-
-  override val mergeExpressions = {
-
-    // count merge
-    val newCount = count.left + count.right
-
-    // average merge
-    val newAvg = ((avg.left * count.left) + (avg.right * count.right)) / newCount
-
-    // update sum of square differences
-    val newMk = {
-      val avgDelta = avg.right - avg.left
-      val mkDelta = (avgDelta * avgDelta) * (count.left * count.right) / newCount
-      mk.left + mk.right + mkDelta
-    }
-
-    Seq(
-      /* count = */ If(IsNull(count.left), count.right,
-                       If(IsNull(count.right), count.left, newCount)),
-      /* avg = */ If(IsNull(avg.left), avg.right,
-                     If(IsNull(avg.right), avg.left, newAvg)),
-      /* mk = */ If(IsNull(mk.left), mk.right,
-                    If(IsNull(mk.right), mk.left, newMk))
-    )
-  }
-
-  override val evaluateExpression = {
-    // when count == 0, return null
-    // when count == 1, return 0
-    // when count >1
-    // stddev_samp = sqrt (mk/(count -1))
-    // stddev_pop = sqrt (mk/count)
-    val varCol =
-      if (isSample) {
-        mk / Cast((count - Cast(Literal(1), resultType)), resultType)
-      } else {
-        mk / count
-      }
-
-    If(EqualTo(count, Cast(Literal(0), resultType)), Cast(Literal(null), resultType),
-      If(EqualTo(count, Cast(Literal(1), resultType)), Cast(Literal(0), resultType),
-        Cast(Sqrt(varCol), resultType)))
-  }
-}
-
-case class Sum(child: Expression) extends DeclarativeAggregate {
-
-  override def children: Seq[Expression] = child :: Nil
-
-  override def nullable: Boolean = true
-
-  // Return data type.
-  override def dataType: DataType = resultType
-
-  // Expected input data type.
-  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
-  // new version at planning time (after analysis phase). For now, NullType is added at here
-  // to make it resolved when we have cases like `select sum(null)`.
-  // We can use our analyzer to cast NullType to the default data type of the NumericType once
-  // we remove the old aggregate functions. Then, we will not need NullType at here.
-  override def inputTypes: Seq[AbstractDataType] =
-    Seq(TypeCollection(LongType, DoubleType, DecimalType, NullType))
-
-  private val resultType = child.dataType match {
-    case DecimalType.Fixed(precision, scale) =>
-      DecimalType.bounded(precision + 10, scale)
-    // TODO: Remove this line once we remove the NullType from inputTypes.
-    case NullType => IntegerType
-    case _ => child.dataType
-  }
-
-  private val sumDataType = resultType
-
-  private val sum = AttributeReference("sum", sumDataType)()
-
-  private val zero = Cast(Literal(0), sumDataType)
-
-  override val aggBufferAttributes = sum :: Nil
-
-  override val initialValues = Seq(
-    /* sum = */ Literal.create(null, sumDataType)
-  )
-
-  override val updateExpressions = Seq(
-    /* sum = */
-    Coalesce(Seq(Add(Coalesce(Seq(sum, zero)), Cast(child, sumDataType)), sum))
-  )
-
-  override val mergeExpressions = {
-    val add = Add(Coalesce(Seq(sum.left, zero)), Cast(sum.right, sumDataType))
-    Seq(
-      /* sum = */
-      Coalesce(Seq(add, sum.left))
-    )
-  }
-
-  override val evaluateExpression = Cast(sum, resultType)
-}
-
-/**
- * Compute Pearson correlation between two expressions.
- * When applied on empty data (i.e., count is zero), it returns NULL.
- *
- * Definition of Pearson correlation can be found at
- * http://en.wikipedia.org/wiki/Pearson_product-moment_correlation_coefficient
- *
- * @param left one of the expressions to compute correlation with.
- * @param right another expression to compute correlation with.
- */
-case class Corr(
-    left: Expression,
-    right: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0)
-  extends ImperativeAggregate {
-
-  def children: Seq[Expression] = Seq(left, right)
-
-  def nullable: Boolean = false
-
-  def dataType: DataType = DoubleType
-
-  override def inputTypes: Seq[AbstractDataType] = Seq(DoubleType, DoubleType)
-
-  def aggBufferSchema: StructType = StructType.fromAttributes(aggBufferAttributes)
-
-  def inputAggBufferAttributes: Seq[AttributeReference] = aggBufferAttributes.map(_.newInstance())
-
-  val aggBufferAttributes: Seq[AttributeReference] = Seq(
-    AttributeReference("xAvg", DoubleType)(),
-    AttributeReference("yAvg", DoubleType)(),
-    AttributeReference("Ck", DoubleType)(),
-    AttributeReference("MkX", DoubleType)(),
-    AttributeReference("MkY", DoubleType)(),
-    AttributeReference("count", LongType)())
-
-  // Local cache of mutableAggBufferOffset(s) that will be used in update and merge
-  private[this] val mutableAggBufferOffsetPlus1 = mutableAggBufferOffset + 1
-  private[this] val mutableAggBufferOffsetPlus2 = mutableAggBufferOffset + 2
-  private[this] val mutableAggBufferOffsetPlus3 = mutableAggBufferOffset + 3
-  private[this] val mutableAggBufferOffsetPlus4 = mutableAggBufferOffset + 4
-  private[this] val mutableAggBufferOffsetPlus5 = mutableAggBufferOffset + 5
-
-  // Local cache of inputAggBufferOffset(s) that will be used in update and merge
-  private[this] val inputAggBufferOffsetPlus1 = inputAggBufferOffset + 1
-  private[this] val inputAggBufferOffsetPlus2 = inputAggBufferOffset + 2
-  private[this] val inputAggBufferOffsetPlus3 = inputAggBufferOffset + 3
-  private[this] val inputAggBufferOffsetPlus4 = inputAggBufferOffset + 4
-  private[this] val inputAggBufferOffsetPlus5 = inputAggBufferOffset + 5
-
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
-    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
-
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
-    copy(inputAggBufferOffset = newInputAggBufferOffset)
-
-  override def initialize(buffer: MutableRow): Unit = {
-    buffer.setDouble(mutableAggBufferOffset, 0.0)
-    buffer.setDouble(mutableAggBufferOffsetPlus1, 0.0)
-    buffer.setDouble(mutableAggBufferOffsetPlus2, 0.0)
-    buffer.setDouble(mutableAggBufferOffsetPlus3, 0.0)
-    buffer.setDouble(mutableAggBufferOffsetPlus4, 0.0)
-    buffer.setLong(mutableAggBufferOffsetPlus5, 0L)
-  }
-
-  override def update(buffer: MutableRow, input: InternalRow): Unit = {
-    val leftEval = left.eval(input)
-    val rightEval = right.eval(input)
-
-    if (leftEval != null && rightEval != null) {
-      val x = leftEval.asInstanceOf[Double]
-      val y = rightEval.asInstanceOf[Double]
-
-      var xAvg = buffer.getDouble(mutableAggBufferOffset)
-      var yAvg = buffer.getDouble(mutableAggBufferOffsetPlus1)
-      var Ck = buffer.getDouble(mutableAggBufferOffsetPlus2)
-      var MkX = buffer.getDouble(mutableAggBufferOffsetPlus3)
-      var MkY = buffer.getDouble(mutableAggBufferOffsetPlus4)
-      var count = buffer.getLong(mutableAggBufferOffsetPlus5)
-
-      val deltaX = x - xAvg
-      val deltaY = y - yAvg
-      count += 1
-      xAvg += deltaX / count
-      yAvg += deltaY / count
-      Ck += deltaX * (y - yAvg)
-      MkX += deltaX * (x - xAvg)
-      MkY += deltaY * (y - yAvg)
-
-      buffer.setDouble(mutableAggBufferOffset, xAvg)
-      buffer.setDouble(mutableAggBufferOffsetPlus1, yAvg)
-      buffer.setDouble(mutableAggBufferOffsetPlus2, Ck)
-      buffer.setDouble(mutableAggBufferOffsetPlus3, MkX)
-      buffer.setDouble(mutableAggBufferOffsetPlus4, MkY)
-      buffer.setLong(mutableAggBufferOffsetPlus5, count)
-    }
-  }
-
-  // Merge counters from other partitions. Formula can be found at:
-  // http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
-  override def merge(buffer1: MutableRow, buffer2: InternalRow): Unit = {
-    val count2 = buffer2.getLong(inputAggBufferOffsetPlus5)
-
-    // We only go to merge two buffers if there is at least one record aggregated in buffer2.
-    // We don't need to check count in buffer1 because if count2 is more than zero, totalCount
-    // is more than zero too, then we won't get a divide by zero exception.
-    if (count2 > 0) {
-      var xAvg = buffer1.getDouble(mutableAggBufferOffset)
-      var yAvg = buffer1.getDouble(mutableAggBufferOffsetPlus1)
-      var Ck = buffer1.getDouble(mutableAggBufferOffsetPlus2)
-      var MkX = buffer1.getDouble(mutableAggBufferOffsetPlus3)
-      var MkY = buffer1.getDouble(mutableAggBufferOffsetPlus4)
-      var count = buffer1.getLong(mutableAggBufferOffsetPlus5)
-
-      val xAvg2 = buffer2.getDouble(inputAggBufferOffset)
-      val yAvg2 = buffer2.getDouble(inputAggBufferOffsetPlus1)
-      val Ck2 = buffer2.getDouble(inputAggBufferOffsetPlus2)
-      val MkX2 = buffer2.getDouble(inputAggBufferOffsetPlus3)
-      val MkY2 = buffer2.getDouble(inputAggBufferOffsetPlus4)
-
-      val totalCount = count + count2
-      val deltaX = xAvg - xAvg2
-      val deltaY = yAvg - yAvg2
-      Ck += Ck2 + deltaX * deltaY * count / totalCount * count2
-      xAvg = (xAvg * count + xAvg2 * count2) / totalCount
-      yAvg = (yAvg * count + yAvg2 * count2) / totalCount
-      MkX += MkX2 + deltaX * deltaX * count / totalCount * count2
-      MkY += MkY2 + deltaY * deltaY * count / totalCount * count2
-      count = totalCount
-
-      buffer1.setDouble(mutableAggBufferOffset, xAvg)
-      buffer1.setDouble(mutableAggBufferOffsetPlus1, yAvg)
-      buffer1.setDouble(mutableAggBufferOffsetPlus2, Ck)
-      buffer1.setDouble(mutableAggBufferOffsetPlus3, MkX)
-      buffer1.setDouble(mutableAggBufferOffsetPlus4, MkY)
-      buffer1.setLong(mutableAggBufferOffsetPlus5, count)
-    }
-  }
-
-  override def eval(buffer: InternalRow): Any = {
-    val count = buffer.getLong(mutableAggBufferOffsetPlus5)
-    if (count > 0) {
-      val Ck = buffer.getDouble(mutableAggBufferOffsetPlus2)
-      val MkX = buffer.getDouble(mutableAggBufferOffsetPlus3)
-      val MkY = buffer.getDouble(mutableAggBufferOffsetPlus4)
-      val corr = Ck / math.sqrt(MkX * MkY)
-      if (corr.isNaN) {
-        null
-      } else {
-        corr
-      }
-    } else {
-      null
-    }
-  }
-}
-
 // scalastyle:off
 /**
  * HyperLogLog++ (HLL++) is a state of the art cardinality estimation algorithm. This class
@@ -1058,310 +432,3 @@ object HyperLogLogPlusPlus {
   )
   // scalastyle:on
 }
-
-/**
- * A central moment is the expected value of a specified power of the deviation of a random
- * variable from the mean. Central moments are often used to characterize the properties of about
- * the shape of a distribution.
- *
- * This class implements online, one-pass algorithms for computing the central moments of a set of
- * points.
- *
- * Behavior:
- *  - null values are ignored
- *  - returns `Double.NaN` when the column contains `Double.NaN` values
- *
- * References:
- *  - Xiangrui Meng.  "Simpler Online Updates for Arbitrary-Order Central Moments."
- *      2015. http://arxiv.org/abs/1510.04923
- *
- * @see [[https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
- *     Algorithms for calculating variance (Wikipedia)]]
- *
- * @param child to compute central moments of.
- */
-abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate with Serializable {
-
-  /**
-   * The central moment order to be computed.
-   */
-  protected def momentOrder: Int
-
-  override def children: Seq[Expression] = Seq(child)
-
-  override def nullable: Boolean = false
-
-  override def dataType: DataType = DoubleType
-
-  // Expected input data type.
-  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
-  // new version at planning time (after analysis phase). For now, NullType is added at here
-  // to make it resolved when we have cases like `select avg(null)`.
-  // We can use our analyzer to cast NullType to the default data type of the NumericType once
-  // we remove the old aggregate functions. Then, we will not need NullType at here.
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
-
-  override def aggBufferSchema: StructType = StructType.fromAttributes(aggBufferAttributes)
-
-  /**
-   * Size of aggregation buffer.
-   */
-  private[this] val bufferSize = 5
-
-  override val aggBufferAttributes: Seq[AttributeReference] = Seq.tabulate(bufferSize) { i =>
-    AttributeReference(s"M$i", DoubleType)()
-  }
-
-  // Note: although this simply copies aggBufferAttributes, this common code can not be placed
-  // in the superclass because that will lead to initialization ordering issues.
-  override val inputAggBufferAttributes: Seq[AttributeReference] =
-    aggBufferAttributes.map(_.newInstance())
-
-  // buffer offsets
-  private[this] val nOffset = mutableAggBufferOffset
-  private[this] val meanOffset = mutableAggBufferOffset + 1
-  private[this] val secondMomentOffset = mutableAggBufferOffset + 2
-  private[this] val thirdMomentOffset = mutableAggBufferOffset + 3
-  private[this] val fourthMomentOffset = mutableAggBufferOffset + 4
-
-  // frequently used values for online updates
-  private[this] var delta = 0.0
-  private[this] var deltaN = 0.0
-  private[this] var delta2 = 0.0
-  private[this] var deltaN2 = 0.0
-  private[this] var n = 0.0
-  private[this] var mean = 0.0
-  private[this] var m2 = 0.0
-  private[this] var m3 = 0.0
-  private[this] var m4 = 0.0
-
-  /**
-   * Initialize all moments to zero.
-   */
-  override def initialize(buffer: MutableRow): Unit = {
-    for (aggIndex <- 0 until bufferSize) {
-      buffer.setDouble(mutableAggBufferOffset + aggIndex, 0.0)
-    }
-  }
-
-  /**
-   * Update the central moments buffer.
-   */
-  override def update(buffer: MutableRow, input: InternalRow): Unit = {
-    val v = Cast(child, DoubleType).eval(input)
-    if (v != null) {
-      val updateValue = v match {
-        case d: Double => d
-      }
-
-      n = buffer.getDouble(nOffset)
-      mean = buffer.getDouble(meanOffset)
-
-      n += 1.0
-      buffer.setDouble(nOffset, n)
-      delta = updateValue - mean
-      deltaN = delta / n
-      mean += deltaN
-      buffer.setDouble(meanOffset, mean)
-
-      if (momentOrder >= 2) {
-        m2 = buffer.getDouble(secondMomentOffset)
-        m2 += delta * (delta - deltaN)
-        buffer.setDouble(secondMomentOffset, m2)
-      }
-
-      if (momentOrder >= 3) {
-        delta2 = delta * delta
-        deltaN2 = deltaN * deltaN
-        m3 = buffer.getDouble(thirdMomentOffset)
-        m3 += -3.0 * deltaN * m2 + delta * (delta2 - deltaN2)
-        buffer.setDouble(thirdMomentOffset, m3)
-      }
-
-      if (momentOrder >= 4) {
-        m4 = buffer.getDouble(fourthMomentOffset)
-        m4 += -4.0 * deltaN * m3 - 6.0 * deltaN2 * m2 +
-          delta * (delta * delta2 - deltaN * deltaN2)
-        buffer.setDouble(fourthMomentOffset, m4)
-      }
-    }
-  }
-
-  /**
-   * Merge two central moment buffers.
-   */
-  override def merge(buffer1: MutableRow, buffer2: InternalRow): Unit = {
-    val n1 = buffer1.getDouble(nOffset)
-    val n2 = buffer2.getDouble(inputAggBufferOffset)
-    val mean1 = buffer1.getDouble(meanOffset)
-    val mean2 = buffer2.getDouble(inputAggBufferOffset + 1)
-
-    var secondMoment1 = 0.0
-    var secondMoment2 = 0.0
-
-    var thirdMoment1 = 0.0
-    var thirdMoment2 = 0.0
-
-    var fourthMoment1 = 0.0
-    var fourthMoment2 = 0.0
-
-    n = n1 + n2
-    buffer1.setDouble(nOffset, n)
-    delta = mean2 - mean1
-    deltaN = if (n == 0.0) 0.0 else delta / n
-    mean = mean1 + deltaN * n2
-    buffer1.setDouble(mutableAggBufferOffset + 1, mean)
-
-    // higher order moments computed according to:
-    // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Higher-order_statistics
-    if (momentOrder >= 2) {
-      secondMoment1 = buffer1.getDouble(secondMomentOffset)
-      secondMoment2 = buffer2.getDouble(inputAggBufferOffset + 2)
-      m2 = secondMoment1 + secondMoment2 + delta * deltaN * n1 * n2
-      buffer1.setDouble(secondMomentOffset, m2)
-    }
-
-    if (momentOrder >= 3) {
-      thirdMoment1 = buffer1.getDouble(thirdMomentOffset)
-      thirdMoment2 = buffer2.getDouble(inputAggBufferOffset + 3)
-      m3 = thirdMoment1 + thirdMoment2 + deltaN * deltaN * delta * n1 * n2 *
-        (n1 - n2) + 3.0 * deltaN * (n1 * secondMoment2 - n2 * secondMoment1)
-      buffer1.setDouble(thirdMomentOffset, m3)
-    }
-
-    if (momentOrder >= 4) {
-      fourthMoment1 = buffer1.getDouble(fourthMomentOffset)
-      fourthMoment2 = buffer2.getDouble(inputAggBufferOffset + 4)
-      m4 = fourthMoment1 + fourthMoment2 + deltaN * deltaN * deltaN * delta * n1 *
-        n2 * (n1 * n1 - n1 * n2 + n2 * n2) + deltaN * deltaN * 6.0 *
-        (n1 * n1 * secondMoment2 + n2 * n2 * secondMoment1) +
-        4.0 * deltaN * (n1 * thirdMoment2 - n2 * thirdMoment1)
-      buffer1.setDouble(fourthMomentOffset, m4)
-    }
-  }
-
-  /**
-   * Compute aggregate statistic from sufficient moments.
-   * @param centralMoments Length `momentOrder + 1` array of central moments (un-normalized)
-   *                       needed to compute the aggregate stat.
-   */
-  def getStatistic(n: Double, mean: Double, centralMoments: Array[Double]): Double
-
-  override final def eval(buffer: InternalRow): Any = {
-    val n = buffer.getDouble(nOffset)
-    val mean = buffer.getDouble(meanOffset)
-    val moments = Array.ofDim[Double](momentOrder + 1)
-    moments(0) = 1.0
-    moments(1) = 0.0
-    if (momentOrder >= 2) {
-      moments(2) = buffer.getDouble(secondMomentOffset)
-    }
-    if (momentOrder >= 3) {
-      moments(3) = buffer.getDouble(thirdMomentOffset)
-    }
-    if (momentOrder >= 4) {
-      moments(4) = buffer.getDouble(fourthMomentOffset)
-    }
-
-    getStatistic(n, mean, moments)
-  }
-}
-
-case class VarianceSamp(child: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0) extends CentralMomentAgg(child) {
-
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
-    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
-
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
-    copy(inputAggBufferOffset = newInputAggBufferOffset)
-
-  override def prettyName: String = "variance_samp"
-
-  override protected val momentOrder = 2
-
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
-    require(moments.length == momentOrder + 1,
-      s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
-
-    if (n == 0.0 || n == 1.0) Double.NaN else moments(2) / (n - 1.0)
-  }
-}
-
-case class VariancePop(child: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0) extends CentralMomentAgg(child) {
-
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
-    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
-
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
-    copy(inputAggBufferOffset = newInputAggBufferOffset)
-
-  override def prettyName: String = "variance_pop"
-
-  override protected val momentOrder = 2
-
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
-    require(moments.length == momentOrder + 1,
-      s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
-
-    if (n == 0.0) Double.NaN else moments(2) / n
-  }
-}
-
-case class Skewness(child: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0) extends CentralMomentAgg(child) {
-
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
-    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
-
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
-    copy(inputAggBufferOffset = newInputAggBufferOffset)
-
-  override def prettyName: String = "skewness"
-
-  override protected val momentOrder = 3
-
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
-    require(moments.length == momentOrder + 1,
-      s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
-    val m2 = moments(2)
-    val m3 = moments(3)
-    if (n == 0.0 || m2 == 0.0) {
-      Double.NaN
-    } else {
-      math.sqrt(n) * m3 / math.sqrt(m2 * m2 * m2)
-    }
-  }
-}
-
-case class Kurtosis(child: Expression,
-    mutableAggBufferOffset: Int = 0,
-    inputAggBufferOffset: Int = 0) extends CentralMomentAgg(child) {
-
-  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
-    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
-
-  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
-    copy(inputAggBufferOffset = newInputAggBufferOffset)
-
-  override def prettyName: String = "kurtosis"
-
-  override protected val momentOrder = 4
-
-  // NOTE: this is the formula for excess kurtosis, which is default for R and SciPy
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
-    require(moments.length == momentOrder + 1,
-      s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
-    val m2 = moments(2)
-    val m4 = moments(4)
-    if (n == 0.0 || m2 == 0.0) {
-      Double.NaN
-    } else {
-      n * m4 / (m2 * m2) - 3.0
-    }
-  }
-}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
new file mode 100644
index 0000000000000..6da39e7143447
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.expressions._
+
+case class Kurtosis(child: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0)
+  extends CentralMomentAgg(child) {
+
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+  override def prettyName: String = "kurtosis"
+
+  override protected val momentOrder = 4
+
+  // NOTE: this is the formula for excess kurtosis, which is default for R and SciPy
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+    require(moments.length == momentOrder + 1,
+      s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
+    val m2 = moments(2)
+    val m4 = moments(4)
+    if (n == 0.0 || m2 == 0.0) {
+      Double.NaN
+    } else {
+      n * m4 / (m2 * m2) - 3.0
+    }
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Last.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Last.scala
new file mode 100644
index 0000000000000..8636bfe8d07aa
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Last.scala
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+/**
+ * Returns the last value of `child` for a group of rows. If the last value of `child`
+ * is `null`, it returns `null` (respecting nulls). Even if [[Last]] is used on a already
+ * sorted column, if we do partial aggregation and final aggregation (when mergeExpression
+ * is used) its result will not be deterministic (unless the input table is sorted and has
+ * a single partition, and we use a single reducer to do the aggregation.).
+ */
+case class Last(child: Expression, ignoreNullsExpr: Expression) extends DeclarativeAggregate {
+
+  def this(child: Expression) = this(child, Literal.create(false, BooleanType))
+
+  private val ignoreNulls: Boolean = ignoreNullsExpr match {
+    case Literal(b: Boolean, BooleanType) => b
+    case _ =>
+      throw new AnalysisException("The second argument of First should be a boolean literal.")
+  }
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  // Last is not a deterministic function.
+  override def deterministic: Boolean = false
+
+  // Return data type.
+  override def dataType: DataType = child.dataType
+
+  // Expected input data type.
+  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
+
+  private val last = AttributeReference("last", child.dataType)()
+
+  override val aggBufferAttributes: Seq[AttributeReference] = last :: Nil
+
+  override val initialValues: Seq[Literal] = Seq(
+    /* last = */ Literal.create(null, child.dataType)
+  )
+
+  override val updateExpressions: Seq[Expression] = {
+    if (ignoreNulls) {
+      Seq(
+        /* last = */ If(IsNull(child), last, child)
+      )
+    } else {
+      Seq(
+        /* last = */ child
+      )
+    }
+  }
+
+  override val mergeExpressions: Seq[Expression] = {
+    if (ignoreNulls) {
+      Seq(
+        /* last = */ If(IsNull(last.right), last.left, last.right)
+      )
+    } else {
+      Seq(
+        /* last = */ last.right
+      )
+    }
+  }
+
+  override val evaluateExpression: AttributeReference = last
+
+  override def toString: String = s"last($child)${if (ignoreNulls) " ignore nulls"}"
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
new file mode 100644
index 0000000000000..b9d75ad452838
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+case class Max(child: Expression) extends DeclarativeAggregate {
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  // Return data type.
+  override def dataType: DataType = child.dataType
+
+  // Expected input data type.
+  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
+
+  private val max = AttributeReference("max", child.dataType)()
+
+  override val aggBufferAttributes: Seq[AttributeReference] = max :: Nil
+
+  override val initialValues: Seq[Literal] = Seq(
+    /* max = */ Literal.create(null, child.dataType)
+  )
+
+  override val updateExpressions: Seq[Expression] = Seq(
+    /* max = */ If(IsNull(child), max, If(IsNull(max), child, Greatest(Seq(max, child))))
+  )
+
+  override val mergeExpressions: Seq[Expression] = {
+    val greatest = Greatest(Seq(max.left, max.right))
+    Seq(
+      /* max = */ If(IsNull(max.right), max.left, If(IsNull(max.left), max.right, greatest))
+    )
+  }
+
+  override val evaluateExpression: AttributeReference = max
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
new file mode 100644
index 0000000000000..5ed9cd348daba
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+
+case class Min(child: Expression) extends DeclarativeAggregate {
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  // Return data type.
+  override def dataType: DataType = child.dataType
+
+  // Expected input data type.
+  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
+
+  private val min = AttributeReference("min", child.dataType)()
+
+  override val aggBufferAttributes: Seq[AttributeReference] = min :: Nil
+
+  override val initialValues: Seq[Expression] = Seq(
+    /* min = */ Literal.create(null, child.dataType)
+  )
+
+  override val updateExpressions: Seq[Expression] = Seq(
+    /* min = */ If(IsNull(child), min, If(IsNull(min), child, Least(Seq(min, child))))
+  )
+
+  override val mergeExpressions: Seq[Expression] = {
+    val least = Least(Seq(min.left, min.right))
+    Seq(
+      /* min = */ If(IsNull(min.right), min.left, If(IsNull(min.left), min.right, least))
+    )
+  }
+
+  override val evaluateExpression: AttributeReference = min
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
new file mode 100644
index 0000000000000..0def7ddfd9d3d
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.expressions._
+
+case class Skewness(child: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0)
+  extends CentralMomentAgg(child) {
+
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+  override def prettyName: String = "skewness"
+
+  override protected val momentOrder = 3
+
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+    require(moments.length == momentOrder + 1,
+      s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
+    val m2 = moments(2)
+    val m3 = moments(3)
+    if (n == 0.0 || m2 == 0.0) {
+      Double.NaN
+    } else {
+      math.sqrt(n) * m3 / math.sqrt(m2 * m2 * m2)
+    }
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
new file mode 100644
index 0000000000000..3f47ffe13cbc8
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
@@ -0,0 +1,134 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+
+// Compute the population standard deviation of a column
+case class StddevPop(child: Expression) extends StddevAgg(child) {
+  override def isSample: Boolean = false
+  override def prettyName: String = "stddev_pop"
+}
+
+
+// Compute the sample standard deviation of a column
+case class StddevSamp(child: Expression) extends StddevAgg(child) {
+  override def isSample: Boolean = true
+  override def prettyName: String = "stddev_samp"
+}
+
+
+// Compute standard deviation based on online algorithm specified here:
+// http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
+abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
+
+  def isSample: Boolean
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  override def dataType: DataType = resultType
+
+  // Expected input data type.
+  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
+  // new version at planning time (after analysis phase). For now, NullType is added at here
+  // to make it resolved when we have cases like `select stddev(null)`.
+  // We can use our analyzer to cast NullType to the default data type of the NumericType once
+  // we remove the old aggregate functions. Then, we will not need NullType at here.
+  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
+
+  private val resultType = DoubleType
+
+  private val count = AttributeReference("count", resultType)()
+  private val avg = AttributeReference("avg", resultType)()
+  private val mk = AttributeReference("mk", resultType)()
+
+  override val aggBufferAttributes = count :: avg :: mk :: Nil
+
+  override val initialValues: Seq[Expression] = Seq(
+    /* count = */ Cast(Literal(0), resultType),
+    /* avg = */ Cast(Literal(0), resultType),
+    /* mk = */ Cast(Literal(0), resultType)
+  )
+
+  override val updateExpressions: Seq[Expression] = {
+    val value = Cast(child, resultType)
+    val newCount = count + Cast(Literal(1), resultType)
+
+    // update average
+    // avg = avg + (value - avg)/count
+    val newAvg = avg + (value - avg) / newCount
+
+    // update sum ofference from mean
+    // Mk = Mk + (value - preAvg) * (value - updatedAvg)
+    val newMk = mk + (value - avg) * (value - newAvg)
+
+    Seq(
+      /* count = */ If(IsNull(child), count, newCount),
+      /* avg = */ If(IsNull(child), avg, newAvg),
+      /* mk = */ If(IsNull(child), mk, newMk)
+    )
+  }
+
+  override val mergeExpressions: Seq[Expression] = {
+
+    // count merge
+    val newCount = count.left + count.right
+
+    // average merge
+    val newAvg = ((avg.left * count.left) + (avg.right * count.right)) / newCount
+
+    // update sum of square differences
+    val newMk = {
+      val avgDelta = avg.right - avg.left
+      val mkDelta = (avgDelta * avgDelta) * (count.left * count.right) / newCount
+      mk.left + mk.right + mkDelta
+    }
+
+    Seq(
+      /* count = */ If(IsNull(count.left), count.right,
+                       If(IsNull(count.right), count.left, newCount)),
+      /* avg = */ If(IsNull(avg.left), avg.right,
+                     If(IsNull(avg.right), avg.left, newAvg)),
+      /* mk = */ If(IsNull(mk.left), mk.right,
+                    If(IsNull(mk.right), mk.left, newMk))
+    )
+  }
+
+  override val evaluateExpression: Expression = {
+    // when count == 0, return null
+    // when count == 1, return 0
+    // when count >1
+    // stddev_samp = sqrt (mk/(count -1))
+    // stddev_pop = sqrt (mk/count)
+    val varCol =
+      if (isSample) {
+        mk / Cast(count - Cast(Literal(1), resultType), resultType)
+      } else {
+        mk / count
+      }
+
+    If(EqualTo(count, Cast(Literal(0), resultType)), Cast(Literal(null), resultType),
+      If(EqualTo(count, Cast(Literal(1), resultType)), Cast(Literal(0), resultType),
+        Cast(Sqrt(varCol), resultType)))
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
new file mode 100644
index 0000000000000..7f8adbc56ad1d
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types._
+
+case class Sum(child: Expression) extends DeclarativeAggregate {
+
+  override def children: Seq[Expression] = child :: Nil
+
+  override def nullable: Boolean = true
+
+  // Return data type.
+  override def dataType: DataType = resultType
+
+  // Expected input data type.
+  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
+  // new version at planning time (after analysis phase). For now, NullType is added at here
+  // to make it resolved when we have cases like `select sum(null)`.
+  // We can use our analyzer to cast NullType to the default data type of the NumericType once
+  // we remove the old aggregate functions. Then, we will not need NullType at here.
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(TypeCollection(LongType, DoubleType, DecimalType, NullType))
+
+  private val resultType = child.dataType match {
+    case DecimalType.Fixed(precision, scale) =>
+      DecimalType.bounded(precision + 10, scale)
+    // TODO: Remove this line once we remove the NullType from inputTypes.
+    case NullType => IntegerType
+    case _ => child.dataType
+  }
+
+  private val sumDataType = resultType
+
+  private val sum = AttributeReference("sum", sumDataType)()
+
+  private val zero = Cast(Literal(0), sumDataType)
+
+  override val aggBufferAttributes = sum :: Nil
+
+  override val initialValues: Seq[Expression] = Seq(
+    /* sum = */ Literal.create(null, sumDataType)
+  )
+
+  override val updateExpressions: Seq[Expression] = Seq(
+    /* sum = */
+    Coalesce(Seq(Add(Coalesce(Seq(sum, zero)), Cast(child, sumDataType)), sum))
+  )
+
+  override val mergeExpressions: Seq[Expression] = {
+    val add = Add(Coalesce(Seq(sum.left, zero)), Cast(sum.right, sumDataType))
+    Seq(
+      /* sum = */
+      Coalesce(Seq(add, sum.left))
+    )
+  }
+
+  override val evaluateExpression: Expression = Cast(sum, resultType)
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/utils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
similarity index 100%
rename from sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/utils.scala
rename to sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala
new file mode 100644
index 0000000000000..ec63534e5290a
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.expressions._
+
+case class VarianceSamp(child: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0)
+  extends CentralMomentAgg(child) {
+
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+  override def prettyName: String = "var_samp"
+
+  override protected val momentOrder = 2
+
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+    require(moments.length == momentOrder + 1,
+      s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
+
+    if (n == 0.0 || n == 1.0) Double.NaN else moments(2) / (n - 1.0)
+  }
+}
+
+case class VariancePop(child: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0)
+  extends CentralMomentAgg(child) {
+
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+  override def prettyName: String = "var_pop"
+
+  override protected val momentOrder = 2
+
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+    require(moments.length == momentOrder + 1,
+      s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
+
+    if (n == 0.0) Double.NaN else moments(2) / n
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
index 89d63abd9f272..3dcf7915d77b3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
@@ -549,7 +549,7 @@ case class SumDistinct(child: Expression) extends UnaryExpression with PartialAg
     case _ =>
       child.dataType
   }
-  override def toString: String = s"SUM(DISTINCT $child)"
+  override def toString: String = s"sum(distinct $child)"
   override def newInstance(): SumDistinctFunction = new SumDistinctFunction(child, this)
 
   override def asPartial: SplitEvaluation = {
@@ -646,7 +646,7 @@ case class First(
 
   override def nullable: Boolean = true
   override def dataType: DataType = child.dataType
-  override def toString: String = s"FIRST(${child}${if (ignoreNulls) " IGNORE NULLS"})"
+  override def toString: String = s"first(${child}${if (ignoreNulls) " ignore nulls"})"
 
   override def asPartial: SplitEvaluation = {
     val partialFirst = Alias(First(child, ignoreNulls), "PartialFirst")()
@@ -707,7 +707,7 @@ case class Last(
   override def references: AttributeSet = child.references
   override def nullable: Boolean = true
   override def dataType: DataType = child.dataType
-  override def toString: String = s"LAST($child)${if (ignoreNulls) " IGNORE NULLS"}"
+  override def toString: String = s"last($child)${if (ignoreNulls) " ignore nulls"}"
 
   override def asPartial: SplitEvaluation = {
     val partialLast = Alias(Last(child, ignoreNulls), "PartialLast")()
@@ -756,7 +756,7 @@ case class Corr(left: Expression, right: Expression)
     extends BinaryExpression with AggregateExpression1 with ImplicitCastInputTypes {
   override def nullable: Boolean = false
   override def dataType: DoubleType.type = DoubleType
-  override def toString: String = s"CORRELATION($left, $right)"
+  override def toString: String = s"corr($left, $right)"
   override def inputTypes: Seq[AbstractDataType] = Seq(DoubleType, DoubleType)
   override def newInstance(): AggregateFunction1 = {
     throw new UnsupportedOperationException(
@@ -788,14 +788,14 @@ abstract class StddevAgg1(child: Expression) extends UnaryExpression with Partia
 // Compute the population standard deviation of a column
 case class StddevPop(child: Expression) extends StddevAgg1(child) {
 
-  override def toString: String = s"STDDEV_POP($child)"
+  override def toString: String = s"stddev_pop($child)"
   override def isSample: Boolean = false
 }
 
 // Compute the sample standard deviation of a column
 case class StddevSamp(child: Expression) extends StddevAgg1(child) {
 
-  override def toString: String = s"STDDEV_SAMP($child)"
+  override def toString: String = s"stddev_samp($child)"
   override def isSample: Boolean = true
 }
 
@@ -1019,8 +1019,6 @@ case class Kurtosis(child: Expression) extends UnaryExpression with AggregateExp
   override def foldable: Boolean = false
 
   override def prettyName: String = "kurtosis"
-
-  override def toString: String = s"KURTOSIS($child)"
 }
 
 // placeholder
@@ -1038,8 +1036,6 @@ case class Skewness(child: Expression) extends UnaryExpression with AggregateExp
   override def foldable: Boolean = false
 
   override def prettyName: String = "skewness"
-
-  override def toString: String = s"SKEWNESS($child)"
 }
 
 // placeholder
@@ -1056,9 +1052,7 @@ case class VariancePop(child: Expression) extends UnaryExpression with Aggregate
 
   override def foldable: Boolean = false
 
-  override def prettyName: String = "variance_pop"
-
-  override def toString: String = s"VAR_POP($child)"
+  override def prettyName: String = "var_pop"
 }
 
 // placeholder
@@ -1075,7 +1069,5 @@ case class VarianceSamp(child: Expression) extends UnaryExpression with Aggregat
 
   override def foldable: Boolean = false
 
-  override def prettyName: String = "variance_samp"
-
-  override def toString: String = s"VAR_SAMP($child)"
+  override def prettyName: String = "var_samp"
 }

From 701fb5052080fa8c0a79ad7c1e65693ccf444787 Mon Sep 17 00:00:00 2001
From: Adam Roberts <aroberts@uk.ibm.com>
Date: Wed, 4 Nov 2015 14:03:31 -0800
Subject: [PATCH 056/518] [SPARK-10949] Update Snappy version to 1.1.2

This is an updated version of #8995 by a-roberts. Original description follows:

Snappy now supports concatenation of serialized streams, this patch contains a version number change and the "does not support" test is now a "supports" test.

Snappy 1.1.2 changelog mentions:

> snappy-java-1.1.2 (22 September 2015)
> This is a backward compatible release for 1.1.x.
> Add AIX (32-bit) support.
> There is no upgrade for the native libraries of the other platforms.

> A major change since 1.1.1 is a support for reading concatenated results of SnappyOutputStream(s)
> snappy-java-1.1.2-RC2 (18 May 2015)
> Fix #107: SnappyOutputStream.close() is not idempotent
> snappy-java-1.1.2-RC1 (13 May 2015)
> SnappyInputStream now supports reading concatenated compressed results of SnappyOutputStream
> There has been no compressed format change since 1.0.5.x. So You can read the compressed results > interchangeablly between these versions.
> Fixes a problem when java.io.tmpdir does not exist.

Closes #8995.

Author: Adam Roberts <aroberts@uk.ibm.com>
Author: Josh Rosen <joshrosen@databricks.com>

Closes #9439 from JoshRosen/update-snappy.
---
 .../org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java  | 4 ++--
 .../main/scala/org/apache/spark/io/CompressionCodec.scala   | 5 +++++
 .../scala/org/apache/spark/io/CompressionCodecSuite.scala   | 6 ++----
 pom.xml                                                     | 2 +-
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
index e19b37864293c..6a0a89e81c321 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
@@ -254,8 +254,8 @@ private long[] mergeSpills(SpillInfo[] spills) throws IOException {
     final CompressionCodec compressionCodec = CompressionCodec$.MODULE$.createCodec(sparkConf);
     final boolean fastMergeEnabled =
       sparkConf.getBoolean("spark.shuffle.unsafe.fastMergeEnabled", true);
-    final boolean fastMergeIsSupported =
-      !compressionEnabled || compressionCodec instanceof LZFCompressionCodec;
+    final boolean fastMergeIsSupported = !compressionEnabled ||
+      CompressionCodec$.MODULE$.supportsConcatenationOfSerializedStreams(compressionCodec);
     try {
       if (spills.length == 0) {
         new FileOutputStream(outputFile).close(); // Create an empty file
diff --git a/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala b/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
index 9dc36704a676d..ca74eedf89be5 100644
--- a/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
+++ b/core/src/main/scala/org/apache/spark/io/CompressionCodec.scala
@@ -47,6 +47,11 @@ trait CompressionCodec {
 private[spark] object CompressionCodec {
 
   private val configKey = "spark.io.compression.codec"
+
+  private[spark] def supportsConcatenationOfSerializedStreams(codec: CompressionCodec): Boolean = {
+    codec.isInstanceOf[SnappyCompressionCodec] || codec.isInstanceOf[LZFCompressionCodec]
+  }
+
   private val shortCompressionCodecNames = Map(
     "lz4" -> classOf[LZ4CompressionCodec].getName,
     "lzf" -> classOf[LZFCompressionCodec].getName,
diff --git a/core/src/test/scala/org/apache/spark/io/CompressionCodecSuite.scala b/core/src/test/scala/org/apache/spark/io/CompressionCodecSuite.scala
index cbdb33c89d0fb..1553ab60bddaa 100644
--- a/core/src/test/scala/org/apache/spark/io/CompressionCodecSuite.scala
+++ b/core/src/test/scala/org/apache/spark/io/CompressionCodecSuite.scala
@@ -100,12 +100,10 @@ class CompressionCodecSuite extends SparkFunSuite {
     testCodec(codec)
   }
 
-  test("snappy does not support concatenation of serialized streams") {
+  test("snappy supports concatenation of serialized streams") {
     val codec = CompressionCodec.createCodec(conf, classOf[SnappyCompressionCodec].getName)
     assert(codec.getClass === classOf[SnappyCompressionCodec])
-    intercept[Exception] {
-      testConcatenationOfSerializedStreams(codec)
-    }
+    testConcatenationOfSerializedStreams(codec)
   }
 
   test("bad compression codec") {
diff --git a/pom.xml b/pom.xml
index 762bfc7282335..f5a3e44fc0a34 100644
--- a/pom.xml
+++ b/pom.xml
@@ -165,7 +165,7 @@
     <jline.groupid>org.scala-lang</jline.groupid>
     <codehaus.jackson.version>1.9.13</codehaus.jackson.version>
     <fasterxml.jackson.version>2.4.4</fasterxml.jackson.version>
-    <snappy.version>1.1.1.7</snappy.version>
+    <snappy.version>1.1.2</snappy.version>
     <netlib.java.version>1.1.2</netlib.java.version>
     <calcite.version>1.2.0-incubating</calcite.version>
     <commons-codec.version>1.10</commons-codec.version>

From 1b6a5d4af9691c3f7f3ebee3146dc13d12a0e047 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Wed, 4 Nov 2015 14:45:02 -0800
Subject: [PATCH 057/518] [SPARK-11493] remove bitset from BytesToBytesMap

Since we have 4 bytes as number of records in the beginning of a page, the address can not be zero, so we do not need the bitset.

For performance concerns, the bitset could help speed up false lookup if the slot is empty (because bitset is smaller than longArray, cache hit rate will be higher). In practice, the map is filled with 35% - 70% (use 50% as average), so only half of the false lookups can benefit of it, all others will pay the cost of load the bitset (still need to access the longArray anyway).

For aggregation, we always need to access the longArray (insert a new key after false lookup), also confirmed by a benchmark.

 For broadcast hash join, there could be a regression, but a simple benchmark showed that it may not (most of lookup are false):

```
sqlContext.range(1<<20).write.parquet("small")
df = sqlContext.read.parquet('small')
for i in range(3):
    t = time.time()
    df2 = sqlContext.range(1<<26).selectExpr("id * 1111111111 % 987654321 as id2")
    df2.join(df, df.id == df2.id2).count()
    print time.time() -t
```

Having bitset (used time in seconds):
```
17.5404241085
10.2758829594
10.5786800385
```
After removing bitset (used time in seconds):
```
21.8939979076
12.4132959843
9.97224712372
```

cc rxin nongli

Author: Davies Liu <davies@databricks.com>

Closes #9452 from davies/remove_bitset.
---
 .../spark/unsafe/map/BytesToBytesMap.java     |  58 +++------
 .../apache/spark/unsafe/bitset/BitSet.java    | 113 ------------------
 .../spark/unsafe/bitset/BitSetSuite.java      |  88 --------------
 3 files changed, 15 insertions(+), 244 deletions(-)
 delete mode 100644 unsafe/src/main/java/org/apache/spark/unsafe/bitset/BitSet.java
 delete mode 100644 unsafe/src/test/java/org/apache/spark/unsafe/bitset/BitSetSuite.java

diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index e36709c6fc849..07241c827c2ae 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -35,7 +35,6 @@
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.array.ByteArrayMethods;
 import org.apache.spark.unsafe.array.LongArray;
-import org.apache.spark.unsafe.bitset.BitSet;
 import org.apache.spark.unsafe.hash.Murmur3_x86_32;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.unsafe.memory.MemoryLocation;
@@ -123,12 +122,6 @@ public final class BytesToBytesMap extends MemoryConsumer {
    */
   private boolean canGrowArray = true;
 
-  /**
-   * A {@link BitSet} used to track location of the map where the key is set.
-   * Size of the bitset should be half of the size of the long array.
-   */
-  @Nullable private BitSet bitset;
-
   private final double loadFactor;
 
   /**
@@ -427,7 +420,6 @@ public Location lookup(Object keyBase, long keyOffset, int keyLength) {
    * This is a thread-safe version of `lookup`, could be used by multiple threads.
    */
   public void safeLookup(Object keyBase, long keyOffset, int keyLength, Location loc) {
-    assert(bitset != null);
     assert(longArray != null);
 
     if (enablePerfMetrics) {
@@ -440,7 +432,7 @@ public void safeLookup(Object keyBase, long keyOffset, int keyLength, Location l
       if (enablePerfMetrics) {
         numProbes++;
       }
-      if (!bitset.isSet(pos)) {
+      if (longArray.get(pos * 2) == 0) {
         // This is a new key.
         loc.with(pos, hashcode, false);
         return;
@@ -644,7 +636,6 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
       assert (!isDefined) : "Can only set value once for a key";
       assert (keyLength % 8 == 0);
       assert (valueLength % 8 == 0);
-      assert(bitset != null);
       assert(longArray != null);
 
       if (numElements == MAX_CAPACITY || !canGrowArray) {
@@ -678,7 +669,6 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
       Platform.putInt(base, offset, Platform.getInt(base, offset) + 1);
       pageCursor += recordLength;
       numElements++;
-      bitset.set(pos);
       final long storedKeyAddress = taskMemoryManager.encodePageNumberAndOffset(
         currentPage, recordOffset);
       longArray.set(pos * 2, storedKeyAddress);
@@ -734,7 +724,6 @@ private void allocate(int capacity) {
     assert (capacity <= MAX_CAPACITY);
     acquireMemory(capacity * 16);
     longArray = new LongArray(MemoryBlock.fromLongArray(new long[capacity * 2]));
-    bitset = new BitSet(MemoryBlock.fromLongArray(new long[capacity / 64]));
 
     this.growthThreshold = (int) (capacity * loadFactor);
     this.mask = capacity - 1;
@@ -749,7 +738,6 @@ public void freeArray() {
       long used = longArray.memoryBlock().size();
       longArray = null;
       releaseMemory(used);
-      bitset = null;
     }
   }
 
@@ -795,9 +783,7 @@ public long getTotalMemoryConsumption() {
     for (MemoryBlock dataPage : dataPages) {
       totalDataPagesSize += dataPage.size();
     }
-    return totalDataPagesSize +
-      ((bitset != null) ? bitset.memoryBlock().size() : 0L) +
-      ((longArray != null) ? longArray.memoryBlock().size() : 0L);
+    return totalDataPagesSize + ((longArray != null) ? longArray.memoryBlock().size() : 0L);
   }
 
   private void updatePeakMemoryUsed() {
@@ -852,7 +838,6 @@ public int getNumDataPages() {
    */
   @VisibleForTesting
   void growAndRehash() {
-    assert(bitset != null);
     assert(longArray != null);
 
     long resizeStartTime = -1;
@@ -861,39 +846,26 @@ void growAndRehash() {
     }
     // Store references to the old data structures to be used when we re-hash
     final LongArray oldLongArray = longArray;
-    final BitSet oldBitSet = bitset;
-    final int oldCapacity = (int) oldBitSet.capacity();
+    final int oldCapacity = (int) oldLongArray.size() / 2;
 
     // Allocate the new data structures
-    try {
-      allocate(Math.min(growthStrategy.nextCapacity(oldCapacity), MAX_CAPACITY));
-    } catch (OutOfMemoryError oom) {
-      longArray = oldLongArray;
-      bitset = oldBitSet;
-      throw oom;
-    }
+    allocate(Math.min(growthStrategy.nextCapacity(oldCapacity), MAX_CAPACITY));
 
     // Re-mask (we don't recompute the hashcode because we stored all 32 bits of it)
-    for (int pos = oldBitSet.nextSetBit(0); pos >= 0; pos = oldBitSet.nextSetBit(pos + 1)) {
-      final long keyPointer = oldLongArray.get(pos * 2);
-      final int hashcode = (int) oldLongArray.get(pos * 2 + 1);
+    for (int i = 0; i < oldLongArray.size(); i += 2) {
+      final long keyPointer = oldLongArray.get(i);
+      if (keyPointer == 0) {
+        continue;
+      }
+      final int hashcode = (int) oldLongArray.get(i + 1);
       int newPos = hashcode & mask;
       int step = 1;
-      boolean keepGoing = true;
-
-      // No need to check for equality here when we insert so this has one less if branch than
-      // the similar code path in addWithoutResize.
-      while (keepGoing) {
-        if (!bitset.isSet(newPos)) {
-          bitset.set(newPos);
-          longArray.set(newPos * 2, keyPointer);
-          longArray.set(newPos * 2 + 1, hashcode);
-          keepGoing = false;
-        } else {
-          newPos = (newPos + step) & mask;
-          step++;
-        }
+      while (longArray.get(newPos * 2) != 0) {
+        newPos = (newPos + step) & mask;
+        step++;
       }
+      longArray.set(newPos * 2, keyPointer);
+      longArray.set(newPos * 2 + 1, hashcode);
     }
     releaseMemory(oldLongArray.memoryBlock().size());
 
diff --git a/unsafe/src/main/java/org/apache/spark/unsafe/bitset/BitSet.java b/unsafe/src/main/java/org/apache/spark/unsafe/bitset/BitSet.java
deleted file mode 100644
index 7c124173b0bbb..0000000000000
--- a/unsafe/src/main/java/org/apache/spark/unsafe/bitset/BitSet.java
+++ /dev/null
@@ -1,113 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.unsafe.bitset;
-
-import org.apache.spark.unsafe.array.LongArray;
-import org.apache.spark.unsafe.memory.MemoryBlock;
-
-/**
- * A fixed size uncompressed bit set backed by a {@link LongArray}.
- *
- * Each bit occupies exactly one bit of storage.
- */
-public final class BitSet {
-
-  /** A long array for the bits. */
-  private final LongArray words;
-
-  /** Length of the long array. */
-  private final int numWords;
-
-  private final Object baseObject;
-  private final long baseOffset;
-
-  /**
-   * Creates a new {@link BitSet} using the specified memory block. Size of the memory block must be
-   * multiple of 8 bytes (i.e. 64 bits).
-   */
-  public BitSet(MemoryBlock memory) {
-    words = new LongArray(memory);
-    assert (words.size() <= Integer.MAX_VALUE);
-    numWords = (int) words.size();
-    baseObject = words.memoryBlock().getBaseObject();
-    baseOffset = words.memoryBlock().getBaseOffset();
-  }
-
-  public MemoryBlock memoryBlock() {
-    return words.memoryBlock();
-  }
-
-  /**
-   * Returns the number of bits in this {@code BitSet}.
-   */
-  public long capacity() {
-    return numWords * 64;
-  }
-
-  /**
-   * Sets the bit at the specified index to {@code true}.
-   */
-  public void set(int index) {
-    assert index < numWords * 64 : "index (" + index + ") should < length (" + numWords * 64 + ")";
-    BitSetMethods.set(baseObject, baseOffset, index);
-  }
-
-  /**
-   * Sets the bit at the specified index to {@code false}.
-   */
-  public void unset(int index) {
-    assert index < numWords * 64 : "index (" + index + ") should < length (" + numWords * 64 + ")";
-    BitSetMethods.unset(baseObject, baseOffset, index);
-  }
-
-  /**
-   * Returns {@code true} if the bit is set at the specified index.
-   */
-  public boolean isSet(int index) {
-    assert index < numWords * 64 : "index (" + index + ") should < length (" + numWords * 64 + ")";
-    return BitSetMethods.isSet(baseObject, baseOffset, index);
-  }
-
-  /**
-   * Returns the index of the first bit that is set to true that occurs on or after the
-   * specified starting index. If no such bit exists then {@code -1} is returned.
-   * <p>
-   * To iterate over the true bits in a BitSet, use the following loop:
-   * <pre>
-   * <code>
-   *  for (long i = bs.nextSetBit(0); i &gt;= 0; i = bs.nextSetBit(i + 1)) {
-   *    // operate on index i here
-   *  }
-   * </code>
-   * </pre>
-   *
-   * @param fromIndex the index to start checking from (inclusive)
-   * @return the index of the next set bit, or -1 if there is no such bit
-   */
-  public int nextSetBit(int fromIndex) {
-    return BitSetMethods.nextSetBit(baseObject, baseOffset, fromIndex, numWords);
-  }
-
-  /**
-   * Returns {@code true} if any bit is set.
-   */
-  public boolean anySet() {
-    return BitSetMethods.anySet(baseObject, baseOffset, numWords);
-  }
-
-}
diff --git a/unsafe/src/test/java/org/apache/spark/unsafe/bitset/BitSetSuite.java b/unsafe/src/test/java/org/apache/spark/unsafe/bitset/BitSetSuite.java
deleted file mode 100644
index 14e38683df4ab..0000000000000
--- a/unsafe/src/test/java/org/apache/spark/unsafe/bitset/BitSetSuite.java
+++ /dev/null
@@ -1,88 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.unsafe.bitset;
-
-import org.junit.Assert;
-import org.junit.Test;
-
-import org.apache.spark.unsafe.memory.MemoryBlock;
-
-public class BitSetSuite {
-
-  private static BitSet createBitSet(int capacity) {
-    Assert.assertEquals(0, capacity % 64);
-    return new BitSet(MemoryBlock.fromLongArray(new long[capacity / 64]));
-  }
-
-  @Test
-  public void basicOps() {
-    BitSet bs = createBitSet(64);
-    Assert.assertEquals(64, bs.capacity());
-
-    // Make sure the bit set starts empty.
-    for (int i = 0; i < bs.capacity(); i++) {
-      Assert.assertFalse(bs.isSet(i));
-    }
-    // another form of asserting that the bit set is empty
-    Assert.assertFalse(bs.anySet());
-
-    // Set every bit and check it.
-    for (int i = 0; i < bs.capacity(); i++) {
-      bs.set(i);
-      Assert.assertTrue(bs.isSet(i));
-    }
-
-    // Unset every bit and check it.
-    for (int i = 0; i < bs.capacity(); i++) {
-      Assert.assertTrue(bs.isSet(i));
-      bs.unset(i);
-      Assert.assertFalse(bs.isSet(i));
-    }
-
-    // Make sure anySet() can detect any set bit
-    bs = createBitSet(256);
-    bs.set(64);
-    Assert.assertTrue(bs.anySet());
-  }
-
-  @Test
-  public void traversal() {
-    BitSet bs = createBitSet(256);
-
-    Assert.assertEquals(-1, bs.nextSetBit(0));
-    Assert.assertEquals(-1, bs.nextSetBit(10));
-    Assert.assertEquals(-1, bs.nextSetBit(64));
-
-    bs.set(10);
-    Assert.assertEquals(10, bs.nextSetBit(0));
-    Assert.assertEquals(10, bs.nextSetBit(1));
-    Assert.assertEquals(10, bs.nextSetBit(10));
-    Assert.assertEquals(-1, bs.nextSetBit(11));
-
-    bs.set(11);
-    Assert.assertEquals(10, bs.nextSetBit(10));
-    Assert.assertEquals(11, bs.nextSetBit(11));
-
-    // Skip a whole word and find it
-    bs.set(190);
-    Assert.assertEquals(190, bs.nextSetBit(12));
-
-    Assert.assertEquals(-1, bs.nextSetBit(191));
-    Assert.assertEquals(-1, bs.nextSetBit(256));
-  }
-}

From 411ff6afb485c9d8cfc667c9346f836f2529ea9f Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Wed, 4 Nov 2015 15:28:19 -0800
Subject: [PATCH 058/518] [SPARK-10028][MLLIB][PYTHON] Add Python API for
 PrefixSpan

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>

Closes #9469 from yu-iskw/SPARK-10028.
---
 .../api/python/PrefixSpanModelWrapper.scala   | 32 +++++++++
 .../mllib/api/python/PythonMLLibAPI.scala     | 23 ++++++-
 python/pyspark/mllib/fpm.py                   | 69 ++++++++++++++++++-
 3 files changed, 122 insertions(+), 2 deletions(-)
 create mode 100644 mllib/src/main/scala/org/apache/spark/mllib/api/python/PrefixSpanModelWrapper.scala

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PrefixSpanModelWrapper.scala b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PrefixSpanModelWrapper.scala
new file mode 100644
index 0000000000000..0027602a04f81
--- /dev/null
+++ b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PrefixSpanModelWrapper.scala
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.api.python
+
+import org.apache.spark.mllib.fpm.PrefixSpanModel
+import org.apache.spark.rdd.RDD
+
+/**
+ * A Wrapper of PrefixSpanModel to provide helper method for Python
+ */
+private[python] class PrefixSpanModelWrapper(model: PrefixSpanModel[Any])
+  extends PrefixSpanModel(model.freqSequences) {
+
+  def getFreqSequences: RDD[Array[Any]] = {
+    SerDe.fromTuple2RDD(model.freqSequences.map(x => (x.javaSequence, x.freq)))
+  }
+}
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
index 21e55938fa7aa..40c41806cdfea 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
@@ -35,7 +35,7 @@ import org.apache.spark.mllib.classification._
 import org.apache.spark.mllib.clustering._
 import org.apache.spark.mllib.evaluation.RankingMetrics
 import org.apache.spark.mllib.feature._
-import org.apache.spark.mllib.fpm.{FPGrowth, FPGrowthModel}
+import org.apache.spark.mllib.fpm.{FPGrowth, FPGrowthModel, PrefixSpan}
 import org.apache.spark.mllib.linalg._
 import org.apache.spark.mllib.linalg.distributed._
 import org.apache.spark.mllib.optimization._
@@ -557,6 +557,27 @@ private[python] class PythonMLLibAPI extends Serializable {
     new FPGrowthModelWrapper(model)
   }
 
+  /**
+   * Java stub for Python mllib PrefixSpan.train().  This stub returns a handle
+   * to the Java object instead of the content of the Java object.  Extra care
+   * needs to be taken in the Python code to ensure it gets freed on exit; see
+   * the Py4J documentation.
+   */
+  def trainPrefixSpanModel(
+      data: JavaRDD[java.util.ArrayList[java.util.ArrayList[Any]]],
+      minSupport: Double,
+      maxPatternLength: Int,
+      localProjDBSize: Int ): PrefixSpanModelWrapper = {
+    val prefixSpan = new PrefixSpan()
+      .setMinSupport(minSupport)
+      .setMaxPatternLength(maxPatternLength)
+      .setMaxLocalProjDBSize(localProjDBSize)
+
+    val trainData = data.rdd.map(_.asScala.toArray.map(_.asScala.toArray))
+    val model = prefixSpan.run(trainData)
+    new PrefixSpanModelWrapper(model)
+  }
+
   /**
    * Java stub for Normalizer.transform()
    */
diff --git a/python/pyspark/mllib/fpm.py b/python/pyspark/mllib/fpm.py
index bdabba9602a8c..2039decc0cb3c 100644
--- a/python/pyspark/mllib/fpm.py
+++ b/python/pyspark/mllib/fpm.py
@@ -23,7 +23,7 @@
 from pyspark.rdd import ignore_unicode_prefix
 from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc, inherit_doc
 
-__all__ = ['FPGrowth', 'FPGrowthModel']
+__all__ = ['FPGrowth', 'FPGrowthModel', 'PrefixSpan', 'PrefixSpanModel']
 
 
 @inherit_doc
@@ -85,6 +85,73 @@ class FreqItemset(namedtuple("FreqItemset", ["items", "freq"])):
         """
 
 
+@inherit_doc
+@ignore_unicode_prefix
+class PrefixSpanModel(JavaModelWrapper):
+    """
+    .. note:: Experimental
+
+    Model fitted by PrefixSpan
+
+    >>> data = [
+    ...    [["a", "b"], ["c"]],
+    ...    [["a"], ["c", "b"], ["a", "b"]],
+    ...    [["a", "b"], ["e"]],
+    ...    [["f"]]]
+    >>> rdd = sc.parallelize(data, 2)
+    >>> model = PrefixSpan.train(rdd)
+    >>> sorted(model.freqSequences().collect())
+    [FreqSequence(sequence=[[u'a']], freq=3), FreqSequence(sequence=[[u'a'], [u'a']], freq=1), ...
+
+    .. versionadded:: 1.6.0
+    """
+
+    @since("1.6.0")
+    def freqSequences(self):
+        """Gets frequence sequences"""
+        return self.call("getFreqSequences").map(lambda x: PrefixSpan.FreqSequence(x[0], x[1]))
+
+
+class PrefixSpan(object):
+    """
+    .. note:: Experimental
+
+    A parallel PrefixSpan algorithm to mine frequent sequential patterns.
+    The PrefixSpan algorithm is described in J. Pei, et al., PrefixSpan:
+    Mining Sequential Patterns Efficiently by Prefix-Projected Pattern Growth
+    ([[http://doi.org/10.1109/ICDE.2001.914830]]).
+
+    .. versionadded:: 1.6.0
+    """
+
+    @classmethod
+    @since("1.6.0")
+    def train(cls, data, minSupport=0.1, maxPatternLength=10, maxLocalProjDBSize=32000000):
+        """
+        Finds the complete set of frequent sequential patterns in the input sequences of itemsets.
+
+        :param data: The input data set, each element contains a sequnce of itemsets.
+        :param minSupport: the minimal support level of the sequential pattern, any pattern appears
+            more than  (minSupport * size-of-the-dataset) times will be output (default: `0.1`)
+        :param maxPatternLength: the maximal length of the sequential pattern, any pattern appears
+            less than maxPatternLength will be output. (default: `10`)
+        :param maxLocalProjDBSize: The maximum number of items (including delimiters used in
+            the internal storage format) allowed in a projected database before local
+            processing. If a projected database exceeds this size, another
+            iteration of distributed prefix growth is run. (default: `32000000`)
+        """
+        model = callMLlibFunc("trainPrefixSpanModel",
+                              data, minSupport, maxPatternLength, maxLocalProjDBSize)
+        return PrefixSpanModel(model)
+
+    class FreqSequence(namedtuple("FreqSequence", ["sequence", "freq"])):
+        """
+        Represents a (sequence, freq) tuple.
+
+        .. versionadded:: 1.6.0
+        """
+
+
 def _test():
     import doctest
     import pyspark.mllib.fpm

From b6e0a5ae6f243139f11c9cbbf18cddd3f25db208 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 4 Nov 2015 16:49:25 -0800
Subject: [PATCH 059/518] [SPARK-11510][SQL] Remove SQL aggregation tests for
 higher order statistics

We have some aggregate function tests in both DataFrameAggregateSuite and SQLQuerySuite. The two have almost the same coverage and we should just remove the SQL one.

Author: Reynold Xin <rxin@databricks.com>

Closes #9475 from rxin/SPARK-11510.
---
 .../spark/sql/DataFrameAggregateSuite.scala   | 97 ++++++-------------
 .../org/apache/spark/sql/SQLQuerySuite.scala  | 77 ---------------
 .../spark/sql/StringFunctionsSuite.scala      |  1 -
 3 files changed, 28 insertions(+), 147 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index b0e2ffaa60687..2e679e7bc4e0a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -83,13 +83,8 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
 
   test("average") {
     checkAnswer(
-      testData2.agg(avg('a)),
-      Row(2.0))
-
-    // Also check mean
-    checkAnswer(
-      testData2.agg(mean('a)),
-      Row(2.0))
+      testData2.agg(avg('a), mean('a)),
+      Row(2.0, 2.0))
 
     checkAnswer(
       testData2.agg(avg('a), sumDistinct('a)), // non-partial
@@ -98,6 +93,7 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
     checkAnswer(
       decimalData.agg(avg('a)),
       Row(new java.math.BigDecimal(2.0)))
+
     checkAnswer(
       decimalData.agg(avg('a), sumDistinct('a)), // non-partial
       Row(new java.math.BigDecimal(2.0), new java.math.BigDecimal(6)) :: Nil)
@@ -168,44 +164,23 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
 
   test("zero count") {
     val emptyTableData = Seq.empty[(Int, Int)].toDF("a", "b")
-    assert(emptyTableData.count() === 0)
-
     checkAnswer(
       emptyTableData.agg(count('a), sumDistinct('a)), // non-partial
       Row(0, null))
   }
 
   test("stddev") {
-    val testData2ADev = math.sqrt(4/5.0)
-
+    val testData2ADev = math.sqrt(4 / 5.0)
     checkAnswer(
-      testData2.agg(stddev('a)),
-      Row(testData2ADev))
-
-    checkAnswer(
-      testData2.agg(stddev_pop('a)),
-      Row(math.sqrt(4/6.0)))
-
-    checkAnswer(
-      testData2.agg(stddev_samp('a)),
-      Row(testData2ADev))
+      testData2.agg(stddev('a), stddev_pop('a), stddev_samp('a)),
+      Row(testData2ADev, math.sqrt(4 / 6.0), testData2ADev))
   }
 
   test("zero stddev") {
     val emptyTableData = Seq.empty[(Int, Int)].toDF("a", "b")
-    assert(emptyTableData.count() == 0)
-
-    checkAnswer(
-    emptyTableData.agg(stddev('a)),
-    Row(null))
-
     checkAnswer(
-    emptyTableData.agg(stddev_pop('a)),
-    Row(null))
-
-    checkAnswer(
-    emptyTableData.agg(stddev_samp('a)),
-    Row(null))
+    emptyTableData.agg(stddev('a), stddev_pop('a), stddev_samp('a)),
+    Row(null, null, null))
   }
 
   test("zero sum") {
@@ -227,6 +202,7 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
 
     val sparkVariance = testData2.agg(variance('a))
     checkAggregatesWithTol(sparkVariance, Row(4.0 / 5.0), absTol)
+
     val sparkVariancePop = testData2.agg(var_pop('a))
     checkAggregatesWithTol(sparkVariancePop, Row(4.0 / 6.0), absTol)
 
@@ -241,52 +217,35 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
   }
 
   test("zero moments") {
-    val emptyTableData = Seq((1, 2)).toDF("a", "b")
-    assert(emptyTableData.count() === 1)
-
-    checkAnswer(
-      emptyTableData.agg(variance('a)),
-      Row(Double.NaN))
-
-    checkAnswer(
-      emptyTableData.agg(var_samp('a)),
-      Row(Double.NaN))
-
+    val input = Seq((1, 2)).toDF("a", "b")
     checkAnswer(
-      emptyTableData.agg(var_pop('a)),
-      Row(0.0))
+      input.agg(variance('a), var_samp('a), var_pop('a), skewness('a), kurtosis('a)),
+      Row(Double.NaN, Double.NaN, 0.0, Double.NaN, Double.NaN))
 
     checkAnswer(
-      emptyTableData.agg(skewness('a)),
-      Row(Double.NaN))
-
-    checkAnswer(
-      emptyTableData.agg(kurtosis('a)),
-      Row(Double.NaN))
+      input.agg(
+        expr("variance(a)"),
+        expr("var_samp(a)"),
+        expr("var_pop(a)"),
+        expr("skewness(a)"),
+        expr("kurtosis(a)")),
+      Row(Double.NaN, Double.NaN, 0.0, Double.NaN, Double.NaN))
   }
 
   test("null moments") {
     val emptyTableData = Seq.empty[(Int, Int)].toDF("a", "b")
-    assert(emptyTableData.count() === 0)
-
-    checkAnswer(
-      emptyTableData.agg(variance('a)),
-      Row(Double.NaN))
-
-    checkAnswer(
-      emptyTableData.agg(var_samp('a)),
-      Row(Double.NaN))
-
-    checkAnswer(
-      emptyTableData.agg(var_pop('a)),
-      Row(Double.NaN))
 
     checkAnswer(
-      emptyTableData.agg(skewness('a)),
-      Row(Double.NaN))
+      emptyTableData.agg(variance('a), var_samp('a), var_pop('a), skewness('a), kurtosis('a)),
+      Row(Double.NaN, Double.NaN, Double.NaN, Double.NaN, Double.NaN))
 
     checkAnswer(
-      emptyTableData.agg(kurtosis('a)),
-      Row(Double.NaN))
+      emptyTableData.agg(
+        expr("variance(a)"),
+        expr("var_samp(a)"),
+        expr("var_pop(a)"),
+        expr("skewness(a)"),
+        expr("kurtosis(a)")),
+      Row(Double.NaN, Double.NaN, Double.NaN, Double.NaN, Double.NaN))
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index 5731a356243e5..3de277a79a52c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -726,83 +726,6 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
     }
   }
 
-  test("stddev") {
-    checkAnswer(
-      sql("SELECT STDDEV(a) FROM testData2"),
-      Row(math.sqrt(4.0 / 5.0))
-    )
-  }
-
-  test("stddev_pop") {
-    checkAnswer(
-      sql("SELECT STDDEV_POP(a) FROM testData2"),
-      Row(math.sqrt(4.0 / 6.0))
-    )
-  }
-
-  test("stddev_samp") {
-    checkAnswer(
-      sql("SELECT STDDEV_SAMP(a) FROM testData2"),
-      Row(math.sqrt(4/5.0))
-    )
-  }
-
-  test("var_samp") {
-    val absTol = 1e-8
-    val sparkAnswer = sql("SELECT VAR_SAMP(a) FROM testData2")
-    val expectedAnswer = Row(4.0 / 5.0)
-    checkAggregatesWithTol(sparkAnswer, expectedAnswer, absTol)
-  }
-
-  test("variance") {
-    val absTol = 1e-8
-    val sparkAnswer = sql("SELECT VARIANCE(a) FROM testData2")
-    val expectedAnswer = Row(0.8)
-    checkAggregatesWithTol(sparkAnswer, expectedAnswer, absTol)
-  }
-
-  test("var_pop") {
-    val absTol = 1e-8
-    val sparkAnswer = sql("SELECT VAR_POP(a) FROM testData2")
-    val expectedAnswer = Row(4.0 / 6.0)
-    checkAggregatesWithTol(sparkAnswer, expectedAnswer, absTol)
-  }
-
-  test("skewness") {
-    val absTol = 1e-8
-    val sparkAnswer = sql("SELECT skewness(a) FROM testData2")
-    val expectedAnswer = Row(0.0)
-    checkAggregatesWithTol(sparkAnswer, expectedAnswer, absTol)
-  }
-
-  test("kurtosis") {
-    val absTol = 1e-8
-    val sparkAnswer = sql("SELECT kurtosis(a) FROM testData2")
-    val expectedAnswer = Row(-1.5)
-    checkAggregatesWithTol(sparkAnswer, expectedAnswer, absTol)
-  }
-
-  test("stddev agg") {
-    checkAnswer(
-      sql("SELECT a, stddev(b), stddev_pop(b), stddev_samp(b) FROM testData2 GROUP BY a"),
-      (1 to 3).map(i => Row(i, math.sqrt(1.0 / 2.0), math.sqrt(1.0 / 4.0), math.sqrt(1.0 / 2.0))))
-  }
-
-  test("variance agg") {
-    val absTol = 1e-8
-    checkAggregatesWithTol(
-      sql("SELECT a, variance(b), var_samp(b), var_pop(b) FROM testData2 GROUP BY a"),
-      (1 to 3).map(i => Row(i, 1.0 / 2.0, 1.0 / 2.0, 1.0 / 4.0)),
-      absTol)
-  }
-
-  test("skewness and kurtosis agg") {
-    val absTol = 1e-8
-    val sparkAnswer = sql("SELECT a, skewness(b), kurtosis(b) FROM testData2 GROUP BY a")
-    val expectedAnswer = (1 to 3).map(i => Row(i, 0.0, -2.0))
-    checkAggregatesWithTol(sparkAnswer, expectedAnswer, absTol)
-  }
-
   test("inner join where, one match per row") {
     checkAnswer(
       sql("SELECT * FROM upperCaseData JOIN lowerCaseData WHERE n = N"),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
index e12e6bea30260..e2090b0a83ce7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala
@@ -19,7 +19,6 @@ package org.apache.spark.sql
 
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
-import org.apache.spark.sql.types.Decimal
 
 
 class StringFunctionsSuite extends QueryTest with SharedSQLContext {

From ce5e6a2849ae860689fa3e7d5aaa12216945ea99 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Wed, 4 Nov 2015 16:58:38 -0800
Subject: [PATCH 060/518] [SPARK-11491] Update build to use Scala 2.10.5

Spark should build against Scala 2.10.5, since that includes a fix for Scaladoc that will fix doc snapshot publishing: https://issues.scala-lang.org/browse/SI-8479

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9450 from JoshRosen/upgrade-to-scala-2.10.5.
---
 LICENSE                            | 10 +++++-----
 dev/audit-release/README.md        |  2 +-
 dev/audit-release/audit_release.py |  2 +-
 docker/spark-test/base/Dockerfile  |  2 +-
 docs/_config.yml                   |  2 +-
 pom.xml                            |  4 ++--
 project/SparkBuild.scala           |  2 +-
 7 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/LICENSE b/LICENSE
index 790476ece15bd..0db2d14465bd3 100644
--- a/LICENSE
+++ b/LICENSE
@@ -250,11 +250,11 @@ The text of each license is also included at licenses/LICENSE-[project].txt.
         (Interpreter classes (all .scala files in repl/src/main/scala
         except for Main.Scala, SparkHelper.scala and ExecutorClassLoader.scala),
         and for SerializableMapWrapper in JavaUtils.scala)
-     (BSD-like) Scala Actors library (org.scala-lang:scala-actors:2.10.4 - http://www.scala-lang.org/)
-     (BSD-like) Scala Compiler (org.scala-lang:scala-compiler:2.10.4 - http://www.scala-lang.org/)
-     (BSD-like) Scala Compiler (org.scala-lang:scala-reflect:2.10.4 - http://www.scala-lang.org/)
-     (BSD-like) Scala Library (org.scala-lang:scala-library:2.10.4 - http://www.scala-lang.org/)
-     (BSD-like) Scalap (org.scala-lang:scalap:2.10.4 - http://www.scala-lang.org/)
+     (BSD-like) Scala Actors library (org.scala-lang:scala-actors:2.10.5 - http://www.scala-lang.org/)
+     (BSD-like) Scala Compiler (org.scala-lang:scala-compiler:2.10.5 - http://www.scala-lang.org/)
+     (BSD-like) Scala Compiler (org.scala-lang:scala-reflect:2.10.5 - http://www.scala-lang.org/)
+     (BSD-like) Scala Library (org.scala-lang:scala-library:2.10.5 - http://www.scala-lang.org/)
+     (BSD-like) Scalap (org.scala-lang:scalap:2.10.5 - http://www.scala-lang.org/)
      (BSD-style) scalacheck (org.scalacheck:scalacheck_2.10:1.10.0 - http://www.scalacheck.org)
      (BSD-style) spire (org.spire-math:spire_2.10:0.7.1 - http://spire-math.org)
      (BSD-style) spire-macros (org.spire-math:spire-macros_2.10:0.7.1 - http://spire-math.org)
diff --git a/dev/audit-release/README.md b/dev/audit-release/README.md
index 38becda0eae92..f72f8c653a265 100644
--- a/dev/audit-release/README.md
+++ b/dev/audit-release/README.md
@@ -4,7 +4,7 @@ run them locally by setting appropriate environment variables.
 
 ```
 $ cd sbt_app_core
-$ SCALA_VERSION=2.10.4 \
+$ SCALA_VERSION=2.10.5 \
   SPARK_VERSION=1.0.0-SNAPSHOT \
   SPARK_RELEASE_REPOSITORY=file:///home/patrick/.ivy2/local \
   sbt run
diff --git a/dev/audit-release/audit_release.py b/dev/audit-release/audit_release.py
index 0b7069f6e116a..27d1dd784ce2e 100755
--- a/dev/audit-release/audit_release.py
+++ b/dev/audit-release/audit_release.py
@@ -35,7 +35,7 @@
 RELEASE_KEY = "XXXXXXXX" # Your 8-digit hex
 RELEASE_REPOSITORY = "https://repository.apache.org/content/repositories/orgapachespark-1033"
 RELEASE_VERSION = "1.1.1"
-SCALA_VERSION = "2.10.4"
+SCALA_VERSION = "2.10.5"
 SCALA_BINARY_VERSION = "2.10"
 
 # Do not set these
diff --git a/docker/spark-test/base/Dockerfile b/docker/spark-test/base/Dockerfile
index 5dbdb8b22a44f..7ba0de603dc7d 100644
--- a/docker/spark-test/base/Dockerfile
+++ b/docker/spark-test/base/Dockerfile
@@ -25,7 +25,7 @@ RUN apt-get update && \
     apt-get install -y less openjdk-7-jre-headless net-tools vim-tiny sudo openssh-server && \
     rm -rf /var/lib/apt/lists/*
 
-ENV SCALA_VERSION 2.10.4
+ENV SCALA_VERSION 2.10.5
 ENV CDH_VERSION cdh4
 ENV SCALA_HOME /opt/scala-$SCALA_VERSION
 ENV SPARK_HOME /opt/spark
diff --git a/docs/_config.yml b/docs/_config.yml
index c59cc465ef89d..2c70b76be8b7a 100644
--- a/docs/_config.yml
+++ b/docs/_config.yml
@@ -17,7 +17,7 @@ include:
 SPARK_VERSION: 1.6.0-SNAPSHOT
 SPARK_VERSION_SHORT: 1.6.0
 SCALA_BINARY_VERSION: "2.10"
-SCALA_VERSION: "2.10.4"
+SCALA_VERSION: "2.10.5"
 MESOS_VERSION: 0.21.0
 SPARK_ISSUE_TRACKER_URL: https://issues.apache.org/jira/browse/SPARK
 SPARK_GITHUB_URL: https://github.com/apache/spark
diff --git a/pom.xml b/pom.xml
index f5a3e44fc0a34..4ed1c0c82dee6 100644
--- a/pom.xml
+++ b/pom.xml
@@ -159,7 +159,7 @@
     <!--  commons-httpclient/commons-httpclient-->
     <httpclient.classic.version>3.1</httpclient.classic.version>
     <commons.math3.version>3.4.1</commons.math3.version>
-    <scala.version>2.10.4</scala.version>
+    <scala.version>2.10.5</scala.version>
     <scala.binary.version>2.10</scala.binary.version>
     <jline.version>${scala.version}</jline.version>
     <jline.groupid>org.scala-lang</jline.groupid>
@@ -2422,7 +2422,7 @@
         <property><name>!scala-2.11</name></property>
       </activation>
       <properties>
-        <scala.version>2.10.4</scala.version>
+        <scala.version>2.10.5</scala.version>
         <scala.binary.version>2.10</scala.binary.version>
         <jline.version>${scala.version}</jline.version>
         <jline.groupid>org.scala-lang</jline.groupid>
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index 766edd9500c30..75c36930decef 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -316,7 +316,7 @@ object OldDeps {
 
   def oldDepsSettings() = Defaults.coreDefaultSettings ++ Seq(
     name := "old-deps",
-    scalaVersion := "2.10.4",
+    scalaVersion := "2.10.5",
     retrieveManaged := true,
     retrievePattern := "[type]s/[artifact](-[revision])(-[classifier]).[ext]",
     libraryDependencies := Seq("spark-streaming-mqtt", "spark-streaming-zeromq",

From a752ddad7fe1d0f01b51f7551ec017ff87e1eea5 Mon Sep 17 00:00:00 2001
From: Zhenhua Wang <wangzhenhua@huawei.com>
Date: Wed, 4 Nov 2015 17:16:00 -0800
Subject: [PATCH 061/518] [SPARK-11398] [SQL] unnecessary def dialectClassName
 in HiveContext, and misleading dialect conf at the start of spark-sql
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

1. def dialectClassName in HiveContext is unnecessary.
In HiveContext, if conf.dialect == "hiveql", getSQLDialect() will return new HiveQLDialect(this);
else it will use super.getSQLDialect(). Then in super.getSQLDialect(), it calls dialectClassName, which is overriden in HiveContext and still return super.dialectClassName.
So we'll never reach the code "classOf[HiveQLDialect].getCanonicalName" of def dialectClassName in HiveContext.

2. When we start bin/spark-sql, the default context is HiveContext, and the corresponding dialect is hiveql.
However, if we type "set spark.sql.dialect;", the result is "sql", which is inconsistent with the actual dialect and is misleading. For example, we can use sql like "create table" which is only allowed in hiveql, but this dialect conf shows it's "sql".
Although this problem will not cause any execution error, it's misleading to spark sql users. Therefore I think we should fix it.
In this pr, while procesing “set spark.sql.dialect” in SetCommand, I use "conf.dialect" instead of "getConf()" for the case of key == SQLConf.DIALECT.key, so that it will return the right dialect conf.

Author: Zhenhua Wang <wangzhenhua@huawei.com>

Closes #9349 from wzhfy/dialect.
---
 .../scala/org/apache/spark/sql/execution/commands.scala    | 6 +++++-
 .../main/scala/org/apache/spark/sql/hive/HiveContext.scala | 6 ------
 .../apache/spark/sql/hive/execution/SQLQuerySuite.scala    | 7 +++++++
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
index 856607615ae87..e5f60b15e7359 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
@@ -156,7 +156,11 @@ case class SetCommand(kv: Option[(String, Option[String])]) extends RunnableComm
       val runFunc = (sqlContext: SQLContext) => {
         val value =
           try {
-            sqlContext.getConf(key)
+            if (key == SQLConf.DIALECT.key) {
+              sqlContext.conf.dialect
+            } else {
+              sqlContext.getConf(key)
+            }
           } catch {
             case _: NoSuchElementException => "<undefined>"
           }
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index 83a81cf5d1fcf..1f5135320326c 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -555,12 +555,6 @@ class HiveContext private[hive](
     override def caseSensitiveAnalysis: Boolean = getConf(SQLConf.CASE_SENSITIVE, false)
   }
 
-  protected[sql] override def dialectClassName = if (conf.dialect == "hiveql") {
-    classOf[HiveQLDialect].getCanonicalName
-  } else {
-    super.dialectClassName
-  }
-
   protected[sql] override def getSQLDialect(): ParserDialect = {
     if (conf.dialect == "hiveql") {
       new HiveQLDialect(this)
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
index fd380641dcc71..af48d478953b4 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
@@ -335,6 +335,13 @@ class SQLQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
     }
   }
 
+  test("SQL dialect at the start of HiveContext") {
+    val hiveContext = new HiveContext(sqlContext.sparkContext)
+    val dialectConf = "spark.sql.dialect"
+    checkAnswer(hiveContext.sql(s"set $dialectConf"), Row(dialectConf, "hiveql"))
+    assert(hiveContext.getSQLDialect().getClass === classOf[HiveQLDialect])
+  }
+
   test("SQL Dialect Switching") {
     assert(getSQLDialect().getClass === classOf[HiveQLDialect])
     setConf("spark.sql.dialect", classOf[MyDialect].getCanonicalName())

From d0b56339625727744e2c30fc2167bc6a457d37f7 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Wed, 4 Nov 2015 17:19:52 -0800
Subject: [PATCH 062/518] [SPARK-11307] Reduce memory consumption of
 OutputCommitCoordinator

OutputCommitCoordinator uses a map in a place where an array would suffice, increasing its memory consumption for result stages with millions of tasks.

This patch replaces that map with an array. The only tricky part of this is reasoning about the range of possible array indexes in order to make sure that we never index out of bounds.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9274 from JoshRosen/SPARK-11307.
---
 .../apache/spark/scheduler/DAGScheduler.scala |  8 +++-
 .../scheduler/OutputCommitCoordinator.scala   | 40 ++++++++++++-------
 .../OutputCommitCoordinatorSuite.scala        |  2 +-
 3 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index 5673fbf2c8fea..a1f0fd05f661a 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -947,7 +947,13 @@ class DAGScheduler(
     // serializable. If tasks are not serializable, a SparkListenerStageCompleted event
     // will be posted, which should always come after a corresponding SparkListenerStageSubmitted
     // event.
-    outputCommitCoordinator.stageStart(stage.id)
+    stage match {
+      case s: ShuffleMapStage =>
+        outputCommitCoordinator.stageStart(stage = s.id, maxPartitionId = s.numPartitions - 1)
+      case s: ResultStage =>
+        outputCommitCoordinator.stageStart(
+          stage = s.id, maxPartitionId = s.rdd.partitions.length - 1)
+    }
     val taskIdToLocations: Map[Int, Seq[TaskLocation]] = try {
       stage match {
         case s: ShuffleMapStage =>
diff --git a/core/src/main/scala/org/apache/spark/scheduler/OutputCommitCoordinator.scala b/core/src/main/scala/org/apache/spark/scheduler/OutputCommitCoordinator.scala
index add0dedc03f44..4d146678174f6 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/OutputCommitCoordinator.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/OutputCommitCoordinator.scala
@@ -47,6 +47,8 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
   private type PartitionId = Int
   private type TaskAttemptNumber = Int
 
+  private val NO_AUTHORIZED_COMMITTER: TaskAttemptNumber = -1
+
   /**
    * Map from active stages's id => partition id => task attempt with exclusive lock on committing
    * output for that partition.
@@ -56,9 +58,7 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
    *
    * Access to this map should be guarded by synchronizing on the OutputCommitCoordinator instance.
    */
-  private val authorizedCommittersByStage: CommittersByStageMap = mutable.Map()
-  private type CommittersByStageMap =
-    mutable.Map[StageId, mutable.Map[PartitionId, TaskAttemptNumber]]
+  private val authorizedCommittersByStage = mutable.Map[StageId, Array[TaskAttemptNumber]]()
 
   /**
    * Returns whether the OutputCommitCoordinator's internal data structures are all empty.
@@ -95,9 +95,21 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
     }
   }
 
-  // Called by DAGScheduler
-  private[scheduler] def stageStart(stage: StageId): Unit = synchronized {
-    authorizedCommittersByStage(stage) = mutable.HashMap[PartitionId, TaskAttemptNumber]()
+  /**
+   * Called by the DAGScheduler when a stage starts.
+   *
+   * @param stage the stage id.
+   * @param maxPartitionId the maximum partition id that could appear in this stage's tasks (i.e.
+   *                       the maximum possible value of `context.partitionId`).
+   */
+  private[scheduler] def stageStart(
+      stage: StageId,
+      maxPartitionId: Int): Unit = {
+    val arr = new Array[TaskAttemptNumber](maxPartitionId + 1)
+    java.util.Arrays.fill(arr, NO_AUTHORIZED_COMMITTER)
+    synchronized {
+      authorizedCommittersByStage(stage) = arr
+    }
   }
 
   // Called by DAGScheduler
@@ -122,10 +134,10 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
         logInfo(s"Task was denied committing, stage: $stage, partition: $partition, " +
           s"attempt: $attemptNumber")
       case otherReason =>
-        if (authorizedCommitters.get(partition).exists(_ == attemptNumber)) {
+        if (authorizedCommitters(partition) == attemptNumber) {
           logDebug(s"Authorized committer (attemptNumber=$attemptNumber, stage=$stage, " +
             s"partition=$partition) failed; clearing lock")
-          authorizedCommitters.remove(partition)
+          authorizedCommitters(partition) = NO_AUTHORIZED_COMMITTER
         }
     }
   }
@@ -145,16 +157,16 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
       attemptNumber: TaskAttemptNumber): Boolean = synchronized {
     authorizedCommittersByStage.get(stage) match {
       case Some(authorizedCommitters) =>
-        authorizedCommitters.get(partition) match {
-          case Some(existingCommitter) =>
-            logDebug(s"Denying attemptNumber=$attemptNumber to commit for stage=$stage, " +
-              s"partition=$partition; existingCommitter = $existingCommitter")
-            false
-          case None =>
+        authorizedCommitters(partition) match {
+          case NO_AUTHORIZED_COMMITTER =>
             logDebug(s"Authorizing attemptNumber=$attemptNumber to commit for stage=$stage, " +
               s"partition=$partition")
             authorizedCommitters(partition) = attemptNumber
             true
+          case existingCommitter =>
+            logDebug(s"Denying attemptNumber=$attemptNumber to commit for stage=$stage, " +
+              s"partition=$partition; existingCommitter = $existingCommitter")
+            false
         }
       case None =>
         logDebug(s"Stage $stage has completed, so not allowing attempt number $attemptNumber of" +
diff --git a/core/src/test/scala/org/apache/spark/scheduler/OutputCommitCoordinatorSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/OutputCommitCoordinatorSuite.scala
index 48456a9cd6e7b..7345508bfe995 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/OutputCommitCoordinatorSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/OutputCommitCoordinatorSuite.scala
@@ -171,7 +171,7 @@ class OutputCommitCoordinatorSuite extends SparkFunSuite with BeforeAndAfter {
     val partition: Int = 2
     val authorizedCommitter: Int = 3
     val nonAuthorizedCommitter: Int = 100
-    outputCommitCoordinator.stageStart(stage)
+    outputCommitCoordinator.stageStart(stage, maxPartitionId = 2)
 
     assert(outputCommitCoordinator.canCommit(stage, partition, authorizedCommitter))
     assert(!outputCommitCoordinator.canCommit(stage, partition, nonAuthorizedCommitter))

From 81498dd5c86ca51d2fb351c8ef52cbb28e6844f4 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Wed, 4 Nov 2015 21:30:21 -0800
Subject: [PATCH 063/518] [SPARK-11425] [SPARK-11486] Improve hybrid
 aggregation

After aggregation, the dataset could be smaller than inputs, so it's better to do hash based aggregation for all inputs, then using sort based aggregation to merge them.

Author: Davies Liu <davies@databricks.com>

Closes #9383 from davies/fix_switch.
---
 .../spark/unsafe/map/BytesToBytesMap.java     |  46 +++--
 .../unsafe/sort/UnsafeExternalSorter.java     |  39 ++--
 .../unsafe/sort/UnsafeInMemorySorter.java     |  15 +-
 .../UnsafeFixedWidthAggregationMap.java       |   9 +-
 .../sql/execution/UnsafeKVExternalSorter.java |  23 ++-
 .../TungstenAggregationIterator.scala         | 171 +++++-------------
 .../UnsafeFixedWidthAggregationMapSuite.scala |  64 +++----
 .../execution/AggregationQuerySuite.scala     |  12 +-
 8 files changed, 165 insertions(+), 214 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 07241c827c2ae..6656fd1d0bc59 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -20,6 +20,7 @@
 import javax.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.Iterator;
 import java.util.LinkedList;
 
@@ -638,7 +639,11 @@ public boolean putNewKey(Object keyBase, long keyOffset, int keyLength,
       assert (valueLength % 8 == 0);
       assert(longArray != null);
 
-      if (numElements == MAX_CAPACITY || !canGrowArray) {
+
+      if (numElements == MAX_CAPACITY
+        // The map could be reused from last spill (because of no enough memory to grow),
+        // then we don't try to grow again if hit the `growthThreshold`.
+        || !canGrowArray && numElements > growthThreshold) {
         return false;
       }
 
@@ -730,25 +735,18 @@ private void allocate(int capacity) {
   }
 
   /**
-   * Free the memory used by longArray.
+   * Free all allocated memory associated with this map, including the storage for keys and values
+   * as well as the hash map array itself.
+   *
+   * This method is idempotent and can be called multiple times.
    */
-  public void freeArray() {
+  public void free() {
     updatePeakMemoryUsed();
     if (longArray != null) {
       long used = longArray.memoryBlock().size();
       longArray = null;
       releaseMemory(used);
     }
-  }
-
-  /**
-   * Free all allocated memory associated with this map, including the storage for keys and values
-   * as well as the hash map array itself.
-   *
-   * This method is idempotent and can be called multiple times.
-   */
-  public void free() {
-    freeArray();
     Iterator<MemoryBlock> dataPagesIterator = dataPages.iterator();
     while (dataPagesIterator.hasNext()) {
       MemoryBlock dataPage = dataPagesIterator.next();
@@ -833,6 +831,28 @@ public int getNumDataPages() {
     return dataPages.size();
   }
 
+  /**
+   * Returns the underline long[] of longArray.
+   */
+  public long[] getArray() {
+    assert(longArray != null);
+    return (long[]) longArray.memoryBlock().getBaseObject();
+  }
+
+  /**
+   * Reset this map to initialized state.
+   */
+  public void reset() {
+    numElements = 0;
+    Arrays.fill(getArray(), 0);
+    while (dataPages.size() > 0) {
+      MemoryBlock dataPage = dataPages.removeLast();
+      freePage(dataPage);
+    }
+    currentPage = null;
+    pageCursor = 0;
+  }
+
   /**
    * Grows the size of the hash table and re-hash everything.
    */
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 509fb0a044c0c..cba043bc48cc8 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -79,9 +79,13 @@ public static UnsafeExternalSorter createWithExistingInMemorySorter(
       PrefixComparator prefixComparator,
       int initialSize,
       long pageSizeBytes,
-      UnsafeInMemorySorter inMemorySorter) {
-    return new UnsafeExternalSorter(taskMemoryManager, blockManager,
+      UnsafeInMemorySorter inMemorySorter) throws IOException {
+    UnsafeExternalSorter sorter = new UnsafeExternalSorter(taskMemoryManager, blockManager,
       taskContext, recordComparator, prefixComparator, initialSize, pageSizeBytes, inMemorySorter);
+    sorter.spill(Long.MAX_VALUE, sorter);
+    // The external sorter will be used to insert records, in-memory sorter is not needed.
+    sorter.inMemSorter = null;
+    return sorter;
   }
 
   public static UnsafeExternalSorter create(
@@ -124,7 +128,6 @@ private UnsafeExternalSorter(
       acquireMemory(inMemSorter.getMemoryUsage());
     } else {
       this.inMemSorter = existingInMemorySorter;
-      // will acquire after free the map
     }
     this.peakMemoryUsedBytes = getMemoryUsage();
 
@@ -157,12 +160,9 @@ public void closeCurrentPage() {
    */
   @Override
   public long spill(long size, MemoryConsumer trigger) throws IOException {
-    assert(inMemSorter != null);
     if (trigger != this) {
       if (readingIterator != null) {
         return readingIterator.spill();
-      } else {
-
       }
       return 0L; // this should throw exception
     }
@@ -388,25 +388,38 @@ public void insertKVRecord(Object keyBase, long keyOffset, int keyLen,
     inMemSorter.insertRecord(recordAddress, prefix);
   }
 
+  /**
+   * Merges another UnsafeExternalSorters into this one, the other one will be emptied.
+   *
+   * @throws IOException
+   */
+  public void merge(UnsafeExternalSorter other) throws IOException {
+    other.spill();
+    spillWriters.addAll(other.spillWriters);
+    // remove them from `spillWriters`, or the files will be deleted in `cleanupResources`.
+    other.spillWriters.clear();
+    other.cleanupResources();
+  }
+
   /**
    * Returns a sorted iterator. It is the caller's responsibility to call `cleanupResources()`
    * after consuming this iterator.
    */
   public UnsafeSorterIterator getSortedIterator() throws IOException {
-    assert(inMemSorter != null);
-    readingIterator = new SpillableIterator(inMemSorter.getSortedIterator());
-    int numIteratorsToMerge = spillWriters.size() + (readingIterator.hasNext() ? 1 : 0);
     if (spillWriters.isEmpty()) {
+      assert(inMemSorter != null);
+      readingIterator = new SpillableIterator(inMemSorter.getSortedIterator());
       return readingIterator;
     } else {
       final UnsafeSorterSpillMerger spillMerger =
-        new UnsafeSorterSpillMerger(recordComparator, prefixComparator, numIteratorsToMerge);
+        new UnsafeSorterSpillMerger(recordComparator, prefixComparator, spillWriters.size());
       for (UnsafeSorterSpillWriter spillWriter : spillWriters) {
         spillMerger.addSpillIfNotEmpty(spillWriter.getReader(blockManager));
       }
-      spillWriters.clear();
-      spillMerger.addSpillIfNotEmpty(readingIterator);
-
+      if (inMemSorter != null) {
+        readingIterator = new SpillableIterator(inMemSorter.getSortedIterator());
+        spillMerger.addSpillIfNotEmpty(readingIterator);
+      }
       return spillMerger.getSortedIterator();
     }
   }
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
index 1480f0681ed9c..d57213b9b8bfc 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -19,9 +19,9 @@
 
 import java.util.Comparator;
 
+import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.util.collection.Sorter;
-import org.apache.spark.memory.TaskMemoryManager;
 
 /**
  * Sorts records using an AlphaSort-style key-prefix sort. This sort stores pointers to records
@@ -77,13 +77,20 @@ public int compare(RecordPointerAndKeyPrefix r1, RecordPointerAndKeyPrefix r2) {
    */
   private int pos = 0;
 
+  public UnsafeInMemorySorter(
+    final TaskMemoryManager memoryManager,
+    final RecordComparator recordComparator,
+    final PrefixComparator prefixComparator,
+    int initialSize) {
+    this(memoryManager, recordComparator, prefixComparator, new long[initialSize * 2]);
+  }
+
   public UnsafeInMemorySorter(
       final TaskMemoryManager memoryManager,
       final RecordComparator recordComparator,
       final PrefixComparator prefixComparator,
-      int initialSize) {
-    assert (initialSize > 0);
-    this.array = new long[initialSize * 2];
+      long[] array) {
+    this.array = array;
     this.memoryManager = memoryManager;
     this.sorter = new Sorter<>(UnsafeSortDataFormat.INSTANCE);
     this.sortComparator = new SortComparator(recordComparator, prefixComparator, memoryManager);
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java
index d4b6d75b4d981..a2f99d566d471 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java
@@ -236,16 +236,13 @@ public void printPerfMetrics() {
 
   /**
    * Sorts the map's records in place, spill them to disk, and returns an [[UnsafeKVExternalSorter]]
-   * that can be used to insert more records to do external sorting.
    *
-   * The only memory that is allocated is the address/prefix array, 16 bytes per record.
-   *
-   * Note that this destroys the map, and as a result, the map cannot be used anymore after this.
+   * Note that the map will be reset for inserting new records, and the returned sorter can NOT be used
+   * to insert records.
    */
   public UnsafeKVExternalSorter destructAndCreateExternalSorter() throws IOException {
-    UnsafeKVExternalSorter sorter = new UnsafeKVExternalSorter(
+    return new UnsafeKVExternalSorter(
       groupingKeySchema, aggregationBufferSchema,
       SparkEnv.get().blockManager(), map.getPageSizeBytes(), map);
-    return sorter;
   }
 }
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
index 845f2ae6859b7..e2898ef2e2158 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
@@ -83,11 +83,10 @@ public UnsafeKVExternalSorter(
         /* initialSize */ 4096,
         pageSizeBytes);
     } else {
-      // The memory needed for UnsafeInMemorySorter should be less than longArray in map.
-      map.freeArray();
-      // The memory used by UnsafeInMemorySorter will be counted later (end of this block)
+      // During spilling, the array in map will not be used, so we can borrow that and use it
+      // as the underline array for in-memory sorter (it's always large enough).
       final UnsafeInMemorySorter inMemSorter = new UnsafeInMemorySorter(
-        taskMemoryManager, recordComparator, prefixComparator, Math.max(1, map.numElements()));
+        taskMemoryManager, recordComparator, prefixComparator, map.getArray());
 
       // We cannot use the destructive iterator here because we are reusing the existing memory
       // pages in BytesToBytesMap to hold records during sorting.
@@ -123,10 +122,9 @@ public UnsafeKVExternalSorter(
         pageSizeBytes,
         inMemSorter);
 
-      sorter.spill();
-      map.free();
-      // counting the memory used UnsafeInMemorySorter
-      taskMemoryManager.acquireExecutionMemory(inMemSorter.getMemoryUsage(), sorter);
+      // reset the map, so we can re-use it to insert new records. the inMemSorter will not used
+      // anymore, so the underline array could be used by map again.
+      map.reset();
     }
   }
 
@@ -142,6 +140,15 @@ public void insertKV(UnsafeRow key, UnsafeRow value) throws IOException {
       value.getBaseObject(), value.getBaseOffset(), value.getSizeInBytes(), prefix);
   }
 
+  /**
+   * Merges another UnsafeKVExternalSorter into `this`, the other one will be emptied.
+   *
+   * @throws IOException
+   */
+  public void merge(UnsafeKVExternalSorter other) throws IOException {
+    sorter.merge(other.sorter);
+  }
+
   /**
    * Returns a sorted iterator. It is the caller's responsibility to call `cleanupResources()`
    * after consuming this iterator.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala
index 713a4db0cd59b..ce8d592c368ee 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala
@@ -34,14 +34,18 @@ import org.apache.spark.sql.types.StructType
  *
  * This iterator first uses hash-based aggregation to process input rows. It uses
  * a hash map to store groups and their corresponding aggregation buffers. If we
- * this map cannot allocate memory from memory manager,
- * it switches to sort-based aggregation. The process of the switch has the following step:
+ * this map cannot allocate memory from memory manager, it spill the map into disk
+ * and create a new one. After processed all the input, then merge all the spills
+ * together using external sorter, and do sort-based aggregation.
+ *
+ * The process has the following step:
+ *  - Step 0: Do hash-based aggregation.
  *  - Step 1: Sort all entries of the hash map based on values of grouping expressions and
  *            spill them to disk.
- *  - Step 2: Create a external sorter based on the spilled sorted map entries.
- *  - Step 3: Redirect all input rows to the external sorter.
- *  - Step 4: Get a sorted [[KVIterator]] from the external sorter.
- *  - Step 5: Initialize sort-based aggregation.
+ *  - Step 2: Create a external sorter based on the spilled sorted map entries and reset the map.
+ *  - Step 3: Get a sorted [[KVIterator]] from the external sorter.
+ *  - Step 4: Repeat step 0 until no more input.
+ *  - Step 5: Initialize sort-based aggregation on the sorted iterator.
  * Then, this iterator works in the way of sort-based aggregation.
  *
  * The code of this class is organized as follows:
@@ -488,9 +492,10 @@ class TungstenAggregationIterator(
 
   // The function used to read and process input rows. When processing input rows,
   // it first uses hash-based aggregation by putting groups and their buffers in
-  // hashMap. If we could not allocate more memory for the map, we switch to
-  // sort-based aggregation (by calling switchToSortBasedAggregation).
-  private def processInputs(): Unit = {
+  // hashMap. If there is not enough memory, it will multiple hash-maps, spilling
+  // after each becomes full then using sort to merge these spills, finally do sort
+  // based aggregation.
+  private def processInputs(fallbackStartsAt: Int): Unit = {
     if (groupingExpressions.isEmpty) {
       // If there is no grouping expressions, we can just reuse the same buffer over and over again.
       // Note that it would be better to eliminate the hash map entirely in the future.
@@ -502,44 +507,40 @@ class TungstenAggregationIterator(
         processRow(buffer, newInput)
       }
     } else {
-      while (!sortBased && inputIter.hasNext) {
+      var i = 0
+      while (inputIter.hasNext) {
         val newInput = inputIter.next()
         numInputRows += 1
         val groupingKey = groupProjection.apply(newInput)
-        val buffer: UnsafeRow = hashMap.getAggregationBufferFromUnsafeRow(groupingKey)
+        var buffer: UnsafeRow = null
+        if (i < fallbackStartsAt) {
+          buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey)
+        }
         if (buffer == null) {
-          // buffer == null means that we could not allocate more memory.
-          // Now, we need to spill the map and switch to sort-based aggregation.
-          switchToSortBasedAggregation(groupingKey, newInput)
-        } else {
-          processRow(buffer, newInput)
+          val sorter = hashMap.destructAndCreateExternalSorter()
+          if (externalSorter == null) {
+            externalSorter = sorter
+          } else {
+            externalSorter.merge(sorter)
+          }
+          i = 0
+          buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey)
+          if (buffer == null) {
+            // failed to allocate the first page
+            throw new OutOfMemoryError("No enough memory for aggregation")
+          }
         }
+        processRow(buffer, newInput)
+        i += 1
       }
-    }
-  }
 
-  // This function is only used for testing. It basically the same as processInputs except
-  // that it switch to sort-based aggregation after `fallbackStartsAt` input rows have
-  // been processed.
-  private def processInputsWithControlledFallback(fallbackStartsAt: Int): Unit = {
-    var i = 0
-    while (!sortBased && inputIter.hasNext) {
-      val newInput = inputIter.next()
-      numInputRows += 1
-      val groupingKey = groupProjection.apply(newInput)
-      val buffer: UnsafeRow = if (i < fallbackStartsAt) {
-        hashMap.getAggregationBufferFromUnsafeRow(groupingKey)
-      } else {
-        null
-      }
-      if (buffer == null) {
-        // buffer == null means that we could not allocate more memory.
-        // Now, we need to spill the map and switch to sort-based aggregation.
-        switchToSortBasedAggregation(groupingKey, newInput)
-      } else {
-        processRow(buffer, newInput)
+      if (externalSorter != null) {
+        val sorter = hashMap.destructAndCreateExternalSorter()
+        externalSorter.merge(sorter)
+        hashMap.free()
+
+        switchToSortBasedAggregation()
       }
-      i += 1
     }
   }
 
@@ -561,88 +562,8 @@ class TungstenAggregationIterator(
   /**
    * Switch to sort-based aggregation when the hash-based approach is unable to acquire memory.
    */
-  private def switchToSortBasedAggregation(firstKey: UnsafeRow, firstInput: InternalRow): Unit = {
+  private def switchToSortBasedAggregation(): Unit = {
     logInfo("falling back to sort based aggregation.")
-    // Step 1: Get the ExternalSorter containing sorted entries of the map.
-    externalSorter = hashMap.destructAndCreateExternalSorter()
-
-    // Step 2: If we have aggregate function with mode Partial or Complete,
-    // we need to process input rows to get aggregation buffer.
-    // So, later in the sort-based aggregation iterator, we can do merge.
-    // If aggregate functions are with mode Final and PartialMerge,
-    // we just need to project the aggregation buffer from an input row.
-    val needsProcess = aggregationMode match {
-      case (Some(Partial), None) => true
-      case (None, Some(Complete)) => true
-      case (Some(Final), Some(Complete)) => true
-      case _ => false
-    }
-
-    // Note: Since we spill the sorter's contents immediately after creating it, we must insert
-    // something into the sorter here to ensure that we acquire at least a page of memory.
-    // This is done through `externalSorter.insertKV`, which will trigger the page allocation.
-    // Otherwise, children operators may steal the window of opportunity and starve our sorter.
-
-    if (needsProcess) {
-      // First, we create a buffer.
-      val buffer = createNewAggregationBuffer()
-
-      // Process firstKey and firstInput.
-      // Initialize buffer.
-      buffer.copyFrom(initialAggregationBuffer)
-      processRow(buffer, firstInput)
-      externalSorter.insertKV(firstKey, buffer)
-
-      // Process the rest of input rows.
-      while (inputIter.hasNext) {
-        val newInput = inputIter.next()
-        numInputRows += 1
-        val groupingKey = groupProjection.apply(newInput)
-        buffer.copyFrom(initialAggregationBuffer)
-        processRow(buffer, newInput)
-        externalSorter.insertKV(groupingKey, buffer)
-      }
-    } else {
-      // When needsProcess is false, the format of input rows is groupingKey + aggregation buffer.
-      // We need to project the aggregation buffer part from an input row.
-      val buffer = createNewAggregationBuffer()
-      // In principle, we could use `allAggregateFunctions.flatMap(_.inputAggBufferAttributes)` to
-      // extract the aggregation buffer. In practice, however, we extract it positionally by relying
-      // on it being present at the end of the row. The reason for this relates to how the different
-      // aggregates handle input binding.
-      //
-      // ImperativeAggregate uses field numbers and field number offsets to manipulate its buffers,
-      // so its correctness does not rely on attribute bindings. When we fall back to sort-based
-      // aggregation, these field number offsets (mutableAggBufferOffset and inputAggBufferOffset)
-      // need to be updated and any internal state in the aggregate functions themselves must be
-      // reset, so we call withNewMutableAggBufferOffset and withNewInputAggBufferOffset to reset
-      // this state and update the offsets.
-      //
-      // The updated ImperativeAggregate will have different attribute ids for its
-      // aggBufferAttributes and inputAggBufferAttributes. This isn't a problem for the actual
-      // ImperativeAggregate evaluation, but it means that
-      // `allAggregateFunctions.flatMap(_.inputAggBufferAttributes)` will no longer match the
-      // attributes in `originalInputAttributes`, which is why we can't use those attributes here.
-      //
-      // For more details, see the discussion on PR #9038.
-      val bufferExtractor = newMutableProjection(
-        originalInputAttributes.drop(initialInputBufferOffset),
-        originalInputAttributes)()
-      bufferExtractor.target(buffer)
-
-      // Insert firstKey and its buffer.
-      bufferExtractor(firstInput)
-      externalSorter.insertKV(firstKey, buffer)
-
-      // Insert the rest of input rows.
-      while (inputIter.hasNext) {
-        val newInput = inputIter.next()
-        numInputRows += 1
-        val groupingKey = groupProjection.apply(newInput)
-        bufferExtractor(newInput)
-        externalSorter.insertKV(groupingKey, buffer)
-      }
-    }
 
     // Set aggregationMode, processRow, and generateOutput for sort-based aggregation.
     val newAggregationMode = aggregationMode match {
@@ -762,15 +683,7 @@ class TungstenAggregationIterator(
   /**
    * Start processing input rows.
    */
-  testFallbackStartsAt match {
-    case None =>
-      processInputs()
-    case Some(fallbackStartsAt) =>
-      // This is the testing path. processInputsWithControlledFallback is same as processInputs
-      // except that it switches to sort-based aggregation after `fallbackStartsAt` input rows
-      // have been processed.
-      processInputsWithControlledFallback(fallbackStartsAt)
-  }
+  processInputs(testFallbackStartsAt.getOrElse(Int.MaxValue))
 
   // If we did not switch to sort-based aggregation in processInputs,
   // we pre-load the first key-value pair from the map (to make hasNext idempotent).
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
index a38623623a441..7ceaee38d131b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMapSuite.scala
@@ -170,9 +170,6 @@ class UnsafeFixedWidthAggregationMapSuite
   }
 
   testWithMemoryLeakDetection("test external sorting") {
-    // Memory consumption in the beginning of the task.
-    val initialMemoryConsumption = taskMemoryManager.getMemoryConsumptionForThisTask()
-
     val map = new UnsafeFixedWidthAggregationMap(
       emptyAggregationBuffer,
       aggBufferSchema,
@@ -189,35 +186,33 @@ class UnsafeFixedWidthAggregationMapSuite
       buf.setInt(0, keyString.length)
       assert(buf != null)
     }
-
-    // Convert the map into a sorter
     val sorter = map.destructAndCreateExternalSorter()
 
     // Add more keys to the sorter and make sure the results come out sorted.
     val additionalKeys = randomStrings(1024)
-    val keyConverter = UnsafeProjection.create(groupKeySchema)
-    val valueConverter = UnsafeProjection.create(aggBufferSchema)
-
     additionalKeys.zipWithIndex.foreach { case (str, i) =>
-      val k = InternalRow(UTF8String.fromString(str))
-      val v = InternalRow(str.length)
-      sorter.insertKV(keyConverter.apply(k), valueConverter.apply(v))
+      val buf = map.getAggregationBuffer(InternalRow(UTF8String.fromString(str)))
+      buf.setInt(0, str.length)
 
       if ((i % 100) == 0) {
-        memoryManager.markExecutionAsOutOfMemoryOnce()
-        sorter.closeCurrentPage()
+        val sorter2 = map.destructAndCreateExternalSorter()
+        sorter.merge(sorter2)
       }
     }
+    val sorter2 = map.destructAndCreateExternalSorter()
+    sorter.merge(sorter2)
 
     val out = new scala.collection.mutable.ArrayBuffer[String]
     val iter = sorter.sortedIterator()
     while (iter.next()) {
-      assert(iter.getKey.getString(0).length === iter.getValue.getInt(0))
-      out += iter.getKey.getString(0)
+      // At here, we also test if copy is correct.
+      val key = iter.getKey.copy()
+      val value = iter.getValue.copy()
+      assert(key.getString(0).length === value.getInt(0))
+      out += key.getString(0)
     }
 
     assert(out === (keys ++ additionalKeys).sorted)
-
     map.free()
   }
 
@@ -232,25 +227,21 @@ class UnsafeFixedWidthAggregationMapSuite
       PAGE_SIZE_BYTES,
       false // disable perf metrics
     )
-
-    // Convert the map into a sorter
     val sorter = map.destructAndCreateExternalSorter()
 
     // Add more keys to the sorter and make sure the results come out sorted.
     val additionalKeys = randomStrings(1024)
-    val keyConverter = UnsafeProjection.create(groupKeySchema)
-    val valueConverter = UnsafeProjection.create(aggBufferSchema)
-
     additionalKeys.zipWithIndex.foreach { case (str, i) =>
-      val k = InternalRow(UTF8String.fromString(str))
-      val v = InternalRow(str.length)
-      sorter.insertKV(keyConverter.apply(k), valueConverter.apply(v))
+      val buf = map.getAggregationBuffer(InternalRow(UTF8String.fromString(str)))
+      buf.setInt(0, str.length)
 
       if ((i % 100) == 0) {
-        memoryManager.markExecutionAsOutOfMemoryOnce()
-        sorter.closeCurrentPage()
+        val sorter2 = map.destructAndCreateExternalSorter()
+        sorter.merge(sorter2)
       }
     }
+    val sorter2 = map.destructAndCreateExternalSorter()
+    sorter.merge(sorter2)
 
     val out = new scala.collection.mutable.ArrayBuffer[String]
     val iter = sorter.sortedIterator()
@@ -262,16 +253,12 @@ class UnsafeFixedWidthAggregationMapSuite
       out += key.getString(0)
     }
 
-    assert(out === (additionalKeys).sorted)
-
+    assert(out === additionalKeys.sorted)
     map.free()
   }
 
   testWithMemoryLeakDetection("test external sorting with empty records") {
 
-    // Memory consumption in the beginning of the task.
-    val initialMemoryConsumption = taskMemoryManager.getMemoryConsumptionForThisTask()
-
     val map = new UnsafeFixedWidthAggregationMap(
       emptyAggregationBuffer,
       StructType(Nil),
@@ -281,7 +268,6 @@ class UnsafeFixedWidthAggregationMapSuite
       PAGE_SIZE_BYTES,
       false // disable perf metrics
     )
-
     (1 to 10).foreach { i =>
       val buf = map.getAggregationBuffer(UnsafeRow.createFromByteArray(0, 0))
       assert(buf != null)
@@ -292,13 +278,15 @@ class UnsafeFixedWidthAggregationMapSuite
 
     // Add more keys to the sorter and make sure the results come out sorted.
     (1 to 4096).foreach { i =>
-      sorter.insertKV(UnsafeRow.createFromByteArray(0, 0), UnsafeRow.createFromByteArray(0, 0))
+      map.getAggregationBufferFromUnsafeRow(UnsafeRow.createFromByteArray(0, 0))
 
       if ((i % 100) == 0) {
-        memoryManager.markExecutionAsOutOfMemoryOnce()
-        sorter.closeCurrentPage()
+        val sorter2 = map.destructAndCreateExternalSorter()
+        sorter.merge(sorter2)
       }
     }
+    val sorter2 = map.destructAndCreateExternalSorter()
+    sorter.merge(sorter2)
 
     var count = 0
     val iter = sorter.sortedIterator()
@@ -309,9 +297,8 @@ class UnsafeFixedWidthAggregationMapSuite
       count += 1
     }
 
-    // 1 record was from the map and 4096 records were explicitly inserted.
-    assert(count === 4097)
-
+    // 1 record per map, spilled 42 times.
+    assert(count === 42)
     map.free()
   }
 
@@ -345,6 +332,7 @@ class UnsafeFixedWidthAggregationMapSuite
     var sorter: UnsafeKVExternalSorter = null
     try {
       sorter = map.destructAndCreateExternalSorter()
+      map.free()
     } finally {
       if (sorter != null) {
         sorter.cleanupResources()
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 74061db0f28af..ea80060e370e0 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -22,13 +22,12 @@ import scala.collection.JavaConverters._
 import org.apache.spark.SparkException
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.expressions.UnsafeRow
-import org.apache.spark.sql.execution.aggregate
 import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
 import org.apache.spark.sql.functions._
-import org.apache.spark.sql.test.SQLTestUtils
-import org.apache.spark.sql.types._
 import org.apache.spark.sql.hive.aggregate.{MyDoubleAvg, MyDoubleSum}
 import org.apache.spark.sql.hive.test.TestHiveSingleton
+import org.apache.spark.sql.test.SQLTestUtils
+import org.apache.spark.sql.types._
 
 class ScalaAggregateFunction(schema: StructType) extends UserDefinedAggregateFunction {
 
@@ -702,6 +701,13 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
     }
   }
 
+  test("no aggregation function (SPARK-11486)") {
+    val df = sqlContext.range(20).selectExpr("id", "repeat(id, 1) as s")
+      .groupBy("s").count()
+      .groupBy().count()
+    checkAnswer(df, Row(20) :: Nil)
+  }
+
   test("udaf with all data types") {
     val struct =
       StructType(

From 6f81eae24f83df51a99d4bb2629dd7daadc01519 Mon Sep 17 00:00:00 2001
From: Sean Owen <sowen@cloudera.com>
Date: Thu, 5 Nov 2015 09:08:53 +0000
Subject: [PATCH 064/518] [SPARK-11440][CORE][STREAMING][BUILD] Declare rest of
 @Experimental items non-experimental if they've existed since 1.2.0

Remove `Experimental` annotations in core, streaming for items that existed in 1.2.0 or before. The changes are:

* SparkContext
  * binary{Files,Records} : 1.2.0
  * submitJob : 1.0.0
* JavaSparkContext
  * binary{Files,Records} : 1.2.0
* DoubleRDDFunctions, JavaDoubleRDD
  * {mean,sum}Approx : 1.0.0
* PairRDDFunctions, JavaPairRDD
  * sampleByKeyExact : 1.2.0
  * countByKeyApprox : 1.0.0
* PairRDDFunctions
  * countApproxDistinctByKey : 1.1.0
* RDD
  * countApprox, countByValueApprox, countApproxDistinct : 1.0.0
* JavaRDDLike
  * countApprox : 1.0.0
* PythonHadoopUtil.Converter : 1.1.0
* PortableDataStream : 1.2.0 (related to binaryFiles)
* BoundedDouble : 1.0.0
* PartialResult : 1.0.0
* StreamingContext, JavaStreamingContext
  * binaryRecordsStream : 1.2.0
* HiveContext
  * analyze : 1.2.0

Author: Sean Owen <sowen@cloudera.com>

Closes #9396 from srowen/SPARK-11440.
---
 .../src/main/scala/org/apache/spark/SparkContext.scala | 10 +---------
 .../org/apache/spark/api/java/JavaDoubleRDD.scala      |  7 -------
 .../scala/org/apache/spark/api/java/JavaPairRDD.scala  |  9 ---------
 .../scala/org/apache/spark/api/java/JavaRDDLike.scala  |  5 -----
 .../org/apache/spark/api/java/JavaSparkContext.scala   |  7 -------
 .../org/apache/spark/api/python/PythonHadoopUtil.scala |  3 ---
 .../org/apache/spark/input/PortableDataStream.scala    |  2 --
 .../scala/org/apache/spark/partial/BoundedDouble.scala |  4 ----
 .../scala/org/apache/spark/partial/PartialResult.scala |  3 ---
 .../org/apache/spark/rdd/DoubleRDDFunctions.scala      |  4 ----
 .../scala/org/apache/spark/rdd/PairRDDFunctions.scala  |  7 -------
 core/src/main/scala/org/apache/spark/rdd/RDD.scala     |  8 +-------
 .../scala/org/apache/spark/sql/hive/HiveContext.scala  |  2 --
 .../org/apache/spark/streaming/StreamingContext.scala  |  3 ---
 .../streaming/api/java/JavaStreamingContext.scala      |  3 ---
 15 files changed, 2 insertions(+), 75 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index a6857b4c7d882..7421821e2601b 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -45,7 +45,7 @@ import org.apache.hadoop.mapreduce.lib.input.{FileInputFormat => NewFileInputFor
 
 import org.apache.mesos.MesosNativeLibrary
 
-import org.apache.spark.annotation.{DeveloperApi, Experimental}
+import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.deploy.{LocalSparkCluster, SparkHadoopUtil}
 import org.apache.spark.executor.{ExecutorEndpoint, TriggerThreadDump}
@@ -870,8 +870,6 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
   }
 
   /**
-   * :: Experimental ::
-   *
    * Get an RDD for a Hadoop-readable dataset as PortableDataStream for each file
    * (useful for binary data)
    *
@@ -902,7 +900,6 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
    *             list of inputs.
    * @param minPartitions A suggestion value of the minimal splitting number for input data.
    */
-  @Experimental
   def binaryFiles(
       path: String,
       minPartitions: Int = defaultMinPartitions): RDD[(String, PortableDataStream)] = withScope {
@@ -922,8 +919,6 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
   }
 
   /**
-   * :: Experimental ::
-   *
    * Load data from a flat binary file, assuming the length of each record is constant.
    *
    * '''Note:''' We ensure that the byte array for each record in the resulting RDD
@@ -936,7 +931,6 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
    *
    * @return An RDD of data with values, represented as byte arrays
    */
-  @Experimental
   def binaryRecords(
       path: String,
       recordLength: Int,
@@ -1963,10 +1957,8 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
   }
 
   /**
-   * :: Experimental ::
    * Submit a job for execution and return a FutureJob holding the result.
    */
-  @Experimental
   def submitJob[T, U, R](
       rdd: RDD[T],
       processPartition: Iterator[T] => U,
diff --git a/core/src/main/scala/org/apache/spark/api/java/JavaDoubleRDD.scala b/core/src/main/scala/org/apache/spark/api/java/JavaDoubleRDD.scala
index a650df605b92e..c32aefac465bc 100644
--- a/core/src/main/scala/org/apache/spark/api/java/JavaDoubleRDD.scala
+++ b/core/src/main/scala/org/apache/spark/api/java/JavaDoubleRDD.scala
@@ -24,7 +24,6 @@ import scala.reflect.ClassTag
 
 import org.apache.spark.Partitioner
 import org.apache.spark.SparkContext.doubleRDDToDoubleRDDFunctions
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.function.{Function => JFunction}
 import org.apache.spark.partial.{BoundedDouble, PartialResult}
 import org.apache.spark.rdd.RDD
@@ -209,25 +208,19 @@ class JavaDoubleRDD(val srdd: RDD[scala.Double])
     srdd.meanApprox(timeout, confidence)
 
   /**
-   * :: Experimental ::
    * Approximate operation to return the mean within a timeout.
    */
-  @Experimental
   def meanApprox(timeout: Long): PartialResult[BoundedDouble] = srdd.meanApprox(timeout)
 
   /**
-   * :: Experimental ::
    * Approximate operation to return the sum within a timeout.
    */
-  @Experimental
   def sumApprox(timeout: Long, confidence: JDouble): PartialResult[BoundedDouble] =
     srdd.sumApprox(timeout, confidence)
 
   /**
-   * :: Experimental ::
    * Approximate operation to return the sum within a timeout.
    */
-  @Experimental
   def sumApprox(timeout: Long): PartialResult[BoundedDouble] = srdd.sumApprox(timeout)
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/api/java/JavaPairRDD.scala b/core/src/main/scala/org/apache/spark/api/java/JavaPairRDD.scala
index 8344f6368ac48..0b0c6e5bb8cc1 100644
--- a/core/src/main/scala/org/apache/spark/api/java/JavaPairRDD.scala
+++ b/core/src/main/scala/org/apache/spark/api/java/JavaPairRDD.scala
@@ -32,7 +32,6 @@ import org.apache.hadoop.mapreduce.{OutputFormat => NewOutputFormat}
 
 import org.apache.spark.{HashPartitioner, Partitioner}
 import org.apache.spark.Partitioner._
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.JavaSparkContext.fakeClassTag
 import org.apache.spark.api.java.JavaUtils.mapAsSerializableJavaMap
 import org.apache.spark.api.java.function.{Function => JFunction, Function2 => JFunction2, PairFunction}
@@ -159,7 +158,6 @@ class JavaPairRDD[K, V](val rdd: RDD[(K, V)])
     sampleByKey(withReplacement, fractions, Utils.random.nextLong)
 
   /**
-   * ::Experimental::
    * Return a subset of this RDD sampled by key (via stratified sampling) containing exactly
    * math.ceil(numItems * samplingRate) for each stratum (group of pairs with the same key).
    *
@@ -169,14 +167,12 @@ class JavaPairRDD[K, V](val rdd: RDD[(K, V)])
    * additional pass over the RDD to guarantee sample size; when sampling with replacement, we need
    * two additional passes.
    */
-  @Experimental
   def sampleByKeyExact(withReplacement: Boolean,
       fractions: JMap[K, Double],
       seed: Long): JavaPairRDD[K, V] =
     new JavaPairRDD[K, V](rdd.sampleByKeyExact(withReplacement, fractions.asScala, seed))
 
   /**
-   * ::Experimental::
    * Return a subset of this RDD sampled by key (via stratified sampling) containing exactly
    * math.ceil(numItems * samplingRate) for each stratum (group of pairs with the same key).
    *
@@ -188,7 +184,6 @@ class JavaPairRDD[K, V](val rdd: RDD[(K, V)])
    *
    * Use Utils.random.nextLong as the default seed for the random number generator.
    */
-  @Experimental
   def sampleByKeyExact(withReplacement: Boolean, fractions: JMap[K, Double]): JavaPairRDD[K, V] =
     sampleByKeyExact(withReplacement, fractions, Utils.random.nextLong)
 
@@ -300,20 +295,16 @@ class JavaPairRDD[K, V](val rdd: RDD[(K, V)])
   def countByKey(): java.util.Map[K, Long] = mapAsSerializableJavaMap(rdd.countByKey())
 
   /**
-   * :: Experimental ::
    * Approximate version of countByKey that can return a partial result if it does
    * not finish within a timeout.
    */
-  @Experimental
   def countByKeyApprox(timeout: Long): PartialResult[java.util.Map[K, BoundedDouble]] =
     rdd.countByKeyApprox(timeout).map(mapAsSerializableJavaMap)
 
   /**
-   * :: Experimental ::
    * Approximate version of countByKey that can return a partial result if it does
    * not finish within a timeout.
    */
-  @Experimental
   def countByKeyApprox(timeout: Long, confidence: Double = 0.95)
   : PartialResult[java.util.Map[K, BoundedDouble]] =
     rdd.countByKeyApprox(timeout, confidence).map(mapAsSerializableJavaMap)
diff --git a/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala b/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
index fc817cdd6a3f8..871be0b1f39ea 100644
--- a/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
+++ b/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
@@ -28,7 +28,6 @@ import com.google.common.base.Optional
 import org.apache.hadoop.io.compress.CompressionCodec
 
 import org.apache.spark._
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.JavaPairRDD._
 import org.apache.spark.api.java.JavaSparkContext.fakeClassTag
 import org.apache.spark.api.java.JavaUtils.mapAsSerializableJavaMap
@@ -436,20 +435,16 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
   def count(): Long = rdd.count()
 
   /**
-   * :: Experimental ::
    * Approximate version of count() that returns a potentially incomplete result
    * within a timeout, even if not all tasks have finished.
    */
-  @Experimental
   def countApprox(timeout: Long, confidence: Double): PartialResult[BoundedDouble] =
     rdd.countApprox(timeout, confidence)
 
   /**
-   * :: Experimental ::
    * Approximate version of count() that returns a potentially incomplete result
    * within a timeout, even if not all tasks have finished.
    */
-  @Experimental
   def countApprox(timeout: Long): PartialResult[BoundedDouble] =
     rdd.countApprox(timeout)
 
diff --git a/core/src/main/scala/org/apache/spark/api/java/JavaSparkContext.scala b/core/src/main/scala/org/apache/spark/api/java/JavaSparkContext.scala
index 609496ccdfef1..4f54cd69e2175 100644
--- a/core/src/main/scala/org/apache/spark/api/java/JavaSparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/api/java/JavaSparkContext.scala
@@ -33,7 +33,6 @@ import org.apache.hadoop.mapreduce.{InputFormat => NewInputFormat}
 
 import org.apache.spark._
 import org.apache.spark.AccumulatorParam._
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.JavaSparkContext.fakeClassTag
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.{EmptyRDD, HadoopRDD, NewHadoopRDD, RDD}
@@ -266,8 +265,6 @@ class JavaSparkContext(val sc: SparkContext)
     new JavaPairRDD(sc.binaryFiles(path, minPartitions))
 
   /**
-   * :: Experimental ::
-   *
    * Read a directory of binary files from HDFS, a local file system (available on all nodes),
    * or any Hadoop-supported file system URI as a byte array. Each file is read as a single
    * record and returned in a key-value pair, where the key is the path of each file,
@@ -294,19 +291,15 @@ class JavaSparkContext(val sc: SparkContext)
    *
    * @note Small files are preferred; very large files but may cause bad performance.
    */
-  @Experimental
   def binaryFiles(path: String): JavaPairRDD[String, PortableDataStream] =
     new JavaPairRDD(sc.binaryFiles(path, defaultMinPartitions))
 
   /**
-   * :: Experimental ::
-   *
    * Load data from a flat binary file, assuming the length of each record is constant.
    *
    * @param path Directory to the input data files
    * @return An RDD of data with values, represented as byte arrays
    */
-  @Experimental
   def binaryRecords(path: String, recordLength: Int): JavaRDD[Array[Byte]] = {
     new JavaRDD(sc.binaryRecords(path, recordLength))
   }
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonHadoopUtil.scala b/core/src/main/scala/org/apache/spark/api/python/PythonHadoopUtil.scala
index a7dfa1d257cf2..d2beef2a0dd43 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonHadoopUtil.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonHadoopUtil.scala
@@ -24,17 +24,14 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.io._
 
 import org.apache.spark.{Logging, SparkException}
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.rdd.RDD
 import org.apache.spark.util.{SerializableConfiguration, Utils}
 
 /**
- * :: Experimental ::
  * A trait for use with reading custom classes in PySpark. Implement this trait and add custom
  * transformation code by overriding the convert method.
  */
-@Experimental
 trait Converter[T, + U] extends Serializable {
   def convert(obj: T): U
 }
diff --git a/core/src/main/scala/org/apache/spark/input/PortableDataStream.scala b/core/src/main/scala/org/apache/spark/input/PortableDataStream.scala
index e2ffc3b64e5db..33e4ee0215817 100644
--- a/core/src/main/scala/org/apache/spark/input/PortableDataStream.scala
+++ b/core/src/main/scala/org/apache/spark/input/PortableDataStream.scala
@@ -27,7 +27,6 @@ import org.apache.hadoop.fs.Path
 import org.apache.hadoop.mapreduce.{InputSplit, JobContext, RecordReader, TaskAttemptContext}
 import org.apache.hadoop.mapreduce.lib.input.{CombineFileInputFormat, CombineFileRecordReader, CombineFileSplit}
 
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.deploy.SparkHadoopUtil
 
 /**
@@ -129,7 +128,6 @@ private[spark] class StreamInputFormat extends StreamFileInputFormat[PortableDat
  * @note TaskAttemptContext is not serializable resulting in the confBytes construct
  * @note CombineFileSplit is not serializable resulting in the splitBytes construct
  */
-@Experimental
 class PortableDataStream(
     isplit: CombineFileSplit,
     context: TaskAttemptContext,
diff --git a/core/src/main/scala/org/apache/spark/partial/BoundedDouble.scala b/core/src/main/scala/org/apache/spark/partial/BoundedDouble.scala
index aed0353344427..48b9434153172 100644
--- a/core/src/main/scala/org/apache/spark/partial/BoundedDouble.scala
+++ b/core/src/main/scala/org/apache/spark/partial/BoundedDouble.scala
@@ -17,13 +17,9 @@
 
 package org.apache.spark.partial
 
-import org.apache.spark.annotation.Experimental
-
 /**
- * :: Experimental ::
  * A Double value with error bars and associated confidence.
  */
-@Experimental
 class BoundedDouble(val mean: Double, val confidence: Double, val low: Double, val high: Double) {
   override def toString(): String = "[%.3f, %.3f]".format(low, high)
 }
diff --git a/core/src/main/scala/org/apache/spark/partial/PartialResult.scala b/core/src/main/scala/org/apache/spark/partial/PartialResult.scala
index 53c4b32c95ab3..25cb7490aa9c9 100644
--- a/core/src/main/scala/org/apache/spark/partial/PartialResult.scala
+++ b/core/src/main/scala/org/apache/spark/partial/PartialResult.scala
@@ -17,9 +17,6 @@
 
 package org.apache.spark.partial
 
-import org.apache.spark.annotation.Experimental
-
-@Experimental
 class PartialResult[R](initialVal: R, isFinal: Boolean) {
   private var finalValue: Option[R] = if (isFinal) Some(initialVal) else None
   private var failure: Option[Exception] = None
diff --git a/core/src/main/scala/org/apache/spark/rdd/DoubleRDDFunctions.scala b/core/src/main/scala/org/apache/spark/rdd/DoubleRDDFunctions.scala
index 926bce6f15a2a..7fbaadcea3a3b 100644
--- a/core/src/main/scala/org/apache/spark/rdd/DoubleRDDFunctions.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/DoubleRDDFunctions.scala
@@ -74,10 +74,8 @@ class DoubleRDDFunctions(self: RDD[Double]) extends Logging with Serializable {
   }
 
   /**
-   * :: Experimental ::
    * Approximate operation to return the mean within a timeout.
    */
-  @Experimental
   def meanApprox(
       timeout: Long,
       confidence: Double = 0.95): PartialResult[BoundedDouble] = self.withScope {
@@ -87,10 +85,8 @@ class DoubleRDDFunctions(self: RDD[Double]) extends Logging with Serializable {
   }
 
   /**
-   * :: Experimental ::
    * Approximate operation to return the sum within a timeout.
    */
-  @Experimental
   def sumApprox(
       timeout: Long,
       confidence: Double = 0.95): PartialResult[BoundedDouble] = self.withScope {
diff --git a/core/src/main/scala/org/apache/spark/rdd/PairRDDFunctions.scala b/core/src/main/scala/org/apache/spark/rdd/PairRDDFunctions.scala
index a981b63942e6d..c6181902ace6d 100644
--- a/core/src/main/scala/org/apache/spark/rdd/PairRDDFunctions.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/PairRDDFunctions.scala
@@ -274,7 +274,6 @@ class PairRDDFunctions[K, V](self: RDD[(K, V)])
   }
 
   /**
-   * ::Experimental::
    * Return a subset of this RDD sampled by key (via stratified sampling) containing exactly
    * math.ceil(numItems * samplingRate) for each stratum (group of pairs with the same key).
    *
@@ -289,7 +288,6 @@ class PairRDDFunctions[K, V](self: RDD[(K, V)])
    * @param seed seed for the random number generator
    * @return RDD containing the sampled subset
    */
-  @Experimental
   def sampleByKeyExact(
       withReplacement: Boolean,
       fractions: Map[K, Double],
@@ -384,19 +382,15 @@ class PairRDDFunctions[K, V](self: RDD[(K, V)])
   }
 
   /**
-   * :: Experimental ::
    * Approximate version of countByKey that can return a partial result if it does
    * not finish within a timeout.
    */
-  @Experimental
   def countByKeyApprox(timeout: Long, confidence: Double = 0.95)
       : PartialResult[Map[K, BoundedDouble]] = self.withScope {
     self.map(_._1).countByValueApprox(timeout, confidence)
   }
 
   /**
-   * :: Experimental ::
-   *
    * Return approximate number of distinct values for each key in this RDD.
    *
    * The algorithm used is based on streamlib's implementation of "HyperLogLog in Practice:
@@ -413,7 +407,6 @@ class PairRDDFunctions[K, V](self: RDD[(K, V)])
    *           If `sp` equals 0, the sparse representation is skipped.
    * @param partitioner Partitioner to use for the resulting RDD.
    */
-  @Experimental
   def countApproxDistinctByKey(
       p: Int,
       sp: Int,
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index a97bb174438a5..800ef53cbef07 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -31,7 +31,7 @@ import org.apache.hadoop.mapred.TextOutputFormat
 
 import org.apache.spark._
 import org.apache.spark.Partitioner._
-import org.apache.spark.annotation.{DeveloperApi, Experimental}
+import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.api.java.JavaRDD
 import org.apache.spark.partial.BoundedDouble
 import org.apache.spark.partial.CountEvaluator
@@ -1119,11 +1119,9 @@ abstract class RDD[T: ClassTag](
   def count(): Long = sc.runJob(this, Utils.getIteratorSize _).sum
 
   /**
-   * :: Experimental ::
    * Approximate version of count() that returns a potentially incomplete result
    * within a timeout, even if not all tasks have finished.
    */
-  @Experimental
   def countApprox(
       timeout: Long,
       confidence: Double = 0.95): PartialResult[BoundedDouble] = withScope {
@@ -1152,10 +1150,8 @@ abstract class RDD[T: ClassTag](
   }
 
   /**
-   * :: Experimental ::
    * Approximate version of countByValue().
    */
-  @Experimental
   def countByValueApprox(timeout: Long, confidence: Double = 0.95)
       (implicit ord: Ordering[T] = null)
       : PartialResult[Map[T, BoundedDouble]] = withScope {
@@ -1174,7 +1170,6 @@ abstract class RDD[T: ClassTag](
   }
 
   /**
-   * :: Experimental ::
    * Return approximate number of distinct elements in the RDD.
    *
    * The algorithm used is based on streamlib's implementation of "HyperLogLog in Practice:
@@ -1190,7 +1185,6 @@ abstract class RDD[T: ClassTag](
    * @param sp The precision value for the sparse set, between 0 and 32.
    *           If `sp` equals 0, the sparse representation is skipped.
    */
-  @Experimental
   def countApproxDistinct(p: Int, sp: Int): Long = withScope {
     require(p >= 4, s"p ($p) must be >= 4")
     require(sp <= 32, s"sp ($sp) must be <= 32")
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index 1f5135320326c..670d6a78e36e4 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -36,7 +36,6 @@ import org.apache.hadoop.hive.ql.metadata.Table
 import org.apache.hadoop.hive.ql.parse.VariableSubstitution
 import org.apache.hadoop.hive.serde2.io.{DateWritable, TimestampWritable}
 
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.JavaSparkContext
 import org.apache.spark.sql.SQLConf.SQLConfEntry
 import org.apache.spark.sql.SQLConf.SQLConfEntry._
@@ -356,7 +355,6 @@ class HiveContext private[hive](
    *
    * @since 1.2.0
    */
-  @Experimental
   def analyze(tableName: String) {
     val tableIdent = SqlParser.parseTableIdentifier(tableName)
     val relation = EliminateSubQueries(catalog.lookupRelation(tableIdent))
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala b/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala
index 051f53de64cd5..97113835f3bd0 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala
@@ -445,8 +445,6 @@ class StreamingContext private[streaming] (
   }
 
   /**
-   * :: Experimental ::
-   *
    * Create an input stream that monitors a Hadoop-compatible filesystem
    * for new files and reads them as flat binary files, assuming a fixed length per record,
    * generating one byte array per record. Files must be written to the monitored directory
@@ -459,7 +457,6 @@ class StreamingContext private[streaming] (
    * @param directory HDFS directory to monitor for new file
    * @param recordLength length of each record in bytes
    */
-  @Experimental
   def binaryRecordsStream(
       directory: String,
       recordLength: Int): DStream[Array[Byte]] = withNamedScope("binary records stream") {
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingContext.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingContext.scala
index 13f371f29603a..8f21c79a760c1 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingContext.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingContext.scala
@@ -222,8 +222,6 @@ class JavaStreamingContext(val ssc: StreamingContext) extends Closeable {
   }
 
   /**
-   * :: Experimental ::
-   *
    * Create an input stream that monitors a Hadoop-compatible filesystem
    * for new files and reads them as flat binary files with fixed record lengths,
    * yielding byte arrays
@@ -234,7 +232,6 @@ class JavaStreamingContext(val ssc: StreamingContext) extends Closeable {
    * @param directory HDFS directory to monitor for new files
    * @param recordLength The length at which to split the records
    */
-  @Experimental
   def binaryRecordsStream(directory: String, recordLength: Int): JavaDStream[Array[Byte]] = {
     ssc.binaryRecordsStream(directory, recordLength)
   }

From 859dff56eb0f8c63c86e7e900a12340c199e6247 Mon Sep 17 00:00:00 2001
From: Nick Evans <me@nicolasevans.org>
Date: Thu, 5 Nov 2015 09:18:20 +0000
Subject: [PATCH 065/518] [SPARK-11378][STREAMING] make
 StreamingContext.awaitTerminationOrTimeout return properly

This adds a failing test checking that `awaitTerminationOrTimeout` returns the expected value, and then fixes that failing test with the addition of a `return`.

tdas zsxwing

Author: Nick Evans <me@nicolasevans.org>

Closes #9336 from manygrams/fix_await_termination_or_timeout.
---
 python/pyspark/streaming/context.py | 2 +-
 python/pyspark/streaming/tests.py   | 7 +++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/python/pyspark/streaming/context.py b/python/pyspark/streaming/context.py
index 975c75473214a..8be56c9915265 100644
--- a/python/pyspark/streaming/context.py
+++ b/python/pyspark/streaming/context.py
@@ -218,7 +218,7 @@ def awaitTerminationOrTimeout(self, timeout):
 
         @param timeout: time to wait in seconds
         """
-        self._jssc.awaitTerminationOrTimeout(int(timeout * 1000))
+        return self._jssc.awaitTerminationOrTimeout(int(timeout * 1000))
 
     def stop(self, stopSparkContext=True, stopGraceFully=False):
         """
diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index f7fa481d50235..179479625bca4 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -596,6 +596,13 @@ def setupFunc():
         self.ssc = StreamingContext.getActiveOrCreate(None, setupFunc)
         self.assertTrue(self.setupCalled)
 
+    def test_await_termination_or_timeout(self):
+        self._add_input_stream()
+        self.ssc.start()
+        self.assertFalse(self.ssc.awaitTerminationOrTimeout(0.001))
+        self.ssc.stop(False)
+        self.assertTrue(self.ssc.awaitTerminationOrTimeout(0.001))
+
 
 class CheckpointTests(unittest.TestCase):
 

From 7bdc92197cce0edc0110dc9c2158e6e3f42c72ee Mon Sep 17 00:00:00 2001
From: Herman van Hovell <hvanhovell@questtec.nl>
Date: Thu, 5 Nov 2015 09:23:09 +0000
Subject: [PATCH 066/518] [SPARK-11449][CORE] PortableDataStream should be a
 factory

```PortableDataStream``` maintains some internal state. This makes it tricky to reuse a stream (one needs to call ```close``` on both the ```PortableDataStream``` and the ```InputStream``` it produces).

This PR removes all state from ```PortableDataStream``` and effectively turns it into an ```InputStream```/```Array[Byte]``` factory. This makes the user responsible for managing the ```InputStream``` it returns.

cc srowen

Author: Herman van Hovell <hvanhovell@questtec.nl>

Closes #9417 from hvanhovell/SPARK-11449.
---
 .../spark/input/PortableDataStream.scala      | 45 +++++++------------
 1 file changed, 16 insertions(+), 29 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/input/PortableDataStream.scala b/core/src/main/scala/org/apache/spark/input/PortableDataStream.scala
index 33e4ee0215817..280e7a5fe893c 100644
--- a/core/src/main/scala/org/apache/spark/input/PortableDataStream.scala
+++ b/core/src/main/scala/org/apache/spark/input/PortableDataStream.scala
@@ -21,7 +21,7 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, Da
 
 import scala.collection.JavaConverters._
 
-import com.google.common.io.ByteStreams
+import com.google.common.io.{Closeables, ByteStreams}
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.hadoop.mapreduce.{InputSplit, JobContext, RecordReader, TaskAttemptContext}
@@ -82,7 +82,6 @@ private[spark] abstract class StreamBasedRecordReader[T](
     if (!processed) {
       val fileIn = new PortableDataStream(split, context, index)
       value = parseStream(fileIn)
-      fileIn.close() // if it has not been open yet, close does nothing
       key = fileIn.getPath
       processed = true
       true
@@ -134,12 +133,6 @@ class PortableDataStream(
     index: Integer)
   extends Serializable {
 
-  // transient forces file to be reopened after being serialization
-  // it is also used for non-serializable classes
-
-  @transient private var fileIn: DataInputStream = null
-  @transient private var isOpen = false
-
   private val confBytes = {
     val baos = new ByteArrayOutputStream()
     SparkHadoopUtil.get.getConfigurationFromJobContext(context).
@@ -175,40 +168,34 @@ class PortableDataStream(
   }
 
   /**
-   * Create a new DataInputStream from the split and context
+   * Create a new DataInputStream from the split and context. The user of this method is responsible
+   * for closing the stream after usage.
    */
   def open(): DataInputStream = {
-    if (!isOpen) {
-      val pathp = split.getPath(index)
-      val fs = pathp.getFileSystem(conf)
-      fileIn = fs.open(pathp)
-      isOpen = true
-    }
-    fileIn
+    val pathp = split.getPath(index)
+    val fs = pathp.getFileSystem(conf)
+    fs.open(pathp)
   }
 
   /**
    * Read the file as a byte array
    */
   def toArray(): Array[Byte] = {
-    open()
-    val innerBuffer = ByteStreams.toByteArray(fileIn)
-    close()
-    innerBuffer
+    val stream = open()
+    try {
+      ByteStreams.toByteArray(stream)
+    } finally {
+      Closeables.close(stream, true)
+    }
   }
 
   /**
-   * Close the file (if it is currently open)
+   * Closing the PortableDataStream is not needed anymore. The user either can use the
+   * PortableDataStream to get a DataInputStream (which the user needs to close after usage),
+   * or a byte array.
    */
+  @deprecated("Closing the PortableDataStream is not needed anymore.", "1.6.0")
   def close(): Unit = {
-    if (isOpen) {
-      try {
-        fileIn.close()
-        isOpen = false
-      } catch {
-        case ioe: java.io.IOException => // do nothing
-      }
-    }
   }
 
   def getPath(): String = path

From a94671a027c29bacea37f56b95eccb115638abff Mon Sep 17 00:00:00 2001
From: a1singh <a1singh@ucsd.edu>
Date: Thu, 5 Nov 2015 12:51:10 +0000
Subject: [PATCH 067/518] [SPARK-11506][MLLIB] Removed redundant operation in
 Online LDA implementation

In file LDAOptimizer.scala:

line 441: since "idx" was never used, replaced unrequired zipWithIndex.foreach with foreach.

-      nonEmptyDocs.zipWithIndex.foreach { case ((_, termCounts: Vector), idx: Int) =>
+      nonEmptyDocs.foreach { case (_, termCounts: Vector) =>

Author: a1singh <a1singh@ucsd.edu>

Closes #9456 from a1singh/master.
---
 .../scala/org/apache/spark/mllib/clustering/LDAOptimizer.scala  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAOptimizer.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAOptimizer.scala
index 38486e949bbcf..17c0609800e90 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAOptimizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAOptimizer.scala
@@ -438,7 +438,7 @@ final class OnlineLDAOptimizer extends LDAOptimizer {
 
       val stat = BDM.zeros[Double](k, vocabSize)
       var gammaPart = List[BDV[Double]]()
-      nonEmptyDocs.zipWithIndex.foreach { case ((_, termCounts: Vector), idx: Int) =>
+      nonEmptyDocs.foreach { case (_, termCounts: Vector) =>
         val ids: List[Int] = termCounts match {
           case v: DenseVector => (0 until v.size).toList
           case v: SparseVector => v.indices.toList

From 77488fb8e586103ba4c0858b73e1715f1a66671f Mon Sep 17 00:00:00 2001
From: Cheng Lian <lian@databricks.com>
Date: Thu, 5 Nov 2015 23:49:44 +0800
Subject: [PATCH 068/518] [MINOR][SQL] A minor log line fix

`jars` in the log line is an array, so `$jars` doesn't print its content.

Author: Cheng Lian <lian@databricks.com>

Closes #9494 from liancheng/minor.log-fix.
---
 .../src/main/scala/org/apache/spark/sql/hive/HiveContext.scala | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index 670d6a78e36e4..2d72b959af134 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -309,7 +309,8 @@ class HiveContext private[hive](
           .map(_.toURI.toURL)
 
       logInfo(
-        s"Initializing HiveMetastoreConnection version $hiveMetastoreVersion using $jars")
+        s"Initializing HiveMetastoreConnection version $hiveMetastoreVersion " +
+          s"using ${jars.mkString(":")}")
       new IsolatedClientLoader(
         version = metaVersion,
         execJars = jars.toSeq,

From 72634f27e3110fd7f5bfca498752f69d0b1f873c Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 5 Nov 2015 08:59:06 -0800
Subject: [PATCH 069/518] [MINOR][ML][DOC] Rename weights to coefficients in
 user guide

We should use ```coefficients``` rather than ```weights``` in user guide that freshman can get the right conventional name at the outset. mengxr vectorijk

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9493 from yanboliang/docs-coefficients.
---
 docs/ml-linear-methods.md | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/docs/ml-linear-methods.md b/docs/ml-linear-methods.md
index 4e94e2f9c708d..16e2ee71293ae 100644
--- a/docs/ml-linear-methods.md
+++ b/docs/ml-linear-methods.md
@@ -71,8 +71,8 @@ val lr = new LogisticRegression()
 // Fit the model
 val lrModel = lr.fit(training)
 
-// Print the weights and intercept for logistic regression
-println(s"Weights: ${lrModel.weights} Intercept: ${lrModel.intercept}")
+// Print the coefficients and intercept for logistic regression
+println(s"Coefficients: ${lrModel.coefficients} Intercept: ${lrModel.intercept}")
 {% endhighlight %}
 </div>
 
@@ -105,8 +105,8 @@ public class LogisticRegressionWithElasticNetExample {
     // Fit the model
     LogisticRegressionModel lrModel = lr.fit(training);
 
-    // Print the weights and intercept for logistic regression
-    System.out.println("Weights: " + lrModel.weights() + " Intercept: " + lrModel.intercept());
+    // Print the coefficients and intercept for logistic regression
+    System.out.println("Coefficients: " + lrModel.coefficients() + " Intercept: " + lrModel.intercept());
   }
 }
 {% endhighlight %}
@@ -124,8 +124,8 @@ lr = LogisticRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)
 # Fit the model
 lrModel = lr.fit(training)
 
-# Print the weights and intercept for logistic regression
-print("Weights: " + str(lrModel.weights))
+# Print the coefficients and intercept for logistic regression
+print("Coefficients: " + str(lrModel.coefficients))
 print("Intercept: " + str(lrModel.intercept))
 {% endhighlight %}
 </div>
@@ -258,8 +258,8 @@ val lr = new LinearRegression()
 // Fit the model
 val lrModel = lr.fit(training)
 
-// Print the weights and intercept for linear regression
-println(s"Weights: ${lrModel.weights} Intercept: ${lrModel.intercept}")
+// Print the coefficients and intercept for linear regression
+println(s"Coefficients: ${lrModel.coefficients} Intercept: ${lrModel.intercept}")
 
 // Summarize the model over the training set and print out some metrics
 val trainingSummary = lrModel.summary
@@ -302,8 +302,8 @@ public class LinearRegressionWithElasticNetExample {
     // Fit the model
     LinearRegressionModel lrModel = lr.fit(training);
 
-    // Print the weights and intercept for linear regression
-    System.out.println("Weights: " + lrModel.weights() + " Intercept: " + lrModel.intercept());
+    // Print the coefficients and intercept for linear regression
+    System.out.println("Coefficients: " + lrModel.coefficients() + " Intercept: " + lrModel.intercept());
 
     // Summarize the model over the training set and print out some metrics
     LinearRegressionTrainingSummary trainingSummary = lrModel.summary();
@@ -330,8 +330,8 @@ lr = LinearRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)
 # Fit the model
 lrModel = lr.fit(training)
 
-# Print the weights and intercept for linear regression
-print("Weights: " + str(lrModel.weights))
+# Print the coefficients and intercept for linear regression
+print("Coefficients: " + str(lrModel.coefficients))
 print("Intercept: " + str(lrModel.intercept))
 
 # Linear regression model summary is not yet supported in Python.

From 2e86cf1b01ae0ed69f72bf8054330440d432eeb7 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 5 Nov 2015 09:00:03 -0800
Subject: [PATCH 070/518] [SPARK-11527][ML][PYSPARK] PySpark
 AFTSurvivalRegressionModel should expose coefficients/intercept/scale

PySpark ```AFTSurvivalRegressionModel``` should expose coefficients/intercept/scale. mengxr vectorijk

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9492 from yanboliang/spark-11527.
---
 python/pyspark/ml/regression.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/python/pyspark/ml/regression.py b/python/pyspark/ml/regression.py
index ab26616f4a01d..d7b4fd92c3817 100644
--- a/python/pyspark/ml/regression.py
+++ b/python/pyspark/ml/regression.py
@@ -824,6 +824,30 @@ class AFTSurvivalRegressionModel(JavaModel):
     .. versionadded:: 1.6.0
     """
 
+    @property
+    @since("1.6.0")
+    def coefficients(self):
+        """
+        Model coefficients.
+        """
+        return self._call_java("coefficients")
+
+    @property
+    @since("1.6.0")
+    def intercept(self):
+        """
+        Model intercept.
+        """
+        return self._call_java("intercept")
+
+    @property
+    @since("1.6.0")
+    def scale(self):
+        """
+        Model scale paramter.
+        """
+        return self._call_java("scale")
+
     def predictQuantiles(self, features):
         """
         Predicted Quantiles

From a4b5cefcf1a196e6b257e6127d6b43a7e50200ac Mon Sep 17 00:00:00 2001
From: Nishkam Ravi <nishkamravi@gmail.com>
Date: Thu, 5 Nov 2015 09:35:49 -0800
Subject: [PATCH 071/518] [SPARK-11501][CORE][YARN] Propagate spark.rpc config
 to executors

spark.rpc is supposed to be configurable but is not currently (doesn't get propagated to executors because RpcEnv.create is done before driver properties are fetched).

Author: Nishkam Ravi <nishkamravi@gmail.com>

Closes #9460 from nishkamravi2/master_akka.
---
 core/src/main/scala/org/apache/spark/SparkConf.scala | 1 +
 1 file changed, 1 insertion(+)

diff --git a/core/src/main/scala/org/apache/spark/SparkConf.scala b/core/src/main/scala/org/apache/spark/SparkConf.scala
index f023e4b21cb40..19633a3ce6a02 100644
--- a/core/src/main/scala/org/apache/spark/SparkConf.scala
+++ b/core/src/main/scala/org/apache/spark/SparkConf.scala
@@ -629,6 +629,7 @@ private[spark] object SparkConf extends Logging {
     name.startsWith("spark.akka") ||
     (name.startsWith("spark.auth") && name != SecurityManager.SPARK_AUTH_SECRET_CONF) ||
     name.startsWith("spark.ssl") ||
+    name.startsWith("spark.rpc") ||
     isSparkPortConf(name)
   }
 

From b072ff4d1d05fc212cd7036d1897a032a395f0b3 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@oc0558782468.ibm.com>
Date: Thu, 5 Nov 2015 09:41:14 -0800
Subject: [PATCH 072/518] [SPARK-11474][SQL] change fetchSize to fetchsize

In DefaultDataSource.scala, it has
override def createRelation(
sqlContext: SQLContext,
parameters: Map[String, String]): BaseRelation
The parameters is CaseInsensitiveMap.
After this line
parameters.foreach(kv => properties.setProperty(kv._1, kv._2))
properties is set to all lower case key/value pairs and fetchSize becomes fetchsize.
However, in compute method in JDBCRDD, it has
val fetchSize = properties.getProperty("fetchSize", "0").toInt
so fetchSize value is always 0 and never gets set correctly.

Author: Huaxin Gao <huaxing@oc0558782468.ibm.com>

Closes #9473 from huaxingao/spark-11474.
---
 .../apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala  | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
index 730d88b024cb1..018a009fbda6d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
@@ -347,6 +347,7 @@ private[sql] class JDBCRDD(
 
   /**
    * Runs the SQL query against the JDBC driver.
+   *
    */
   override def compute(thePart: Partition, context: TaskContext): Iterator[InternalRow] =
     new Iterator[InternalRow] {
@@ -368,7 +369,7 @@ private[sql] class JDBCRDD(
     val sqlText = s"SELECT $columnList FROM $fqTable $myWhereClause"
     val stmt = conn.prepareStatement(sqlText,
         ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)
-    val fetchSize = properties.getProperty("fetchSize", "0").toInt
+    val fetchSize = properties.getProperty("fetchsize", "0").toInt
     stmt.setFetchSize(fetchSize)
     val rs = stmt.executeQuery()
 

From 9da7ceed81b0afce7deb8f39f3a6d565d401a391 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 5 Nov 2015 09:56:18 -0800
Subject: [PATCH 073/518] [SPARK-11473][ML] R-like summary statistics with
 intercept for OLS via normal equation solver

Follow up [SPARK-9836](https://issues.apache.org/jira/browse/SPARK-9836), we should also support summary statistics for ```intercept```.

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9485 from yanboliang/spark-11473.
---
 .../spark/ml/optim/WeightedLeastSquares.scala | 35 ++++++++++---------
 .../ml/regression/LinearRegression.scala      | 22 +++++++-----
 .../ml/regression/LinearRegressionSuite.scala | 16 ++++-----
 python/pyspark/ml/regression.py               | 16 ++++-----
 4 files changed, 48 insertions(+), 41 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/optim/WeightedLeastSquares.scala b/mllib/src/main/scala/org/apache/spark/ml/optim/WeightedLeastSquares.scala
index e612a2122ed62..8617722ae542f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/optim/WeightedLeastSquares.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/optim/WeightedLeastSquares.scala
@@ -75,7 +75,7 @@ private[ml] class WeightedLeastSquares(
     val summary = instances.treeAggregate(new Aggregator)(_.add(_), _.merge(_))
     summary.validate()
     logInfo(s"Number of instances: ${summary.count}.")
-    val k = summary.k
+    val k = if (fitIntercept) summary.k + 1 else summary.k
     val triK = summary.triK
     val wSum = summary.wSum
     val bBar = summary.bBar
@@ -86,14 +86,6 @@ private[ml] class WeightedLeastSquares(
     val aaBar = summary.aaBar
     val aaValues = aaBar.values
 
-    if (fitIntercept) {
-      // shift centers
-      // A^T A - aBar aBar^T
-      BLAS.spr(-1.0, aBar, aaValues)
-      // A^T b - bBar aBar
-      BLAS.axpy(-bBar, aBar, abBar)
-    }
-
     // add regularization to diagonals
     var i = 0
     var j = 2
@@ -111,21 +103,32 @@ private[ml] class WeightedLeastSquares(
       j += 1
     }
 
-    val x = new DenseVector(CholeskyDecomposition.solve(aaBar.values, abBar.values))
+    val aa = if (fitIntercept) {
+      Array.concat(aaBar.values, aBar.values, Array(1.0))
+    } else {
+      aaBar.values
+    }
+    val ab = if (fitIntercept) {
+      Array.concat(abBar.values, Array(bBar))
+    } else {
+      abBar.values
+    }
+
+    val x = CholeskyDecomposition.solve(aa, ab)
+
+    val aaInv = CholeskyDecomposition.inverse(aa, k)
 
-    val aaInv = CholeskyDecomposition.inverse(aaBar.values, k)
     // aaInv is a packed upper triangular matrix, here we get all elements on diagonal
     val diagInvAtWA = new DenseVector((1 to k).map { i =>
       aaInv(i + (i - 1) * i / 2 - 1) / wSum }.toArray)
 
-    // compute intercept
-    val intercept = if (fitIntercept) {
-      bBar - BLAS.dot(aBar, x)
+    val (coefficients, intercept) = if (fitIntercept) {
+      (new DenseVector(x.slice(0, x.length - 1)), x.last)
     } else {
-      0.0
+      (new DenseVector(x), 0.0)
     }
 
-    new WeightedLeastSquaresModel(x, intercept, diagInvAtWA)
+    new WeightedLeastSquaresModel(coefficients, intercept, diagInvAtWA)
   }
 }
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index c51e30483ab3d..6638313818703 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -511,8 +511,7 @@ class LinearRegressionSummary private[regression] (
   }
 
   /**
-   * Standard error of estimated coefficients.
-   * Note that standard error of estimated intercept is not supported currently.
+   * Standard error of estimated coefficients and intercept.
    */
   lazy val coefficientStandardErrors: Array[Double] = {
     if (diagInvAtWA.length == 1 && diagInvAtWA(0) == 0) {
@@ -532,21 +531,26 @@ class LinearRegressionSummary private[regression] (
     }
   }
 
-  /** T-statistic of estimated coefficients.
-    * Note that t-statistic of estimated intercept is not supported currently.
-    */
+  /**
+   * T-statistic of estimated coefficients and intercept.
+   */
   lazy val tValues: Array[Double] = {
     if (diagInvAtWA.length == 1 && diagInvAtWA(0) == 0) {
       throw new UnsupportedOperationException(
         "No t-statistic available for this LinearRegressionModel")
     } else {
-      model.coefficients.toArray.zip(coefficientStandardErrors).map { x => x._1 / x._2 }
+      val estimate = if (model.getFitIntercept) {
+        Array.concat(model.coefficients.toArray, Array(model.intercept))
+      } else {
+        model.coefficients.toArray
+      }
+      estimate.zip(coefficientStandardErrors).map { x => x._1 / x._2 }
     }
   }
 
-  /** Two-sided p-value of estimated coefficients.
-    * Note that p-value of estimated intercept is not supported currently.
-    */
+  /**
+   * Two-sided p-value of estimated coefficients and intercept.
+   */
   lazy val pValues: Array[Double] = {
     if (diagInvAtWA.length == 1 && diagInvAtWA(0) == 0) {
       throw new UnsupportedOperationException(
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
index fbf83e8922861..a1d86fe8fedad 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
@@ -621,13 +621,13 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
         assert(model.summary.objectiveHistory.length == 1)
         assert(model.summary.objectiveHistory(0) == 0.0)
         val devianceResidualsR = Array(-0.35566, 0.34504)
-        val seCoefR = Array(0.0011756, 0.0009032)
-        val tValsR = Array(3998, 7971)
-        val pValsR = Array(0, 0)
+        val seCoefR = Array(0.0011756, 0.0009032, 0.0018489)
+        val tValsR = Array(3998, 7971, 3407)
+        val pValsR = Array(0, 0, 0)
         model.summary.devianceResiduals.zip(devianceResidualsR).foreach { x =>
-          assert(x._1 ~== x._2 absTol 1E-3) }
+          assert(x._1 ~== x._2 absTol 1E-5) }
         model.summary.coefficientStandardErrors.zip(seCoefR).foreach{ x =>
-          assert(x._1 ~== x._2 absTol 1E-3) }
+          assert(x._1 ~== x._2 absTol 1E-5) }
         model.summary.tValues.map(_.round).zip(tValsR).foreach{ x => assert(x._1 === x._2) }
         model.summary.pValues.map(_.round).zip(pValsR).foreach{ x => assert(x._1 === x._2) }
       }
@@ -789,9 +789,9 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
     val coefficientsR = Vectors.dense(Array(6.080, -0.600))
     val interceptR = 18.080
     val devianceResidualsR = Array(-1.358, 1.920)
-    val seCoefR = Array(5.556, 1.960)
-    val tValsR = Array(1.094, -0.306)
-    val pValsR = Array(0.471, 0.811)
+    val seCoefR = Array(5.556, 1.960, 9.608)
+    val tValsR = Array(1.094, -0.306, 1.882)
+    val pValsR = Array(0.471, 0.811, 0.311)
 
     assert(model.coefficients ~== coefficientsR absTol 1E-3)
     assert(model.intercept ~== interceptR absTol 1E-3)
diff --git a/python/pyspark/ml/regression.py b/python/pyspark/ml/regression.py
index d7b4fd92c3817..7648bf13266bf 100644
--- a/python/pyspark/ml/regression.py
+++ b/python/pyspark/ml/regression.py
@@ -55,15 +55,15 @@ class LinearRegression(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPrediction
     >>> lr = LinearRegression(maxIter=5, regParam=0.0, solver="normal")
     >>> model = lr.fit(df)
     >>> test0 = sqlContext.createDataFrame([(Vectors.dense(-1.0),)], ["features"])
-    >>> model.transform(test0).head().prediction
-    -1.0
-    >>> model.weights
-    DenseVector([1.0])
-    >>> model.intercept
-    0.0
+    >>> abs(model.transform(test0).head().prediction - (-1.0)) < 0.001
+    True
+    >>> abs(model.coefficients[0] - 1.0) < 0.001
+    True
+    >>> abs(model.intercept - 0.0) < 0.001
+    True
     >>> test1 = sqlContext.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"])
-    >>> model.transform(test1).head().prediction
-    1.0
+    >>> abs(model.transform(test1).head().prediction - 1.0) < 0.001
+    True
     >>> lr.setParams("vector")
     Traceback (most recent call last):
         ...

From c76865c6220e3e7b2a266bbc4935567ef55303d8 Mon Sep 17 00:00:00 2001
From: Srinivasa Reddy Vundela <vsr@cloudera.com>
Date: Thu, 5 Nov 2015 11:30:44 -0800
Subject: [PATCH 074/518] [SPARK-11484][WEBUI] Using proxyBase set by spark AM

Use the proxyBase set by the AM, if not found then use env. This is to fix the issue if somebody accidentally set APPLICATION_WEB_PROXY_BASE to wrong proxyBase

Author: Srinivasa Reddy Vundela <vsr@cloudera.com>

Closes #9448 from vundela/master.
---
 .../src/main/scala/org/apache/spark/ui/UIUtils.scala | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/ui/UIUtils.scala b/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
index 68a9f912a5d2c..25dcb604d9e5f 100644
--- a/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
+++ b/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
@@ -143,14 +143,10 @@ private[spark] object UIUtils extends Logging {
 
   // Yarn has to go through a proxy so the base uri is provided and has to be on all links
   def uiRoot: String = {
-    if (System.getenv("APPLICATION_WEB_PROXY_BASE") != null) {
-      System.getenv("APPLICATION_WEB_PROXY_BASE")
-    } else if (System.getProperty("spark.ui.proxyBase") != null) {
-      System.getProperty("spark.ui.proxyBase")
-    }
-    else {
-      ""
-    }
+    // SPARK-11484 - Use the proxyBase set by the AM, if not found then use env.
+    sys.props.get("spark.ui.proxyBase")
+      .orElse(sys.env.get("APPLICATION_WEB_PROXY_BASE"))
+      .getOrElse("")
   }
 
   def prependBaseUri(basePath: String = "", resource: String = ""): String = {

From 6b87acd6649a3390b5c2c4fcb61e58d125d0d87c Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 5 Nov 2015 11:58:13 -0800
Subject: [PATCH 075/518] [SPARK-11513][SQL] Remove implicit conversion from
 LogicalPlan to DataFrame

This internal implicit conversion has been a source of confusion for a lot of new developers.

Author: Reynold Xin <rxin@databricks.com>

Closes #9479 from rxin/SPARK-11513.
---
 .../org/apache/spark/sql/DataFrame.scala      | 123 +++++++++++-------
 .../scala/org/apache/spark/sql/Dataset.scala  |   5 +-
 2 files changed, 78 insertions(+), 50 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index d3a2249d7006c..6336dee7be6a3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -147,14 +147,6 @@ class DataFrame private[sql](
       queryExecution.analyzed
   }
 
-  /**
-   * An implicit conversion function internal to this class for us to avoid doing
-   * "new DataFrame(...)" everywhere.
-   */
-  @inline private implicit def logicalPlanToDataFrame(logicalPlan: LogicalPlan): DataFrame = {
-    new DataFrame(sqlContext, logicalPlan)
-  }
-
   protected[sql] def resolve(colName: String): NamedExpression = {
     queryExecution.analyzed.resolveQuoted(colName, sqlContext.analyzer.resolver).getOrElse {
       throw new AnalysisException(
@@ -235,7 +227,7 @@ class DataFrame private[sql](
     // For Data that has more than "numRows" records
     if (hasMoreData) {
       val rowsString = if (numRows == 1) "row" else "rows"
-      sb.append(s"only showing top $numRows ${rowsString}\n")
+      sb.append(s"only showing top $numRows $rowsString\n")
     }
 
     sb.toString()
@@ -332,7 +324,7 @@ class DataFrame private[sql](
    */
   def explain(extended: Boolean): Unit = {
     val explain = ExplainCommand(queryExecution.logical, extended = extended)
-    explain.queryExecution.executedPlan.executeCollect().foreach {
+    withPlan(explain).queryExecution.executedPlan.executeCollect().foreach {
       // scalastyle:off println
       r => println(r.getString(0))
       // scalastyle:on println
@@ -370,7 +362,7 @@ class DataFrame private[sql](
    * @group action
    * @since 1.3.0
    */
-  def show(numRows: Int): Unit = show(numRows, true)
+  def show(numRows: Int): Unit = show(numRows, truncate = true)
 
   /**
    * Displays the top 20 rows of [[DataFrame]] in a tabular form. Strings more than 20 characters
@@ -445,7 +437,7 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def join(right: DataFrame): DataFrame = {
+  def join(right: DataFrame): DataFrame = withPlan {
     Join(logicalPlan, right.logicalPlan, joinType = Inner, None)
   }
 
@@ -520,21 +512,25 @@ class DataFrame private[sql](
       Join(logicalPlan, right.logicalPlan, joinType = Inner, None)).analyzed.asInstanceOf[Join]
 
     // Project only one of the join columns.
-    val joinedCols = usingColumns.map(col => joined.right.resolve(col))
+    val joinedCols = usingColumns.map(col => withPlan(joined.right).resolve(col))
     val condition = usingColumns.map { col =>
-      catalyst.expressions.EqualTo(joined.left.resolve(col), joined.right.resolve(col))
+      catalyst.expressions.EqualTo(
+        withPlan(joined.left).resolve(col),
+        withPlan(joined.right).resolve(col))
     }.reduceLeftOption[catalyst.expressions.BinaryExpression] { (cond, eqTo) =>
       catalyst.expressions.And(cond, eqTo)
     }
 
-    Project(
-      joined.output.filterNot(joinedCols.contains(_)),
-      Join(
-        joined.left,
-        joined.right,
-        joinType = JoinType(joinType),
-        condition)
-    )
+    withPlan {
+      Project(
+        joined.output.filterNot(joinedCols.contains(_)),
+        Join(
+          joined.left,
+          joined.right,
+          joinType = JoinType(joinType),
+          condition)
+      )
+    }
   }
 
   /**
@@ -581,19 +577,20 @@ class DataFrame private[sql](
 
     // Trigger analysis so in the case of self-join, the analyzer will clone the plan.
     // After the cloning, left and right side will have distinct expression ids.
-    val plan = Join(logicalPlan, right.logicalPlan, JoinType(joinType), Some(joinExprs.expr))
+    val plan = withPlan(
+      Join(logicalPlan, right.logicalPlan, JoinType(joinType), Some(joinExprs.expr)))
       .queryExecution.analyzed.asInstanceOf[Join]
 
     // If auto self join alias is disabled, return the plan.
     if (!sqlContext.conf.dataFrameSelfJoinAutoResolveAmbiguity) {
-      return plan
+      return withPlan(plan)
     }
 
     // If left/right have no output set intersection, return the plan.
-    val lanalyzed = this.logicalPlan.queryExecution.analyzed
-    val ranalyzed = right.logicalPlan.queryExecution.analyzed
+    val lanalyzed = withPlan(this.logicalPlan).queryExecution.analyzed
+    val ranalyzed = withPlan(right.logicalPlan).queryExecution.analyzed
     if (lanalyzed.outputSet.intersect(ranalyzed.outputSet).isEmpty) {
-      return plan
+      return withPlan(plan)
     }
 
     // Otherwise, find the trivially true predicates and automatically resolves them to both sides.
@@ -602,9 +599,14 @@ class DataFrame private[sql](
     val cond = plan.condition.map { _.transform {
       case catalyst.expressions.EqualTo(a: AttributeReference, b: AttributeReference)
           if a.sameRef(b) =>
-        catalyst.expressions.EqualTo(plan.left.resolve(a.name), plan.right.resolve(b.name))
+        catalyst.expressions.EqualTo(
+          withPlan(plan.left).resolve(a.name),
+          withPlan(plan.right).resolve(b.name))
     }}
-    plan.copy(condition = cond)
+
+    withPlan {
+      plan.copy(condition = cond)
+    }
   }
 
   /**
@@ -707,7 +709,9 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def as(alias: String): DataFrame = Subquery(alias, logicalPlan)
+  def as(alias: String): DataFrame = withPlan {
+    Subquery(alias, logicalPlan)
+  }
 
   /**
    * (Scala-specific) Returns a new [[DataFrame]] with an alias set.
@@ -739,7 +743,7 @@ class DataFrame private[sql](
    * @since 1.3.0
    */
   @scala.annotation.varargs
-  def select(cols: Column*): DataFrame = {
+  def select(cols: Column*): DataFrame = withPlan {
     val namedExpressions = cols.map {
       // Wrap UnresolvedAttribute with UnresolvedAlias, as when we resolve UnresolvedAttribute, we
       // will remove intermediate Alias for ExtractValue chain, and we need to alias it again to
@@ -798,7 +802,9 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def filter(condition: Column): DataFrame = Filter(condition.expr, logicalPlan)
+  def filter(condition: Column): DataFrame = withPlan {
+    Filter(condition.expr, logicalPlan)
+  }
 
   /**
    * Filters rows using the given SQL expression.
@@ -1039,7 +1045,9 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def limit(n: Int): DataFrame = Limit(Literal(n), logicalPlan)
+  def limit(n: Int): DataFrame = withPlan {
+    Limit(Literal(n), logicalPlan)
+  }
 
   /**
    * Returns a new [[DataFrame]] containing union of rows in this frame and another frame.
@@ -1047,7 +1055,9 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def unionAll(other: DataFrame): DataFrame = Union(logicalPlan, other.logicalPlan)
+  def unionAll(other: DataFrame): DataFrame = withPlan {
+    Union(logicalPlan, other.logicalPlan)
+  }
 
   /**
    * Returns a new [[DataFrame]] containing rows only in both this frame and another frame.
@@ -1055,7 +1065,9 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def intersect(other: DataFrame): DataFrame = Intersect(logicalPlan, other.logicalPlan)
+  def intersect(other: DataFrame): DataFrame = withPlan {
+    Intersect(logicalPlan, other.logicalPlan)
+  }
 
   /**
    * Returns a new [[DataFrame]] containing rows in this frame but not in another frame.
@@ -1063,7 +1075,9 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def except(other: DataFrame): DataFrame = Except(logicalPlan, other.logicalPlan)
+  def except(other: DataFrame): DataFrame = withPlan {
+    Except(logicalPlan, other.logicalPlan)
+  }
 
   /**
    * Returns a new [[DataFrame]] by sampling a fraction of rows.
@@ -1074,7 +1088,7 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def sample(withReplacement: Boolean, fraction: Double, seed: Long): DataFrame = {
+  def sample(withReplacement: Boolean, fraction: Double, seed: Long): DataFrame = withPlan {
     Sample(0.0, fraction, withReplacement, seed, logicalPlan)
   }
 
@@ -1102,7 +1116,7 @@ class DataFrame private[sql](
     val sum = weights.sum
     val normalizedCumWeights = weights.map(_ / sum).scanLeft(0.0d)(_ + _)
     normalizedCumWeights.sliding(2).map { x =>
-      new DataFrame(sqlContext, Sample(x(0), x(1), false, seed, logicalPlan))
+      new DataFrame(sqlContext, Sample(x(0), x(1), withReplacement = false, seed, logicalPlan))
     }.toArray
   }
 
@@ -1162,8 +1176,10 @@ class DataFrame private[sql](
       f.andThen(_.map(convert(_).asInstanceOf[InternalRow]))
     val generator = UserDefinedGenerator(elementTypes, rowFunction, input.map(_.expr))
 
-    Generate(generator, join = true, outer = false,
-      qualifier = None, generatorOutput = Nil, logicalPlan)
+    withPlan {
+      Generate(generator, join = true, outer = false,
+        qualifier = None, generatorOutput = Nil, logicalPlan)
+    }
   }
 
   /**
@@ -1190,8 +1206,10 @@ class DataFrame private[sql](
     }
     val generator = UserDefinedGenerator(elementTypes, rowFunction, apply(inputColumn).expr :: Nil)
 
-    Generate(generator, join = true, outer = false,
-      qualifier = None, generatorOutput = Nil, logicalPlan)
+    withPlan {
+      Generate(generator, join = true, outer = false,
+        qualifier = None, generatorOutput = Nil, logicalPlan)
+    }
   }
 
   /////////////////////////////////////////////////////////////////////////////
@@ -1309,7 +1327,7 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.4.0
    */
-  def dropDuplicates(colNames: Seq[String]): DataFrame = {
+  def dropDuplicates(colNames: Seq[String]): DataFrame = withPlan {
     val groupCols = colNames.map(resolve)
     val groupColExprIds = groupCols.map(_.exprId)
     val aggCols = logicalPlan.output.map { attr =>
@@ -1355,7 +1373,7 @@ class DataFrame private[sql](
    * @since 1.3.1
    */
   @scala.annotation.varargs
-  def describe(cols: String*): DataFrame = {
+  def describe(cols: String*): DataFrame = withPlan {
 
     // The list of summary statistics to compute, in the form of expressions.
     val statistics = List[(String, Expression => Expression)](
@@ -1505,7 +1523,7 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def repartition(numPartitions: Int): DataFrame = {
+  def repartition(numPartitions: Int): DataFrame = withPlan {
     Repartition(numPartitions, shuffle = true, logicalPlan)
   }
 
@@ -1519,7 +1537,7 @@ class DataFrame private[sql](
    * @since 1.6.0
    */
   @scala.annotation.varargs
-  def repartition(numPartitions: Int, partitionExprs: Column*): DataFrame = {
+  def repartition(numPartitions: Int, partitionExprs: Column*): DataFrame = withPlan {
     RepartitionByExpression(partitionExprs.map(_.expr), logicalPlan, Some(numPartitions))
   }
 
@@ -1533,7 +1551,7 @@ class DataFrame private[sql](
    * @since 1.6.0
    */
   @scala.annotation.varargs
-  def repartition(partitionExprs: Column*): DataFrame = {
+  def repartition(partitionExprs: Column*): DataFrame = withPlan {
     RepartitionByExpression(partitionExprs.map(_.expr), logicalPlan, numPartitions = None)
   }
 
@@ -1545,7 +1563,7 @@ class DataFrame private[sql](
    * @group rdd
    * @since 1.4.0
    */
-  def coalesce(numPartitions: Int): DataFrame = {
+  def coalesce(numPartitions: Int): DataFrame = withPlan {
     Repartition(numPartitions, shuffle = false, logicalPlan)
   }
 
@@ -2066,7 +2084,14 @@ class DataFrame private[sql](
           SortOrder(expr, Ascending)
       }
     }
-    Sort(sortOrder, global = global, logicalPlan)
+    withPlan {
+      Sort(sortOrder, global = global, logicalPlan)
+    }
+  }
+
+  /** A convenient function to wrap a logical plan and produce a DataFrame. */
+  @inline private def withPlan(logicalPlan: => LogicalPlan): DataFrame = {
+    new DataFrame(sqlContext, logicalPlan)
   }
 
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 7b75aeec4cf3a..500227e93a472 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -107,13 +107,16 @@ class Dataset[T] private(
    * strongly typed objects that Dataset operations work on, a Dataframe returns generic [[Row]]
    * objects that allow fields to be accessed by ordinal or name.
    */
+  // This is declared with parentheses to prevent the Scala compiler from treating
+  // `ds.toDF("1")` as invoking this toDF and then apply on the returned DataFrame.
   def toDF(): DataFrame = DataFrame(sqlContext, logicalPlan)
 
-
   /**
    * Returns this Dataset.
    * @since 1.6.0
    */
+  // This is declared with parentheses to prevent the Scala compiler from treating
+  // `ds.toDS("1")` as invoking this toDF and then apply on the returned Dataset.
   def toDS(): Dataset[T] = this
 
   /**

From f80f7b69a3f81d0ea879a31c769d17ffbbac74aa Mon Sep 17 00:00:00 2001
From: "Ehsan M.Kermani" <ehsanmo1367@gmail.com>
Date: Thu, 5 Nov 2015 12:11:57 -0800
Subject: [PATCH 076/518] [SPARK-10265][DOCUMENTATION, ML] Fixed @Since
 annotation to ml.regression

Here is my first commit.

Author: Ehsan M.Kermani <ehsanmo1367@gmail.com>

Closes #8728 from ehsanmok/SinceAnn.
---
 .../ml/regression/DecisionTreeRegressor.scala | 20 +++++++++--
 .../spark/ml/regression/GBTRegressor.scala    | 33 ++++++++++++++++---
 .../ml/regression/IsotonicRegression.scala    | 26 +++++++++++++--
 .../ml/regression/LinearRegression.scala      | 28 ++++++++++++++--
 .../ml/regression/RandomForestRegressor.scala | 30 ++++++++++++++---
 5 files changed, 119 insertions(+), 18 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
index 88b79a4eb82be..04420fc6e8251 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.ml.regression
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.{PredictionModel, Predictor}
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.tree.{DecisionTreeModel, DecisionTreeParams, Node, TreeRegressorParams}
@@ -36,30 +36,39 @@ import org.apache.spark.sql.DataFrame
  * for regression.
  * It supports both continuous and categorical features.
  */
+@Since("1.4.0")
 @Experimental
-final class DecisionTreeRegressor(override val uid: String)
+final class DecisionTreeRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: String)
   extends Predictor[Vector, DecisionTreeRegressor, DecisionTreeRegressionModel]
   with DecisionTreeParams with TreeRegressorParams {
 
+  @Since("1.4.0")
   def this() = this(Identifiable.randomUID("dtr"))
 
   // Override parameter setters from parent trait for Java API compatibility.
-
+  @Since("1.4.0")
   override def setMaxDepth(value: Int): this.type = super.setMaxDepth(value)
 
+  @Since("1.4.0")
   override def setMaxBins(value: Int): this.type = super.setMaxBins(value)
 
+  @Since("1.4.0")
   override def setMinInstancesPerNode(value: Int): this.type =
     super.setMinInstancesPerNode(value)
 
+  @Since("1.4.0")
   override def setMinInfoGain(value: Double): this.type = super.setMinInfoGain(value)
 
+  @Since("1.4.0")
   override def setMaxMemoryInMB(value: Int): this.type = super.setMaxMemoryInMB(value)
 
+  @Since("1.4.0")
   override def setCacheNodeIds(value: Boolean): this.type = super.setCacheNodeIds(value)
 
+  @Since("1.4.0")
   override def setCheckpointInterval(value: Int): this.type = super.setCheckpointInterval(value)
 
+  @Since("1.4.0")
   override def setImpurity(value: String): this.type = super.setImpurity(value)
 
   override protected def train(dataset: DataFrame): DecisionTreeRegressionModel = {
@@ -78,9 +87,11 @@ final class DecisionTreeRegressor(override val uid: String)
       subsamplingRate = 1.0)
   }
 
+  @Since("1.4.0")
   override def copy(extra: ParamMap): DecisionTreeRegressor = defaultCopy(extra)
 }
 
+@Since("1.4.0")
 @Experimental
 object DecisionTreeRegressor {
   /** Accessor for supported impurities: variance */
@@ -93,6 +104,7 @@ object DecisionTreeRegressor {
  * It supports both continuous and categorical features.
  * @param rootNode  Root of the decision tree
  */
+@Since("1.4.0")
 @Experimental
 final class DecisionTreeRegressionModel private[ml] (
     override val uid: String,
@@ -115,10 +127,12 @@ final class DecisionTreeRegressionModel private[ml] (
     rootNode.predictImpl(features).prediction
   }
 
+  @Since("1.4.0")
   override def copy(extra: ParamMap): DecisionTreeRegressionModel = {
     copyValues(new DecisionTreeRegressionModel(uid, rootNode, numFeatures), extra).setParent(parent)
   }
 
+  @Since("1.4.0")
   override def toString: String = {
     s"DecisionTreeRegressionModel (uid=$uid) of depth $depth with $numNodes nodes"
   }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
index 65b5b3e0727df..07144cc7cfbd7 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
@@ -20,7 +20,7 @@ package org.apache.spark.ml.regression
 import com.github.fommil.netlib.BLAS.{getInstance => blas}
 
 import org.apache.spark.Logging
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.{PredictionModel, Predictor}
 import org.apache.spark.ml.param.{Param, ParamMap}
 import org.apache.spark.ml.tree.{DecisionTreeModel, GBTParams, TreeEnsembleModel, TreeRegressorParams}
@@ -42,54 +42,65 @@ import org.apache.spark.sql.types.DoubleType
  * learning algorithm for regression.
  * It supports both continuous and categorical features.
  */
+@Since("1.4.0")
 @Experimental
-final class GBTRegressor(override val uid: String)
+final class GBTRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: String)
   extends Predictor[Vector, GBTRegressor, GBTRegressionModel]
   with GBTParams with TreeRegressorParams with Logging {
 
+  @Since("1.4.0")
   def this() = this(Identifiable.randomUID("gbtr"))
 
   // Override parameter setters from parent trait for Java API compatibility.
 
   // Parameters from TreeRegressorParams:
-
+  @Since("1.4.0")
   override def setMaxDepth(value: Int): this.type = super.setMaxDepth(value)
 
+  @Since("1.4.0")
   override def setMaxBins(value: Int): this.type = super.setMaxBins(value)
 
+  @Since("1.4.0")
   override def setMinInstancesPerNode(value: Int): this.type =
     super.setMinInstancesPerNode(value)
 
+  @Since("1.4.0")
   override def setMinInfoGain(value: Double): this.type = super.setMinInfoGain(value)
 
+  @Since("1.4.0")
   override def setMaxMemoryInMB(value: Int): this.type = super.setMaxMemoryInMB(value)
 
+  @Since("1.4.0")
   override def setCacheNodeIds(value: Boolean): this.type = super.setCacheNodeIds(value)
 
+  @Since("1.4.0")
   override def setCheckpointInterval(value: Int): this.type = super.setCheckpointInterval(value)
 
   /**
    * The impurity setting is ignored for GBT models.
    * Individual trees are built using impurity "Variance."
    */
+  @Since("1.4.0")
   override def setImpurity(value: String): this.type = {
     logWarning("GBTRegressor.setImpurity should NOT be used")
     this
   }
 
   // Parameters from TreeEnsembleParams:
-
+  @Since("1.4.0")
   override def setSubsamplingRate(value: Double): this.type = super.setSubsamplingRate(value)
 
+  @Since("1.4.0")
   override def setSeed(value: Long): this.type = {
     logWarning("The 'seed' parameter is currently ignored by Gradient Boosting.")
     super.setSeed(value)
   }
 
   // Parameters from GBTParams:
-
+  @Since("1.4.0")
   override def setMaxIter(value: Int): this.type = super.setMaxIter(value)
 
+  @Since("1.4.0")
   override def setStepSize(value: Double): this.type = super.setStepSize(value)
 
   // Parameters for GBTRegressor:
@@ -100,6 +111,7 @@ final class GBTRegressor(override val uid: String)
    * (default = squared)
    * @group param
    */
+  @Since("1.4.0")
   val lossType: Param[String] = new Param[String](this, "lossType", "Loss function which GBT" +
     " tries to minimize (case-insensitive). Supported options:" +
     s" ${GBTRegressor.supportedLossTypes.mkString(", ")}",
@@ -108,9 +120,11 @@ final class GBTRegressor(override val uid: String)
   setDefault(lossType -> "squared")
 
   /** @group setParam */
+  @Since("1.4.0")
   def setLossType(value: String): this.type = set(lossType, value)
 
   /** @group getParam */
+  @Since("1.4.0")
   def getLossType: String = $(lossType).toLowerCase
 
   /** (private[ml]) Convert new loss to old loss. */
@@ -135,13 +149,16 @@ final class GBTRegressor(override val uid: String)
     GBTRegressionModel.fromOld(oldModel, this, categoricalFeatures, numFeatures)
   }
 
+  @Since("1.4.0")
   override def copy(extra: ParamMap): GBTRegressor = defaultCopy(extra)
 }
 
+@Since("1.4.0")
 @Experimental
 object GBTRegressor {
   // The losses below should be lowercase.
   /** Accessor for supported loss settings: squared (L2), absolute (L1) */
+  @Since("1.4.0")
   final val supportedLossTypes: Array[String] = Array("squared", "absolute").map(_.toLowerCase)
 }
 
@@ -154,6 +171,7 @@ object GBTRegressor {
  * @param _trees  Decision trees in the ensemble.
  * @param _treeWeights  Weights for the decision trees in the ensemble.
  */
+@Since("1.4.0")
 @Experimental
 final class GBTRegressionModel private[ml](
     override val uid: String,
@@ -172,11 +190,14 @@ final class GBTRegressionModel private[ml](
    * @param _trees  Decision trees in the ensemble.
    * @param _treeWeights  Weights for the decision trees in the ensemble.
    */
+  @Since("1.4.0")
   def this(uid: String, _trees: Array[DecisionTreeRegressionModel], _treeWeights: Array[Double]) =
     this(uid, _trees, _treeWeights, -1)
 
+  @Since("1.4.0")
   override def trees: Array[DecisionTreeModel] = _trees.asInstanceOf[Array[DecisionTreeModel]]
 
+  @Since("1.4.0")
   override def treeWeights: Array[Double] = _treeWeights
 
   override protected def transformImpl(dataset: DataFrame): DataFrame = {
@@ -194,11 +215,13 @@ final class GBTRegressionModel private[ml](
     blas.ddot(numTrees, treePredictions, 1, _treeWeights, 1)
   }
 
+  @Since("1.4.0")
   override def copy(extra: ParamMap): GBTRegressionModel = {
     copyValues(new GBTRegressionModel(uid, _trees, _treeWeights, numFeatures),
       extra).setParent(parent)
   }
 
+  @Since("1.4.0")
   override def toString: String = {
     s"GBTRegressionModel (uid=$uid) with $numTrees trees"
   }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
index f4a17c8f9a582..a1fe01b047108 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
@@ -18,7 +18,7 @@
 package org.apache.spark.ml.regression
 
 import org.apache.spark.Logging
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared.{HasFeaturesCol, HasLabelCol, HasPredictionCol, HasWeightCol}
@@ -124,32 +124,42 @@ private[regression] trait IsotonicRegressionBase extends Params with HasFeatures
  *
  * Uses [[org.apache.spark.mllib.regression.IsotonicRegression]].
  */
+@Since("1.5.0")
 @Experimental
-class IsotonicRegression(override val uid: String) extends Estimator[IsotonicRegressionModel]
-  with IsotonicRegressionBase {
+class IsotonicRegression @Since("1.5.0") (@Since("1.5.0") override val uid: String)
+  extends Estimator[IsotonicRegressionModel] with IsotonicRegressionBase {
 
+  @Since("1.5.0")
   def this() = this(Identifiable.randomUID("isoReg"))
 
   /** @group setParam */
+  @Since("1.5.0")
   def setLabelCol(value: String): this.type = set(labelCol, value)
 
   /** @group setParam */
+  @Since("1.5.0")
   def setFeaturesCol(value: String): this.type = set(featuresCol, value)
 
   /** @group setParam */
+  @Since("1.5.0")
   def setPredictionCol(value: String): this.type = set(predictionCol, value)
 
   /** @group setParam */
+  @Since("1.5.0")
   def setIsotonic(value: Boolean): this.type = set(isotonic, value)
 
   /** @group setParam */
+  @Since("1.5.0")
   def setWeightCol(value: String): this.type = set(weightCol, value)
 
   /** @group setParam */
+  @Since("1.5.0")
   def setFeatureIndex(value: Int): this.type = set(featureIndex, value)
 
+  @Since("1.5.0")
   override def copy(extra: ParamMap): IsotonicRegression = defaultCopy(extra)
 
+  @Since("1.5.0")
   override def fit(dataset: DataFrame): IsotonicRegressionModel = {
     validateAndTransformSchema(dataset.schema, fitting = true)
     // Extract columns from data.  If dataset is persisted, do not persist oldDataset.
@@ -163,6 +173,7 @@ class IsotonicRegression(override val uid: String) extends Estimator[IsotonicReg
     copyValues(new IsotonicRegressionModel(uid, oldModel).setParent(this))
   }
 
+  @Since("1.5.0")
   override def transformSchema(schema: StructType): StructType = {
     validateAndTransformSchema(schema, fitting = true)
   }
@@ -178,6 +189,7 @@ class IsotonicRegression(override val uid: String) extends Estimator[IsotonicReg
  * @param oldModel A [[org.apache.spark.mllib.regression.IsotonicRegressionModel]]
  *                 model trained by [[org.apache.spark.mllib.regression.IsotonicRegression]].
  */
+@Since("1.5.0")
 @Experimental
 class IsotonicRegressionModel private[ml] (
     override val uid: String,
@@ -185,27 +197,34 @@ class IsotonicRegressionModel private[ml] (
   extends Model[IsotonicRegressionModel] with IsotonicRegressionBase {
 
   /** @group setParam */
+  @Since("1.5.0")
   def setFeaturesCol(value: String): this.type = set(featuresCol, value)
 
   /** @group setParam */
+  @Since("1.5.0")
   def setPredictionCol(value: String): this.type = set(predictionCol, value)
 
   /** @group setParam */
+  @Since("1.5.0")
   def setFeatureIndex(value: Int): this.type = set(featureIndex, value)
 
   /** Boundaries in increasing order for which predictions are known. */
+  @Since("1.5.0")
   def boundaries: Vector = Vectors.dense(oldModel.boundaries)
 
   /**
    * Predictions associated with the boundaries at the same index, monotone because of isotonic
    * regression.
    */
+  @Since("1.5.0")
   def predictions: Vector = Vectors.dense(oldModel.predictions)
 
+  @Since("1.5.0")
   override def copy(extra: ParamMap): IsotonicRegressionModel = {
     copyValues(new IsotonicRegressionModel(uid, oldModel), extra).setParent(parent)
   }
 
+  @Since("1.5.0")
   override def transform(dataset: DataFrame): DataFrame = {
     val predict = dataset.schema($(featuresCol)).dataType match {
       case DoubleType =>
@@ -217,6 +236,7 @@ class IsotonicRegressionModel private[ml] (
     dataset.withColumn($(predictionCol), predict(col($(featuresCol))))
   }
 
+  @Since("1.5.0")
   override def transformSchema(schema: StructType): StructType = {
     validateAndTransformSchema(schema, fitting = false)
   }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index 6638313818703..913140e581983 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -24,9 +24,9 @@ import breeze.optimize.{CachedDiffFunction, DiffFunction, LBFGS => BreezeLBFGS,
 import breeze.stats.distributions.StudentsT
 
 import org.apache.spark.{Logging, SparkException}
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.ml.feature.Instance
 import org.apache.spark.ml.optim.WeightedLeastSquares
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.PredictorParams
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.param.shared._
@@ -61,11 +61,13 @@ private[regression] trait LinearRegressionParams extends PredictorParams
  *  - L1 (Lasso)
  *  - L2 + L1 (elastic net)
  */
+@Since("1.3.0")
 @Experimental
-class LinearRegression(override val uid: String)
+class LinearRegression @Since("1.3.0") (@Since("1.3.0") override val uid: String)
   extends Regressor[Vector, LinearRegression, LinearRegressionModel]
   with LinearRegressionParams with Logging {
 
+  @Since("1.4.0")
   def this() = this(Identifiable.randomUID("linReg"))
 
   /**
@@ -73,6 +75,7 @@ class LinearRegression(override val uid: String)
    * Default is 0.0.
    * @group setParam
    */
+  @Since("1.3.0")
   def setRegParam(value: Double): this.type = set(regParam, value)
   setDefault(regParam -> 0.0)
 
@@ -81,6 +84,7 @@ class LinearRegression(override val uid: String)
    * Default is true.
    * @group setParam
    */
+  @Since("1.5.0")
   def setFitIntercept(value: Boolean): this.type = set(fitIntercept, value)
   setDefault(fitIntercept -> true)
 
@@ -93,6 +97,7 @@ class LinearRegression(override val uid: String)
    * Default is true.
    * @group setParam
    */
+  @Since("1.5.0")
   def setStandardization(value: Boolean): this.type = set(standardization, value)
   setDefault(standardization -> true)
 
@@ -103,6 +108,7 @@ class LinearRegression(override val uid: String)
    * Default is 0.0 which is an L2 penalty.
    * @group setParam
    */
+  @Since("1.4.0")
   def setElasticNetParam(value: Double): this.type = set(elasticNetParam, value)
   setDefault(elasticNetParam -> 0.0)
 
@@ -111,6 +117,7 @@ class LinearRegression(override val uid: String)
    * Default is 100.
    * @group setParam
    */
+  @Since("1.3.0")
   def setMaxIter(value: Int): this.type = set(maxIter, value)
   setDefault(maxIter -> 100)
 
@@ -120,6 +127,7 @@ class LinearRegression(override val uid: String)
    * Default is 1E-6.
    * @group setParam
    */
+  @Since("1.4.0")
   def setTol(value: Double): this.type = set(tol, value)
   setDefault(tol -> 1E-6)
 
@@ -129,6 +137,7 @@ class LinearRegression(override val uid: String)
    * Default is empty, so all instances have weight one.
    * @group setParam
    */
+  @Since("1.6.0")
   def setWeightCol(value: String): this.type = set(weightCol, value)
   setDefault(weightCol -> "")
 
@@ -139,6 +148,7 @@ class LinearRegression(override val uid: String)
    * selected automatically.
    * @group setParam
    */
+  @Since("1.6.0")
   def setSolver(value: String): this.type = set(solver, value)
   setDefault(solver -> "auto")
 
@@ -329,6 +339,7 @@ class LinearRegression(override val uid: String)
     model.setSummary(trainingSummary)
   }
 
+  @Since("1.4.0")
   override def copy(extra: ParamMap): LinearRegression = defaultCopy(extra)
 }
 
@@ -336,6 +347,7 @@ class LinearRegression(override val uid: String)
  * :: Experimental ::
  * Model produced by [[LinearRegression]].
  */
+@Since("1.3.0")
 @Experimental
 class LinearRegressionModel private[ml] (
     override val uid: String,
@@ -355,6 +367,7 @@ class LinearRegressionModel private[ml] (
    * Gets summary (e.g. residuals, mse, r-squared ) of model on training set. An exception is
    * thrown if `trainingSummary == None`.
    */
+  @Since("1.5.0")
   def summary: LinearRegressionTrainingSummary = trainingSummary match {
     case Some(summ) => summ
     case None =>
@@ -369,6 +382,7 @@ class LinearRegressionModel private[ml] (
   }
 
   /** Indicates whether a training summary exists for this model instance. */
+  @Since("1.5.0")
   def hasSummary: Boolean = trainingSummary.isDefined
 
   /**
@@ -402,6 +416,7 @@ class LinearRegressionModel private[ml] (
     dot(features, coefficients) + intercept
   }
 
+  @Since("1.4.0")
   override def copy(extra: ParamMap): LinearRegressionModel = {
     val newModel = copyValues(new LinearRegressionModel(uid, coefficients, intercept), extra)
     if (trainingSummary.isDefined) newModel.setSummary(trainingSummary.get)
@@ -416,6 +431,7 @@ class LinearRegressionModel private[ml] (
  * @param predictions predictions outputted by the model's `transform` method.
  * @param objectiveHistory objective function (scaled loss + regularization) at each iteration.
  */
+@Since("1.5.0")
 @Experimental
 class LinearRegressionTrainingSummary private[regression] (
     predictions: DataFrame,
@@ -428,6 +444,7 @@ class LinearRegressionTrainingSummary private[regression] (
   extends LinearRegressionSummary(predictions, predictionCol, labelCol, model, diagInvAtWA) {
 
   /** Number of training iterations until termination */
+  @Since("1.5.0")
   val totalIterations = objectiveHistory.length
 
 }
@@ -437,6 +454,7 @@ class LinearRegressionTrainingSummary private[regression] (
  * Linear regression results evaluated on a dataset.
  * @param predictions predictions outputted by the model's `transform` method.
  */
+@Since("1.5.0")
 @Experimental
 class LinearRegressionSummary private[regression] (
     @transient val predictions: DataFrame,
@@ -455,33 +473,39 @@ class LinearRegressionSummary private[regression] (
    * explainedVariance = 1 - variance(y - \hat{y}) / variance(y)
    * Reference: [[http://en.wikipedia.org/wiki/Explained_variation]]
    */
+  @Since("1.5.0")
   val explainedVariance: Double = metrics.explainedVariance
 
   /**
    * Returns the mean absolute error, which is a risk function corresponding to the
    * expected value of the absolute error loss or l1-norm loss.
    */
+  @Since("1.5.0")
   val meanAbsoluteError: Double = metrics.meanAbsoluteError
 
   /**
    * Returns the mean squared error, which is a risk function corresponding to the
    * expected value of the squared error loss or quadratic loss.
    */
+  @Since("1.5.0")
   val meanSquaredError: Double = metrics.meanSquaredError
 
   /**
    * Returns the root mean squared error, which is defined as the square root of
    * the mean squared error.
    */
+  @Since("1.5.0")
   val rootMeanSquaredError: Double = metrics.rootMeanSquaredError
 
   /**
    * Returns R^2^, the coefficient of determination.
    * Reference: [[http://en.wikipedia.org/wiki/Coefficient_of_determination]]
    */
+  @Since("1.5.0")
   val r2: Double = metrics.r2
 
   /** Residuals (label - predicted value) */
+  @Since("1.5.0")
   @transient lazy val residuals: DataFrame = {
     val t = udf { (pred: Double, label: Double) => label - pred }
     predictions.select(t(col(predictionCol), col(labelCol)).as("residuals"))
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala
index 64fc17247cce6..71e40b513ee0a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.ml.regression
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.{PredictionModel, Predictor}
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.tree.{DecisionTreeModel, RandomForestParams, TreeEnsembleModel, TreeRegressorParams}
@@ -37,44 +37,55 @@ import org.apache.spark.sql.functions._
  * [[http://en.wikipedia.org/wiki/Random_forest  Random Forest]] learning algorithm for regression.
  * It supports both continuous and categorical features.
  */
+@Since("1.4.0")
 @Experimental
-final class RandomForestRegressor(override val uid: String)
+final class RandomForestRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: String)
   extends Predictor[Vector, RandomForestRegressor, RandomForestRegressionModel]
   with RandomForestParams with TreeRegressorParams {
 
+  @Since("1.4.0")
   def this() = this(Identifiable.randomUID("rfr"))
 
   // Override parameter setters from parent trait for Java API compatibility.
 
   // Parameters from TreeRegressorParams:
-
+  @Since("1.4.0")
   override def setMaxDepth(value: Int): this.type = super.setMaxDepth(value)
 
+  @Since("1.4.0")
   override def setMaxBins(value: Int): this.type = super.setMaxBins(value)
 
+  @Since("1.4.0")
   override def setMinInstancesPerNode(value: Int): this.type =
     super.setMinInstancesPerNode(value)
 
+  @Since("1.4.0")
   override def setMinInfoGain(value: Double): this.type = super.setMinInfoGain(value)
 
+  @Since("1.4.0")
   override def setMaxMemoryInMB(value: Int): this.type = super.setMaxMemoryInMB(value)
 
+  @Since("1.4.0")
   override def setCacheNodeIds(value: Boolean): this.type = super.setCacheNodeIds(value)
 
+  @Since("1.4.0")
   override def setCheckpointInterval(value: Int): this.type = super.setCheckpointInterval(value)
 
+  @Since("1.4.0")
   override def setImpurity(value: String): this.type = super.setImpurity(value)
 
   // Parameters from TreeEnsembleParams:
-
+  @Since("1.4.0")
   override def setSubsamplingRate(value: Double): this.type = super.setSubsamplingRate(value)
 
+  @Since("1.4.0")
   override def setSeed(value: Long): this.type = super.setSeed(value)
 
   // Parameters from RandomForestParams:
-
+  @Since("1.4.0")
   override def setNumTrees(value: Int): this.type = super.setNumTrees(value)
 
+  @Since("1.4.0")
   override def setFeatureSubsetStrategy(value: String): this.type =
     super.setFeatureSubsetStrategy(value)
 
@@ -91,15 +102,19 @@ final class RandomForestRegressor(override val uid: String)
     new RandomForestRegressionModel(trees, numFeatures)
   }
 
+  @Since("1.4.0")
   override def copy(extra: ParamMap): RandomForestRegressor = defaultCopy(extra)
 }
 
+@Since("1.4.0")
 @Experimental
 object RandomForestRegressor {
   /** Accessor for supported impurity settings: variance */
+  @Since("1.4.0")
   final val supportedImpurities: Array[String] = TreeRegressorParams.supportedImpurities
 
   /** Accessor for supported featureSubsetStrategy settings: auto, all, onethird, sqrt, log2 */
+  @Since("1.4.0")
   final val supportedFeatureSubsetStrategies: Array[String] =
     RandomForestParams.supportedFeatureSubsetStrategies
 }
@@ -111,6 +126,7 @@ object RandomForestRegressor {
  * @param _trees  Decision trees in the ensemble.
  * @param numFeatures  Number of features used by this model
  */
+@Since("1.4.0")
 @Experimental
 final class RandomForestRegressionModel private[ml] (
     override val uid: String,
@@ -128,11 +144,13 @@ final class RandomForestRegressionModel private[ml] (
   private[ml] def this(trees: Array[DecisionTreeRegressionModel], numFeatures: Int) =
     this(Identifiable.randomUID("rfr"), trees, numFeatures)
 
+  @Since("1.4.0")
   override def trees: Array[DecisionTreeModel] = _trees.asInstanceOf[Array[DecisionTreeModel]]
 
   // Note: We may add support for weights (based on tree performance) later on.
   private lazy val _treeWeights: Array[Double] = Array.fill[Double](numTrees)(1.0)
 
+  @Since("1.4.0")
   override def treeWeights: Array[Double] = _treeWeights
 
   override protected def transformImpl(dataset: DataFrame): DataFrame = {
@@ -150,10 +168,12 @@ final class RandomForestRegressionModel private[ml] (
     _trees.map(_.rootNode.predictImpl(features).prediction).sum / numTrees
   }
 
+  @Since("1.4.0")
   override def copy(extra: ParamMap): RandomForestRegressionModel = {
     copyValues(new RandomForestRegressionModel(uid, _trees, numFeatures), extra).setParent(parent)
   }
 
+  @Since("1.4.0")
   override def toString: String = {
     s"RandomForestRegressionModel (uid=$uid) with $numTrees trees"
   }

From 14ee0f5726f96e2c4c28ac328d43fd85a0630b48 Mon Sep 17 00:00:00 2001
From: Travis Hegner <thegner@trilliumit.com>
Date: Thu, 5 Nov 2015 12:35:23 -0800
Subject: [PATCH 077/518] [SPARK-10648] Oracle dialect to handle nonspecific
 numeric types

This is the alternative/agreed upon solution to PR #8780.

Creating an OracleDialect to handle the nonspecific numeric types that can be defined in oracle.

Author: Travis Hegner <thegner@trilliumit.com>

Closes #9495 from travishegner/OracleDialect.
---
 .../apache/spark/sql/jdbc/JdbcDialects.scala  | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
index 88ae83957a708..f9a6a09b6270d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
@@ -139,6 +139,7 @@ object JdbcDialects {
   registerDialect(DB2Dialect)
   registerDialect(MsSqlServerDialect)
   registerDialect(DerbyDialect)
+  registerDialect(OracleDialect)
 
 
   /**
@@ -315,3 +316,27 @@ case object DerbyDialect extends JdbcDialect {
 
 }
 
+/**
+ * :: DeveloperApi ::
+ * Default Oracle dialect, mapping a nonspecific numeric type to a general decimal type.
+ */
+@DeveloperApi
+case object OracleDialect extends JdbcDialect {
+  override def canHandle(url: String): Boolean = url.startsWith("jdbc:oracle")
+  override def getCatalystType(
+      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
+    // Handle NUMBER fields that have no precision/scale in special way
+    // because JDBC ResultSetMetaData converts this to 0 procision and -127 scale
+    // For more details, please see
+    // https://github.com/apache/spark/pull/8780#issuecomment-145598968
+    // and
+    // https://github.com/apache/spark/pull/8780#issuecomment-144541760
+    if (sqlType == Types.NUMERIC && size == 0) {
+      // This is sub-optimal as we have to pick a precision/scale in advance whereas the data
+      //  in Oracle is allowed to have different precision/scale for each value.
+      Some(DecimalType(DecimalType.MAX_PRECISION, 10))
+    } else {
+      None
+    }
+  }
+}

From 8a5314efd19fb8f8a194a373fd994b954cc1fd47 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 5 Nov 2015 13:34:36 -0800
Subject: [PATCH 078/518] [SPARK-11532][SQL] Remove implicit conversion from
 Expression to Column

Author: Reynold Xin <rxin@databricks.com>

Closes #9500 from rxin/SPARK-11532.
---
 .../scala/org/apache/spark/sql/Column.scala   | 118 ++++++++++--------
 1 file changed, 66 insertions(+), 52 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index c73f696962de5..c32c93897ce0b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -68,7 +68,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
   })
 
   /** Creates a column based on the given expression. */
-  implicit private def exprToColumn(newExpr: Expression): Column = new Column(newExpr)
+  private def withExpr(newExpr: Expression): Column = new Column(newExpr)
 
   override def toString: String = expr.prettyString
 
@@ -99,7 +99,9 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.4.0
    */
-  def apply(extraction: Any): Column = UnresolvedExtractValue(expr, lit(extraction).expr)
+  def apply(extraction: Any): Column = withExpr {
+    UnresolvedExtractValue(expr, lit(extraction).expr)
+  }
 
   /**
    * Unary minus, i.e. negate the expression.
@@ -115,7 +117,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def unary_- : Column = UnaryMinus(expr)
+  def unary_- : Column = withExpr { UnaryMinus(expr) }
 
   /**
    * Inversion of boolean expression, i.e. NOT.
@@ -131,7 +133,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def unary_! : Column = Not(expr)
+  def unary_! : Column = withExpr { Not(expr) }
 
   /**
    * Equality test.
@@ -147,7 +149,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def === (other: Any): Column = {
+  def === (other: Any): Column = withExpr {
     val right = lit(other).expr
     if (this.expr == right) {
       logWarning(
@@ -188,7 +190,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def !== (other: Any): Column = Not(EqualTo(expr, lit(other).expr))
+  def !== (other: Any): Column = withExpr{ Not(EqualTo(expr, lit(other).expr)) }
 
   /**
    * Inequality test.
@@ -205,7 +207,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group java_expr_ops
    * @since 1.3.0
    */
-  def notEqual(other: Any): Column = Not(EqualTo(expr, lit(other).expr))
+  def notEqual(other: Any): Column = withExpr { Not(EqualTo(expr, lit(other).expr)) }
 
   /**
    * Greater than.
@@ -221,7 +223,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def > (other: Any): Column = GreaterThan(expr, lit(other).expr)
+  def > (other: Any): Column = withExpr { GreaterThan(expr, lit(other).expr) }
 
   /**
    * Greater than.
@@ -252,7 +254,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def < (other: Any): Column = LessThan(expr, lit(other).expr)
+  def < (other: Any): Column = withExpr { LessThan(expr, lit(other).expr) }
 
   /**
    * Less than.
@@ -282,7 +284,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def <= (other: Any): Column = LessThanOrEqual(expr, lit(other).expr)
+  def <= (other: Any): Column = withExpr { LessThanOrEqual(expr, lit(other).expr) }
 
   /**
    * Less than or equal to.
@@ -312,7 +314,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def >= (other: Any): Column = GreaterThanOrEqual(expr, lit(other).expr)
+  def >= (other: Any): Column = withExpr { GreaterThanOrEqual(expr, lit(other).expr) }
 
   /**
    * Greater than or equal to an expression.
@@ -335,7 +337,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def <=> (other: Any): Column = EqualNullSafe(expr, lit(other).expr)
+  def <=> (other: Any): Column = withExpr { EqualNullSafe(expr, lit(other).expr) }
 
   /**
    * Equality test that is safe for null values.
@@ -368,7 +370,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    */
   def when(condition: Column, value: Any): Column = this.expr match {
     case CaseWhen(branches: Seq[Expression]) =>
-      CaseWhen(branches ++ Seq(lit(condition).expr, lit(value).expr))
+      withExpr { CaseWhen(branches ++ Seq(lit(condition).expr, lit(value).expr)) }
     case _ =>
       throw new IllegalArgumentException(
         "when() can only be applied on a Column previously generated by when() function")
@@ -398,7 +400,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
   def otherwise(value: Any): Column = this.expr match {
     case CaseWhen(branches: Seq[Expression]) =>
       if (branches.size % 2 == 0) {
-        CaseWhen(branches :+ lit(value).expr)
+        withExpr { CaseWhen(branches :+ lit(value).expr) }
       } else {
         throw new IllegalArgumentException(
           "otherwise() can only be applied once on a Column previously generated by when()")
@@ -424,7 +426,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.5.0
    */
-  def isNaN: Column = IsNaN(expr)
+  def isNaN: Column = withExpr { IsNaN(expr) }
 
   /**
    * True if the current expression is null.
@@ -432,7 +434,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def isNull: Column = IsNull(expr)
+  def isNull: Column = withExpr { IsNull(expr) }
 
   /**
    * True if the current expression is NOT null.
@@ -440,7 +442,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def isNotNull: Column = IsNotNull(expr)
+  def isNotNull: Column = withExpr { IsNotNull(expr) }
 
   /**
    * Boolean OR.
@@ -455,7 +457,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def || (other: Any): Column = Or(expr, lit(other).expr)
+  def || (other: Any): Column = withExpr { Or(expr, lit(other).expr) }
 
   /**
    * Boolean OR.
@@ -485,7 +487,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def && (other: Any): Column = And(expr, lit(other).expr)
+  def && (other: Any): Column = withExpr { And(expr, lit(other).expr) }
 
   /**
    * Boolean AND.
@@ -515,7 +517,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def + (other: Any): Column = Add(expr, lit(other).expr)
+  def + (other: Any): Column = withExpr { Add(expr, lit(other).expr) }
 
   /**
    * Sum of this expression and another expression.
@@ -545,7 +547,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def - (other: Any): Column = Subtract(expr, lit(other).expr)
+  def - (other: Any): Column = withExpr { Subtract(expr, lit(other).expr) }
 
   /**
    * Subtraction. Subtract the other expression from this expression.
@@ -575,7 +577,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def * (other: Any): Column = Multiply(expr, lit(other).expr)
+  def * (other: Any): Column = withExpr { Multiply(expr, lit(other).expr) }
 
   /**
    * Multiplication of this expression and another expression.
@@ -605,7 +607,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def / (other: Any): Column = Divide(expr, lit(other).expr)
+  def / (other: Any): Column = withExpr { Divide(expr, lit(other).expr) }
 
   /**
    * Division this expression by another expression.
@@ -628,7 +630,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def % (other: Any): Column = Remainder(expr, lit(other).expr)
+  def % (other: Any): Column = withExpr { Remainder(expr, lit(other).expr) }
 
   /**
    * Modulo (a.k.a. remainder) expression.
@@ -657,7 +659,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @since 1.5.0
    */
   @scala.annotation.varargs
-  def isin(list: Any*): Column = In(expr, list.map(lit(_).expr))
+  def isin(list: Any*): Column = withExpr { In(expr, list.map(lit(_).expr)) }
 
   /**
    * SQL like expression.
@@ -665,7 +667,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def like(literal: String): Column = Like(expr, lit(literal).expr)
+  def like(literal: String): Column = withExpr { Like(expr, lit(literal).expr) }
 
   /**
    * SQL RLIKE expression (LIKE with Regex).
@@ -673,7 +675,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def rlike(literal: String): Column = RLike(expr, lit(literal).expr)
+  def rlike(literal: String): Column = withExpr { RLike(expr, lit(literal).expr) }
 
   /**
    * An expression that gets an item at position `ordinal` out of an array,
@@ -682,7 +684,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def getItem(key: Any): Column = UnresolvedExtractValue(expr, Literal(key))
+  def getItem(key: Any): Column = withExpr { UnresolvedExtractValue(expr, Literal(key)) }
 
   /**
    * An expression that gets a field by name in a [[StructType]].
@@ -690,7 +692,9 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def getField(fieldName: String): Column = UnresolvedExtractValue(expr, Literal(fieldName))
+  def getField(fieldName: String): Column = withExpr {
+    UnresolvedExtractValue(expr, Literal(fieldName))
+  }
 
   /**
    * An expression that returns a substring.
@@ -700,7 +704,9 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def substr(startPos: Column, len: Column): Column = Substring(expr, startPos.expr, len.expr)
+  def substr(startPos: Column, len: Column): Column = withExpr {
+    Substring(expr, startPos.expr, len.expr)
+  }
 
   /**
    * An expression that returns a substring.
@@ -710,7 +716,9 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def substr(startPos: Int, len: Int): Column = Substring(expr, lit(startPos).expr, lit(len).expr)
+  def substr(startPos: Int, len: Int): Column = withExpr {
+    Substring(expr, lit(startPos).expr, lit(len).expr)
+  }
 
   /**
    * Contains the other element.
@@ -718,7 +726,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def contains(other: Any): Column = Contains(expr, lit(other).expr)
+  def contains(other: Any): Column = withExpr { Contains(expr, lit(other).expr) }
 
   /**
    * String starts with.
@@ -726,7 +734,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def startsWith(other: Column): Column = StartsWith(expr, lit(other).expr)
+  def startsWith(other: Column): Column = withExpr { StartsWith(expr, lit(other).expr) }
 
   /**
    * String starts with another string literal.
@@ -742,7 +750,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def endsWith(other: Column): Column = EndsWith(expr, lit(other).expr)
+  def endsWith(other: Column): Column = withExpr { EndsWith(expr, lit(other).expr) }
 
   /**
    * String ends with another string literal.
@@ -777,9 +785,11 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def as(alias: String): Column = expr match {
-    case ne: NamedExpression => Alias(expr, alias)(explicitMetadata = Some(ne.metadata))
-    case other => Alias(other, alias)()
+  def as(alias: String): Column = withExpr {
+    expr match {
+      case ne: NamedExpression => Alias(expr, alias)(explicitMetadata = Some(ne.metadata))
+      case other => Alias(other, alias)()
+    }
   }
 
   /**
@@ -792,7 +802,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.4.0
    */
-  def as(aliases: Seq[String]): Column = MultiAlias(expr, aliases)
+  def as(aliases: Seq[String]): Column = withExpr { MultiAlias(expr, aliases) }
 
   /**
    * Assigns the given aliases to the results of a table generating function.
@@ -804,7 +814,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.4.0
    */
-  def as(aliases: Array[String]): Column = MultiAlias(expr, aliases)
+  def as(aliases: Array[String]): Column = withExpr { MultiAlias(expr, aliases) }
 
   /**
    * Gives the column an alias.
@@ -819,9 +829,11 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def as(alias: Symbol): Column = expr match {
-    case ne: NamedExpression => Alias(expr, alias.name)(explicitMetadata = Some(ne.metadata))
-    case other => Alias(other, alias.name)()
+  def as(alias: Symbol): Column = withExpr {
+    expr match {
+      case ne: NamedExpression => Alias(expr, alias.name)(explicitMetadata = Some(ne.metadata))
+      case other => Alias(other, alias.name)()
+    }
   }
 
   /**
@@ -834,7 +846,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def as(alias: String, metadata: Metadata): Column = {
+  def as(alias: String, metadata: Metadata): Column = withExpr {
     Alias(expr, alias)(explicitMetadata = Some(metadata))
   }
 
@@ -852,10 +864,12 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def cast(to: DataType): Column = expr match {
-    // keeps the name of expression if possible when do cast.
-    case ne: NamedExpression => UnresolvedAlias(Cast(expr, to))
-    case _ => Cast(expr, to)
+  def cast(to: DataType): Column = withExpr {
+    expr match {
+      // keeps the name of expression if possible when do cast.
+      case ne: NamedExpression => UnresolvedAlias(Cast(expr, to))
+      case _ => Cast(expr, to)
+    }
   }
 
   /**
@@ -885,7 +899,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def desc: Column = SortOrder(expr, Descending)
+  def desc: Column = withExpr { SortOrder(expr, Descending) }
 
   /**
    * Returns an ordering used in sorting.
@@ -900,7 +914,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.3.0
    */
-  def asc: Column = SortOrder(expr, Ascending)
+  def asc: Column = withExpr { SortOrder(expr, Ascending) }
 
   /**
    * Prints the expression to the console for debugging purpose.
@@ -927,7 +941,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.4.0
    */
-  def bitwiseOR(other: Any): Column = BitwiseOr(expr, lit(other).expr)
+  def bitwiseOR(other: Any): Column = withExpr { BitwiseOr(expr, lit(other).expr) }
 
   /**
    * Compute bitwise AND of this expression with another expression.
@@ -938,7 +952,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.4.0
    */
-  def bitwiseAND(other: Any): Column = BitwiseAnd(expr, lit(other).expr)
+  def bitwiseAND(other: Any): Column = withExpr { BitwiseAnd(expr, lit(other).expr) }
 
   /**
    * Compute bitwise XOR of this expression with another expression.
@@ -949,7 +963,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * @group expr_ops
    * @since 1.4.0
    */
-  def bitwiseXOR(other: Any): Column = BitwiseXor(expr, lit(other).expr)
+  def bitwiseXOR(other: Any): Column = withExpr { BitwiseXor(expr, lit(other).expr) }
 
   /**
    * Define a windowing column.

From b9455d1f1810e1e3f472014f665ad3ad3122bcc0 Mon Sep 17 00:00:00 2001
From: adrian555 <wzhuang@us.ibm.com>
Date: Thu, 5 Nov 2015 14:47:38 -0800
Subject: [PATCH 079/518] [SPARK-11260][SPARKR] with() function support

Author: adrian555 <wzhuang@us.ibm.com>
Author: Adrian Zhuang <adrian555@users.noreply.github.com>

Closes #9443 from adrian555/with.
---
 R/pkg/NAMESPACE                  |  1 +
 R/pkg/R/DataFrame.R              | 30 ++++++++++++++++++++++++------
 R/pkg/R/generics.R               |  4 ++++
 R/pkg/R/utils.R                  | 13 +++++++++++++
 R/pkg/inst/tests/test_sparkSQL.R |  9 +++++++++
 5 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/R/pkg/NAMESPACE b/R/pkg/NAMESPACE
index cd9537a2655f0..56b8ed0bf271b 100644
--- a/R/pkg/NAMESPACE
+++ b/R/pkg/NAMESPACE
@@ -83,6 +83,7 @@ exportMethods("arrange",
               "unique",
               "unpersist",
               "where",
+              "with",
               "withColumn",
               "withColumnRenamed",
               "write.df")
diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index df5bc8137187b..44ce9414da5cf 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -2126,11 +2126,29 @@ setMethod("as.data.frame",
 setMethod("attach",
           signature(what = "DataFrame"),
           function(what, pos = 2, name = deparse(substitute(what)), warn.conflicts = TRUE) {
-            cols <- columns(what)
-            stopifnot(length(cols) > 0)
-            newEnv <- new.env()
-            for (i in 1:length(cols)) {
-              assign(x = cols[i], value = what[, cols[i]], envir = newEnv)
-            }
+            newEnv <- assignNewEnv(what)
             attach(newEnv, pos = pos, name = name, warn.conflicts = warn.conflicts)
           })
+
+#' Evaluate a R expression in an environment constructed from a DataFrame
+#' with() allows access to columns of a DataFrame by simply referring to
+#' their name. It appends every column of a DataFrame into a new
+#' environment. Then, the given expression is evaluated in this new
+#' environment.
+#'
+#' @rdname with
+#' @title Evaluate a R expression in an environment constructed from a DataFrame
+#' @param data (DataFrame) DataFrame to use for constructing an environment.
+#' @param expr (expression) Expression to evaluate.
+#' @param ... arguments to be passed to future methods.
+#' @examples
+#' \dontrun{
+#' with(irisDf, nrow(Sepal_Width))
+#' }
+#' @seealso \link{attach}
+setMethod("with",
+          signature(data = "DataFrame"),
+          function(data, expr, ...) {
+            newEnv <- assignNewEnv(data)
+            eval(substitute(expr), envir = newEnv, enclos = newEnv)
+          })
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 0b35340e48e42..083d37fee28a4 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -1043,3 +1043,7 @@ setGeneric("as.data.frame")
 #' @rdname attach
 #' @export
 setGeneric("attach")
+
+#' @rdname with
+#' @export
+setGeneric("with")
diff --git a/R/pkg/R/utils.R b/R/pkg/R/utils.R
index 0b9e2957fe9a5..db3b2c4bbd799 100644
--- a/R/pkg/R/utils.R
+++ b/R/pkg/R/utils.R
@@ -623,3 +623,16 @@ convertNamedListToEnv <- function(namedList) {
   }
   env
 }
+
+# Assign a new environment for attach() and with() methods
+assignNewEnv <- function(data) {
+  stopifnot(class(data) == "DataFrame")
+  cols <- columns(data)
+  stopifnot(length(cols) > 0)
+
+  env <- new.env()
+  for (i in 1:length(cols)) {
+    assign(x = cols[i], value = data[, cols[i]], envir = env)
+  }
+  env
+}
\ No newline at end of file
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index b4a4d03b2643b..816315b1e4e13 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -1494,6 +1494,15 @@ test_that("attach() on a DataFrame", {
   expect_error(age)
 })
 
+test_that("with() on a DataFrame", {
+  df <- createDataFrame(sqlContext, iris)
+  expect_error(Sepal_Length)
+  sum1 <- with(df, list(summary(Sepal_Length), summary(Sepal_Width)))
+  expect_equal(collect(sum1[[1]])[1, "Sepal_Length"], "150")
+  sum2 <- with(df, distinct(Sepal_Length))
+  expect_equal(nrow(sum2), 35)
+})
+
 unlink(parquetPath)
 unlink(jsonPath)
 unlink(jsonPathNa)

From d9e30c59cede7f57786bb19e64ba422eda43bdcb Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Thu, 5 Nov 2015 14:53:16 -0800
Subject: [PATCH 080/518] [SPARK-10656][SQL] completely support special chars
 in DataFrame

the main problem is: we interpret column name with special handling of `.` for DataFrame. This enables us to write something like `df("a.b")` to get the field `b` of `a`. However, we don't need this feature in `DataFrame.apply("*")` or `DataFrame.withColumnRenamed`. In these 2 cases, the column name is the final name already, we don't need extra process to interpret it.

The solution is simple, use `queryExecution.analyzed.output` to get resolved column directly, instead of using `DataFrame.resolve`.

close https://github.com/apache/spark/pull/8811

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9462 from cloud-fan/special-chars.
---
 .../scala/org/apache/spark/sql/DataFrame.scala   | 16 ++++++++++------
 .../org/apache/spark/sql/DataFrameSuite.scala    |  6 ++++++
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 6336dee7be6a3..f2d4db5550273 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -698,7 +698,7 @@ class DataFrame private[sql](
    */
   def col(colName: String): Column = colName match {
     case "*" =>
-      Column(ResolvedStar(schema.fieldNames.map(resolve)))
+      Column(ResolvedStar(queryExecution.analyzed.output))
     case _ =>
       val expr = resolve(colName)
       Column(expr)
@@ -1259,13 +1259,17 @@ class DataFrame private[sql](
    */
   def withColumnRenamed(existingName: String, newName: String): DataFrame = {
     val resolver = sqlContext.analyzer.resolver
-    val shouldRename = schema.exists(f => resolver(f.name, existingName))
+    val output = queryExecution.analyzed.output
+    val shouldRename = output.exists(f => resolver(f.name, existingName))
     if (shouldRename) {
-      val colNames = schema.map { field =>
-        val name = field.name
-        if (resolver(name, existingName)) Column(name).as(newName) else Column(name)
+      val columns = output.map { col =>
+        if (resolver(col.name, existingName)) {
+          Column(col).as(newName)
+        } else {
+          Column(col)
+        }
       }
-      select(colNames : _*)
+      select(columns : _*)
     } else {
       this
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index 84a616d0b9081..f3a7aa280367a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -1128,4 +1128,10 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
       }
     }
   }
+
+  test("SPARK-10656: completely support special chars") {
+    val df = Seq(1 -> "a").toDF("i_$.a", "d^'a.")
+    checkAnswer(df.select(df("*")), Row(1, "a"))
+    checkAnswer(df.withColumnRenamed("d^'a.", "a"), Row(1, "a"))
+  }
 }

From b6974f8fed1726a381636e996834111a8e7ced8d Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 5 Nov 2015 15:34:05 -0800
Subject: [PATCH 081/518] [SPARK-11536][SQL] Remove the internal implicit
 conversion from Expression to Column in functions.scala

Author: Reynold Xin <rxin@databricks.com>

Closes #9505 from rxin/SPARK-11536.
---
 .../org/apache/spark/sql/functions.scala      | 580 +++++++++---------
 1 file changed, 299 insertions(+), 281 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index c70c965a9b04c..04627589886a8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -51,7 +51,7 @@ import org.apache.spark.util.Utils
 object functions {
 // scalastyle:on
 
-  private[this] implicit def toColumn(expr: Expression): Column = Column(expr)
+  private def withExpr(expr: Expression): Column = Column(expr)
 
   /**
    * Returns a [[Column]] based on the given column name.
@@ -128,7 +128,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def approxCountDistinct(e: Column): Column = ApproxCountDistinct(e.expr)
+  def approxCountDistinct(e: Column): Column = withExpr { ApproxCountDistinct(e.expr) }
 
   /**
    * Aggregate function: returns the approximate number of distinct items in a group.
@@ -144,7 +144,9 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def approxCountDistinct(e: Column, rsd: Double): Column = ApproxCountDistinct(e.expr, rsd)
+  def approxCountDistinct(e: Column, rsd: Double): Column = withExpr {
+    ApproxCountDistinct(e.expr, rsd)
+  }
 
   /**
    * Aggregate function: returns the approximate number of distinct items in a group.
@@ -162,7 +164,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def avg(e: Column): Column = Average(e.expr)
+  def avg(e: Column): Column = withExpr { Average(e.expr) }
 
   /**
    * Aggregate function: returns the average of the values in a group.
@@ -178,8 +180,9 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def corr(column1: Column, column2: Column): Column =
+  def corr(column1: Column, column2: Column): Column = withExpr {
     Corr(column1.expr, column2.expr)
+  }
 
   /**
    * Aggregate function: returns the Pearson Correlation Coefficient for two columns.
@@ -187,8 +190,9 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def corr(columnName1: String, columnName2: String): Column =
+  def corr(columnName1: String, columnName2: String): Column = {
     corr(Column(columnName1), Column(columnName2))
+  }
 
   /**
    * Aggregate function: returns the number of items in a group.
@@ -196,10 +200,12 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def count(e: Column): Column = e.expr match {
-    // Turn count(*) into count(1)
-    case s: Star => Count(Literal(1))
-    case _ => Count(e.expr)
+  def count(e: Column): Column = withExpr {
+    e.expr match {
+      // Turn count(*) into count(1)
+      case s: Star => Count(Literal(1))
+      case _ => Count(e.expr)
+    }
   }
 
   /**
@@ -217,8 +223,9 @@ object functions {
    * @since 1.3.0
    */
   @scala.annotation.varargs
-  def countDistinct(expr: Column, exprs: Column*): Column =
+  def countDistinct(expr: Column, exprs: Column*): Column = withExpr {
     CountDistinct((expr +: exprs).map(_.expr))
+  }
 
   /**
    * Aggregate function: returns the number of distinct items in a group.
@@ -236,7 +243,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def first(e: Column): Column = First(e.expr)
+  def first(e: Column): Column = withExpr { First(e.expr) }
 
   /**
    * Aggregate function: returns the first value of a column in a group.
@@ -252,7 +259,7 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def kurtosis(e: Column): Column = Kurtosis(e.expr)
+  def kurtosis(e: Column): Column = withExpr { Kurtosis(e.expr) }
 
   /**
    * Aggregate function: returns the last value in a group.
@@ -260,7 +267,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def last(e: Column): Column = Last(e.expr)
+  def last(e: Column): Column = withExpr { Last(e.expr) }
 
   /**
    * Aggregate function: returns the last value of the column in a group.
@@ -276,7 +283,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def max(e: Column): Column = Max(e.expr)
+  def max(e: Column): Column = withExpr { Max(e.expr) }
 
   /**
    * Aggregate function: returns the maximum value of the column in a group.
@@ -310,7 +317,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def min(e: Column): Column = Min(e.expr)
+  def min(e: Column): Column = withExpr { Min(e.expr) }
 
   /**
    * Aggregate function: returns the minimum value of the column in a group.
@@ -326,7 +333,7 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def skewness(e: Column): Column = Skewness(e.expr)
+  def skewness(e: Column): Column = withExpr { Skewness(e.expr) }
 
   /**
    * Aggregate function: alias for [[stddev_samp]].
@@ -334,7 +341,7 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def stddev(e: Column): Column = StddevSamp(e.expr)
+  def stddev(e: Column): Column = withExpr { StddevSamp(e.expr) }
 
   /**
    * Aggregate function: returns the unbiased sample standard deviation of
@@ -343,7 +350,7 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def stddev_samp(e: Column): Column = StddevSamp(e.expr)
+  def stddev_samp(e: Column): Column = withExpr { StddevSamp(e.expr) }
 
   /**
    * Aggregate function: returns the population standard deviation of
@@ -352,7 +359,7 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def stddev_pop(e: Column): Column = StddevPop(e.expr)
+  def stddev_pop(e: Column): Column = withExpr { StddevPop(e.expr) }
 
   /**
    * Aggregate function: returns the sum of all values in the expression.
@@ -360,7 +367,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def sum(e: Column): Column = Sum(e.expr)
+  def sum(e: Column): Column = withExpr { Sum(e.expr) }
 
   /**
    * Aggregate function: returns the sum of all values in the given column.
@@ -376,7 +383,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def sumDistinct(e: Column): Column = SumDistinct(e.expr)
+  def sumDistinct(e: Column): Column = withExpr { SumDistinct(e.expr) }
 
   /**
    * Aggregate function: returns the sum of distinct values in the expression.
@@ -392,7 +399,7 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def variance(e: Column): Column = VarianceSamp(e.expr)
+  def variance(e: Column): Column = withExpr { VarianceSamp(e.expr) }
 
   /**
    * Aggregate function: returns the unbiased variance of the values in a group.
@@ -400,7 +407,7 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def var_samp(e: Column): Column = VarianceSamp(e.expr)
+  def var_samp(e: Column): Column = withExpr { VarianceSamp(e.expr) }
 
   /**
    * Aggregate function: returns the population variance of the values in a group.
@@ -408,7 +415,7 @@ object functions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def var_pop(e: Column): Column = VariancePop(e.expr)
+  def var_pop(e: Column): Column = withExpr { VariancePop(e.expr) }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
   // Window functions
@@ -429,9 +436,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def cumeDist(): Column = {
-    UnresolvedWindowFunction("cume_dist", Nil)
-  }
+  def cumeDist(): Column = withExpr { UnresolvedWindowFunction("cume_dist", Nil) }
 
   /**
    * Window function: returns the rank of rows within a window partition, without any gaps.
@@ -446,9 +451,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def denseRank(): Column = {
-    UnresolvedWindowFunction("dense_rank", Nil)
-  }
+  def denseRank(): Column = withExpr { UnresolvedWindowFunction("dense_rank", Nil) }
 
   /**
    * Window function: returns the value that is `offset` rows before the current row, and
@@ -460,9 +463,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def lag(e: Column, offset: Int): Column = {
-    lag(e, offset, null)
-  }
+  def lag(e: Column, offset: Int): Column = lag(e, offset, null)
 
   /**
    * Window function: returns the value that is `offset` rows before the current row, and
@@ -474,9 +475,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def lag(columnName: String, offset: Int): Column = {
-    lag(columnName, offset, null)
-  }
+  def lag(columnName: String, offset: Int): Column = lag(columnName, offset, null)
 
   /**
    * Window function: returns the value that is `offset` rows before the current row, and
@@ -502,7 +501,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def lag(e: Column, offset: Int, defaultValue: Any): Column = {
+  def lag(e: Column, offset: Int, defaultValue: Any): Column = withExpr {
     UnresolvedWindowFunction("lag", e.expr :: Literal(offset) :: Literal(defaultValue) :: Nil)
   }
 
@@ -516,9 +515,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def lead(columnName: String, offset: Int): Column = {
-    lead(columnName, offset, null)
-  }
+  def lead(columnName: String, offset: Int): Column = { lead(columnName, offset, null) }
 
   /**
    * Window function: returns the value that is `offset` rows after the current row, and
@@ -530,9 +527,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def lead(e: Column, offset: Int): Column = {
-    lead(e, offset, null)
-  }
+  def lead(e: Column, offset: Int): Column = { lead(e, offset, null) }
 
   /**
    * Window function: returns the value that is `offset` rows after the current row, and
@@ -558,7 +553,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def lead(e: Column, offset: Int, defaultValue: Any): Column = {
+  def lead(e: Column, offset: Int, defaultValue: Any): Column = withExpr {
     UnresolvedWindowFunction("lead", e.expr :: Literal(offset) :: Literal(defaultValue) :: Nil)
   }
 
@@ -572,9 +567,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def ntile(n: Int): Column = {
-    UnresolvedWindowFunction("ntile", lit(n).expr :: Nil)
-  }
+  def ntile(n: Int): Column = withExpr { UnresolvedWindowFunction("ntile", lit(n).expr :: Nil) }
 
   /**
    * Window function: returns the relative rank (i.e. percentile) of rows within a window partition.
@@ -589,9 +582,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def percentRank(): Column = {
-    UnresolvedWindowFunction("percent_rank", Nil)
-  }
+  def percentRank(): Column = withExpr { UnresolvedWindowFunction("percent_rank", Nil) }
 
   /**
    * Window function: returns the rank of rows within a window partition.
@@ -606,9 +597,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def rank(): Column = {
-    UnresolvedWindowFunction("rank", Nil)
-  }
+  def rank(): Column = withExpr { UnresolvedWindowFunction("rank", Nil) }
 
   /**
    * Window function: returns a sequential number starting at 1 within a window partition.
@@ -618,9 +607,7 @@ object functions {
    * @group window_funcs
    * @since 1.4.0
    */
-  def rowNumber(): Column = {
-    UnresolvedWindowFunction("row_number", Nil)
-  }
+  def rowNumber(): Column = withExpr { UnresolvedWindowFunction("row_number", Nil) }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
   // Non-aggregate functions
@@ -632,7 +619,7 @@ object functions {
    * @group normal_funcs
    * @since 1.3.0
    */
-  def abs(e: Column): Column = Abs(e.expr)
+  def abs(e: Column): Column = withExpr { Abs(e.expr) }
 
   /**
    * Creates a new array column. The input columns must all have the same data type.
@@ -641,7 +628,7 @@ object functions {
    * @since 1.4.0
    */
   @scala.annotation.varargs
-  def array(cols: Column*): Column = CreateArray(cols.map(_.expr))
+  def array(cols: Column*): Column = withExpr { CreateArray(cols.map(_.expr)) }
 
   /**
    * Creates a new array column. The input columns must all have the same data type.
@@ -679,14 +666,14 @@ object functions {
    * @since 1.3.0
    */
   @scala.annotation.varargs
-  def coalesce(e: Column*): Column = Coalesce(e.map(_.expr))
+  def coalesce(e: Column*): Column = withExpr { Coalesce(e.map(_.expr)) }
 
   /**
    * Creates a string column for the file name of the current Spark task.
    *
    * @group normal_funcs
    */
-  def inputFileName(): Column = InputFileName()
+  def inputFileName(): Column = withExpr { InputFileName() }
 
   /**
    * Return true iff the column is NaN.
@@ -694,7 +681,7 @@ object functions {
    * @group normal_funcs
    * @since 1.5.0
    */
-  def isNaN(e: Column): Column = IsNaN(e.expr)
+  def isNaN(e: Column): Column = withExpr { IsNaN(e.expr) }
 
   /**
    * A column expression that generates monotonically increasing 64-bit integers.
@@ -711,7 +698,7 @@ object functions {
    * @group normal_funcs
    * @since 1.4.0
    */
-  def monotonicallyIncreasingId(): Column = MonotonicallyIncreasingID()
+  def monotonicallyIncreasingId(): Column = withExpr { MonotonicallyIncreasingID() }
 
   /**
    * Returns col1 if it is not NaN, or col2 if col1 is NaN.
@@ -721,7 +708,7 @@ object functions {
    * @group normal_funcs
    * @since 1.5.0
    */
-  def nanvl(col1: Column, col2: Column): Column = NaNvl(col1.expr, col2.expr)
+  def nanvl(col1: Column, col2: Column): Column = withExpr { NaNvl(col1.expr, col2.expr) }
 
   /**
    * Unary minus, i.e. negate the expression.
@@ -760,7 +747,7 @@ object functions {
    * @group normal_funcs
    * @since 1.4.0
    */
-  def rand(seed: Long): Column = Rand(seed)
+  def rand(seed: Long): Column = withExpr { Rand(seed) }
 
   /**
    * Generate a random column with i.i.d. samples from U[0.0, 1.0].
@@ -776,7 +763,7 @@ object functions {
    * @group normal_funcs
    * @since 1.4.0
    */
-  def randn(seed: Long): Column = Randn(seed)
+  def randn(seed: Long): Column = withExpr { Randn(seed) }
 
   /**
    * Generate a column with i.i.d. samples from the standard normal distribution.
@@ -794,7 +781,7 @@ object functions {
    * @group normal_funcs
    * @since 1.4.0
    */
-  def sparkPartitionId(): Column = SparkPartitionID()
+  def sparkPartitionId(): Column = withExpr { SparkPartitionID() }
 
   /**
    * Computes the square root of the specified float value.
@@ -802,7 +789,7 @@ object functions {
    * @group math_funcs
    * @since 1.3.0
    */
-  def sqrt(e: Column): Column = Sqrt(e.expr)
+  def sqrt(e: Column): Column = withExpr { Sqrt(e.expr) }
 
   /**
    * Computes the square root of the specified float value.
@@ -823,9 +810,7 @@ object functions {
    * @since 1.4.0
    */
   @scala.annotation.varargs
-  def struct(cols: Column*): Column = {
-    CreateStruct(cols.map(_.expr))
-  }
+  def struct(cols: Column*): Column = withExpr { CreateStruct(cols.map(_.expr)) }
 
   /**
    * Creates a new struct column that composes multiple input columns.
@@ -858,7 +843,7 @@ object functions {
    * @group normal_funcs
    * @since 1.4.0
    */
-  def when(condition: Column, value: Any): Column = {
+  def when(condition: Column, value: Any): Column = withExpr {
     CaseWhen(Seq(condition.expr, lit(value).expr))
   }
 
@@ -868,7 +853,7 @@ object functions {
    * @group normal_funcs
    * @since 1.4.0
    */
-  def bitwiseNOT(e: Column): Column = BitwiseNot(e.expr)
+  def bitwiseNOT(e: Column): Column = withExpr { BitwiseNot(e.expr) }
 
   /**
    * Parses the expression string into the column that it represents, similar to
@@ -893,7 +878,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def acos(e: Column): Column = Acos(e.expr)
+  def acos(e: Column): Column = withExpr { Acos(e.expr) }
 
   /**
    * Computes the cosine inverse of the given column; the returned angle is in the range
@@ -911,7 +896,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def asin(e: Column): Column = Asin(e.expr)
+  def asin(e: Column): Column = withExpr { Asin(e.expr) }
 
   /**
    * Computes the sine inverse of the given column; the returned angle is in the range
@@ -928,7 +913,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def atan(e: Column): Column = Atan(e.expr)
+  def atan(e: Column): Column = withExpr { Atan(e.expr) }
 
   /**
    * Computes the tangent inverse of the given column.
@@ -945,7 +930,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def atan2(l: Column, r: Column): Column = Atan2(l.expr, r.expr)
+  def atan2(l: Column, r: Column): Column = withExpr { Atan2(l.expr, r.expr) }
 
   /**
    * Returns the angle theta from the conversion of rectangular coordinates (x, y) to
@@ -982,7 +967,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def atan2(l: Column, r: Double): Column = atan2(l, lit(r).expr)
+  def atan2(l: Column, r: Double): Column = atan2(l, lit(r))
 
   /**
    * Returns the angle theta from the conversion of rectangular coordinates (x, y) to
@@ -1000,7 +985,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def atan2(l: Double, r: Column): Column = atan2(lit(l).expr, r)
+  def atan2(l: Double, r: Column): Column = atan2(lit(l), r)
 
   /**
    * Returns the angle theta from the conversion of rectangular coordinates (x, y) to
@@ -1018,7 +1003,7 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def bin(e: Column): Column = Bin(e.expr)
+  def bin(e: Column): Column = withExpr { Bin(e.expr) }
 
   /**
    * An expression that returns the string representation of the binary value of the given long
@@ -1035,7 +1020,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def cbrt(e: Column): Column = Cbrt(e.expr)
+  def cbrt(e: Column): Column = withExpr { Cbrt(e.expr) }
 
   /**
    * Computes the cube-root of the given column.
@@ -1051,7 +1036,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def ceil(e: Column): Column = Ceil(e.expr)
+  def ceil(e: Column): Column = withExpr { Ceil(e.expr) }
 
   /**
    * Computes the ceiling of the given column.
@@ -1067,8 +1052,9 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def conv(num: Column, fromBase: Int, toBase: Int): Column =
+  def conv(num: Column, fromBase: Int, toBase: Int): Column = withExpr {
     Conv(num.expr, lit(fromBase).expr, lit(toBase).expr)
+  }
 
   /**
    * Computes the cosine of the given value.
@@ -1076,7 +1062,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def cos(e: Column): Column = Cos(e.expr)
+  def cos(e: Column): Column = withExpr { Cos(e.expr) }
 
   /**
    * Computes the cosine of the given column.
@@ -1092,7 +1078,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def cosh(e: Column): Column = Cosh(e.expr)
+  def cosh(e: Column): Column = withExpr { Cosh(e.expr) }
 
   /**
    * Computes the hyperbolic cosine of the given column.
@@ -1108,7 +1094,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def exp(e: Column): Column = Exp(e.expr)
+  def exp(e: Column): Column = withExpr { Exp(e.expr) }
 
   /**
    * Computes the exponential of the given column.
@@ -1124,7 +1110,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def expm1(e: Column): Column = Expm1(e.expr)
+  def expm1(e: Column): Column = withExpr { Expm1(e.expr) }
 
   /**
    * Computes the exponential of the given column.
@@ -1140,7 +1126,7 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def factorial(e: Column): Column = Factorial(e.expr)
+  def factorial(e: Column): Column = withExpr { Factorial(e.expr) }
 
   /**
    * Computes the floor of the given value.
@@ -1148,7 +1134,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def floor(e: Column): Column = Floor(e.expr)
+  def floor(e: Column): Column = withExpr { Floor(e.expr) }
 
   /**
    * Computes the floor of the given column.
@@ -1166,7 +1152,7 @@ object functions {
    * @since 1.5.0
    */
   @scala.annotation.varargs
-  def greatest(exprs: Column*): Column = {
+  def greatest(exprs: Column*): Column = withExpr {
     require(exprs.length > 1, "greatest requires at least 2 arguments.")
     Greatest(exprs.map(_.expr))
   }
@@ -1189,7 +1175,7 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def hex(column: Column): Column = Hex(column.expr)
+  def hex(column: Column): Column = withExpr { Hex(column.expr) }
 
   /**
    * Inverse of hex. Interprets each pair of characters as a hexadecimal number
@@ -1198,7 +1184,7 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def unhex(column: Column): Column = Unhex(column.expr)
+  def unhex(column: Column): Column = withExpr { Unhex(column.expr) }
 
   /**
    * Computes `sqrt(a^2^ + b^2^)` without intermediate overflow or underflow.
@@ -1206,7 +1192,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def hypot(l: Column, r: Column): Column = Hypot(l.expr, r.expr)
+  def hypot(l: Column, r: Column): Column = withExpr { Hypot(l.expr, r.expr) }
 
   /**
    * Computes `sqrt(a^2^ + b^2^)` without intermediate overflow or underflow.
@@ -1239,7 +1225,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def hypot(l: Column, r: Double): Column = hypot(l, lit(r).expr)
+  def hypot(l: Column, r: Double): Column = hypot(l, lit(r))
 
   /**
    * Computes `sqrt(a^2^ + b^2^)` without intermediate overflow or underflow.
@@ -1255,7 +1241,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def hypot(l: Double, r: Column): Column = hypot(lit(l).expr, r)
+  def hypot(l: Double, r: Column): Column = hypot(lit(l), r)
 
   /**
    * Computes `sqrt(a^2^ + b^2^)` without intermediate overflow or underflow.
@@ -1273,7 +1259,7 @@ object functions {
    * @since 1.5.0
    */
   @scala.annotation.varargs
-  def least(exprs: Column*): Column = {
+  def least(exprs: Column*): Column = withExpr {
     require(exprs.length > 1, "least requires at least 2 arguments.")
     Least(exprs.map(_.expr))
   }
@@ -1296,7 +1282,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def log(e: Column): Column = Log(e.expr)
+  def log(e: Column): Column = withExpr { Log(e.expr) }
 
   /**
    * Computes the natural logarithm of the given column.
@@ -1312,7 +1298,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def log(base: Double, a: Column): Column = Logarithm(lit(base).expr, a.expr)
+  def log(base: Double, a: Column): Column = withExpr { Logarithm(lit(base).expr, a.expr) }
 
   /**
    * Returns the first argument-base logarithm of the second argument.
@@ -1328,7 +1314,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def log10(e: Column): Column = Log10(e.expr)
+  def log10(e: Column): Column = withExpr { Log10(e.expr) }
 
   /**
    * Computes the logarithm of the given value in base 10.
@@ -1344,7 +1330,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def log1p(e: Column): Column = Log1p(e.expr)
+  def log1p(e: Column): Column = withExpr { Log1p(e.expr) }
 
   /**
    * Computes the natural logarithm of the given column plus one.
@@ -1360,7 +1346,7 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def log2(expr: Column): Column = Log2(expr.expr)
+  def log2(expr: Column): Column = withExpr { Log2(expr.expr) }
 
   /**
    * Computes the logarithm of the given value in base 2.
@@ -1376,7 +1362,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def pow(l: Column, r: Column): Column = Pow(l.expr, r.expr)
+  def pow(l: Column, r: Column): Column = withExpr { Pow(l.expr, r.expr) }
 
   /**
    * Returns the value of the first argument raised to the power of the second argument.
@@ -1408,7 +1394,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def pow(l: Column, r: Double): Column = pow(l, lit(r).expr)
+  def pow(l: Column, r: Double): Column = pow(l, lit(r))
 
   /**
    * Returns the value of the first argument raised to the power of the second argument.
@@ -1424,7 +1410,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def pow(l: Double, r: Column): Column = pow(lit(l).expr, r)
+  def pow(l: Double, r: Column): Column = pow(lit(l), r)
 
   /**
    * Returns the value of the first argument raised to the power of the second argument.
@@ -1440,7 +1426,9 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def pmod(dividend: Column, divisor: Column): Column = Pmod(dividend.expr, divisor.expr)
+  def pmod(dividend: Column, divisor: Column): Column = withExpr {
+    Pmod(dividend.expr, divisor.expr)
+  }
 
   /**
    * Returns the double value that is closest in value to the argument and
@@ -1449,7 +1437,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def rint(e: Column): Column = Rint(e.expr)
+  def rint(e: Column): Column = withExpr { Rint(e.expr) }
 
   /**
    * Returns the double value that is closest in value to the argument and
@@ -1466,7 +1454,7 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def round(e: Column): Column = round(e.expr, 0)
+  def round(e: Column): Column = round(e, 0)
 
   /**
    * Round the value of `e` to `scale` decimal places if `scale` >= 0
@@ -1475,7 +1463,7 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def round(e: Column, scale: Int): Column = Round(e.expr, Literal(scale))
+  def round(e: Column, scale: Int): Column = withExpr { Round(e.expr, Literal(scale)) }
 
   /**
    * Shift the the given value numBits left. If the given value is a long value, this function
@@ -1484,7 +1472,7 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def shiftLeft(e: Column, numBits: Int): Column = ShiftLeft(e.expr, lit(numBits).expr)
+  def shiftLeft(e: Column, numBits: Int): Column = withExpr { ShiftLeft(e.expr, lit(numBits).expr) }
 
   /**
    * Shift the the given value numBits right. If the given value is a long value, it will return
@@ -1493,7 +1481,9 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def shiftRight(e: Column, numBits: Int): Column = ShiftRight(e.expr, lit(numBits).expr)
+  def shiftRight(e: Column, numBits: Int): Column = withExpr {
+    ShiftRight(e.expr, lit(numBits).expr)
+  }
 
   /**
    * Unsigned shift the the given value numBits right. If the given value is a long value,
@@ -1502,8 +1492,9 @@ object functions {
    * @group math_funcs
    * @since 1.5.0
    */
-  def shiftRightUnsigned(e: Column, numBits: Int): Column =
+  def shiftRightUnsigned(e: Column, numBits: Int): Column = withExpr {
     ShiftRightUnsigned(e.expr, lit(numBits).expr)
+  }
 
   /**
    * Computes the signum of the given value.
@@ -1511,7 +1502,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def signum(e: Column): Column = Signum(e.expr)
+  def signum(e: Column): Column = withExpr { Signum(e.expr) }
 
   /**
    * Computes the signum of the given column.
@@ -1527,7 +1518,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def sin(e: Column): Column = Sin(e.expr)
+  def sin(e: Column): Column = withExpr { Sin(e.expr) }
 
   /**
    * Computes the sine of the given column.
@@ -1543,7 +1534,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def sinh(e: Column): Column = Sinh(e.expr)
+  def sinh(e: Column): Column = withExpr { Sinh(e.expr) }
 
   /**
    * Computes the hyperbolic sine of the given column.
@@ -1559,7 +1550,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def tan(e: Column): Column = Tan(e.expr)
+  def tan(e: Column): Column = withExpr { Tan(e.expr) }
 
   /**
    * Computes the tangent of the given column.
@@ -1575,7 +1566,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def tanh(e: Column): Column = Tanh(e.expr)
+  def tanh(e: Column): Column = withExpr { Tanh(e.expr) }
 
   /**
    * Computes the hyperbolic tangent of the given column.
@@ -1591,7 +1582,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def toDegrees(e: Column): Column = ToDegrees(e.expr)
+  def toDegrees(e: Column): Column = withExpr { ToDegrees(e.expr) }
 
   /**
    * Converts an angle measured in radians to an approximately equivalent angle measured in degrees.
@@ -1607,7 +1598,7 @@ object functions {
    * @group math_funcs
    * @since 1.4.0
    */
-  def toRadians(e: Column): Column = ToRadians(e.expr)
+  def toRadians(e: Column): Column = withExpr { ToRadians(e.expr) }
 
   /**
    * Converts an angle measured in degrees to an approximately equivalent angle measured in radians.
@@ -1628,7 +1619,7 @@ object functions {
    * @group misc_funcs
    * @since 1.5.0
    */
-  def md5(e: Column): Column = Md5(e.expr)
+  def md5(e: Column): Column = withExpr { Md5(e.expr) }
 
   /**
    * Calculates the SHA-1 digest of a binary column and returns the value
@@ -1637,7 +1628,7 @@ object functions {
    * @group misc_funcs
    * @since 1.5.0
    */
-  def sha1(e: Column): Column = Sha1(e.expr)
+  def sha1(e: Column): Column = withExpr { Sha1(e.expr) }
 
   /**
    * Calculates the SHA-2 family of hash functions of a binary column and
@@ -1652,7 +1643,7 @@ object functions {
   def sha2(e: Column, numBits: Int): Column = {
     require(Seq(0, 224, 256, 384, 512).contains(numBits),
       s"numBits $numBits is not in the permitted values (0, 224, 256, 384, 512)")
-    Sha2(e.expr, lit(numBits).expr)
+    withExpr { Sha2(e.expr, lit(numBits).expr) }
   }
 
   /**
@@ -1662,7 +1653,7 @@ object functions {
    * @group misc_funcs
    * @since 1.5.0
    */
-  def crc32(e: Column): Column = Crc32(e.expr)
+  def crc32(e: Column): Column = withExpr { Crc32(e.expr) }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
   // String functions
@@ -1675,7 +1666,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def ascii(e: Column): Column = Ascii(e.expr)
+  def ascii(e: Column): Column = withExpr { Ascii(e.expr) }
 
   /**
    * Computes the BASE64 encoding of a binary column and returns it as a string column.
@@ -1684,7 +1675,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def base64(e: Column): Column = Base64(e.expr)
+  def base64(e: Column): Column = withExpr { Base64(e.expr) }
 
   /**
    * Concatenates multiple input string columns together into a single string column.
@@ -1693,7 +1684,7 @@ object functions {
    * @since 1.5.0
    */
   @scala.annotation.varargs
-  def concat(exprs: Column*): Column = Concat(exprs.map(_.expr))
+  def concat(exprs: Column*): Column = withExpr { Concat(exprs.map(_.expr)) }
 
   /**
    * Concatenates multiple input string columns together into a single string column,
@@ -1703,7 +1694,7 @@ object functions {
    * @since 1.5.0
    */
   @scala.annotation.varargs
-  def concat_ws(sep: String, exprs: Column*): Column = {
+  def concat_ws(sep: String, exprs: Column*): Column = withExpr {
     ConcatWs(Literal.create(sep, StringType) +: exprs.map(_.expr))
   }
 
@@ -1715,7 +1706,9 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def decode(value: Column, charset: String): Column = Decode(value.expr, lit(charset).expr)
+  def decode(value: Column, charset: String): Column = withExpr {
+    Decode(value.expr, lit(charset).expr)
+  }
 
   /**
    * Computes the first argument into a binary from a string using the provided character set
@@ -1725,7 +1718,9 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def encode(value: Column, charset: String): Column = Encode(value.expr, lit(charset).expr)
+  def encode(value: Column, charset: String): Column = withExpr {
+    Encode(value.expr, lit(charset).expr)
+  }
 
   /**
    * Formats numeric column x to a format like '#,###,###.##', rounded to d decimal places,
@@ -1737,7 +1732,9 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def format_number(x: Column, d: Int): Column = FormatNumber(x.expr, lit(d).expr)
+  def format_number(x: Column, d: Int): Column = withExpr {
+    FormatNumber(x.expr, lit(d).expr)
+  }
 
   /**
    * Formats the arguments in printf-style and returns the result as a string column.
@@ -1746,7 +1743,7 @@ object functions {
    * @since 1.5.0
    */
   @scala.annotation.varargs
-  def format_string(format: String, arguments: Column*): Column = {
+  def format_string(format: String, arguments: Column*): Column = withExpr {
     FormatString((lit(format) +: arguments).map(_.expr): _*)
   }
 
@@ -1759,7 +1756,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def initcap(e: Column): Column = InitCap(e.expr)
+  def initcap(e: Column): Column = withExpr { InitCap(e.expr) }
 
   /**
    * Locate the position of the first occurrence of substr column in the given string.
@@ -1771,7 +1768,9 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def instr(str: Column, substring: String): Column = StringInstr(str.expr, lit(substring).expr)
+  def instr(str: Column, substring: String): Column = withExpr {
+    StringInstr(str.expr, lit(substring).expr)
+  }
 
   /**
    * Computes the length of a given string or binary column.
@@ -1779,7 +1778,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def length(e: Column): Column = Length(e.expr)
+  def length(e: Column): Column = withExpr { Length(e.expr) }
 
   /**
    * Converts a string column to lower case.
@@ -1787,14 +1786,14 @@ object functions {
    * @group string_funcs
    * @since 1.3.0
    */
-  def lower(e: Column): Column = Lower(e.expr)
+  def lower(e: Column): Column = withExpr { Lower(e.expr) }
 
   /**
    * Computes the Levenshtein distance of the two given string columns.
    * @group string_funcs
    * @since 1.5.0
    */
-  def levenshtein(l: Column, r: Column): Column = Levenshtein(l.expr, r.expr)
+  def levenshtein(l: Column, r: Column): Column = withExpr { Levenshtein(l.expr, r.expr) }
 
   /**
    * Locate the position of the first occurrence of substr.
@@ -1804,7 +1803,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def locate(substr: String, str: Column): Column = {
+  def locate(substr: String, str: Column): Column = withExpr {
     new StringLocate(lit(substr).expr, str.expr)
   }
 
@@ -1817,7 +1816,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def locate(substr: String, str: Column, pos: Int): Column = {
+  def locate(substr: String, str: Column, pos: Int): Column = withExpr {
     StringLocate(lit(substr).expr, str.expr, lit(pos).expr)
   }
 
@@ -1827,7 +1826,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def lpad(str: Column, len: Int, pad: String): Column = {
+  def lpad(str: Column, len: Int, pad: String): Column = withExpr {
     StringLPad(str.expr, lit(len).expr, lit(pad).expr)
   }
 
@@ -1837,7 +1836,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def ltrim(e: Column): Column = StringTrimLeft(e.expr)
+  def ltrim(e: Column): Column = withExpr {StringTrimLeft(e.expr) }
 
   /**
    * Extract a specific(idx) group identified by a java regex, from the specified string column.
@@ -1845,7 +1844,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def regexp_extract(e: Column, exp: String, groupIdx: Int): Column = {
+  def regexp_extract(e: Column, exp: String, groupIdx: Int): Column = withExpr {
     RegExpExtract(e.expr, lit(exp).expr, lit(groupIdx).expr)
   }
 
@@ -1855,7 +1854,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def regexp_replace(e: Column, pattern: String, replacement: String): Column = {
+  def regexp_replace(e: Column, pattern: String, replacement: String): Column = withExpr {
     RegExpReplace(e.expr, lit(pattern).expr, lit(replacement).expr)
   }
 
@@ -1866,7 +1865,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def unbase64(e: Column): Column = UnBase64(e.expr)
+  def unbase64(e: Column): Column = withExpr { UnBase64(e.expr) }
 
   /**
    * Right-padded with pad to a length of len.
@@ -1874,7 +1873,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def rpad(str: Column, len: Int, pad: String): Column = {
+  def rpad(str: Column, len: Int, pad: String): Column = withExpr {
     StringRPad(str.expr, lit(len).expr, lit(pad).expr)
   }
 
@@ -1884,7 +1883,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def repeat(str: Column, n: Int): Column = {
+  def repeat(str: Column, n: Int): Column = withExpr {
     StringRepeat(str.expr, lit(n).expr)
   }
 
@@ -1894,9 +1893,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def reverse(str: Column): Column = {
-    StringReverse(str.expr)
-  }
+  def reverse(str: Column): Column = withExpr { StringReverse(str.expr) }
 
   /**
    * Trim the spaces from right end for the specified string value.
@@ -1904,7 +1901,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def rtrim(e: Column): Column = StringTrimRight(e.expr)
+  def rtrim(e: Column): Column = withExpr { StringTrimRight(e.expr) }
 
   /**
    * * Return the soundex code for the specified expression.
@@ -1912,7 +1909,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def soundex(e: Column): Column = SoundEx(e.expr)
+  def soundex(e: Column): Column = withExpr { SoundEx(e.expr) }
 
   /**
    * Splits str around pattern (pattern is a regular expression).
@@ -1921,7 +1918,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def split(str: Column, pattern: String): Column = {
+  def split(str: Column, pattern: String): Column = withExpr {
     StringSplit(str.expr, lit(pattern).expr)
   }
 
@@ -1933,8 +1930,9 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def substring(str: Column, pos: Int, len: Int): Column =
+  def substring(str: Column, pos: Int, len: Int): Column = withExpr {
     Substring(str.expr, lit(pos).expr, lit(len).expr)
+  }
 
   /**
    * Returns the substring from string str before count occurrences of the delimiter delim.
@@ -1944,8 +1942,9 @@ object functions {
    *
    * @group string_funcs
    */
-  def substring_index(str: Column, delim: String, count: Int): Column =
+  def substring_index(str: Column, delim: String, count: Int): Column = withExpr {
     SubstringIndex(str.expr, lit(delim).expr, lit(count).expr)
+  }
 
   /**
    * Translate any character in the src by a character in replaceString.
@@ -1956,8 +1955,9 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def translate(src: Column, matchingString: String, replaceString: String): Column =
+  def translate(src: Column, matchingString: String, replaceString: String): Column = withExpr {
     StringTranslate(src.expr, lit(matchingString).expr, lit(replaceString).expr)
+  }
 
   /**
    * Trim the spaces from both ends for the specified string column.
@@ -1965,7 +1965,7 @@ object functions {
    * @group string_funcs
    * @since 1.5.0
    */
-  def trim(e: Column): Column = StringTrim(e.expr)
+  def trim(e: Column): Column = withExpr { StringTrim(e.expr) }
 
   /**
    * Converts a string column to upper case.
@@ -1973,7 +1973,7 @@ object functions {
    * @group string_funcs
    * @since 1.3.0
    */
-  def upper(e: Column): Column = Upper(e.expr)
+  def upper(e: Column): Column = withExpr { Upper(e.expr) }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
   // DateTime functions
@@ -1985,8 +1985,9 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def add_months(startDate: Column, numMonths: Int): Column =
+  def add_months(startDate: Column, numMonths: Int): Column = withExpr {
     AddMonths(startDate.expr, Literal(numMonths))
+  }
 
   /**
    * Returns the current date as a date column.
@@ -1994,7 +1995,7 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def current_date(): Column = CurrentDate()
+  def current_date(): Column = withExpr { CurrentDate() }
 
   /**
    * Returns the current timestamp as a timestamp column.
@@ -2002,7 +2003,7 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def current_timestamp(): Column = CurrentTimestamp()
+  def current_timestamp(): Column = withExpr { CurrentTimestamp() }
 
   /**
    * Converts a date/timestamp/string to a value of string in the format specified by the date
@@ -2017,71 +2018,72 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def date_format(dateExpr: Column, format: String): Column =
+  def date_format(dateExpr: Column, format: String): Column = withExpr {
     DateFormatClass(dateExpr.expr, Literal(format))
+  }
 
   /**
    * Returns the date that is `days` days after `start`
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def date_add(start: Column, days: Int): Column = DateAdd(start.expr, Literal(days))
+  def date_add(start: Column, days: Int): Column = withExpr { DateAdd(start.expr, Literal(days)) }
 
   /**
    * Returns the date that is `days` days before `start`
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def date_sub(start: Column, days: Int): Column = DateSub(start.expr, Literal(days))
+  def date_sub(start: Column, days: Int): Column = withExpr { DateSub(start.expr, Literal(days)) }
 
   /**
    * Returns the number of days from `start` to `end`.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def datediff(end: Column, start: Column): Column = DateDiff(end.expr, start.expr)
+  def datediff(end: Column, start: Column): Column = withExpr { DateDiff(end.expr, start.expr) }
 
   /**
    * Extracts the year as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def year(e: Column): Column = Year(e.expr)
+  def year(e: Column): Column = withExpr { Year(e.expr) }
 
   /**
    * Extracts the quarter as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def quarter(e: Column): Column = Quarter(e.expr)
+  def quarter(e: Column): Column = withExpr { Quarter(e.expr) }
 
   /**
    * Extracts the month as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def month(e: Column): Column = Month(e.expr)
+  def month(e: Column): Column = withExpr { Month(e.expr) }
 
   /**
    * Extracts the day of the month as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def dayofmonth(e: Column): Column = DayOfMonth(e.expr)
+  def dayofmonth(e: Column): Column = withExpr { DayOfMonth(e.expr) }
 
   /**
    * Extracts the day of the year as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def dayofyear(e: Column): Column = DayOfYear(e.expr)
+  def dayofyear(e: Column): Column = withExpr { DayOfYear(e.expr) }
 
   /**
    * Extracts the hours as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def hour(e: Column): Column = Hour(e.expr)
+  def hour(e: Column): Column = withExpr { Hour(e.expr) }
 
   /**
    * Given a date column, returns the last day of the month which the given date belongs to.
@@ -2091,21 +2093,23 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def last_day(e: Column): Column = LastDay(e.expr)
+  def last_day(e: Column): Column = withExpr { LastDay(e.expr) }
 
   /**
    * Extracts the minutes as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def minute(e: Column): Column = Minute(e.expr)
+  def minute(e: Column): Column = withExpr { Minute(e.expr) }
 
   /*
    * Returns number of months between dates `date1` and `date2`.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def months_between(date1: Column, date2: Column): Column = MonthsBetween(date1.expr, date2.expr)
+  def months_between(date1: Column, date2: Column): Column = withExpr {
+    MonthsBetween(date1.expr, date2.expr)
+  }
 
   /**
    * Given a date column, returns the first date which is later than the value of the date column
@@ -2120,21 +2124,23 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def next_day(date: Column, dayOfWeek: String): Column = NextDay(date.expr, lit(dayOfWeek).expr)
+  def next_day(date: Column, dayOfWeek: String): Column = withExpr {
+    NextDay(date.expr, lit(dayOfWeek).expr)
+  }
 
   /**
    * Extracts the seconds as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def second(e: Column): Column = Second(e.expr)
+  def second(e: Column): Column = withExpr { Second(e.expr) }
 
   /**
    * Extracts the week number as an integer from a given date/timestamp/string.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def weekofyear(e: Column): Column = WeekOfYear(e.expr)
+  def weekofyear(e: Column): Column = withExpr { WeekOfYear(e.expr) }
 
   /**
    * Converts the number of seconds from unix epoch (1970-01-01 00:00:00 UTC) to a string
@@ -2143,7 +2149,9 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def from_unixtime(ut: Column): Column = FromUnixTime(ut.expr, Literal("yyyy-MM-dd HH:mm:ss"))
+  def from_unixtime(ut: Column): Column = withExpr {
+    FromUnixTime(ut.expr, Literal("yyyy-MM-dd HH:mm:ss"))
+  }
 
   /**
    * Converts the number of seconds from unix epoch (1970-01-01 00:00:00 UTC) to a string
@@ -2152,14 +2160,18 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def from_unixtime(ut: Column, f: String): Column = FromUnixTime(ut.expr, Literal(f))
+  def from_unixtime(ut: Column, f: String): Column = withExpr {
+    FromUnixTime(ut.expr, Literal(f))
+  }
 
   /**
    * Gets current Unix timestamp in seconds.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def unix_timestamp(): Column = UnixTimestamp(CurrentTimestamp(), Literal("yyyy-MM-dd HH:mm:ss"))
+  def unix_timestamp(): Column = withExpr {
+    UnixTimestamp(CurrentTimestamp(), Literal("yyyy-MM-dd HH:mm:ss"))
+  }
 
   /**
    * Converts time string in format yyyy-MM-dd HH:mm:ss to Unix timestamp (in seconds),
@@ -2167,7 +2179,9 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def unix_timestamp(s: Column): Column = UnixTimestamp(s.expr, Literal("yyyy-MM-dd HH:mm:ss"))
+  def unix_timestamp(s: Column): Column = withExpr {
+    UnixTimestamp(s.expr, Literal("yyyy-MM-dd HH:mm:ss"))
+  }
 
   /**
    * Convert time string with given pattern
@@ -2176,7 +2190,7 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def unix_timestamp(s: Column, p: String): Column = UnixTimestamp(s.expr, Literal(p))
+  def unix_timestamp(s: Column, p: String): Column = withExpr {UnixTimestamp(s.expr, Literal(p)) }
 
   /**
    * Converts the column into DateType.
@@ -2184,7 +2198,7 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def to_date(e: Column): Column = ToDate(e.expr)
+  def to_date(e: Column): Column = withExpr { ToDate(e.expr) }
 
   /**
    * Returns date truncated to the unit specified by the format.
@@ -2195,22 +2209,27 @@ object functions {
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def trunc(date: Column, format: String): Column = TruncDate(date.expr, Literal(format))
+  def trunc(date: Column, format: String): Column = withExpr {
+    TruncDate(date.expr, Literal(format))
+  }
 
   /**
    * Assumes given timestamp is UTC and converts to given timezone.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def from_utc_timestamp(ts: Column, tz: String): Column =
-    FromUTCTimestamp(ts.expr, Literal(tz).expr)
+  def from_utc_timestamp(ts: Column, tz: String): Column = withExpr {
+    FromUTCTimestamp(ts.expr, Literal(tz))
+  }
 
   /**
    * Assumes given timestamp is in given timezone and converts to UTC.
    * @group datetime_funcs
    * @since 1.5.0
    */
-  def to_utc_timestamp(ts: Column, tz: String): Column = ToUTCTimestamp(ts.expr, Literal(tz).expr)
+  def to_utc_timestamp(ts: Column, tz: String): Column = withExpr {
+    ToUTCTimestamp(ts.expr, Literal(tz))
+  }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
   // Collection functions
@@ -2221,8 +2240,9 @@ object functions {
    * @group collection_funcs
    * @since 1.5.0
    */
-  def array_contains(column: Column, value: Any): Column =
+  def array_contains(column: Column, value: Any): Column = withExpr {
     ArrayContains(column.expr, Literal(value))
+  }
 
   /**
    * Creates a new row for each element in the given array or map column.
@@ -2230,7 +2250,7 @@ object functions {
    * @group collection_funcs
    * @since 1.3.0
    */
-  def explode(e: Column): Column = Explode(e.expr)
+  def explode(e: Column): Column = withExpr { Explode(e.expr) }
 
   /**
    * Returns length of array or map.
@@ -2238,7 +2258,7 @@ object functions {
    * @group collection_funcs
    * @since 1.5.0
    */
-  def size(e: Column): Column = Size(e.expr)
+  def size(e: Column): Column = withExpr { Size(e.expr) }
 
   /**
    * Sorts the input array for the given column in ascending order,
@@ -2256,7 +2276,7 @@ object functions {
    * @group collection_funcs
    * @since 1.5.0
    */
-  def sort_array(e: Column, asc: Boolean): Column = SortArray(e.expr, lit(asc).expr)
+  def sort_array(e: Column, asc: Boolean): Column = withExpr { SortArray(e.expr, lit(asc).expr) }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
   //////////////////////////////////////////////////////////////////////////////////////////////
@@ -2296,11 +2316,10 @@ object functions {
      * @deprecated As of 1.5.0, since it's redundant with udf()
      */
     @deprecated("Use udf", "1.5.0")
-    def callUDF(f: Function$x[$fTypes], returnType: DataType${if (args.length > 0) ", " + args else ""}): Column = {
+    def callUDF(f: Function$x[$fTypes], returnType: DataType${if (args.length > 0) ", " + args else ""}): Column = withExpr {
       ScalaUDF(f, returnType, Seq($argsInUDF))
     }""")
   }
-  }
   */
   /**
    * Defines a user-defined function of 0 arguments as user-defined function (UDF).
@@ -2435,147 +2454,146 @@ object functions {
   }
 
   //////////////////////////////////////////////////////////////////////////////////////////////////
-
   /**
-   * Call a Scala function of 0 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 0 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function0[_], returnType: DataType): Column = {
+  def callUDF(f: Function0[_], returnType: DataType): Column = withExpr {
     ScalaUDF(f, returnType, Seq())
   }
 
   /**
-   * Call a Scala function of 1 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 1 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function1[_, _], returnType: DataType, arg1: Column): Column = {
+  def callUDF(f: Function1[_, _], returnType: DataType, arg1: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr))
   }
 
   /**
-   * Call a Scala function of 2 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 2 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function2[_, _, _], returnType: DataType, arg1: Column, arg2: Column): Column = {
+  def callUDF(f: Function2[_, _, _], returnType: DataType, arg1: Column, arg2: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr))
   }
 
   /**
-   * Call a Scala function of 3 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 3 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function3[_, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column): Column = {
+  def callUDF(f: Function3[_, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr))
   }
 
   /**
-   * Call a Scala function of 4 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 4 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function4[_, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column): Column = {
+  def callUDF(f: Function4[_, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr))
   }
 
   /**
-   * Call a Scala function of 5 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 5 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function5[_, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column): Column = {
+  def callUDF(f: Function5[_, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr))
   }
 
   /**
-   * Call a Scala function of 6 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 6 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function6[_, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column): Column = {
+  def callUDF(f: Function6[_, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr))
   }
 
   /**
-   * Call a Scala function of 7 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 7 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function7[_, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column): Column = {
+  def callUDF(f: Function7[_, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr, arg7.expr))
   }
 
   /**
-   * Call a Scala function of 8 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 8 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function8[_, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column): Column = {
+  def callUDF(f: Function8[_, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr, arg7.expr, arg8.expr))
   }
 
   /**
-   * Call a Scala function of 9 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 9 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function9[_, _, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column, arg9: Column): Column = {
+  def callUDF(f: Function9[_, _, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column, arg9: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr, arg7.expr, arg8.expr, arg9.expr))
   }
 
   /**
-   * Call a Scala function of 10 arguments as user-defined function (UDF). This requires
-   * you to specify the return data type.
-   *
-   * @group udf_funcs
-   * @since 1.3.0
-   * @deprecated As of 1.5.0, since it's redundant with udf()
-   */
+    * Call a Scala function of 10 arguments as user-defined function (UDF). This requires
+    * you to specify the return data type.
+    *
+    * @group udf_funcs
+    * @since 1.3.0
+    * @deprecated As of 1.5.0, since it's redundant with udf()
+    */
   @deprecated("Use udf", "1.5.0")
-  def callUDF(f: Function10[_, _, _, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column, arg9: Column, arg10: Column): Column = {
+  def callUDF(f: Function10[_, _, _, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column, arg9: Column, arg10: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr, arg7.expr, arg8.expr, arg9.expr, arg10.expr))
   }
 
@@ -2597,7 +2615,7 @@ object functions {
    * @since 1.5.0
    */
   @scala.annotation.varargs
-  def callUDF(udfName: String, cols: Column*): Column = {
+  def callUDF(udfName: String, cols: Column*): Column = withExpr {
     UnresolvedFunction(udfName, cols.map(_.expr), isDistinct = false)
   }
 
@@ -2618,7 +2636,7 @@ object functions {
    * @deprecated As of 1.5.0, since it was not coherent to have two functions callUdf and callUDF
    */
   @deprecated("Use callUDF", "1.5.0")
-  def callUdf(udfName: String, cols: Column*): Column = {
+  def callUdf(udfName: String, cols: Column*): Column = withExpr {
     // Note: we avoid using closures here because on file systems that are case-insensitive, the
     // compiled class file for the closure here will conflict with the one in callUDF (upper case).
     val exprs = new Array[Expression](cols.size)

From 244010624200eddea6dfd1b2c89f40be45212e96 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Thu, 5 Nov 2015 16:34:10 -0800
Subject: [PATCH 082/518] [SPARK-11542] [SPARKR] fix glm with long fomular

Because deparse() will break the long string into multiple lines, the deserialization will fail

Author: Davies Liu <davies@databricks.com>

Closes #9510 from davies/fix_glm.
---
 R/pkg/R/mllib.R               |  3 ++-
 R/pkg/inst/tests/test_mllib.R | 12 ++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/R/pkg/R/mllib.R b/R/pkg/R/mllib.R
index 60bfadb8e7503..b0d73dd93a79d 100644
--- a/R/pkg/R/mllib.R
+++ b/R/pkg/R/mllib.R
@@ -48,8 +48,9 @@ setMethod("glm", signature(formula = "formula", family = "ANY", data = "DataFram
           function(formula, family = c("gaussian", "binomial"), data, lambda = 0, alpha = 0,
             standardize = TRUE, solver = "auto") {
             family <- match.arg(family)
+            formula <- paste(deparse(formula), collapse="")
             model <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
-                                 "fitRModelFormula", deparse(formula), data@sdf, family, lambda,
+                                 "fitRModelFormula", formula, data@sdf, family, lambda,
                                  alpha, standardize, solver)
             return(new("PipelineModel", model = model))
           })
diff --git a/R/pkg/inst/tests/test_mllib.R b/R/pkg/inst/tests/test_mllib.R
index 032cfef061fd3..4761e285a2479 100644
--- a/R/pkg/inst/tests/test_mllib.R
+++ b/R/pkg/inst/tests/test_mllib.R
@@ -33,6 +33,18 @@ test_that("glm and predict", {
   expect_equal(typeof(take(select(prediction, "prediction"), 1)$prediction), "double")
 })
 
+test_that("glm should work with long formula", {
+  training <- createDataFrame(sqlContext, iris)
+  training$LongLongLongLongLongName <- training$Sepal_Width
+  training$VeryLongLongLongLonLongName <- training$Sepal_Length
+  training$AnotherLongLongLongLongName <- training$Species
+  model <- glm(LongLongLongLongLongName ~ VeryLongLongLongLonLongName + AnotherLongLongLongLongName,
+               data = training)
+  vals <- collect(select(predict(model, training), "prediction"))
+  rVals <- predict(glm(Sepal.Width ~ Sepal.Length + Species, data = iris), iris)
+  expect_true(all(abs(rVals - vals) < 1e-6), rVals - vals)
+})
+
 test_that("predictions match with native glm", {
   training <- createDataFrame(sqlContext, iris)
   model <- glm(Sepal_Width ~ Sepal_Length + Species, data = training)

From 07414afac9a100ede1dee5f3d45a657802c8bd2a Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Thu, 5 Nov 2015 17:02:22 -0800
Subject: [PATCH 083/518] [SPARK-11537] [SQL] fix negative
 hours/minutes/seconds

Currently, if the Timestamp is before epoch (1970/01/01), the hours, minutes and seconds will be negative (also rounding up).

Author: Davies Liu <davies@databricks.com>

Closes #9502 from davies/neg_hour.
---
 .../sql/catalyst/util/DateTimeUtils.scala     | 23 ++++++++++++-------
 .../catalyst/util/DateTimeUtilsSuite.scala    | 13 +++++++++++
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
index 781ed1688a327..f5fff90e5a542 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
@@ -392,29 +392,36 @@ object DateTimeUtils {
     Some((c.getTimeInMillis / MILLIS_PER_DAY).toInt)
   }
 
+  /**
+   * Returns the microseconds since year zero (-17999) from microseconds since epoch.
+   */
+  def absoluteMicroSecond(microsec: SQLTimestamp): SQLTimestamp = {
+    microsec + toYearZero * MICROS_PER_DAY
+  }
+
   /**
    * Returns the hour value of a given timestamp value. The timestamp is expressed in microseconds.
    */
-  def getHours(timestamp: SQLTimestamp): Int = {
-    val localTs = (timestamp / 1000) + defaultTimeZone.getOffset(timestamp / 1000)
-    ((localTs / 1000 / 3600) % 24).toInt
+  def getHours(microsec: SQLTimestamp): Int = {
+    val localTs = absoluteMicroSecond(microsec) + defaultTimeZone.getOffset(microsec / 1000) * 1000L
+    ((localTs / MICROS_PER_SECOND / 3600) % 24).toInt
   }
 
   /**
    * Returns the minute value of a given timestamp value. The timestamp is expressed in
    * microseconds.
    */
-  def getMinutes(timestamp: SQLTimestamp): Int = {
-    val localTs = (timestamp / 1000) + defaultTimeZone.getOffset(timestamp / 1000)
-    ((localTs / 1000 / 60) % 60).toInt
+  def getMinutes(microsec: SQLTimestamp): Int = {
+    val localTs = absoluteMicroSecond(microsec) + defaultTimeZone.getOffset(microsec / 1000) * 1000L
+    ((localTs / MICROS_PER_SECOND / 60) % 60).toInt
   }
 
   /**
    * Returns the second value of a given timestamp value. The timestamp is expressed in
    * microseconds.
    */
-  def getSeconds(timestamp: SQLTimestamp): Int = {
-    ((timestamp / 1000 / 1000) % 60).toInt
+  def getSeconds(microsec: SQLTimestamp): Int = {
+    ((absoluteMicroSecond(microsec) / MICROS_PER_SECOND) % 60).toInt
   }
 
   private[this] def isLeapYear(year: Int): Boolean = {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
index 46335941b62d6..64d15e6b910c1 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
@@ -358,6 +358,19 @@ class DateTimeUtilsSuite extends SparkFunSuite {
     assert(getSeconds(c.getTimeInMillis * 1000) === 9)
   }
 
+  test("hours / miniute / seconds") {
+    Seq(Timestamp.valueOf("2015-06-11 10:12:35.789"),
+      Timestamp.valueOf("2015-06-11 20:13:40.789"),
+      Timestamp.valueOf("1900-06-11 12:14:50.789"),
+      Timestamp.valueOf("1700-02-28 12:14:50.123456")).foreach { t =>
+      val us = fromJavaTimestamp(t)
+      assert(toJavaTimestamp(us) === t)
+      assert(getHours(us) === t.getHours)
+      assert(getMinutes(us) === t.getMinutes)
+      assert(getSeconds(us) === t.getSeconds)
+    }
+  }
+
   test("get day in year") {
     val c = Calendar.getInstance()
     c.set(2015, 2, 18, 0, 0, 0)

From 6091e91fca58078a0f1d9c35d68c0ae7205a534c Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 5 Nov 2015 17:10:35 -0800
Subject: [PATCH 084/518] Revert "[SPARK-11469][SQL] Allow users to define
 nondeterministic udfs."

This reverts commit 9cf56c96b7d02a14175d40b336da14c2e1c88339.
---
 project/MimaExcludes.scala                    |  47 -----
 .../sql/catalyst/expressions/ScalaUDF.scala   |   7 +-
 .../apache/spark/sql/UDFRegistration.scala    | 164 ++++++++----------
 .../spark/sql/UserDefinedFunction.scala       |  13 +-
 .../scala/org/apache/spark/sql/UDFSuite.scala | 105 -----------
 .../datasources/parquet/ParquetIOSuite.scala  |   4 +-
 6 files changed, 78 insertions(+), 262 deletions(-)

diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 90dc947d4e588..40f5c9fec8bb8 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -114,53 +114,6 @@ object MimaExcludes {
           "org.apache.spark.rdd.MapPartitionsWithPreparationRDD"),
         ProblemFilters.exclude[MissingClassProblem](
           "org.apache.spark.rdd.MapPartitionsWithPreparationRDD$")
-      ) ++ Seq(
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$2"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$3"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$4"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$5"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$6"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$7"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$8"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$9"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$10"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$11"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$12"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$13"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$14"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$15"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$16"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$17"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$18"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$19"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$20"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$21"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$22"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$23"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.sql.UDFRegistration.org$apache$spark$sql$UDFRegistration$$builder$24")
       ) ++ Seq(
         // SPARK-11485
         ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.DataFrameHolder.df")
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
index a04af7f1dd877..11c7950c0613b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
@@ -30,18 +30,13 @@ case class ScalaUDF(
     function: AnyRef,
     dataType: DataType,
     children: Seq[Expression],
-    inputTypes: Seq[DataType] = Nil,
-    isDeterministic: Boolean = true)
+    inputTypes: Seq[DataType] = Nil)
   extends Expression with ImplicitCastInputTypes with CodegenFallback {
 
   override def nullable: Boolean = true
 
   override def toString: String = s"UDF(${children.mkString(",")})"
 
-  override def foldable: Boolean = deterministic && children.forall(_.foldable)
-
-  override def deterministic: Boolean = isDeterministic && children.forall(_.deterministic)
-
   // scalastyle:off
 
   /** This method has been generated by this script
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala b/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
index f5b95e13e47bc..fc4d0938c533a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
@@ -58,10 +58,8 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
    * Register a user-defined aggregate function (UDAF).
    *
    * @param name the name of the UDAF.
-   * @param udaf the UDAF that needs to be registered.
+   * @param udaf the UDAF needs to be registered.
    * @return the registered UDAF.
-   *
-   * @since 1.5.0
    */
   def register(
       name: String,
@@ -71,22 +69,6 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     udaf
   }
 
-  /**
-   * Register a user-defined function (UDF).
-   *
-   * @param name the name of the UDF.
-   * @param udf the UDF that needs to be registered.
-   * @return the registered UDF.
-   *
-   * @since 1.6.0
-   */
-  def register(
-      name: String,
-      udf: UserDefinedFunction): UserDefinedFunction = {
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
-  }
-
   // scalastyle:off
 
   /* register 0-22 were generated by this script
@@ -104,9 +86,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
         def register[$typeTags](name: String, func: Function$x[$types]): UserDefinedFunction = {
           val dataType = ScalaReflection.schemaFor[RT].dataType
           val inputTypes = Try($inputTypes).getOrElse(Nil)
-          val udf = UserDefinedFunction(func, dataType, inputTypes)
-          functionRegistry.registerFunction(name, udf.builder)
-          udf
+          def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+          functionRegistry.registerFunction(name, builder)
+          UserDefinedFunction(func, dataType)
         }""")
     }
 
@@ -136,9 +118,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag](name: String, func: Function0[RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -149,9 +131,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag](name: String, func: Function1[A1, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -162,9 +144,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag](name: String, func: Function2[A1, A2, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -175,9 +157,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag](name: String, func: Function3[A1, A2, A3, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -188,9 +170,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag](name: String, func: Function4[A1, A2, A3, A4, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -201,9 +183,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag](name: String, func: Function5[A1, A2, A3, A4, A5, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -214,9 +196,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag](name: String, func: Function6[A1, A2, A3, A4, A5, A6, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -227,9 +209,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag](name: String, func: Function7[A1, A2, A3, A4, A5, A6, A7, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -240,9 +222,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag](name: String, func: Function8[A1, A2, A3, A4, A5, A6, A7, A8, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -253,9 +235,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag](name: String, func: Function9[A1, A2, A3, A4, A5, A6, A7, A8, A9, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -266,9 +248,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag](name: String, func: Function10[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -279,9 +261,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag](name: String, func: Function11[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -292,9 +274,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag](name: String, func: Function12[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -305,9 +287,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag](name: String, func: Function13[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -318,9 +300,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag](name: String, func: Function14[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -331,9 +313,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag](name: String, func: Function15[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -344,9 +326,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag](name: String, func: Function16[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -357,9 +339,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag](name: String, func: Function17[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -370,9 +352,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag](name: String, func: Function18[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -383,9 +365,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag, A19: TypeTag](name: String, func: Function19[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, A19, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -396,9 +378,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag, A19: TypeTag, A20: TypeTag](name: String, func: Function20[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, A19, A20, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -409,9 +391,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag, A19: TypeTag, A20: TypeTag, A21: TypeTag](name: String, func: Function21[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, A19, A20, A21, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: ScalaReflection.schemaFor[A21].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   /**
@@ -422,9 +404,9 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
   def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag, A5: TypeTag, A6: TypeTag, A7: TypeTag, A8: TypeTag, A9: TypeTag, A10: TypeTag, A11: TypeTag, A12: TypeTag, A13: TypeTag, A14: TypeTag, A15: TypeTag, A16: TypeTag, A17: TypeTag, A18: TypeTag, A19: TypeTag, A20: TypeTag, A21: TypeTag, A22: TypeTag](name: String, func: Function22[A1, A2, A3, A4, A5, A6, A7, A8, A9, A10, A11, A12, A13, A14, A15, A16, A17, A18, A19, A20, A21, A22, RT]): UserDefinedFunction = {
     val dataType = ScalaReflection.schemaFor[RT].dataType
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: ScalaReflection.schemaFor[A21].dataType :: ScalaReflection.schemaFor[A22].dataType :: Nil).getOrElse(Nil)
-    val udf = UserDefinedFunction(func, dataType, inputTypes)
-    functionRegistry.registerFunction(name, udf.builder)
-    udf
+    def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
+    functionRegistry.registerFunction(name, builder)
+    UserDefinedFunction(func, dataType)
   }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/UserDefinedFunction.scala b/sql/core/src/main/scala/org/apache/spark/sql/UserDefinedFunction.scala
index 1319391db5375..0f8cd280b5acb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/UserDefinedFunction.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/UserDefinedFunction.scala
@@ -44,20 +44,11 @@ import org.apache.spark.sql.types.DataType
 case class UserDefinedFunction protected[sql] (
     f: AnyRef,
     dataType: DataType,
-    inputTypes: Seq[DataType] = Nil,
-    deterministic: Boolean = true) {
+    inputTypes: Seq[DataType] = Nil) {
 
   def apply(exprs: Column*): Column = {
-    Column(ScalaUDF(f, dataType, exprs.map(_.expr), inputTypes, deterministic))
+    Column(ScalaUDF(f, dataType, exprs.map(_.expr), inputTypes))
   }
-
-  protected[sql] def builder: Seq[Expression] => ScalaUDF = {
-    (exprs: Seq[Expression]) =>
-      ScalaUDF(f, dataType, exprs, inputTypes, deterministic)
-  }
-
-  def nondeterministic: UserDefinedFunction =
-    UserDefinedFunction(f, dataType, inputTypes, deterministic = false)
 }
 
 /**
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
index 6e510f0b8aff4..e0435a0dba6ad 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql
 
-import org.apache.spark.sql.catalyst.expressions.ScalaUDF
-import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.test.SQLTestData._
 
@@ -193,107 +191,4 @@ class UDFSuite extends QueryTest with SharedSQLContext {
     // pass a decimal to intExpected.
     assert(sql("SELECT intExpected(1.0)").head().getInt(0) === 1)
   }
-
-  private def checkNumUDFs(df: DataFrame, expectedNumUDFs: Int): Unit = {
-    val udfs = df.queryExecution.optimizedPlan.collect {
-      case p: logical.Project => p.projectList.flatMap {
-        case e => e.collect {
-          case udf: ScalaUDF => udf
-        }
-      }
-    }.flatten
-    assert(udfs.length === expectedNumUDFs)
-  }
-
-  test("foldable udf") {
-    import org.apache.spark.sql.functions._
-
-    val myUDF = udf((x: Int) => x + 1)
-
-    {
-      val df = sql("SELECT 1 as a")
-        .select(col("a"), myUDF(col("a")).as("b"))
-        .select(col("a"), col("b"), myUDF(col("b")).as("c"))
-      checkNumUDFs(df, 0)
-      checkAnswer(df, Row(1, 2, 3))
-    }
-  }
-
-  test("nondeterministic udf: using UDFRegistration") {
-    import org.apache.spark.sql.functions._
-
-    val myUDF = sqlContext.udf.register("plusOne1", (x: Int) => x + 1)
-    sqlContext.udf.register("plusOne2", myUDF.nondeterministic)
-
-    {
-      val df = sqlContext.range(1, 2).select(col("id").as("a"))
-        .select(col("a"), myUDF(col("a")).as("b"))
-        .select(col("a"), col("b"), myUDF(col("b")).as("c"))
-      checkNumUDFs(df, 3)
-      checkAnswer(df, Row(1, 2, 3))
-    }
-
-    {
-      val df = sqlContext.range(1, 2).select(col("id").as("a"))
-        .select(col("a"), callUDF("plusOne1", col("a")).as("b"))
-        .select(col("a"), col("b"), callUDF("plusOne1", col("b")).as("c"))
-      checkNumUDFs(df, 3)
-      checkAnswer(df, Row(1, 2, 3))
-    }
-
-    {
-      val df = sqlContext.range(1, 2).select(col("id").as("a"))
-        .select(col("a"), myUDF.nondeterministic(col("a")).as("b"))
-        .select(col("a"), col("b"), myUDF.nondeterministic(col("b")).as("c"))
-      checkNumUDFs(df, 2)
-      checkAnswer(df, Row(1, 2, 3))
-    }
-
-    {
-      val df = sqlContext.range(1, 2).select(col("id").as("a"))
-        .select(col("a"), callUDF("plusOne2", col("a")).as("b"))
-        .select(col("a"), col("b"), callUDF("plusOne2", col("b")).as("c"))
-      checkNumUDFs(df, 2)
-      checkAnswer(df, Row(1, 2, 3))
-    }
-  }
-
-  test("nondeterministic udf: using udf function") {
-    import org.apache.spark.sql.functions._
-
-    val myUDF = udf((x: Int) => x + 1)
-
-    {
-      val df = sqlContext.range(1, 2).select(col("id").as("a"))
-        .select(col("a"), myUDF(col("a")).as("b"))
-        .select(col("a"), col("b"), myUDF(col("b")).as("c"))
-      checkNumUDFs(df, 3)
-      checkAnswer(df, Row(1, 2, 3))
-    }
-
-    {
-      val df = sqlContext.range(1, 2).select(col("id").as("a"))
-        .select(col("a"), myUDF.nondeterministic(col("a")).as("b"))
-        .select(col("a"), col("b"), myUDF.nondeterministic(col("b")).as("c"))
-      checkNumUDFs(df, 2)
-      checkAnswer(df, Row(1, 2, 3))
-    }
-
-    {
-      // nondeterministicUDF will not be foldable.
-      val df = sql("SELECT 1 as a")
-        .select(col("a"), myUDF.nondeterministic(col("a")).as("b"))
-        .select(col("a"), col("b"), myUDF.nondeterministic(col("b")).as("c"))
-      checkNumUDFs(df, 2)
-      checkAnswer(df, Row(1, 2, 3))
-    }
-  }
-
-  test("override a registered udf") {
-    sqlContext.udf.register("intExpected", (x: Int) => x)
-    assert(sql("SELECT intExpected(1.0)").head().getInt(0) === 1)
-
-    sqlContext.udf.register("intExpected", (x: Int) => x + 1)
-    assert(sql("SELECT intExpected(1.0)").head().getInt(0) === 2)
-  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index f14b2886a9ecb..72744799897be 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -381,7 +381,7 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
       sqlContext.udf.register("div0", (x: Int) => x / 0)
       withTempPath { dir =>
         intercept[org.apache.spark.SparkException] {
-          sqlContext.range(1, 2).selectExpr("div0(id) as a").write.parquet(dir.getCanonicalPath)
+          sqlContext.sql("select div0(1)").write.parquet(dir.getCanonicalPath)
         }
         val path = new Path(dir.getCanonicalPath, "_temporary")
         val fs = path.getFileSystem(hadoopConfiguration)
@@ -405,7 +405,7 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
       sqlContext.udf.register("div0", (x: Int) => x / 0)
       withTempPath { dir =>
         intercept[org.apache.spark.SparkException] {
-          sqlContext.range(1, 2).selectExpr("div0(id) as a").write.parquet(dir.getCanonicalPath)
+          sqlContext.sql("select div0(1)").write.parquet(dir.getCanonicalPath)
         }
         val path = new Path(dir.getCanonicalPath, "_temporary")
         val fs = path.getFileSystem(hadoopConfiguration)

From 8fa8c8375d7015a0332aa9ee613d7c6b6d62bae7 Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Thu, 5 Nov 2015 17:59:01 -0800
Subject: [PATCH 085/518] [SPARK-11514][ML] Pass random seed to spark.ml
 DecisionTree*

cc jkbradley

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>

Closes #9486 from yu-iskw/SPARK-11514.
---
 .../ml/classification/DecisionTreeClassifier.scala    |  4 +++-
 .../spark/ml/regression/DecisionTreeRegressor.scala   |  4 +++-
 .../scala/org/apache/spark/ml/tree/treeParams.scala   | 11 ++++++-----
 .../classification/DecisionTreeClassifierSuite.scala  |  1 +
 .../ml/regression/DecisionTreeRegressorSuite.scala    |  1 +
 5 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
index b0157f7ce24ec..c478aea44ace8 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
@@ -62,6 +62,8 @@ final class DecisionTreeClassifier(override val uid: String)
 
   override def setImpurity(value: String): this.type = super.setImpurity(value)
 
+  override def setSeed(value: Long): this.type = super.setSeed(value)
+
   override protected def train(dataset: DataFrame): DecisionTreeClassificationModel = {
     val categoricalFeatures: Map[Int, Int] =
       MetadataUtils.getCategoricalFeatures(dataset.schema($(featuresCol)))
@@ -75,7 +77,7 @@ final class DecisionTreeClassifier(override val uid: String)
     val oldDataset: RDD[LabeledPoint] = extractLabeledPoints(dataset)
     val strategy = getOldStrategy(categoricalFeatures, numClasses)
     val trees = RandomForest.run(oldDataset, strategy, numTrees = 1, featureSubsetStrategy = "all",
-      seed = 0L, parentUID = Some(uid))
+      seed = $(seed), parentUID = Some(uid))
     trees.head.asInstanceOf[DecisionTreeClassificationModel]
   }
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
index 04420fc6e8251..477030d9ea3ee 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
@@ -71,13 +71,15 @@ final class DecisionTreeRegressor @Since("1.4.0") (@Since("1.4.0") override val
   @Since("1.4.0")
   override def setImpurity(value: String): this.type = super.setImpurity(value)
 
+  override def setSeed(value: Long): this.type = super.setSeed(value)
+
   override protected def train(dataset: DataFrame): DecisionTreeRegressionModel = {
     val categoricalFeatures: Map[Int, Int] =
       MetadataUtils.getCategoricalFeatures(dataset.schema($(featuresCol)))
     val oldDataset: RDD[LabeledPoint] = extractLabeledPoints(dataset)
     val strategy = getOldStrategy(categoricalFeatures)
     val trees = RandomForest.run(oldDataset, strategy, numTrees = 1, featureSubsetStrategy = "all",
-      seed = 0L, parentUID = Some(uid))
+      seed = $(seed), parentUID = Some(uid))
     trees.head.asInstanceOf[DecisionTreeRegressionModel]
   }
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala
index 281ba6eeffa92..1da97db9277d8 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala
@@ -29,7 +29,8 @@ import org.apache.spark.mllib.tree.loss.{Loss => OldLoss}
  *
  * Note: Marked as private and DeveloperApi since this may be made public in the future.
  */
-private[ml] trait DecisionTreeParams extends PredictorParams with HasCheckpointInterval {
+private[ml] trait DecisionTreeParams extends PredictorParams
+  with HasCheckpointInterval with HasSeed {
 
   /**
    * Maximum depth of the tree (>= 0).
@@ -123,6 +124,9 @@ private[ml] trait DecisionTreeParams extends PredictorParams with HasCheckpointI
   /** @group getParam */
   final def getMinInfoGain: Double = $(minInfoGain)
 
+  /** @group setParam */
+  def setSeed(value: Long): this.type = set(seed, value)
+
   /** @group expertSetParam */
   def setMaxMemoryInMB(value: Int): this.type = set(maxMemoryInMB, value)
 
@@ -257,7 +261,7 @@ private[ml] object TreeRegressorParams {
  *
  * Note: Marked as private and DeveloperApi since this may be made public in the future.
  */
-private[ml] trait TreeEnsembleParams extends DecisionTreeParams with HasSeed {
+private[ml] trait TreeEnsembleParams extends DecisionTreeParams {
 
   /**
    * Fraction of the training data used for learning each decision tree, in range (0, 1].
@@ -276,9 +280,6 @@ private[ml] trait TreeEnsembleParams extends DecisionTreeParams with HasSeed {
   /** @group getParam */
   final def getSubsamplingRate: Double = $(subsamplingRate)
 
-  /** @group setParam */
-  def setSeed(value: Long): this.type = set(seed, value)
-
   /**
    * Create a Strategy instance to use with the old API.
    * NOTE: The caller should set impurity and seed.
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
index 815f6fd997584..92b8f84144ab0 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
@@ -72,6 +72,7 @@ class DecisionTreeClassifierSuite extends SparkFunSuite with MLlibTestSparkConte
       .setImpurity("gini")
       .setMaxDepth(2)
       .setMaxBins(100)
+      .setSeed(1)
     val categoricalFeatures = Map(0 -> 3, 1-> 3)
     val numClasses = 2
     compareAPIs(categoricalDataPointsRDD, dt, categoricalFeatures, numClasses)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
index 868fb8eecb8bb..e0d5afa7a7e97 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
@@ -49,6 +49,7 @@ class DecisionTreeRegressorSuite extends SparkFunSuite with MLlibTestSparkContex
       .setImpurity("variance")
       .setMaxDepth(2)
       .setMaxBins(100)
+      .setSeed(1)
     val categoricalFeatures = Map(0 -> 3, 1-> 3)
     compareAPIs(categoricalDataPointsRDD, dt, categoricalFeatures)
   }

From 468ad0ae874d5cf55712ee976faf77f19c937ccb Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Thu, 5 Nov 2015 18:03:12 -0800
Subject: [PATCH 086/518] [SPARK-11457][STREAMING][YARN] Fix incorrect AM proxy
 filter conf recovery from checkpoint

Currently Yarn AM proxy filter configuration is recovered from checkpoint file when Spark Streaming application is restarted, which will lead to some unwanted behaviors:

1. Wrong RM address if RM is redeployed from failure.
2. Wrong proxyBase, since app id is updated, old app id for proxyBase is wrong.

So instead of recovering from checkpoint file, these configurations should be reloaded each time when app started.

This problem only exists in Yarn cluster mode, for Yarn client mode, these configurations will be updated with RPC message `AddWebUIFilter`.

Please help to review tdas harishreedharan vanzin , thanks a lot.

Author: jerryshao <sshao@hortonworks.com>

Closes #9412 from jerryshao/SPARK-11457.
---
 .../org/apache/spark/streaming/Checkpoint.scala     | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala b/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala
index b7de6dde61c63..0cd55d9aec2cd 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala
@@ -55,7 +55,8 @@ class Checkpoint(ssc: StreamingContext, val checkpointTime: Time)
       "spark.driver.port",
       "spark.master",
       "spark.yarn.keytab",
-      "spark.yarn.principal")
+      "spark.yarn.principal",
+      "spark.ui.filters")
 
     val newSparkConf = new SparkConf(loadDefaults = false).setAll(sparkConfPairs)
       .remove("spark.driver.host")
@@ -66,6 +67,16 @@ class Checkpoint(ssc: StreamingContext, val checkpointTime: Time)
         newSparkConf.set(prop, value)
       }
     }
+
+    // Add Yarn proxy filter specific configurations to the recovered SparkConf
+    val filter = "org.apache.hadoop.yarn.server.webproxy.amfilter.AmIpFilter"
+    val filterPrefix = s"spark.$filter.param."
+    newReloadConf.getAll.foreach { case (k, v) =>
+      if (k.startsWith(filterPrefix) && k.length > filterPrefix.length) {
+        newSparkConf.set(k, v)
+      }
+    }
+
     newSparkConf
   }
 

From 5e31db70bb783656ba042863fcd3c223e17a8f81 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Thu, 5 Nov 2015 18:05:58 -0800
Subject: [PATCH 087/518] [SPARK-11538][BUILD] Force guava 14 in sbt build.

sbt's version resolution code always picks the most recent version, and we
don't want that for guava.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9508 from vanzin/SPARK-11538.
---
 project/SparkBuild.scala | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index 75c36930decef..b75ed13a78c68 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -207,7 +207,8 @@ object SparkBuild extends PomBuild {
   // Note ordering of these settings matter.
   /* Enable shared settings on all projects */
   (allProjects ++ optionallyEnabledProjects ++ assemblyProjects ++ Seq(spark, tools))
-    .foreach(enable(sharedSettings ++ ExcludedDependencies.settings ++ Revolver.settings))
+    .foreach(enable(sharedSettings ++ DependencyOverrides.settings ++
+      ExcludedDependencies.settings ++ Revolver.settings))
 
   /* Enable tests settings for all projects except examples, assembly and tools */
   (allProjects ++ optionallyEnabledProjects).foreach(enable(TestSettings.settings))
@@ -291,6 +292,14 @@ object Flume {
   lazy val settings = sbtavro.SbtAvro.avroSettings
 }
 
+/**
+ * Overrides to work around sbt's dependency resolution being different from Maven's.
+ */
+object DependencyOverrides {
+  lazy val settings = Seq(
+    dependencyOverrides += "com.google.guava" % "guava" % "14.0.1")
+}
+
 /**
   This excludes library dependencies in sbt, which are specified in maven but are
   not needed by sbt build.

From 3cc2c053b5d68c747a30bd58cf388b87b1922f13 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 5 Nov 2015 18:12:54 -0800
Subject: [PATCH 088/518] [SPARK-11540][SQL] API audit for
 QueryExecutionListener.

Author: Reynold Xin <rxin@databricks.com>

Closes #9509 from rxin/SPARK-11540.
---
 .../spark/sql/execution/QueryExecution.scala  |  30 +++---
 .../sql/util/QueryExecutionListener.scala     | 101 ++++++++++--------
 2 files changed, 72 insertions(+), 59 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index fc9174549e642..c2142d03f422b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.execution
 
+import com.google.common.annotations.VisibleForTesting
+
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.SQLContext
 import org.apache.spark.sql.catalyst.InternalRow
@@ -25,31 +27,33 @@ import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 /**
  * The primary workflow for executing relational queries using Spark.  Designed to allow easy
  * access to the intermediate phases of query execution for developers.
+ *
+ * While this is not a public class, we should avoid changing the function names for the sake of
+ * changing them, because a lot of developers use the feature for debugging.
  */
 class QueryExecution(val sqlContext: SQLContext, val logical: LogicalPlan) {
-  val analyzer = sqlContext.analyzer
-  val optimizer = sqlContext.optimizer
-  val planner = sqlContext.planner
-  val cacheManager = sqlContext.cacheManager
-  val prepareForExecution = sqlContext.prepareForExecution
 
-  def assertAnalyzed(): Unit = analyzer.checkAnalysis(analyzed)
+  @VisibleForTesting
+  def assertAnalyzed(): Unit = sqlContext.analyzer.checkAnalysis(analyzed)
+
+  lazy val analyzed: LogicalPlan = sqlContext.analyzer.execute(logical)
 
-  lazy val analyzed: LogicalPlan = analyzer.execute(logical)
   lazy val withCachedData: LogicalPlan = {
     assertAnalyzed()
-    cacheManager.useCachedData(analyzed)
+    sqlContext.cacheManager.useCachedData(analyzed)
   }
-  lazy val optimizedPlan: LogicalPlan = optimizer.execute(withCachedData)
+
+  lazy val optimizedPlan: LogicalPlan = sqlContext.optimizer.execute(withCachedData)
 
   // TODO: Don't just pick the first one...
   lazy val sparkPlan: SparkPlan = {
     SparkPlan.currentContext.set(sqlContext)
-    planner.plan(optimizedPlan).next()
+    sqlContext.planner.plan(optimizedPlan).next()
   }
+
   // executedPlan should not be used to initialize any SparkPlan. It should be
   // only used for execution.
-  lazy val executedPlan: SparkPlan = prepareForExecution.execute(sparkPlan)
+  lazy val executedPlan: SparkPlan = sqlContext.prepareForExecution.execute(sparkPlan)
 
   /** Internal version of the RDD. Avoids copies and has no schema */
   lazy val toRdd: RDD[InternalRow] = executedPlan.execute()
@@ -57,11 +61,11 @@ class QueryExecution(val sqlContext: SQLContext, val logical: LogicalPlan) {
   protected def stringOrError[A](f: => A): String =
     try f.toString catch { case e: Throwable => e.toString }
 
-  def simpleString: String =
+  def simpleString: String = {
     s"""== Physical Plan ==
        |${stringOrError(executedPlan)}
       """.stripMargin.trim
-
+  }
 
   override def toString: String = {
     def output =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/util/QueryExecutionListener.scala b/sql/core/src/main/scala/org/apache/spark/sql/util/QueryExecutionListener.scala
index 909a8abd225b8..ac432e2baa3c0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/util/QueryExecutionListener.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/util/QueryExecutionListener.scala
@@ -19,36 +19,38 @@ package org.apache.spark.sql.util
 
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import scala.collection.mutable.ListBuffer
+import scala.util.control.NonFatal
 
-import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.Logging
+import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.sql.execution.QueryExecution
 
 
 /**
+ * :: Experimental ::
  * The interface of query execution listener that can be used to analyze execution metrics.
  *
- * Note that implementations should guarantee thread-safety as they will be used in a non
- * thread-safe way.
+ * Note that implementations should guarantee thread-safety as they can be invoked by
+ * multiple different threads.
  */
 @Experimental
 trait QueryExecutionListener {
 
   /**
    * A callback function that will be called when a query executed successfully.
-   * Implementations should guarantee thread-safe.
+   * Note that this can be invoked by multiple different threads.
    *
-   * @param funcName the name of the action that triggered this query.
+   * @param funcName name of the action that triggered this query.
    * @param qe the QueryExecution object that carries detail information like logical plan,
    *           physical plan, etc.
-   * @param duration the execution time for this query in nanoseconds.
+   * @param durationNs the execution time for this query in nanoseconds.
    */
   @DeveloperApi
-  def onSuccess(funcName: String, qe: QueryExecution, duration: Long)
+  def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit
 
   /**
    * A callback function that will be called when a query execution failed.
-   * Implementations should guarantee thread-safe.
+   * Note that this can be invoked by multiple different threads.
    *
    * @param funcName the name of the action that triggered this query.
    * @param qe the QueryExecution object that carries detail information like logical plan,
@@ -56,34 +58,20 @@ trait QueryExecutionListener {
    * @param exception the exception that failed this query.
    */
   @DeveloperApi
-  def onFailure(funcName: String, qe: QueryExecution, exception: Exception)
+  def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit
 }
 
-@Experimental
-class ExecutionListenerManager extends Logging {
-  private[this] val listeners = ListBuffer.empty[QueryExecutionListener]
-  private[this] val lock = new ReentrantReadWriteLock()
-
-  /** Acquires a read lock on the cache for the duration of `f`. */
-  private def readLock[A](f: => A): A = {
-    val rl = lock.readLock()
-    rl.lock()
-    try f finally {
-      rl.unlock()
-    }
-  }
 
-  /** Acquires a write lock on the cache for the duration of `f`. */
-  private def writeLock[A](f: => A): A = {
-    val wl = lock.writeLock()
-    wl.lock()
-    try f finally {
-      wl.unlock()
-    }
-  }
+/**
+ * :: Experimental ::
+ *
+ * Manager for [[QueryExecutionListener]]. See [[org.apache.spark.sql.SQLContext.listenerManager]].
+ */
+@Experimental
+class ExecutionListenerManager private[sql] () extends Logging {
 
   /**
-   * Registers the specified QueryExecutionListener.
+   * Registers the specified [[QueryExecutionListener]].
    */
   @DeveloperApi
   def register(listener: QueryExecutionListener): Unit = writeLock {
@@ -91,7 +79,7 @@ class ExecutionListenerManager extends Logging {
   }
 
   /**
-   * Unregisters the specified QueryExecutionListener.
+   * Unregisters the specified [[QueryExecutionListener]].
    */
   @DeveloperApi
   def unregister(listener: QueryExecutionListener): Unit = writeLock {
@@ -99,38 +87,59 @@ class ExecutionListenerManager extends Logging {
   }
 
   /**
-   * clears out all registered QueryExecutionListeners.
+   * Removes all the registered [[QueryExecutionListener]].
    */
   @DeveloperApi
   def clear(): Unit = writeLock {
     listeners.clear()
   }
 
-  private[sql] def onSuccess(
-      funcName: String,
-      qe: QueryExecution,
-      duration: Long): Unit = readLock {
-    withErrorHandling { listener =>
-      listener.onSuccess(funcName, qe, duration)
+  private[sql] def onSuccess(funcName: String, qe: QueryExecution, duration: Long): Unit = {
+    readLock {
+      withErrorHandling { listener =>
+        listener.onSuccess(funcName, qe, duration)
+      }
     }
   }
 
-  private[sql] def onFailure(
-      funcName: String,
-      qe: QueryExecution,
-      exception: Exception): Unit = readLock {
-    withErrorHandling { listener =>
-      listener.onFailure(funcName, qe, exception)
+  private[sql] def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = {
+    readLock {
+      withErrorHandling { listener =>
+        listener.onFailure(funcName, qe, exception)
+      }
     }
   }
 
+  private[this] val listeners = ListBuffer.empty[QueryExecutionListener]
+
+  /** A lock to prevent updating the list of listeners while we are traversing through them. */
+  private[this] val lock = new ReentrantReadWriteLock()
+
   private def withErrorHandling(f: QueryExecutionListener => Unit): Unit = {
     for (listener <- listeners) {
       try {
         f(listener)
       } catch {
-        case e: Exception => logWarning("error executing query execution listener", e)
+        case NonFatal(e) => logWarning("Error executing query execution listener", e)
       }
     }
   }
+
+  /** Acquires a read lock on the cache for the duration of `f`. */
+  private def readLock[A](f: => A): A = {
+    val rl = lock.readLock()
+    rl.lock()
+    try f finally {
+      rl.unlock()
+    }
+  }
+
+  /** Acquires a write lock on the cache for the duration of `f`. */
+  private def writeLock[A](f: => A): A = {
+    val wl = lock.writeLock()
+    wl.lock()
+    try f finally {
+      wl.unlock()
+    }
+  }
 }

From eec74ba8bde7f9446cc38e687bda103e85669d35 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Thu, 5 Nov 2015 19:02:18 -0800
Subject: [PATCH 089/518] [SPARK-7542][SQL] Support off-heap index/sort buffer

This brings the support of off-heap memory for array inside BytesToBytesMap and InMemorySorter, then we could allocate all the memory from off-heap for execution.

Closes #8068

Author: Davies Liu <davies@databricks.com>

Closes #9477 from davies/unsafe_timsort.
---
 .../apache/spark/memory/MemoryConsumer.java   | 36 +++++-----
 .../spark/memory/TaskMemoryManager.java       |  6 +-
 .../shuffle/sort/ShuffleExternalSorter.java   | 26 +++----
 .../shuffle/sort/ShuffleInMemorySorter.java   | 67 ++++++++++---------
 .../shuffle/sort/ShuffleSortDataFormat.java   | 38 +++++++----
 .../spark/unsafe/map/BytesToBytesMap.java     | 18 +++--
 .../unsafe/sort/UnsafeExternalSorter.java     | 28 +++-----
 .../unsafe/sort/UnsafeInMemorySorter.java     | 66 +++++++++++-------
 .../unsafe/sort/UnsafeSortDataFormat.java     | 47 +++++++------
 .../spark/memory/TaskMemoryManagerSuite.java  | 23 -------
 .../spark/memory/TestMemoryConsumer.java      | 45 +++++++++++++
 .../sort/ShuffleInMemorySorterSuite.java      | 16 +++--
 .../sort/UnsafeExternalSorterSuite.java       |  1 -
 .../sort/UnsafeInMemorySorterSuite.java       | 12 ++--
 .../sql/execution/UnsafeKVExternalSorter.java |  3 +-
 .../apache/spark/unsafe/array/LongArray.java  | 18 ++++-
 .../spark/unsafe/array/LongArraySuite.java    |  4 ++
 17 files changed, 265 insertions(+), 189 deletions(-)
 create mode 100644 core/src/test/java/org/apache/spark/memory/TestMemoryConsumer.java

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index 008799cc77395..8fbdb72832adf 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -20,6 +20,7 @@
 
 import java.io.IOException;
 
+import org.apache.spark.unsafe.array.LongArray;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 
 
@@ -28,9 +29,9 @@
  */
 public abstract class MemoryConsumer {
 
-  private final TaskMemoryManager taskMemoryManager;
+  protected final TaskMemoryManager taskMemoryManager;
   private final long pageSize;
-  private long used;
+  protected long used;
 
   protected MemoryConsumer(TaskMemoryManager taskMemoryManager, long pageSize) {
     this.taskMemoryManager = taskMemoryManager;
@@ -74,26 +75,29 @@ public void spill() throws IOException {
   public abstract long spill(long size, MemoryConsumer trigger) throws IOException;
 
   /**
-   * Acquire `size` bytes memory.
-   *
-   * If there is not enough memory, throws OutOfMemoryError.
+   * Allocates a LongArray of `size`.
    */
-  protected void acquireMemory(long size) {
-    long got = taskMemoryManager.acquireExecutionMemory(size, this);
-    if (got < size) {
-      taskMemoryManager.releaseExecutionMemory(got, this);
+  public LongArray allocateArray(long size) {
+    long required = size * 8L;
+    MemoryBlock page = taskMemoryManager.allocatePage(required, this);
+    if (page == null || page.size() < required) {
+      long got = 0;
+      if (page != null) {
+        got = page.size();
+        taskMemoryManager.freePage(page, this);
+      }
       taskMemoryManager.showMemoryUsage();
-      throw new OutOfMemoryError("Could not acquire " + size + " bytes of memory, got " + got);
+      throw new OutOfMemoryError("Unable to acquire " + required + " bytes of memory, got " + got);
     }
-    used += got;
+    used += required;
+    return new LongArray(page);
   }
 
   /**
-   * Release `size` bytes memory.
+   * Frees a LongArray.
    */
-  protected void releaseMemory(long size) {
-    used -= size;
-    taskMemoryManager.releaseExecutionMemory(size, this);
+  public void freeArray(LongArray array) {
+    freePage(array.memoryBlock());
   }
 
   /**
@@ -109,7 +113,7 @@ protected MemoryBlock allocatePage(long required) {
       long got = 0;
       if (page != null) {
         got = page.size();
-        freePage(page);
+        taskMemoryManager.freePage(page, this);
       }
       taskMemoryManager.showMemoryUsage();
       throw new OutOfMemoryError("Unable to acquire " + required + " bytes of memory, got " + got);
diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 4230575446d31..6440f9c0f30de 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -137,7 +137,7 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
       if (got < required) {
         // Call spill() on other consumers to release memory
         for (MemoryConsumer c: consumers) {
-          if (c != null && c != consumer && c.getUsed() > 0) {
+          if (c != consumer && c.getUsed() > 0) {
             try {
               long released = c.spill(required - got, consumer);
               if (released > 0) {
@@ -173,7 +173,9 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
         }
       }
 
-      consumers.add(consumer);
+      if (consumer != null) {
+        consumers.add(consumer);
+      }
       logger.debug("Task {} acquire {} for {}", taskAttemptId, Utils.bytesToString(got), consumer);
       return got;
     }
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
index 400d8520019b9..9affff80143d7 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleExternalSorter.java
@@ -39,6 +39,7 @@
 import org.apache.spark.storage.DiskBlockObjectWriter;
 import org.apache.spark.storage.TempShuffleBlockId;
 import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.array.LongArray;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.util.Utils;
 
@@ -114,8 +115,7 @@ public ShuffleExternalSorter(
     this.numElementsForSpillThreshold =
       conf.getLong("spark.shuffle.spill.numElementsForceSpillThreshold", Long.MAX_VALUE);
     this.writeMetrics = writeMetrics;
-    acquireMemory(initialSize * 8L);
-    this.inMemSorter = new ShuffleInMemorySorter(initialSize);
+    this.inMemSorter = new ShuffleInMemorySorter(this, initialSize);
     this.peakMemoryUsedBytes = getMemoryUsage();
   }
 
@@ -301,9 +301,8 @@ private long freeMemory() {
   public void cleanupResources() {
     freeMemory();
     if (inMemSorter != null) {
-      long sorterMemoryUsage = inMemSorter.getMemoryUsage();
+      inMemSorter.free();
       inMemSorter = null;
-      releaseMemory(sorterMemoryUsage);
     }
     for (SpillInfo spill : spills) {
       if (spill.file.exists() && !spill.file.delete()) {
@@ -321,9 +320,10 @@ private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       long used = inMemSorter.getMemoryUsage();
-      long needed = used + inMemSorter.getMemoryToExpand();
+      LongArray array;
       try {
-        acquireMemory(needed);  // could trigger spilling
+        // could trigger spilling
+        array = allocateArray(used / 8 * 2);
       } catch (OutOfMemoryError e) {
         // should have trigger spilling
         assert(inMemSorter.hasSpaceForAnotherRecord());
@@ -331,16 +331,9 @@ private void growPointerArrayIfNecessary() throws IOException {
       }
       // check if spilling is triggered or not
       if (inMemSorter.hasSpaceForAnotherRecord()) {
-        releaseMemory(needed);
+        freeArray(array);
       } else {
-        try {
-          inMemSorter.expandPointerArray();
-          releaseMemory(used);
-        } catch (OutOfMemoryError oom) {
-          // Just in case that JVM had run out of memory
-          releaseMemory(needed);
-          spill();
-        }
+        inMemSorter.expandPointerArray(array);
       }
     }
   }
@@ -404,9 +397,8 @@ public SpillInfo[] closeAndGetSpills() throws IOException {
         // Do not count the final file towards the spill count.
         writeSortedFile(true);
         freeMemory();
-        long sorterMemoryUsage = inMemSorter.getMemoryUsage();
+        inMemSorter.free();
         inMemSorter = null;
-        releaseMemory(sorterMemoryUsage);
       }
       return spills.toArray(new SpillInfo[spills.size()]);
     } catch (IOException e) {
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
index e630575d1ae19..58ad88e1ed87b 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorter.java
@@ -19,11 +19,14 @@
 
 import java.util.Comparator;
 
+import org.apache.spark.memory.MemoryConsumer;
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.array.LongArray;
 import org.apache.spark.util.collection.Sorter;
 
 final class ShuffleInMemorySorter {
 
-  private final Sorter<PackedRecordPointer, long[]> sorter;
+  private final Sorter<PackedRecordPointer, LongArray> sorter;
   private static final class SortComparator implements Comparator<PackedRecordPointer> {
     @Override
     public int compare(PackedRecordPointer left, PackedRecordPointer right) {
@@ -32,24 +35,34 @@ public int compare(PackedRecordPointer left, PackedRecordPointer right) {
   }
   private static final SortComparator SORT_COMPARATOR = new SortComparator();
 
+  private final MemoryConsumer consumer;
+
   /**
    * An array of record pointers and partition ids that have been encoded by
    * {@link PackedRecordPointer}. The sort operates on this array instead of directly manipulating
    * records.
    */
-  private long[] array;
+  private LongArray array;
 
   /**
    * The position in the pointer array where new records can be inserted.
    */
   private int pos = 0;
 
-  public ShuffleInMemorySorter(int initialSize) {
+  public ShuffleInMemorySorter(MemoryConsumer consumer, int initialSize) {
+    this.consumer = consumer;
     assert (initialSize > 0);
-    this.array = new long[initialSize];
+    this.array = consumer.allocateArray(initialSize);
     this.sorter = new Sorter<>(ShuffleSortDataFormat.INSTANCE);
   }
 
+  public void free() {
+    if (array != null) {
+      consumer.freeArray(array);
+      array = null;
+    }
+  }
+
   public int numRecords() {
     return pos;
   }
@@ -58,30 +71,25 @@ public void reset() {
     pos = 0;
   }
 
-  private int newLength() {
-    // Guard against overflow:
-    return array.length <= Integer.MAX_VALUE / 2 ? (array.length * 2) : Integer.MAX_VALUE;
-  }
-
-  /**
-   * Returns the memory needed to expand
-   */
-  public long getMemoryToExpand() {
-    return ((long) (newLength() - array.length)) * 8;
-  }
-
-  public void expandPointerArray() {
-    final long[] oldArray = array;
-    array = new long[newLength()];
-    System.arraycopy(oldArray, 0, array, 0, oldArray.length);
+  public void expandPointerArray(LongArray newArray) {
+    assert(newArray.size() > array.size());
+    Platform.copyMemory(
+      array.getBaseObject(),
+      array.getBaseOffset(),
+      newArray.getBaseObject(),
+      newArray.getBaseOffset(),
+      array.size() * 8L
+    );
+    consumer.freeArray(array);
+    array = newArray;
   }
 
   public boolean hasSpaceForAnotherRecord() {
-    return pos < array.length;
+    return pos < array.size();
   }
 
   public long getMemoryUsage() {
-    return array.length * 8L;
+    return array.size() * 8L;
   }
 
   /**
@@ -96,14 +104,9 @@ public long getMemoryUsage() {
    */
   public void insertRecord(long recordPointer, int partitionId) {
     if (!hasSpaceForAnotherRecord()) {
-      if (array.length == Integer.MAX_VALUE) {
-        throw new IllegalStateException("Sort pointer array has reached maximum size");
-      } else {
-        expandPointerArray();
-      }
+      expandPointerArray(consumer.allocateArray(array.size() * 2));
     }
-    array[pos] =
-        PackedRecordPointer.packPointer(recordPointer, partitionId);
+    array.set(pos, PackedRecordPointer.packPointer(recordPointer, partitionId));
     pos++;
   }
 
@@ -112,12 +115,12 @@ public void insertRecord(long recordPointer, int partitionId) {
    */
   public static final class ShuffleSorterIterator {
 
-    private final long[] pointerArray;
+    private final LongArray pointerArray;
     private final int numRecords;
     final PackedRecordPointer packedRecordPointer = new PackedRecordPointer();
     private int position = 0;
 
-    public ShuffleSorterIterator(int numRecords, long[] pointerArray) {
+    public ShuffleSorterIterator(int numRecords, LongArray pointerArray) {
       this.numRecords = numRecords;
       this.pointerArray = pointerArray;
     }
@@ -127,7 +130,7 @@ public boolean hasNext() {
     }
 
     public void loadNext() {
-      packedRecordPointer.set(pointerArray[position]);
+      packedRecordPointer.set(pointerArray.get(position));
       position++;
     }
   }
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleSortDataFormat.java b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleSortDataFormat.java
index 8a1e5aec6ff0e..8f4e3229976dc 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleSortDataFormat.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/ShuffleSortDataFormat.java
@@ -17,16 +17,19 @@
 
 package org.apache.spark.shuffle.sort;
 
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.array.LongArray;
+import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.util.collection.SortDataFormat;
 
-final class ShuffleSortDataFormat extends SortDataFormat<PackedRecordPointer, long[]> {
+final class ShuffleSortDataFormat extends SortDataFormat<PackedRecordPointer, LongArray> {
 
   public static final ShuffleSortDataFormat INSTANCE = new ShuffleSortDataFormat();
 
   private ShuffleSortDataFormat() { }
 
   @Override
-  public PackedRecordPointer getKey(long[] data, int pos) {
+  public PackedRecordPointer getKey(LongArray data, int pos) {
     // Since we re-use keys, this method shouldn't be called.
     throw new UnsupportedOperationException();
   }
@@ -37,31 +40,38 @@ public PackedRecordPointer newKey() {
   }
 
   @Override
-  public PackedRecordPointer getKey(long[] data, int pos, PackedRecordPointer reuse) {
-    reuse.set(data[pos]);
+  public PackedRecordPointer getKey(LongArray data, int pos, PackedRecordPointer reuse) {
+    reuse.set(data.get(pos));
     return reuse;
   }
 
   @Override
-  public void swap(long[] data, int pos0, int pos1) {
-    final long temp = data[pos0];
-    data[pos0] = data[pos1];
-    data[pos1] = temp;
+  public void swap(LongArray data, int pos0, int pos1) {
+    final long temp = data.get(pos0);
+    data.set(pos0, data.get(pos1));
+    data.set(pos1, temp);
   }
 
   @Override
-  public void copyElement(long[] src, int srcPos, long[] dst, int dstPos) {
-    dst[dstPos] = src[srcPos];
+  public void copyElement(LongArray src, int srcPos, LongArray dst, int dstPos) {
+    dst.set(dstPos, src.get(srcPos));
   }
 
   @Override
-  public void copyRange(long[] src, int srcPos, long[] dst, int dstPos, int length) {
-    System.arraycopy(src, srcPos, dst, dstPos, length);
+  public void copyRange(LongArray src, int srcPos, LongArray dst, int dstPos, int length) {
+    Platform.copyMemory(
+      src.getBaseObject(),
+      src.getBaseOffset() + srcPos * 8,
+      dst.getBaseObject(),
+      dst.getBaseOffset() + dstPos * 8,
+      length * 8
+    );
   }
 
   @Override
-  public long[] allocate(int length) {
-    return new long[length];
+  public LongArray allocate(int length) {
+    // This buffer is used temporary (usually small), so it's fine to allocated from JVM heap.
+    return new LongArray(MemoryBlock.fromLongArray(new long[length]));
   }
 
 }
diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 6656fd1d0bc59..04694dc54418c 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -20,7 +20,6 @@
 import javax.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.Iterator;
 import java.util.LinkedList;
 
@@ -724,11 +723,10 @@ public long spill(long size, MemoryConsumer trigger) throws IOException {
    */
   private void allocate(int capacity) {
     assert (capacity >= 0);
-    // The capacity needs to be divisible by 64 so that our bit set can be sized properly
     capacity = Math.max((int) Math.min(MAX_CAPACITY, ByteArrayMethods.nextPowerOf2(capacity)), 64);
     assert (capacity <= MAX_CAPACITY);
-    acquireMemory(capacity * 16);
-    longArray = new LongArray(MemoryBlock.fromLongArray(new long[capacity * 2]));
+    longArray = allocateArray(capacity * 2);
+    longArray.zeroOut();
 
     this.growthThreshold = (int) (capacity * loadFactor);
     this.mask = capacity - 1;
@@ -743,9 +741,8 @@ private void allocate(int capacity) {
   public void free() {
     updatePeakMemoryUsed();
     if (longArray != null) {
-      long used = longArray.memoryBlock().size();
+      freeArray(longArray);
       longArray = null;
-      releaseMemory(used);
     }
     Iterator<MemoryBlock> dataPagesIterator = dataPages.iterator();
     while (dataPagesIterator.hasNext()) {
@@ -834,9 +831,9 @@ public int getNumDataPages() {
   /**
    * Returns the underline long[] of longArray.
    */
-  public long[] getArray() {
+  public LongArray getArray() {
     assert(longArray != null);
-    return (long[]) longArray.memoryBlock().getBaseObject();
+    return longArray;
   }
 
   /**
@@ -844,7 +841,8 @@ public long[] getArray() {
    */
   public void reset() {
     numElements = 0;
-    Arrays.fill(getArray(), 0);
+    longArray.zeroOut();
+
     while (dataPages.size() > 0) {
       MemoryBlock dataPage = dataPages.removeLast();
       freePage(dataPage);
@@ -887,7 +885,7 @@ void growAndRehash() {
       longArray.set(newPos * 2, keyPointer);
       longArray.set(newPos * 2 + 1, hashcode);
     }
-    releaseMemory(oldLongArray.memoryBlock().size());
+    freeArray(oldLongArray);
 
     if (enablePerfMetrics) {
       timeSpentResizingNs += System.nanoTime() - resizeStartTime;
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index cba043bc48cc8..9a7b2ad06cab6 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -32,6 +32,7 @@
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.storage.BlockManager;
 import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.array.LongArray;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.util.TaskCompletionListener;
 import org.apache.spark.util.Utils;
@@ -123,9 +124,8 @@ private UnsafeExternalSorter(
     this.writeMetrics = new ShuffleWriteMetrics();
 
     if (existingInMemorySorter == null) {
-      this.inMemSorter =
-        new UnsafeInMemorySorter(taskMemoryManager, recordComparator, prefixComparator, initialSize);
-      acquireMemory(inMemSorter.getMemoryUsage());
+      this.inMemSorter = new UnsafeInMemorySorter(
+        this, taskMemoryManager, recordComparator, prefixComparator, initialSize);
     } else {
       this.inMemSorter = existingInMemorySorter;
     }
@@ -277,9 +277,8 @@ public void cleanupResources() {
       deleteSpillFiles();
       freeMemory();
       if (inMemSorter != null) {
-        long used = inMemSorter.getMemoryUsage();
+        inMemSorter.free();
         inMemSorter = null;
-        releaseMemory(used);
       }
     }
   }
@@ -293,9 +292,10 @@ private void growPointerArrayIfNecessary() throws IOException {
     assert(inMemSorter != null);
     if (!inMemSorter.hasSpaceForAnotherRecord()) {
       long used = inMemSorter.getMemoryUsage();
-      long needed = used + inMemSorter.getMemoryToExpand();
+      LongArray array;
       try {
-        acquireMemory(needed);  // could trigger spilling
+        // could trigger spilling
+        array = allocateArray(used / 8 * 2);
       } catch (OutOfMemoryError e) {
         // should have trigger spilling
         assert(inMemSorter.hasSpaceForAnotherRecord());
@@ -303,16 +303,9 @@ private void growPointerArrayIfNecessary() throws IOException {
       }
       // check if spilling is triggered or not
       if (inMemSorter.hasSpaceForAnotherRecord()) {
-        releaseMemory(needed);
+        freeArray(array);
       } else {
-        try {
-          inMemSorter.expandPointerArray();
-          releaseMemory(used);
-        } catch (OutOfMemoryError oom) {
-          // Just in case that JVM had run out of memory
-          releaseMemory(needed);
-          spill();
-        }
+        inMemSorter.expandPointerArray(array);
       }
     }
   }
@@ -498,9 +491,8 @@ public void loadNext() throws IOException {
           nextUpstream = null;
 
           assert(inMemSorter != null);
-          long used = inMemSorter.getMemoryUsage();
+          inMemSorter.free();
           inMemSorter = null;
-          releaseMemory(used);
         }
         numRecords--;
         upstream.loadNext();
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
index d57213b9b8bfc..a218ad4623f46 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -19,8 +19,10 @@
 
 import java.util.Comparator;
 
+import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.array.LongArray;
 import org.apache.spark.util.collection.Sorter;
 
 /**
@@ -62,15 +64,16 @@ public int compare(RecordPointerAndKeyPrefix r1, RecordPointerAndKeyPrefix r2) {
     }
   }
 
+  private final MemoryConsumer consumer;
   private final TaskMemoryManager memoryManager;
-  private final Sorter<RecordPointerAndKeyPrefix, long[]> sorter;
+  private final Sorter<RecordPointerAndKeyPrefix, LongArray> sorter;
   private final Comparator<RecordPointerAndKeyPrefix> sortComparator;
 
   /**
    * Within this buffer, position {@code 2 * i} holds a pointer pointer to the record at
    * index {@code i}, while position {@code 2 * i + 1} in the array holds an 8-byte key prefix.
    */
-  private long[] array;
+  private LongArray array;
 
   /**
    * The position in the sort buffer where new records can be inserted.
@@ -78,22 +81,33 @@ public int compare(RecordPointerAndKeyPrefix r1, RecordPointerAndKeyPrefix r2) {
   private int pos = 0;
 
   public UnsafeInMemorySorter(
+    final MemoryConsumer consumer,
     final TaskMemoryManager memoryManager,
     final RecordComparator recordComparator,
     final PrefixComparator prefixComparator,
     int initialSize) {
-    this(memoryManager, recordComparator, prefixComparator, new long[initialSize * 2]);
+    this(consumer, memoryManager, recordComparator, prefixComparator,
+      consumer.allocateArray(initialSize * 2));
   }
 
   public UnsafeInMemorySorter(
+    final MemoryConsumer consumer,
       final TaskMemoryManager memoryManager,
       final RecordComparator recordComparator,
       final PrefixComparator prefixComparator,
-      long[] array) {
-    this.array = array;
+      LongArray array) {
+    this.consumer = consumer;
     this.memoryManager = memoryManager;
     this.sorter = new Sorter<>(UnsafeSortDataFormat.INSTANCE);
     this.sortComparator = new SortComparator(recordComparator, prefixComparator, memoryManager);
+    this.array = array;
+  }
+
+  /**
+   * Free the memory used by pointer array.
+   */
+  public void free() {
+    consumer.freeArray(array);
   }
 
   public void reset() {
@@ -107,26 +121,26 @@ public int numRecords() {
     return pos / 2;
   }
 
-  private int newLength() {
-    return array.length < Integer.MAX_VALUE / 2 ? (array.length * 2) : Integer.MAX_VALUE;
-  }
-
-  public long getMemoryToExpand() {
-    return (long) (newLength() - array.length) * 8L;
-  }
-
   public long getMemoryUsage() {
-    return array.length * 8L;
+    return array.size() * 8L;
   }
 
   public boolean hasSpaceForAnotherRecord() {
-    return pos + 2 <= array.length;
+    return pos + 2 <= array.size();
   }
 
-  public void expandPointerArray() {
-    final long[] oldArray = array;
-    array = new long[newLength()];
-    System.arraycopy(oldArray, 0, array, 0, oldArray.length);
+  public void expandPointerArray(LongArray newArray) {
+    if (newArray.size() < array.size()) {
+      throw new OutOfMemoryError("Not enough memory to grow pointer array");
+    }
+    Platform.copyMemory(
+      array.getBaseObject(),
+      array.getBaseOffset(),
+      newArray.getBaseObject(),
+      newArray.getBaseOffset(),
+      array.size() * 8L);
+    consumer.freeArray(array);
+    array = newArray;
   }
 
   /**
@@ -138,11 +152,11 @@ public void expandPointerArray() {
    */
   public void insertRecord(long recordPointer, long keyPrefix) {
     if (!hasSpaceForAnotherRecord()) {
-      expandPointerArray();
+      expandPointerArray(consumer.allocateArray(array.size() * 2));
     }
-    array[pos] = recordPointer;
+    array.set(pos, recordPointer);
     pos++;
-    array[pos] = keyPrefix;
+    array.set(pos, keyPrefix);
     pos++;
   }
 
@@ -150,7 +164,7 @@ public static final class SortedIterator extends UnsafeSorterIterator {
 
     private final TaskMemoryManager memoryManager;
     private final int sortBufferInsertPosition;
-    private final long[] sortBuffer;
+    private final LongArray sortBuffer;
     private int position = 0;
     private Object baseObject;
     private long baseOffset;
@@ -160,7 +174,7 @@ public static final class SortedIterator extends UnsafeSorterIterator {
     private SortedIterator(
         TaskMemoryManager memoryManager,
         int sortBufferInsertPosition,
-        long[] sortBuffer) {
+        LongArray sortBuffer) {
       this.memoryManager = memoryManager;
       this.sortBufferInsertPosition = sortBufferInsertPosition;
       this.sortBuffer = sortBuffer;
@@ -188,11 +202,11 @@ public int numRecordsLeft() {
     @Override
     public void loadNext() {
       // This pointer points to a 4-byte record length, followed by the record's bytes
-      final long recordPointer = sortBuffer[position];
+      final long recordPointer = sortBuffer.get(position);
       baseObject = memoryManager.getPage(recordPointer);
       baseOffset = memoryManager.getOffsetInPage(recordPointer) + 4;  // Skip over record length
       recordLength = Platform.getInt(baseObject, baseOffset - 4);
-      keyPrefix = sortBuffer[position + 1];
+      keyPrefix = sortBuffer.get(position + 1);
       position += 2;
     }
 
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSortDataFormat.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSortDataFormat.java
index d09c728a7a638..d3137f5f31c25 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSortDataFormat.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSortDataFormat.java
@@ -17,6 +17,9 @@
 
 package org.apache.spark.util.collection.unsafe.sort;
 
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.array.LongArray;
+import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.util.collection.SortDataFormat;
 
 /**
@@ -26,14 +29,14 @@
  * Within each long[] buffer, position {@code 2 * i} holds a pointer pointer to the record at
  * index {@code i}, while position {@code 2 * i + 1} in the array holds an 8-byte key prefix.
  */
-final class UnsafeSortDataFormat extends SortDataFormat<RecordPointerAndKeyPrefix, long[]> {
+final class UnsafeSortDataFormat extends SortDataFormat<RecordPointerAndKeyPrefix, LongArray> {
 
   public static final UnsafeSortDataFormat INSTANCE = new UnsafeSortDataFormat();
 
   private UnsafeSortDataFormat() { }
 
   @Override
-  public RecordPointerAndKeyPrefix getKey(long[] data, int pos) {
+  public RecordPointerAndKeyPrefix getKey(LongArray data, int pos) {
     // Since we re-use keys, this method shouldn't be called.
     throw new UnsupportedOperationException();
   }
@@ -44,37 +47,43 @@ public RecordPointerAndKeyPrefix newKey() {
   }
 
   @Override
-  public RecordPointerAndKeyPrefix getKey(long[] data, int pos, RecordPointerAndKeyPrefix reuse) {
-    reuse.recordPointer = data[pos * 2];
-    reuse.keyPrefix = data[pos * 2 + 1];
+  public RecordPointerAndKeyPrefix getKey(LongArray data, int pos, RecordPointerAndKeyPrefix reuse) {
+    reuse.recordPointer = data.get(pos * 2);
+    reuse.keyPrefix = data.get(pos * 2 + 1);
     return reuse;
   }
 
   @Override
-  public void swap(long[] data, int pos0, int pos1) {
-    long tempPointer = data[pos0 * 2];
-    long tempKeyPrefix = data[pos0 * 2 + 1];
-    data[pos0 * 2] = data[pos1 * 2];
-    data[pos0 * 2 + 1] = data[pos1 * 2 + 1];
-    data[pos1 * 2] = tempPointer;
-    data[pos1 * 2 + 1] = tempKeyPrefix;
+  public void swap(LongArray data, int pos0, int pos1) {
+    long tempPointer = data.get(pos0 * 2);
+    long tempKeyPrefix = data.get(pos0 * 2 + 1);
+    data.set(pos0 * 2, data.get(pos1 * 2));
+    data.set(pos0 * 2 + 1, data.get(pos1 * 2 + 1));
+    data.set(pos1 * 2, tempPointer);
+    data.set(pos1 * 2 + 1, tempKeyPrefix);
   }
 
   @Override
-  public void copyElement(long[] src, int srcPos, long[] dst, int dstPos) {
-    dst[dstPos * 2] = src[srcPos * 2];
-    dst[dstPos * 2 + 1] = src[srcPos * 2 + 1];
+  public void copyElement(LongArray src, int srcPos, LongArray dst, int dstPos) {
+    dst.set(dstPos * 2, src.get(srcPos * 2));
+    dst.set(dstPos * 2 + 1, src.get(srcPos * 2 + 1));
   }
 
   @Override
-  public void copyRange(long[] src, int srcPos, long[] dst, int dstPos, int length) {
-    System.arraycopy(src, srcPos * 2, dst, dstPos * 2, length * 2);
+  public void copyRange(LongArray src, int srcPos, LongArray dst, int dstPos, int length) {
+    Platform.copyMemory(
+      src.getBaseObject(),
+      src.getBaseOffset() + srcPos * 16,
+      dst.getBaseObject(),
+      dst.getBaseOffset() + dstPos * 16,
+      length * 16);
   }
 
   @Override
-  public long[] allocate(int length) {
+  public LongArray allocate(int length) {
     assert (length < Integer.MAX_VALUE / 2) : "Length " + length + " is too large";
-    return new long[length * 2];
+    // This is used as temporary buffer, it's fine to allocate from JVM heap.
+    return new LongArray(MemoryBlock.fromLongArray(new long[length * 2]));
   }
 
 }
diff --git a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
index dab7b0592cb4e..c731317395612 100644
--- a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
+++ b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
@@ -17,8 +17,6 @@
 
 package org.apache.spark.memory;
 
-import java.io.IOException;
-
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -27,27 +25,6 @@
 
 public class TaskMemoryManagerSuite {
 
-  class TestMemoryConsumer extends MemoryConsumer {
-    TestMemoryConsumer(TaskMemoryManager memoryManager) {
-      super(memoryManager);
-    }
-
-    @Override
-    public long spill(long size, MemoryConsumer trigger) throws IOException {
-      long used = getUsed();
-      releaseMemory(used);
-      return used;
-    }
-
-    void use(long size) {
-      acquireMemory(size);
-    }
-
-    void free(long size) {
-      releaseMemory(size);
-    }
-  }
-
   @Test
   public void leakedPageMemoryIsDetected() {
     final TaskMemoryManager manager = new TaskMemoryManager(
diff --git a/core/src/test/java/org/apache/spark/memory/TestMemoryConsumer.java b/core/src/test/java/org/apache/spark/memory/TestMemoryConsumer.java
new file mode 100644
index 0000000000000..8ae3642738509
--- /dev/null
+++ b/core/src/test/java/org/apache/spark/memory/TestMemoryConsumer.java
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.memory;
+
+import java.io.IOException;
+
+public class TestMemoryConsumer extends MemoryConsumer {
+  public TestMemoryConsumer(TaskMemoryManager memoryManager) {
+    super(memoryManager);
+  }
+
+  @Override
+  public long spill(long size, MemoryConsumer trigger) throws IOException {
+    long used = getUsed();
+    free(used);
+    return used;
+  }
+
+  void use(long size) {
+    long got = taskMemoryManager.acquireExecutionMemory(size, this);
+    used += got;
+  }
+
+  void free(long size) {
+    used -= size;
+    taskMemoryManager.releaseExecutionMemory(size, this);
+  }
+}
+
+
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
index 2293b1bbc113e..faa5a863ee630 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/ShuffleInMemorySorterSuite.java
@@ -25,13 +25,19 @@
 
 import org.apache.spark.HashPartitioner;
 import org.apache.spark.SparkConf;
-import org.apache.spark.unsafe.Platform;
+import org.apache.spark.memory.TaskMemoryManager;
+import org.apache.spark.memory.TestMemoryConsumer;
 import org.apache.spark.memory.TestMemoryManager;
+import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.memory.MemoryBlock;
-import org.apache.spark.memory.TaskMemoryManager;
 
 public class ShuffleInMemorySorterSuite {
 
+  final TestMemoryManager memoryManager =
+    new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false"));
+  final TaskMemoryManager taskMemoryManager = new TaskMemoryManager(memoryManager, 0);
+  final TestMemoryConsumer consumer = new TestMemoryConsumer(taskMemoryManager);
+
   private static String getStringFromDataPage(Object baseObject, long baseOffset, int strLength) {
     final byte[] strBytes = new byte[strLength];
     Platform.copyMemory(baseObject, baseOffset, strBytes, Platform.BYTE_ARRAY_OFFSET, strLength);
@@ -40,7 +46,7 @@ private static String getStringFromDataPage(Object baseObject, long baseOffset,
 
   @Test
   public void testSortingEmptyInput() {
-    final ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(100);
+    final ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(consumer, 100);
     final ShuffleInMemorySorter.ShuffleSorterIterator iter = sorter.getSortedIterator();
     assert(!iter.hasNext());
   }
@@ -63,7 +69,7 @@ public void testBasicSorting() throws Exception {
       new TaskMemoryManager(new TestMemoryManager(conf), 0);
     final MemoryBlock dataPage = memoryManager.allocatePage(2048, null);
     final Object baseObject = dataPage.getBaseObject();
-    final ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(4);
+    final ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(consumer, 4);
     final HashPartitioner hashPartitioner = new HashPartitioner(4);
 
     // Write the records into the data page and store pointers into the sorter
@@ -104,7 +110,7 @@ public void testBasicSorting() throws Exception {
 
   @Test
   public void testSortingManyNumbers() throws Exception {
-    ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(4);
+    ShuffleInMemorySorter sorter = new ShuffleInMemorySorter(consumer, 4);
     int[] numbersToSort = new int[128000];
     Random random = new Random(16);
     for (int i = 0; i < numbersToSort.length; i++) {
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
index cfead0e5924b8..11c3a7be38875 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
@@ -390,7 +390,6 @@ public void testPeakMemoryUsed() throws Exception {
       for (int i = 0; i < numRecordsPerPage * 10; i++) {
         insertNumber(sorter, i);
         newPeakMemory = sorter.getPeakMemoryUsedBytes();
-        // The first page is pre-allocated on instantiation
         if (i % numRecordsPerPage == 0) {
           // We allocated a new page for this record, so peak memory should change
           assertEquals(previousPeakMemory + pageSizeBytes, newPeakMemory);
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
index 642f6585f8a15..a203a09648ac0 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorterSuite.java
@@ -23,6 +23,7 @@
 
 import org.apache.spark.HashPartitioner;
 import org.apache.spark.SparkConf;
+import org.apache.spark.memory.TestMemoryConsumer;
 import org.apache.spark.memory.TestMemoryManager;
 import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.unsafe.Platform;
@@ -44,9 +45,11 @@ private static String getStringFromDataPage(Object baseObject, long baseOffset,
 
   @Test
   public void testSortingEmptyInput() {
-    final UnsafeInMemorySorter sorter = new UnsafeInMemorySorter(
-      new TaskMemoryManager(
-        new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0),
+    final TaskMemoryManager memoryManager = new TaskMemoryManager(
+      new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
+    final TestMemoryConsumer consumer = new TestMemoryConsumer(memoryManager);
+    final UnsafeInMemorySorter sorter = new UnsafeInMemorySorter(consumer,
+      memoryManager,
       mock(RecordComparator.class),
       mock(PrefixComparator.class),
       100);
@@ -69,6 +72,7 @@ public void testSortingOnlyByIntegerPrefix() throws Exception {
     };
     final TaskMemoryManager memoryManager = new TaskMemoryManager(
       new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
+    final TestMemoryConsumer consumer = new TestMemoryConsumer(memoryManager);
     final MemoryBlock dataPage = memoryManager.allocatePage(2048, null);
     final Object baseObject = dataPage.getBaseObject();
     // Write the records into the data page:
@@ -102,7 +106,7 @@ public int compare(long prefix1, long prefix2) {
         return (int) prefix1 - (int) prefix2;
       }
     };
-    UnsafeInMemorySorter sorter = new UnsafeInMemorySorter(memoryManager, recordComparator,
+    UnsafeInMemorySorter sorter = new UnsafeInMemorySorter(consumer, memoryManager, recordComparator,
       prefixComparator, dataToSort.length);
     // Given a page of records, insert those records into the sorter one-by-one:
     position = dataPage.getBaseOffset();
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
index e2898ef2e2158..8c9b9c85e37fc 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeKVExternalSorter.java
@@ -85,8 +85,9 @@ public UnsafeKVExternalSorter(
     } else {
       // During spilling, the array in map will not be used, so we can borrow that and use it
       // as the underline array for in-memory sorter (it's always large enough).
+      // Since we will not grow the array, it's fine to pass `null` as consumer.
       final UnsafeInMemorySorter inMemSorter = new UnsafeInMemorySorter(
-        taskMemoryManager, recordComparator, prefixComparator, map.getArray());
+        null, taskMemoryManager, recordComparator, prefixComparator, map.getArray());
 
       // We cannot use the destructive iterator here because we are reusing the existing memory
       // pages in BytesToBytesMap to hold records during sorting.
diff --git a/unsafe/src/main/java/org/apache/spark/unsafe/array/LongArray.java b/unsafe/src/main/java/org/apache/spark/unsafe/array/LongArray.java
index 74105050e4191..1a3cdff638264 100644
--- a/unsafe/src/main/java/org/apache/spark/unsafe/array/LongArray.java
+++ b/unsafe/src/main/java/org/apache/spark/unsafe/array/LongArray.java
@@ -39,7 +39,6 @@ public final class LongArray {
   private final long length;
 
   public LongArray(MemoryBlock memory) {
-    assert memory.size() % WIDTH == 0 : "Memory not aligned (" + memory.size() + ")";
     assert memory.size() < (long) Integer.MAX_VALUE * 8: "Array size > 4 billion elements";
     this.memory = memory;
     this.baseObj = memory.getBaseObject();
@@ -51,6 +50,14 @@ public MemoryBlock memoryBlock() {
     return memory;
   }
 
+  public Object getBaseObject() {
+    return baseObj;
+  }
+
+  public long getBaseOffset() {
+    return baseOffset;
+  }
+
   /**
    * Returns the number of elements this array can hold.
    */
@@ -58,6 +65,15 @@ public long size() {
     return length;
   }
 
+  /**
+   * Fill this all with 0L.
+   */
+  public void zeroOut() {
+    for (long off = baseOffset; off < baseOffset + length * WIDTH; off += WIDTH) {
+      Platform.putLong(baseObj, off, 0);
+    }
+  }
+
   /**
    * Sets the value at position {@code index}.
    */
diff --git a/unsafe/src/test/java/org/apache/spark/unsafe/array/LongArraySuite.java b/unsafe/src/test/java/org/apache/spark/unsafe/array/LongArraySuite.java
index 5974cf91ff993..fb8e53b3348f3 100644
--- a/unsafe/src/test/java/org/apache/spark/unsafe/array/LongArraySuite.java
+++ b/unsafe/src/test/java/org/apache/spark/unsafe/array/LongArraySuite.java
@@ -34,5 +34,9 @@ public void basicTest() {
     Assert.assertEquals(2, arr.size());
     Assert.assertEquals(1L, arr.get(0));
     Assert.assertEquals(3L, arr.get(1));
+
+    arr.zeroOut();
+    Assert.assertEquals(0L, arr.get(0));
+    Assert.assertEquals(0L, arr.get(1));
   }
 }

From 363a476c3fefb0263e63fd24df0b2779a64f79ec Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Thu, 5 Nov 2015 21:42:32 -0800
Subject: [PATCH 090/518] [SPARK-11528] [SQL] Typed aggregations for Datasets

This PR adds the ability to do typed SQL aggregations.  We will likely also want to provide an interface to allow users to do aggregations on objects, but this is deferred to another PR.

```scala
val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
ds.groupBy(_._1).agg(sum("_2").as[Int]).collect()

res0: Array(("a", 30), ("b", 3), ("c", 1))
```

Author: Michael Armbrust <michael@databricks.com>

Closes #9499 from marmbrus/dataset-agg.
---
 .../expressions/namedExpressions.scala        |  4 +
 .../scala/org/apache/spark/sql/Dataset.scala  |  2 +-
 .../org/apache/spark/sql/GroupedDataset.scala | 93 ++++++++++++++++++-
 .../org/apache/spark/sql/DatasetSuite.scala   | 36 +++++++
 4 files changed, 132 insertions(+), 3 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
index 8957df0be6814..9ab5c299d0f55 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
@@ -254,6 +254,10 @@ case class AttributeReference(
   }
 
   override def toString: String = s"$name#${exprId.id}$typeSuffix"
+
+  // Since the expression id is not in the first constructor it is missing from the default
+  // tree string.
+  override def simpleString: String = s"$name#${exprId.id}: ${dataType.simpleString}"
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 500227e93a472..4bca9c3b3fe54 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -55,7 +55,7 @@ import org.apache.spark.sql.types.StructType
  * @since 1.6.0
  */
 @Experimental
-class Dataset[T] private(
+class Dataset[T] private[sql](
     @transient val sqlContext: SQLContext,
     @transient val queryExecution: QueryExecution,
     unresolvedEncoder: Encoder[T]) extends Serializable {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 96d6e9dd548e5..b8fc373dffcf5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -17,16 +17,25 @@
 
 package org.apache.spark.sql
 
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttribute}
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor, Encoder}
-import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression, Alias, Attribute}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.QueryExecution
 
 /**
+ * :: Experimental ::
  * A [[Dataset]] has been logically grouped by a user specified grouping key.  Users should not
  * construct a [[GroupedDataset]] directly, but should instead call `groupBy` on an existing
  * [[Dataset]].
+ *
+ * COMPATIBILITY NOTE: Long term we plan to make [[GroupedDataset)]] extend `GroupedData`.  However,
+ * making this change to the class hierarchy would break some function signatures. As such, this
+ * class should be considered a preview of the final API.  Changes will be made to the interface
+ * after Spark 1.6.
  */
+@Experimental
 class GroupedDataset[K, T] private[sql](
     private val kEncoder: Encoder[K],
     private val tEncoder: Encoder[T],
@@ -35,7 +44,7 @@ class GroupedDataset[K, T] private[sql](
     private val groupingAttributes: Seq[Attribute]) extends Serializable {
 
   private implicit val kEnc = kEncoder match {
-    case e: ExpressionEncoder[K] => e.resolve(groupingAttributes)
+    case e: ExpressionEncoder[K] => e.unbind(groupingAttributes).resolve(groupingAttributes)
     case other =>
       throw new UnsupportedOperationException("Only expression encoders are currently supported")
   }
@@ -46,9 +55,16 @@ class GroupedDataset[K, T] private[sql](
       throw new UnsupportedOperationException("Only expression encoders are currently supported")
   }
 
+  /** Encoders for built in aggregations. */
+  private implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder[Long](flat = true)
+
   private def logicalPlan = queryExecution.analyzed
   private def sqlContext = queryExecution.sqlContext
 
+  private def groupedData =
+    new GroupedData(
+      new DataFrame(sqlContext, logicalPlan), groupingAttributes, GroupedData.GroupByType)
+
   /**
    * Returns a new [[GroupedDataset]] where the type of the key has been mapped to the specified
    * type. The mapping of key columns to the type follows the same rules as `as` on [[Dataset]].
@@ -88,6 +104,79 @@ class GroupedDataset[K, T] private[sql](
       MapGroups(f, groupingAttributes, logicalPlan))
   }
 
+  // To ensure valid overloading.
+  protected def agg(expr: Column, exprs: Column*): DataFrame =
+    groupedData.agg(expr, exprs: _*)
+
+  /**
+   * Internal helper function for building typed aggregations that return tuples.  For simplicity
+   * and code reuse, we do this without the help of the type system and then use helper functions
+   * that cast appropriately for the user facing interface.
+   * TODO: does not handle aggrecations that return nonflat results,
+   */
+  protected def aggUntyped(columns: TypedColumn[_]*): Dataset[_] = {
+    val aliases = (groupingAttributes ++ columns.map(_.expr)).map {
+      case u: UnresolvedAttribute => UnresolvedAlias(u)
+      case expr: NamedExpression => expr
+      case expr: Expression => Alias(expr, expr.prettyString)()
+    }
+
+    val unresolvedPlan = Aggregate(groupingAttributes, aliases, logicalPlan)
+    val execution = new QueryExecution(sqlContext, unresolvedPlan)
+
+    val columnEncoders = columns.map(_.encoder.asInstanceOf[ExpressionEncoder[_]])
+
+    // Rebind the encoders to the nested schema that will be produced by the aggregation.
+    val encoders = (kEnc +: columnEncoders).zip(execution.analyzed.output).map {
+      case (e: ExpressionEncoder[_], a) if !e.flat =>
+        e.nested(a).resolve(execution.analyzed.output)
+      case (e, a) =>
+        e.unbind(a :: Nil).resolve(execution.analyzed.output)
+    }
+    new Dataset(sqlContext, execution, ExpressionEncoder.tuple(encoders))
+  }
+
+  /**
+   * Computes the given aggregation, returning a [[Dataset]] of tuples for each unique key
+   * and the result of computing this aggregation over all elements in the group.
+   */
+  def agg[A1](col1: TypedColumn[A1]): Dataset[(K, A1)] =
+    aggUntyped(col1).asInstanceOf[Dataset[(K, A1)]]
+
+  /**
+   * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
+   * and the result of computing these aggregations over all elements in the group.
+   */
+  def agg[A1, A2](col1: TypedColumn[A1], col2: TypedColumn[A2]): Dataset[(K, A1, A2)] =
+    aggUntyped(col1, col2).asInstanceOf[Dataset[(K, A1, A2)]]
+
+  /**
+   * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
+   * and the result of computing these aggregations over all elements in the group.
+   */
+  def agg[A1, A2, A3](
+      col1: TypedColumn[A1],
+      col2: TypedColumn[A2],
+      col3: TypedColumn[A3]): Dataset[(K, A1, A2, A3)] =
+    aggUntyped(col1, col2, col3).asInstanceOf[Dataset[(K, A1, A2, A3)]]
+
+  /**
+   * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
+   * and the result of computing these aggregations over all elements in the group.
+   */
+  def agg[A1, A2, A3, A4](
+      col1: TypedColumn[A1],
+      col2: TypedColumn[A2],
+      col3: TypedColumn[A3],
+      col4: TypedColumn[A4]): Dataset[(K, A1, A2, A3, A4)] =
+    aggUntyped(col1, col2, col3, col4).asInstanceOf[Dataset[(K, A1, A2, A3, A4)]]
+
+  /**
+   * Returns a [[Dataset]] that contains a tuple with each key and the number of items present
+   * for that key.
+   */
+  def count(): Dataset[(K, Long)] = agg(functions.count("*").as[Long])
+
   /**
    * Applies the given function to each cogrouped data.  For each unique group, the function will
    * be passed the grouping key and 2 iterators containing all elements in the group from
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 3e9b621cfd67f..d61e17edc64ed 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -258,6 +258,42 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       (ClassData("a", 1), 30), (ClassData("b", 1), 3), (ClassData("c", 1), 1))
   }
 
+  test("typed aggregation: expr") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+
+    checkAnswer(
+      ds.groupBy(_._1).agg(sum("_2").as[Int]),
+      ("a", 30), ("b", 3), ("c", 1))
+  }
+
+  test("typed aggregation: expr, expr") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+
+    checkAnswer(
+      ds.groupBy(_._1).agg(sum("_2").as[Int], sum($"_2" + 1).as[Long]),
+      ("a", 30, 32L), ("b", 3, 5L), ("c", 1, 2L))
+  }
+
+  test("typed aggregation: expr, expr, expr") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+
+    checkAnswer(
+      ds.groupBy(_._1).agg(sum("_2").as[Int], sum($"_2" + 1).as[Long], count("*").as[Long]),
+      ("a", 30, 32L, 2L), ("b", 3, 5L, 2L), ("c", 1, 2L, 1L))
+  }
+
+  test("typed aggregation: expr, expr, expr, expr") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+
+    checkAnswer(
+      ds.groupBy(_._1).agg(
+        sum("_2").as[Int],
+        sum($"_2" + 1).as[Long],
+        count("*").as[Long],
+        avg("_2").as[Double]),
+      ("a", 30, 32L, 2L, 15.0), ("b", 3, 5L, 2L, 1.5), ("c", 1, 2L, 1L, 1.0))
+  }
+
   test("cogroup") {
     val ds1 = Seq(1 -> "a", 3 -> "abc", 5 -> "hello", 3 -> "foo").toDS()
     val ds2 = Seq(2 -> "q", 3 -> "w", 5 -> "e", 5 -> "r").toDS()

From bc5d6c03893a9bd340d6b94d3550e25648412241 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 5 Nov 2015 22:03:26 -0800
Subject: [PATCH 091/518] [SPARK-11541][SQL] Break JdbcDialects.scala into
 multiple files and mark various dialects as private.

Author: Reynold Xin <rxin@databricks.com>

Closes #9511 from rxin/SPARK-11541.
---
 project/MimaExcludes.scala                    |  19 +-
 .../org/apache/spark/sql/GroupedData.scala    |   2 +-
 .../spark/sql/jdbc/AggregatedDialect.scala    |  44 ++++
 .../apache/spark/sql/jdbc/DB2Dialect.scala    |  32 +++
 .../apache/spark/sql/jdbc/DerbyDialect.scala  |  44 ++++
 .../apache/spark/sql/jdbc/JdbcDialects.scala  | 190 +-----------------
 .../spark/sql/jdbc/MsSqlServerDialect.scala   |  41 ++++
 .../apache/spark/sql/jdbc/MySQLDialect.scala  |  48 +++++
 .../apache/spark/sql/jdbc/OracleDialect.scala |  45 +++++
 .../spark/sql/jdbc/PostgresDialect.scala      |  54 +++++
 10 files changed, 332 insertions(+), 187 deletions(-)
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/jdbc/DB2Dialect.scala
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/jdbc/DerbyDialect.scala
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/jdbc/MsSqlServerDialect.scala
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/jdbc/OracleDialect.scala
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala

diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 40f5c9fec8bb8..dacef911e397e 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -116,7 +116,24 @@ object MimaExcludes {
           "org.apache.spark.rdd.MapPartitionsWithPreparationRDD$")
       ) ++ Seq(
         // SPARK-11485
-        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.DataFrameHolder.df")
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.DataFrameHolder.df"),
+        // SPARK-11541 mark various JDBC dialects as private
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.NoopDialect.productElement"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.NoopDialect.productArity"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.NoopDialect.canEqual"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.NoopDialect.productIterator"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.NoopDialect.productPrefix"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.NoopDialect.toString"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.NoopDialect.hashCode"),
+        ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.jdbc.PostgresDialect$"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.productElement"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.productArity"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.canEqual"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.productIterator"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.productPrefix"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.toString"),
+        ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.hashCode"),
+        ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.jdbc.NoopDialect$")
       )
     case v if v.startsWith("1.5") =>
       Seq(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
index 7cf66b65c8722..f9eab5c2e965b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
@@ -40,7 +40,7 @@ import org.apache.spark.sql.types.NumericType
 class GroupedData protected[sql](
     df: DataFrame,
     groupingExprs: Seq[Expression],
-    private val groupType: GroupedData.GroupType) {
+    groupType: GroupedData.GroupType) {
 
   private[this] def toDF(aggExprs: Seq[Expression]): DataFrame = {
     val aggregates = if (df.sqlContext.conf.dataFrameRetainGroupColumns) {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala
new file mode 100644
index 0000000000000..467d8d62d1b7f
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/AggregatedDialect.scala
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import org.apache.spark.sql.types.{DataType, MetadataBuilder}
+
+/**
+ * AggregatedDialect can unify multiple dialects into one virtual Dialect.
+ * Dialects are tried in order, and the first dialect that does not return a
+ * neutral element will will.
+ *
+ * @param dialects List of dialects.
+ */
+private class AggregatedDialect(dialects: List[JdbcDialect]) extends JdbcDialect {
+
+  require(dialects.nonEmpty)
+
+  override def canHandle(url : String): Boolean =
+    dialects.map(_.canHandle(url)).reduce(_ && _)
+
+  override def getCatalystType(
+      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
+    dialects.flatMap(_.getCatalystType(sqlType, typeName, size, md)).headOption
+  }
+
+  override def getJDBCType(dt: DataType): Option[JdbcType] = {
+    dialects.flatMap(_.getJDBCType(dt)).headOption
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/DB2Dialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/DB2Dialect.scala
new file mode 100644
index 0000000000000..b1cb0e55026be
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/DB2Dialect.scala
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import org.apache.spark.sql.types.{BooleanType, StringType, DataType}
+
+
+private object DB2Dialect extends JdbcDialect {
+
+  override def canHandle(url: String): Boolean = url.startsWith("jdbc:db2")
+
+  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
+    case StringType => Option(JdbcType("CLOB", java.sql.Types.CLOB))
+    case BooleanType => Option(JdbcType("CHAR(1)", java.sql.Types.CHAR))
+    case _ => None
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/DerbyDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/DerbyDialect.scala
new file mode 100644
index 0000000000000..84f68e779c38c
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/DerbyDialect.scala
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import java.sql.Types
+
+import org.apache.spark.sql.types._
+
+
+private object DerbyDialect extends JdbcDialect {
+
+  override def canHandle(url: String): Boolean = url.startsWith("jdbc:derby")
+
+  override def getCatalystType(
+      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
+    if (sqlType == Types.REAL) Option(FloatType) else None
+  }
+
+  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
+    case StringType => Option(JdbcType("CLOB", java.sql.Types.CLOB))
+    case ByteType => Option(JdbcType("SMALLINT", java.sql.Types.SMALLINT))
+    case ShortType => Option(JdbcType("SMALLINT", java.sql.Types.SMALLINT))
+    case BooleanType => Option(JdbcType("BOOLEAN", java.sql.Types.BOOLEAN))
+    // 31 is the maximum precision and 5 is the default scale for a Derby DECIMAL
+    case t: DecimalType if t.precision > 31 =>
+      Option(JdbcType("DECIMAL(31,5)", java.sql.Types.DECIMAL))
+    case _ => None
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
index f9a6a09b6270d..14bfea4e3e287 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql.jdbc
 
-import java.sql.Types
-
 import org.apache.spark.sql.types._
 import org.apache.spark.annotation.DeveloperApi
 
@@ -115,11 +113,10 @@ abstract class JdbcDialect {
 @DeveloperApi
 object JdbcDialects {
 
-  private var dialects = List[JdbcDialect]()
-
   /**
    * Register a dialect for use on all new matching jdbc [[org.apache.spark.sql.DataFrame]].
    * Readding an existing dialect will cause a move-to-front.
+   *
    * @param dialect The new dialect.
    */
   def registerDialect(dialect: JdbcDialect) : Unit = {
@@ -128,12 +125,15 @@ object JdbcDialects {
 
   /**
    * Unregister a dialect. Does nothing if the dialect is not registered.
+   *
    * @param dialect The jdbc dialect.
    */
   def unregisterDialect(dialect : JdbcDialect) : Unit = {
     dialects = dialects.filterNot(_ == dialect)
   }
 
+  private[this] var dialects = List[JdbcDialect]()
+
   registerDialect(MySQLDialect)
   registerDialect(PostgresDialect)
   registerDialect(DB2Dialect)
@@ -141,7 +141,6 @@ object JdbcDialects {
   registerDialect(DerbyDialect)
   registerDialect(OracleDialect)
 
-
   /**
    * Fetch the JdbcDialect class corresponding to a given database url.
    */
@@ -156,187 +155,8 @@ object JdbcDialects {
 }
 
 /**
- * :: DeveloperApi ::
- * AggregatedDialect can unify multiple dialects into one virtual Dialect.
- * Dialects are tried in order, and the first dialect that does not return a
- * neutral element will will.
- * @param dialects List of dialects.
- */
-@DeveloperApi
-class AggregatedDialect(dialects: List[JdbcDialect]) extends JdbcDialect {
-
-  require(dialects.nonEmpty)
-
-  override def canHandle(url : String): Boolean =
-    dialects.map(_.canHandle(url)).reduce(_ && _)
-
-  override def getCatalystType(
-      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
-    dialects.flatMap(_.getCatalystType(sqlType, typeName, size, md)).headOption
-  }
-
-  override def getJDBCType(dt: DataType): Option[JdbcType] = {
-    dialects.flatMap(_.getJDBCType(dt)).headOption
-  }
-}
-
-/**
- * :: DeveloperApi ::
  * NOOP dialect object, always returning the neutral element.
  */
-@DeveloperApi
-case object NoopDialect extends JdbcDialect {
+private object NoopDialect extends JdbcDialect {
   override def canHandle(url : String): Boolean = true
 }
-
-/**
- * :: DeveloperApi ::
- * Default postgres dialect, mapping bit/cidr/inet on read and string/binary/boolean on write.
- */
-@DeveloperApi
-case object PostgresDialect extends JdbcDialect {
-  override def canHandle(url: String): Boolean = url.startsWith("jdbc:postgresql")
-  override def getCatalystType(
-      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
-    if (sqlType == Types.BIT && typeName.equals("bit") && size != 1) {
-      Option(BinaryType)
-    } else if (sqlType == Types.OTHER && typeName.equals("cidr")) {
-      Option(StringType)
-    } else if (sqlType == Types.OTHER && typeName.equals("inet")) {
-      Option(StringType)
-    } else if (sqlType == Types.OTHER && typeName.equals("json")) {
-      Option(StringType)
-    } else if (sqlType == Types.OTHER && typeName.equals("jsonb")) {
-      Option(StringType)
-    } else None
-  }
-
-  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
-    case StringType => Some(JdbcType("TEXT", java.sql.Types.CHAR))
-    case BinaryType => Some(JdbcType("BYTEA", java.sql.Types.BINARY))
-    case BooleanType => Some(JdbcType("BOOLEAN", java.sql.Types.BOOLEAN))
-    case _ => None
-  }
-
-  override def getTableExistsQuery(table: String): String = {
-    s"SELECT 1 FROM $table LIMIT 1"
-  }
-
-}
-
-/**
- * :: DeveloperApi ::
- * Default mysql dialect to read bit/bitsets correctly.
- */
-@DeveloperApi
-case object MySQLDialect extends JdbcDialect {
-  override def canHandle(url : String): Boolean = url.startsWith("jdbc:mysql")
-  override def getCatalystType(
-      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
-    if (sqlType == Types.VARBINARY && typeName.equals("BIT") && size != 1) {
-      // This could instead be a BinaryType if we'd rather return bit-vectors of up to 64 bits as
-      // byte arrays instead of longs.
-      md.putLong("binarylong", 1)
-      Option(LongType)
-    } else if (sqlType == Types.BIT && typeName.equals("TINYINT")) {
-      Option(BooleanType)
-    } else None
-  }
-
-  override def quoteIdentifier(colName: String): String = {
-    s"`$colName`"
-  }
-
-  override def getTableExistsQuery(table: String): String = {
-    s"SELECT 1 FROM $table LIMIT 1"
-  }
-}
-
-/**
- * :: DeveloperApi ::
- * Default DB2 dialect, mapping string/boolean on write to valid DB2 types.
- * By default string, and boolean gets mapped to db2 invalid types TEXT, and BIT(1).
- */
-@DeveloperApi
-case object DB2Dialect extends JdbcDialect {
-
-  override def canHandle(url: String): Boolean = url.startsWith("jdbc:db2")
-
-  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
-    case StringType => Some(JdbcType("CLOB", java.sql.Types.CLOB))
-    case BooleanType => Some(JdbcType("CHAR(1)", java.sql.Types.CHAR))
-    case _ => None
-  }
-}
-
-/**
- * :: DeveloperApi ::
- * Default Microsoft SQL Server dialect, mapping the datetimeoffset types to a String on read.
- */
-@DeveloperApi
-case object MsSqlServerDialect extends JdbcDialect {
-  override def canHandle(url: String): Boolean = url.startsWith("jdbc:sqlserver")
-  override def getCatalystType(
-      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
-    if (typeName.contains("datetimeoffset")) {
-      // String is recommend by Microsoft SQL Server for datetimeoffset types in non-MS clients
-      Option(StringType)
-    } else None
-  }
-
-  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
-    case TimestampType => Some(JdbcType("DATETIME", java.sql.Types.TIMESTAMP))
-    case _ => None
-  }
-}
-
-/**
- * :: DeveloperApi ::
- * Default Apache Derby dialect, mapping real on read
- * and string/byte/short/boolean/decimal on write.
- */
-@DeveloperApi
-case object DerbyDialect extends JdbcDialect {
-  override def canHandle(url: String): Boolean = url.startsWith("jdbc:derby")
-  override def getCatalystType(
-      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
-    if (sqlType == Types.REAL) Option(FloatType) else None
-  }
-
-  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
-    case StringType => Some(JdbcType("CLOB", java.sql.Types.CLOB))
-    case ByteType => Some(JdbcType("SMALLINT", java.sql.Types.SMALLINT))
-    case ShortType => Some(JdbcType("SMALLINT", java.sql.Types.SMALLINT))
-    case BooleanType => Some(JdbcType("BOOLEAN", java.sql.Types.BOOLEAN))
-    // 31 is the maximum precision and 5 is the default scale for a Derby DECIMAL
-    case (t: DecimalType) if (t.precision > 31) =>
-      Some(JdbcType("DECIMAL(31,5)", java.sql.Types.DECIMAL))
-    case _ => None
-  }
-
-}
-
-/**
- * :: DeveloperApi ::
- * Default Oracle dialect, mapping a nonspecific numeric type to a general decimal type.
- */
-@DeveloperApi
-case object OracleDialect extends JdbcDialect {
-  override def canHandle(url: String): Boolean = url.startsWith("jdbc:oracle")
-  override def getCatalystType(
-      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
-    // Handle NUMBER fields that have no precision/scale in special way
-    // because JDBC ResultSetMetaData converts this to 0 procision and -127 scale
-    // For more details, please see
-    // https://github.com/apache/spark/pull/8780#issuecomment-145598968
-    // and
-    // https://github.com/apache/spark/pull/8780#issuecomment-144541760
-    if (sqlType == Types.NUMERIC && size == 0) {
-      // This is sub-optimal as we have to pick a precision/scale in advance whereas the data
-      //  in Oracle is allowed to have different precision/scale for each value.
-      Some(DecimalType(DecimalType.MAX_PRECISION, 10))
-    } else {
-      None
-    }
-  }
-}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MsSqlServerDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MsSqlServerDialect.scala
new file mode 100644
index 0000000000000..3eb722b070d5d
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MsSqlServerDialect.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import org.apache.spark.sql.types._
+
+
+private object MsSqlServerDialect extends JdbcDialect {
+
+  override def canHandle(url: String): Boolean = url.startsWith("jdbc:sqlserver")
+
+  override def getCatalystType(
+      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
+    if (typeName.contains("datetimeoffset")) {
+      // String is recommend by Microsoft SQL Server for datetimeoffset types in non-MS clients
+      Option(StringType)
+    } else {
+      None
+    }
+  }
+
+  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
+    case TimestampType => Some(JdbcType("DATETIME", java.sql.Types.TIMESTAMP))
+    case _ => None
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala
new file mode 100644
index 0000000000000..da413ed1f08b5
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala
@@ -0,0 +1,48 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import java.sql.Types
+
+import org.apache.spark.sql.types.{BooleanType, LongType, DataType, MetadataBuilder}
+
+
+private case object MySQLDialect extends JdbcDialect {
+
+  override def canHandle(url : String): Boolean = url.startsWith("jdbc:mysql")
+
+  override def getCatalystType(
+      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
+    if (sqlType == Types.VARBINARY && typeName.equals("BIT") && size != 1) {
+      // This could instead be a BinaryType if we'd rather return bit-vectors of up to 64 bits as
+      // byte arrays instead of longs.
+      md.putLong("binarylong", 1)
+      Option(LongType)
+    } else if (sqlType == Types.BIT && typeName.equals("TINYINT")) {
+      Option(BooleanType)
+    } else None
+  }
+
+  override def quoteIdentifier(colName: String): String = {
+    s"`$colName`"
+  }
+
+  override def getTableExistsQuery(table: String): String = {
+    s"SELECT 1 FROM $table LIMIT 1"
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/OracleDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/OracleDialect.scala
new file mode 100644
index 0000000000000..4165c382689f9
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/OracleDialect.scala
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import java.sql.Types
+
+import org.apache.spark.sql.types._
+
+
+private case object OracleDialect extends JdbcDialect {
+
+  override def canHandle(url: String): Boolean = url.startsWith("jdbc:oracle")
+
+  override def getCatalystType(
+      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
+    // Handle NUMBER fields that have no precision/scale in special way
+    // because JDBC ResultSetMetaData converts this to 0 procision and -127 scale
+    // For more details, please see
+    // https://github.com/apache/spark/pull/8780#issuecomment-145598968
+    // and
+    // https://github.com/apache/spark/pull/8780#issuecomment-144541760
+    if (sqlType == Types.NUMERIC && size == 0) {
+      // This is sub-optimal as we have to pick a precision/scale in advance whereas the data
+      //  in Oracle is allowed to have different precision/scale for each value.
+      Option(DecimalType(DecimalType.MAX_PRECISION, 10))
+    } else {
+      None
+    }
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
new file mode 100644
index 0000000000000..e701a7fcd9e16
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import java.sql.Types
+
+import org.apache.spark.sql.types._
+
+
+private object PostgresDialect extends JdbcDialect {
+
+  override def canHandle(url: String): Boolean = url.startsWith("jdbc:postgresql")
+
+  override def getCatalystType(
+      sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
+    if (sqlType == Types.BIT && typeName.equals("bit") && size != 1) {
+      Option(BinaryType)
+    } else if (sqlType == Types.OTHER && typeName.equals("cidr")) {
+      Option(StringType)
+    } else if (sqlType == Types.OTHER && typeName.equals("inet")) {
+      Option(StringType)
+    } else if (sqlType == Types.OTHER && typeName.equals("json")) {
+      Option(StringType)
+    } else if (sqlType == Types.OTHER && typeName.equals("jsonb")) {
+      Option(StringType)
+    } else None
+  }
+
+  override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
+    case StringType => Some(JdbcType("TEXT", java.sql.Types.CHAR))
+    case BinaryType => Some(JdbcType("BYTEA", java.sql.Types.BINARY))
+    case BooleanType => Some(JdbcType("BOOLEAN", java.sql.Types.BOOLEAN))
+    case _ => None
+  }
+
+  override def getTableExistsQuery(table: String): String = {
+    s"SELECT 1 FROM $table LIMIT 1"
+  }
+}

From 253e87e8ab8717ffef40a6d0d376b1add155ef90 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Fri, 6 Nov 2015 06:38:49 -0800
Subject: [PATCH 092/518] [SPARK-11453][SQL][FOLLOW-UP] remove DecimalLit

A cleanup for https://github.com/apache/spark/pull/9085.

The `DecimalLit` is very similar to `FloatLit`, we can just keep one of them.
Also added low level unit test at `SqlParserSuite`

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9482 from cloud-fan/parser.
---
 .../sql/catalyst/AbstractSparkSQLParser.scala | 23 ++++++++-----------
 .../apache/spark/sql/catalyst/SqlParser.scala | 20 ++++------------
 .../spark/sql/catalyst/SqlParserSuite.scala   | 21 +++++++++++++++++
 3 files changed, 35 insertions(+), 29 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/AbstractSparkSQLParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/AbstractSparkSQLParser.scala
index 04ac4f20c66ec..bdc52c08acb66 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/AbstractSparkSQLParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/AbstractSparkSQLParser.scala
@@ -78,10 +78,6 @@ private[sql] abstract class AbstractSparkSQLParser
 }
 
 class SqlLexical extends StdLexical {
-  case class FloatLit(chars: String) extends Token {
-    override def toString: String = chars
-  }
-
   case class DecimalLit(chars: String) extends Token {
     override def toString: String = chars
   }
@@ -106,17 +102,16 @@ class SqlLexical extends StdLexical {
   }
 
   override lazy val token: Parser[Token] =
-    ( rep1(digit) ~ ('.' ~> digit.*).? ~ (exp ~> sign.? ~ rep1(digit)) ^^ {
-        case i ~ None ~ (sig ~ rest) =>
-          DecimalLit(i.mkString + "e" + sig.mkString + rest.mkString)
-        case i ~ Some(d) ~ (sig ~ rest) =>
-          DecimalLit(i.mkString + "." + d.mkString + "e" + sig.mkString + rest.mkString)
-      }
+    ( rep1(digit) ~ scientificNotation ^^ { case i ~ s => DecimalLit(i.mkString + s) }
+    | '.' ~> (rep1(digit) ~ scientificNotation) ^^
+      { case i ~ s => DecimalLit("0." + i.mkString + s) }
+    | rep1(digit) ~ ('.' ~> digit.*) ~ scientificNotation ^^
+      { case i1 ~ i2 ~ s => DecimalLit(i1.mkString + "." + i2.mkString + s) }
     | digit.* ~ identChar ~ (identChar | digit).* ^^
       { case first ~ middle ~ rest => processIdent((first ++ (middle :: rest)).mkString) }
     | rep1(digit) ~ ('.' ~> digit.*).? ^^ {
         case i ~ None => NumericLit(i.mkString)
-        case i ~ Some(d) => FloatLit(i.mkString + "." + d.mkString)
+        case i ~ Some(d) => DecimalLit(i.mkString + "." + d.mkString)
       }
     | '\'' ~> chrExcept('\'', '\n', EofCh).* <~ '\'' ^^
       { case chars => StringLit(chars mkString "") }
@@ -133,8 +128,10 @@ class SqlLexical extends StdLexical {
 
   override def identChar: Parser[Elem] = letter | elem('_')
 
-  private lazy val sign: Parser[Elem] = elem("s", c => c == '+' || c == '-')
-  private lazy val exp: Parser[Elem] = elem("e", c => c == 'E' || c == 'e')
+  private lazy val scientificNotation: Parser[String] =
+    (elem('e') | elem('E')) ~> (elem('+') | elem('-')).? ~ rep1(digit) ^^ {
+      case s ~ rest => "e" + s.mkString + rest.mkString
+    }
 
   override def whitespace: Parser[Any] =
     ( whitespaceChar
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
index 440e9e28fa783..cd717c09f8e5e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
@@ -334,27 +334,15 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
 
   protected lazy val numericLiteral: Parser[Literal] =
     ( integral  ^^ { case i => Literal(toNarrowestIntegerType(i)) }
-    | sign.? ~ unsignedFloat ^^ {
-      case s ~ f => Literal(toDecimalOrDouble(s.getOrElse("") + f))
-    }
-    | sign.? ~ unsignedDecimal ^^ {
-      case s ~ d => Literal(toDecimalOrDouble(s.getOrElse("") + d))
-    }
+    | sign.? ~ unsignedFloat ^^
+      { case s ~ f => Literal(toDecimalOrDouble(s.getOrElse("") + f)) }
     )
 
   protected lazy val unsignedFloat: Parser[String] =
     ( "." ~> numericLit ^^ { u => "0." + u }
-    | elem("decimal", _.isInstanceOf[lexical.FloatLit]) ^^ (_.chars)
+    | elem("decimal", _.isInstanceOf[lexical.DecimalLit]) ^^ (_.chars)
     )
 
-  protected lazy val unsignedDecimal: Parser[String] =
-    ( "." ~> decimalLit ^^ { u => "0." + u }
-    | elem("scientific_notation", _.isInstanceOf[lexical.DecimalLit]) ^^ (_.chars)
-    )
-
-  def decimalLit: Parser[String] =
-    elem("scientific_notation", _.isInstanceOf[lexical.DecimalLit]) ^^ (_.chars)
-
   protected lazy val sign: Parser[String] = ("+" | "-")
 
   protected lazy val integral: Parser[String] =
@@ -477,7 +465,7 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
 
   protected lazy val baseExpression: Parser[Expression] =
     ( "*" ^^^ UnresolvedStar(None)
-    | (ident <~ "."). + <~ "*" ^^ { case target => UnresolvedStar(Option(target))}
+    | rep1(ident <~ ".") <~ "*" ^^ { case target => UnresolvedStar(Option(target))}
     | primary
    )
 
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/SqlParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/SqlParserSuite.scala
index ea28bfa021bed..9ff893b84775b 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/SqlParserSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/SqlParserSuite.scala
@@ -126,4 +126,25 @@ class SqlParserSuite extends PlanTest {
     checkSingleUnit("13.123456789", "second")
     checkSingleUnit("-13.123456789", "second")
   }
+
+  test("support scientific notation") {
+    def assertRight(input: String, output: Double): Unit = {
+      val parsed = SqlParser.parse("SELECT " + input)
+      val expected = Project(
+        UnresolvedAlias(
+          Literal(output)
+        ) :: Nil,
+        OneRowRelation)
+      comparePlans(parsed, expected)
+    }
+
+    assertRight("9.0e1", 90)
+    assertRight(".9e+2", 90)
+    assertRight("0.9e+2", 90)
+    assertRight("900e-1", 90)
+    assertRight("900.0E-1", 90)
+    assertRight("9.e+1", 90)
+
+    intercept[RuntimeException](SqlParser.parse("SELECT .e3"))
+  }
 }

From cf69ce136590fea51843bc54f44f0f45c7d0ac36 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Fri, 6 Nov 2015 14:51:53 +0000
Subject: [PATCH 093/518] [SPARK-11511][STREAMING] Fix NPE when an InputDStream
 is not used

Just ignored `InputDStream`s that have null `rememberDuration` in `DStreamGraph.getMaxInputStreamRememberDuration`.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9476 from zsxwing/SPARK-11511.
---
 .../apache/spark/streaming/DStreamGraph.scala    |  3 ++-
 .../spark/streaming/StreamingContextSuite.scala  | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/DStreamGraph.scala b/streaming/src/main/scala/org/apache/spark/streaming/DStreamGraph.scala
index 1b0b7890b3b00..7829f5e887995 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/DStreamGraph.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/DStreamGraph.scala
@@ -167,7 +167,8 @@ final private[streaming] class DStreamGraph extends Serializable with Logging {
    * safe remember duration which can be used to perform cleanup operations.
    */
   def getMaxInputStreamRememberDuration(): Duration = {
-    inputStreams.map { _.rememberDuration }.maxBy { _.milliseconds }
+    // If an InputDStream is not used, its `rememberDuration` will be null and we can ignore them
+    inputStreams.map(_.rememberDuration).filter(_ != null).maxBy(_.milliseconds)
   }
 
   @throws(classOf[IOException])
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/StreamingContextSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/StreamingContextSuite.scala
index c7a877142b374..860fac29c0ee0 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/StreamingContextSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/StreamingContextSuite.scala
@@ -780,6 +780,22 @@ class StreamingContextSuite extends SparkFunSuite with BeforeAndAfter with Timeo
       "Please don't use queueStream when checkpointing is enabled."))
   }
 
+  test("Creating an InputDStream but not using it should not crash") {
+    ssc = new StreamingContext(master, appName, batchDuration)
+    val input1 = addInputStream(ssc)
+    val input2 = addInputStream(ssc)
+    val output = new TestOutputStream(input2)
+    output.register()
+    val batchCount = new BatchCounter(ssc)
+    ssc.start()
+    // Just wait for completing 2 batches to make sure it triggers
+    // `DStream.getMaxInputStreamRememberDuration`
+    batchCount.waitUntilBatchesCompleted(2, 10000)
+    // Throw the exception if crash
+    ssc.awaitTerminationOrTimeout(1)
+    ssc.stop()
+  }
+
   def addInputStream(s: StreamingContext): DStream[Int] = {
     val input = (1 to 100).map(i => 1 to i)
     val inputStream = new TestInputStream(s, input, 1)

From 574141a29835ce78d68c97bb54336cf4fd3c39d3 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Fri, 6 Nov 2015 10:52:04 -0800
Subject: [PATCH 094/518] [SPARK-9162] [SQL] Implement code generation for
 ScalaUDF

JIRA: https://issues.apache.org/jira/browse/SPARK-9162

Currently ScalaUDF extends CodegenFallback and doesn't provide code generation implementation. This path implements code generation for ScalaUDF.

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #9270 from viirya/scalaudf-codegen.
---
 .../sql/catalyst/expressions/ScalaUDF.scala   | 85 ++++++++++++++++++-
 .../scala/org/apache/spark/sql/UDFSuite.scala | 41 +++++++++
 2 files changed, 124 insertions(+), 2 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
index 11c7950c0613b..3388cc20a9803 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.expressions
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.CatalystTypeConverters
-import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.types.DataType
 
 /**
@@ -31,7 +31,7 @@ case class ScalaUDF(
     dataType: DataType,
     children: Seq[Expression],
     inputTypes: Seq[DataType] = Nil)
-  extends Expression with ImplicitCastInputTypes with CodegenFallback {
+  extends Expression with ImplicitCastInputTypes {
 
   override def nullable: Boolean = true
 
@@ -60,6 +60,10 @@ case class ScalaUDF(
 
   */
 
+  // Accessors used in genCode
+  def userDefinedFunc(): AnyRef = function
+  def getChildren(): Seq[Expression] = children
+
   private[this] val f = children.size match {
     case 0 =>
       val func = function.asInstanceOf[() => Any]
@@ -960,6 +964,83 @@ case class ScalaUDF(
   }
 
   // scalastyle:on
+
+  // Generate codes used to convert the arguments to Scala type for user-defined funtions
+  private[this] def genCodeForConverter(ctx: CodeGenContext, index: Int): String = {
+    val converterClassName = classOf[Any => Any].getName
+    val typeConvertersClassName = CatalystTypeConverters.getClass.getName + ".MODULE$"
+    val expressionClassName = classOf[Expression].getName
+    val scalaUDFClassName = classOf[ScalaUDF].getName
+
+    val converterTerm = ctx.freshName("converter")
+    val expressionIdx = ctx.references.size - 1
+    ctx.addMutableState(converterClassName, converterTerm,
+      s"this.$converterTerm = ($converterClassName)$typeConvertersClassName" +
+        s".createToScalaConverter(((${expressionClassName})((($scalaUDFClassName)" +
+          s"expressions[$expressionIdx]).getChildren().apply($index))).dataType());")
+    converterTerm
+  }
+
+  override def genCode(
+      ctx: CodeGenContext,
+      ev: GeneratedExpressionCode): String = {
+
+    ctx.references += this
+
+    val scalaUDFClassName = classOf[ScalaUDF].getName
+    val converterClassName = classOf[Any => Any].getName
+    val typeConvertersClassName = CatalystTypeConverters.getClass.getName + ".MODULE$"
+    val expressionClassName = classOf[Expression].getName
+
+    // Generate codes used to convert the returned value of user-defined functions to Catalyst type
+    val catalystConverterTerm = ctx.freshName("catalystConverter")
+    val catalystConverterTermIdx = ctx.references.size - 1
+    ctx.addMutableState(converterClassName, catalystConverterTerm,
+      s"this.$catalystConverterTerm = ($converterClassName)$typeConvertersClassName" +
+        s".createToCatalystConverter((($scalaUDFClassName)expressions" +
+          s"[$catalystConverterTermIdx]).dataType());")
+
+    val resultTerm = ctx.freshName("result")
+
+    // This must be called before children expressions' codegen
+    // because ctx.references is used in genCodeForConverter
+    val converterTerms = (0 until children.size).map(genCodeForConverter(ctx, _))
+
+    // Initialize user-defined function
+    val funcClassName = s"scala.Function${children.size}"
+
+    val funcTerm = ctx.freshName("udf")
+    val funcExpressionIdx = ctx.references.size - 1
+    ctx.addMutableState(funcClassName, funcTerm,
+      s"this.$funcTerm = ($funcClassName)((($scalaUDFClassName)expressions" +
+        s"[$funcExpressionIdx]).userDefinedFunc());")
+
+    // codegen for children expressions
+    val evals = children.map(_.gen(ctx))
+
+    // Generate the codes for expressions and calling user-defined function
+    // We need to get the boxedType of dataType's javaType here. Because for the dataType
+    // such as IntegerType, its javaType is `int` and the returned type of user-defined
+    // function is Object. Trying to convert an Object to `int` will cause casting exception.
+    val evalCode = evals.map(_.code).mkString
+    val funcArguments = converterTerms.zip(evals).map {
+      case (converter, eval) => s"$converter.apply(${eval.value})"
+    }.mkString(",")
+    val callFunc = s"${ctx.boxedType(ctx.javaType(dataType))} $resultTerm = " +
+      s"(${ctx.boxedType(ctx.javaType(dataType))})${catalystConverterTerm}" +
+        s".apply($funcTerm.apply($funcArguments));"
+
+    evalCode + s"""
+      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+      Boolean ${ev.isNull};
+
+      $callFunc
+
+      ${ev.value} = $resultTerm;
+      ${ev.isNull} = $resultTerm == null;
+    """
+  }
+
   private[this] val converter = CatalystTypeConverters.createToCatalystConverter(dataType)
   override def eval(input: InternalRow): Any = converter(f(input))
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
index e0435a0dba6ad..9837fa6bdb357 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
@@ -191,4 +191,45 @@ class UDFSuite extends QueryTest with SharedSQLContext {
     // pass a decimal to intExpected.
     assert(sql("SELECT intExpected(1.0)").head().getInt(0) === 1)
   }
+
+  test("udf in different types") {
+    sqlContext.udf.register("testDataFunc", (n: Int, s: String) => { (n, s) })
+    sqlContext.udf.register("decimalDataFunc",
+      (a: java.math.BigDecimal, b: java.math.BigDecimal) => { (a, b) })
+    sqlContext.udf.register("binaryDataFunc", (a: Array[Byte], b: Int) => { (a, b) })
+    sqlContext.udf.register("arrayDataFunc",
+      (data: Seq[Int], nestedData: Seq[Seq[Int]]) => { (data, nestedData) })
+    sqlContext.udf.register("mapDataFunc",
+      (data: scala.collection.Map[Int, String]) => { data })
+    sqlContext.udf.register("complexDataFunc",
+      (m: Map[String, Int], a: Seq[Int], b: Boolean) => { (m, a, b) } )
+
+    checkAnswer(
+      sql("SELECT tmp.t.* FROM (SELECT testDataFunc(key, value) AS t from testData) tmp").toDF(),
+      testData)
+    checkAnswer(
+      sql("""
+           | SELECT tmp.t.* FROM
+           | (SELECT decimalDataFunc(a, b) AS t FROM decimalData) tmp
+          """.stripMargin).toDF(), decimalData)
+    checkAnswer(
+      sql("""
+           | SELECT tmp.t.* FROM
+           | (SELECT binaryDataFunc(a, b) AS t FROM binaryData) tmp
+          """.stripMargin).toDF(), binaryData)
+    checkAnswer(
+      sql("""
+           | SELECT tmp.t.* FROM
+           | (SELECT arrayDataFunc(data, nestedData) AS t FROM arrayData) tmp
+          """.stripMargin).toDF(), arrayData.toDF())
+    checkAnswer(
+      sql("""
+           | SELECT mapDataFunc(data) AS t FROM mapData
+          """.stripMargin).toDF(), mapData.toDF())
+    checkAnswer(
+      sql("""
+           | SELECT tmp.t.* FROM
+           | (SELECT complexDataFunc(m, a, b) AS t FROM complexData) tmp
+          """.stripMargin).toDF(), complexData.select("m", "a", "b"))
+  }
 }

From c048929c6a9f7ce57f384037cd6c0bf5751c447a Mon Sep 17 00:00:00 2001
From: Cheng Lian <lian@databricks.com>
Date: Fri, 6 Nov 2015 11:11:36 -0800
Subject: [PATCH 095/518] [SPARK-10978][SQL][FOLLOW-UP] More comprehensive
 tests for PR #9399

This PR adds test cases that test various column pruning and filter push-down cases.

Author: Cheng Lian <lian@databricks.com>

Closes #9468 from liancheng/spark-10978.follow-up.
---
 .../spark/sql/sources/FilteredScanSuite.scala |  21 +-
 .../SimpleTextHadoopFsRelationSuite.scala     | 335 ++++++++++++++++--
 .../sql/sources/SimpleTextRelation.scala      |  11 +
 3 files changed, 321 insertions(+), 46 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
index 7541e723029bf..2cad964e55b2b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
@@ -17,16 +17,15 @@
 
 package org.apache.spark.sql.sources
 
-import org.apache.spark.sql.execution.datasources.LogicalRelation
-
 import scala.language.existentials
 
 import org.apache.spark.rdd.RDD
-import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.sql._
+import org.apache.spark.sql.catalyst.expressions.PredicateHelper
+import org.apache.spark.sql.execution.datasources.LogicalRelation
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
-
+import org.apache.spark.unsafe.types.UTF8String
 
 class FilteredScanSource extends RelationProvider {
   override def createRelation(
@@ -130,7 +129,7 @@ object ColumnsRequired {
   var set: Set[String] = Set.empty
 }
 
-class FilteredScanSuite extends DataSourceTest with SharedSQLContext {
+class FilteredScanSuite extends DataSourceTest with SharedSQLContext with PredicateHelper {
   protected override lazy val sql = caseInsensitiveContext.sql _
 
   override def beforeAll(): Unit = {
@@ -144,9 +143,6 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext {
         |  to '10'
         |)
       """.stripMargin)
-
-    // UDF for testing filter push-down
-    caseInsensitiveContext.udf.register("udf_gt3", (_: Int) > 3)
   }
 
   sqlTest(
@@ -276,14 +272,15 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext {
   testPushDown("SELECT c FROM oneToTenFiltered WHERE c = 'aaaaaAAAAA'", 1, Set("c"))
   testPushDown("SELECT c FROM oneToTenFiltered WHERE c IN ('aaaaaAAAAA', 'foo')", 1, Set("c"))
 
-  // Columns only referenced by UDF filter must be required, as UDF filters can't be pushed down.
-  testPushDown("SELECT c FROM oneToTenFiltered WHERE udf_gt3(A)", 10, Set("a", "c"))
+  // Filters referencing multiple columns are not convertible, all referenced columns must be
+  // required.
+  testPushDown("SELECT c FROM oneToTenFiltered WHERE A + b > 9", 10, Set("a", "b", "c"))
 
-  // A query with an unconvertible filter, an unhandled filter, and a handled filter.
+  // A query with an inconvertible filter, an unhandled filter, and a handled filter.
   testPushDown(
     """SELECT a
       |  FROM oneToTenFiltered
-      | WHERE udf_gt3(b)
+      | WHERE a + b > 9
       |   AND b < 16
       |   AND c IN ('bbbbbBBBBB', 'cccccCCCCC', 'dddddDDDDD', 'foo')
     """.stripMargin.split("\n").map(_.trim).mkString(" "), 3, Set("a", "b"))
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
index d945408341fc9..9251a69f31a47 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
@@ -17,15 +17,21 @@
 
 package org.apache.spark.sql.sources
 
+import java.io.File
+
 import org.apache.hadoop.fs.Path
+
+import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.deploy.SparkHadoopUtil
-import org.apache.spark.sql.Row
-import org.apache.spark.sql.catalyst.CatalystTypeConverters
-import org.apache.spark.sql.execution.PhysicalRDD
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, PredicateHelper}
+import org.apache.spark.sql.execution.{LogicalRDD, PhysicalRDD}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
+import org.apache.spark.sql.{Column, DataFrame, Row, execution}
+import org.apache.spark.util.Utils
 
-class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest {
+class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest with PredicateHelper {
   import testImplicits._
 
   override val dataSourceName: String = classOf[SimpleTextSource].getCanonicalName
@@ -70,43 +76,304 @@ class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest {
     }
   }
 
-  private val writer = testDF.write.option("dataSchema", dataSchema.json).format(dataSourceName)
-  private val reader = sqlContext.read.option("dataSchema", dataSchema.json).format(dataSourceName)
-
-  test("unhandledFilters") {
-    withTempPath { dir =>
-
-      val path = dir.getCanonicalPath
-      writer.save(s"$path/p=0")
-      writer.save(s"$path/p=1")
-
-      val isOdd = udf((_: Int) % 2 == 1)
-      val df = reader.load(path)
-        .filter(
-          // This filter is inconvertible
-          isOdd('a) &&
-            // This filter is convertible but unhandled
-            'a > 1 &&
-            // This filter is convertible and handled
-            'b > "val_1" &&
-            // This filter references a partiiton column, won't be pushed down
-            'p === 1
-        ).select('a, 'p)
-      val rawScan = df.queryExecution.executedPlan collect {
+  private var tempPath: File = _
+
+  private var partitionedDF: DataFrame = _
+
+  private val partitionedDataSchema: StructType = StructType('a.int :: 'b.int :: 'c.string :: Nil)
+
+  protected override def beforeAll(): Unit = {
+    this.tempPath = Utils.createTempDir()
+
+    val df = sqlContext.range(10).select(
+      'id cast IntegerType as 'a,
+      ('id cast IntegerType) * 2 as 'b,
+      concat(lit("val_"), 'id) as 'c
+    )
+
+    partitionedWriter(df).save(s"${tempPath.getCanonicalPath}/p=0")
+    partitionedWriter(df).save(s"${tempPath.getCanonicalPath}/p=1")
+
+    partitionedDF = partitionedReader.load(tempPath.getCanonicalPath)
+  }
+
+  override protected def afterAll(): Unit = {
+    Utils.deleteRecursively(tempPath)
+  }
+
+  private def partitionedWriter(df: DataFrame) =
+    df.write.option("dataSchema", partitionedDataSchema.json).format(dataSourceName)
+
+  private def partitionedReader =
+    sqlContext.read.option("dataSchema", partitionedDataSchema.json).format(dataSourceName)
+
+  /**
+   * Constructs test cases that test column pruning and filter push-down.
+   *
+   * For filter push-down, the following filters are not pushed-down.
+   *
+   * 1. Partitioning filters don't participate filter push-down, they are handled separately in
+   *    `DataSourceStrategy`
+   *
+   * 2. Catalyst filter `Expression`s that cannot be converted to data source `Filter`s are not
+   *    pushed down (e.g. UDF and filters referencing multiple columns).
+   *
+   * 3. Catalyst filter `Expression`s that can be converted to data source `Filter`s but cannot be
+   *    handled by the underlying data source are not pushed down (e.g. returned from
+   *    `BaseRelation.unhandledFilters()`).
+   *
+   *    Note that for [[SimpleTextRelation]], all data source [[Filter]]s other than [[GreaterThan]]
+   *    are unhandled.  We made this assumption in [[SimpleTextRelation.unhandledFilters()]] only
+   *    for testing purposes.
+   *
+   * @param projections Projection list of the query
+   * @param filter Filter condition of the query
+   * @param requiredColumns Expected names of required columns
+   * @param pushedFilters Expected data source [[Filter]]s that are pushed down
+   * @param inconvertibleFilters Expected Catalyst filter [[Expression]]s that cannot be converted
+   *        to data source [[Filter]]s
+   * @param unhandledFilters Expected Catalyst flter [[Expression]]s that can be converted to data
+   *        source [[Filter]]s but cannot be handled by the data source relation
+   * @param partitioningFilters Expected Catalyst filter [[Expression]]s that reference partition
+   *        columns
+   * @param expectedRawScanAnswer Expected query result of the raw table scan returned by the data
+   *        source relation
+   * @param expectedAnswer Expected query result of the full query
+   */
+  def testPruningAndFiltering(
+      projections: Seq[Column],
+      filter: Column,
+      requiredColumns: Seq[String],
+      pushedFilters: Seq[Filter],
+      inconvertibleFilters: Seq[Column],
+      unhandledFilters: Seq[Column],
+      partitioningFilters: Seq[Column])(
+      expectedRawScanAnswer: => Seq[Row])(
+      expectedAnswer: => Seq[Row]): Unit = {
+    test(s"pruning and filtering: df.select(${projections.mkString(", ")}).where($filter)") {
+      val df = partitionedDF.where(filter).select(projections: _*)
+      val queryExecution = df.queryExecution
+      val executedPlan = queryExecution.executedPlan
+
+      val rawScan = executedPlan.collect {
         case p: PhysicalRDD => p
       } match {
-        case Seq(p) => p
+        case Seq(scan) => scan
+        case _ => fail(s"More than one PhysicalRDD found\n$queryExecution")
       }
 
-      val outputSchema = new StructType().add("a", IntegerType).add("p", IntegerType)
+      markup("Checking raw scan answer")
+      checkAnswer(
+        DataFrame(sqlContext, LogicalRDD(rawScan.output, rawScan.rdd)(sqlContext)),
+        expectedRawScanAnswer)
 
-      assertResult(Set((2, 1), (3, 1))) {
-        rawScan.execute().collect()
-          .map { CatalystTypeConverters.convertToScala(_, outputSchema) }
-          .map { case Row(a, p) => (a, p) }.toSet
+      markup("Checking full query answer")
+      checkAnswer(df, expectedAnswer)
+
+      markup("Checking required columns")
+      assert(requiredColumns === SimpleTextRelation.requiredColumns)
+
+      val nonPushedFilters = {
+        val boundFilters = executedPlan.collect {
+          case f: execution.Filter => f
+        } match {
+          case Nil => Nil
+          case Seq(f) => splitConjunctivePredicates(f.condition)
+          case _ => fail(s"More than one PhysicalRDD found\n$queryExecution")
+        }
+
+        // Unbound these bound filters so that we can easily compare them with expected results.
+        boundFilters.map {
+          _.transform { case a: AttributeReference => UnresolvedAttribute(a.name) }
+        }.toSet
       }
 
-      checkAnswer(df, Row(3, 1))
+      markup("Checking pushed filters")
+      assert(SimpleTextRelation.pushedFilters === pushedFilters.toSet)
+
+      val expectedInconvertibleFilters = inconvertibleFilters.map(_.expr).toSet
+      val expectedUnhandledFilters = unhandledFilters.map(_.expr).toSet
+      val expectedPartitioningFilters = partitioningFilters.map(_.expr).toSet
+
+      markup("Checking unhandled and inconvertible filters")
+      assert(expectedInconvertibleFilters ++ expectedUnhandledFilters === nonPushedFilters)
+
+      markup("Checking partitioning filters")
+      val actualPartitioningFilters = splitConjunctivePredicates(filter.expr).filter {
+        _.references.contains(UnresolvedAttribute("p"))
+      }.toSet
+
+      // Partitioning filters are handled separately and don't participate filter push-down. So they
+      // shouldn't be part of non-pushed filters.
+      assert(expectedPartitioningFilters.intersect(nonPushedFilters).isEmpty)
+      assert(expectedPartitioningFilters === actualPartitioningFilters)
     }
   }
+
+  testPruningAndFiltering(
+    projections = Seq('*),
+    filter = 'p > 0,
+    requiredColumns = Seq("a", "b", "c"),
+    pushedFilters = Nil,
+    inconvertibleFilters = Nil,
+    unhandledFilters = Nil,
+    partitioningFilters = Seq('p > 0)
+  ) {
+    Seq(
+      Row(0, 0, "val_0", 1),
+      Row(1, 2, "val_1", 1),
+      Row(2, 4, "val_2", 1),
+      Row(3, 6, "val_3", 1),
+      Row(4, 8, "val_4", 1),
+      Row(5, 10, "val_5", 1),
+      Row(6, 12, "val_6", 1),
+      Row(7, 14, "val_7", 1),
+      Row(8, 16, "val_8", 1),
+      Row(9, 18, "val_9", 1))
+  } {
+    Seq(
+      Row(0, 0, "val_0", 1),
+      Row(1, 2, "val_1", 1),
+      Row(2, 4, "val_2", 1),
+      Row(3, 6, "val_3", 1),
+      Row(4, 8, "val_4", 1),
+      Row(5, 10, "val_5", 1),
+      Row(6, 12, "val_6", 1),
+      Row(7, 14, "val_7", 1),
+      Row(8, 16, "val_8", 1),
+      Row(9, 18, "val_9", 1))
+  }
+
+  testPruningAndFiltering(
+    projections = Seq('c, 'p),
+    filter = 'a < 3 && 'p > 0,
+    requiredColumns = Seq("c", "a"),
+    pushedFilters = Nil,
+    inconvertibleFilters = Nil,
+    unhandledFilters = Seq('a < 3),
+    partitioningFilters = Seq('p > 0)
+  ) {
+    Seq(
+      Row("val_0", 1, 0),
+      Row("val_1", 1, 1),
+      Row("val_2", 1, 2),
+      Row("val_3", 1, 3),
+      Row("val_4", 1, 4),
+      Row("val_5", 1, 5),
+      Row("val_6", 1, 6),
+      Row("val_7", 1, 7),
+      Row("val_8", 1, 8),
+      Row("val_9", 1, 9))
+  } {
+    Seq(
+      Row("val_0", 1),
+      Row("val_1", 1),
+      Row("val_2", 1))
+  }
+
+  testPruningAndFiltering(
+    projections = Seq('*),
+    filter = 'a > 8,
+    requiredColumns = Seq("a", "b", "c"),
+    pushedFilters = Seq(GreaterThan("a", 8)),
+    inconvertibleFilters = Nil,
+    unhandledFilters = Nil,
+    partitioningFilters = Nil
+  ) {
+    Seq(
+      Row(9, 18, "val_9", 0),
+      Row(9, 18, "val_9", 1))
+  } {
+    Seq(
+      Row(9, 18, "val_9", 0),
+      Row(9, 18, "val_9", 1))
+  }
+
+  testPruningAndFiltering(
+    projections = Seq('b, 'p),
+    filter = 'a > 8,
+    requiredColumns = Seq("b"),
+    pushedFilters = Seq(GreaterThan("a", 8)),
+    inconvertibleFilters = Nil,
+    unhandledFilters = Nil,
+    partitioningFilters = Nil
+  ) {
+    Seq(
+      Row(18, 0),
+      Row(18, 1))
+  } {
+    Seq(
+      Row(18, 0),
+      Row(18, 1))
+  }
+
+  testPruningAndFiltering(
+    projections = Seq('b, 'p),
+    filter = 'a > 8 && 'p > 0,
+    requiredColumns = Seq("b"),
+    pushedFilters = Seq(GreaterThan("a", 8)),
+    inconvertibleFilters = Nil,
+    unhandledFilters = Nil,
+    partitioningFilters = Seq('p > 0)
+  ) {
+    Seq(
+      Row(18, 1))
+  } {
+    Seq(
+      Row(18, 1))
+  }
+
+  testPruningAndFiltering(
+    projections = Seq('b, 'p),
+    filter = 'c > "val_7" && 'b < 18 && 'p > 0,
+    requiredColumns = Seq("b"),
+    pushedFilters = Seq(GreaterThan("c", "val_7")),
+    inconvertibleFilters = Nil,
+    unhandledFilters = Seq('b < 18),
+    partitioningFilters = Seq('p > 0)
+  ) {
+    Seq(
+      Row(16, 1),
+      Row(18, 1))
+  } {
+    Seq(
+      Row(16, 1))
+  }
+
+  testPruningAndFiltering(
+    projections = Seq('b, 'p),
+    filter = 'a % 2 === 0 && 'c > "val_7" && 'b < 18 && 'p > 0,
+    requiredColumns = Seq("b", "a"),
+    pushedFilters = Seq(GreaterThan("c", "val_7")),
+    inconvertibleFilters = Seq('a % 2 === 0),
+    unhandledFilters = Seq('b < 18),
+    partitioningFilters = Seq('p > 0)
+  ) {
+    Seq(
+      Row(16, 1, 8),
+      Row(18, 1, 9))
+  } {
+    Seq(
+      Row(16, 1))
+  }
+
+  testPruningAndFiltering(
+    projections = Seq('b, 'p),
+    filter = 'a > 7 && 'a < 9,
+    requiredColumns = Seq("b", "a"),
+    pushedFilters = Seq(GreaterThan("a", 7)),
+    inconvertibleFilters = Nil,
+    unhandledFilters = Seq('a < 9),
+    partitioningFilters = Nil
+  ) {
+    Seq(
+      Row(16, 0, 8),
+      Row(16, 1, 8),
+      Row(18, 0, 9),
+      Row(18, 1, 9))
+  } {
+    Seq(
+      Row(16, 0),
+      Row(16, 1))
+  }
 }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
index da09e1b00ae48..bdc48a383bbbf 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
@@ -128,6 +128,9 @@ class SimpleTextRelation(
       filters: Array[Filter],
       inputFiles: Array[FileStatus]): RDD[Row] = {
 
+    SimpleTextRelation.requiredColumns = requiredColumns
+    SimpleTextRelation.pushedFilters = filters.toSet
+
     val fields = this.dataSchema.map(_.dataType)
     val inputAttributes = this.dataSchema.toAttributes
     val outputAttributes = requiredColumns.flatMap(name => inputAttributes.find(_.name == name))
@@ -191,6 +194,14 @@ class SimpleTextRelation(
   }
 }
 
+object SimpleTextRelation {
+  // Used to test column pruning
+  var requiredColumns: Seq[String] = Nil
+
+  // Used to test filter push-down
+  var pushedFilters: Set[Filter] = Set.empty
+}
+
 /**
  * A simple example [[HadoopFsRelationProvider]].
  */

From 8211aab0793cf64202b99be4f31bb8a9ae77050d Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Fri, 6 Nov 2015 11:13:51 -0800
Subject: [PATCH 096/518] [SPARK-9858][SQL] Add an ExchangeCoordinator to
 estimate the number of post-shuffle partitions for aggregates and joins
 (follow-up)

https://issues.apache.org/jira/browse/SPARK-9858

This PR is the follow-up work of https://github.com/apache/spark/pull/9276. It addresses JoshRosen's comments.

Author: Yin Huai <yhuai@databricks.com>

Closes #9453 from yhuai/numReducer-followUp.
---
 .../plans/physical/partitioning.scala         |   8 -
 .../apache/spark/sql/execution/Exchange.scala |  40 +++--
 .../sql/execution/ExchangeCoordinator.scala   |  31 ++--
 .../apache/spark/sql/CachedTableSuite.scala   | 150 ++++++++++++++----
 4 files changed, 167 insertions(+), 62 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index 9312c8123e92e..86b9417477ba3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -165,11 +165,6 @@ sealed trait Partitioning {
    * produced by `A` could have also been produced by `B`.
    */
   def guarantees(other: Partitioning): Boolean = this == other
-
-  def withNumPartitions(newNumPartitions: Int): Partitioning = {
-    throw new IllegalStateException(
-      s"It is not allowed to call withNumPartitions method of a ${this.getClass.getSimpleName}")
-  }
 }
 
 object Partitioning {
@@ -254,9 +249,6 @@ case class HashPartitioning(expressions: Seq[Expression], numPartitions: Int)
     case _ => false
   }
 
-  override def withNumPartitions(newNumPartitions: Int): HashPartitioning = {
-    HashPartitioning(expressions, newNumPartitions)
-  }
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
index 0f72ec6cc107a..a4ce328c1a9eb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
@@ -242,7 +242,7 @@ case class Exchange(
     // update the number of post-shuffle partitions.
     specifiedPartitionStartIndices.foreach { indices =>
       assert(newPartitioning.isInstanceOf[HashPartitioning])
-      newPartitioning = newPartitioning.withNumPartitions(indices.length)
+      newPartitioning = UnknownPartitioning(indices.length)
     }
     new ShuffledRowRDD(shuffleDependency, specifiedPartitionStartIndices)
   }
@@ -262,7 +262,7 @@ case class Exchange(
 
 object Exchange {
   def apply(newPartitioning: Partitioning, child: SparkPlan): Exchange = {
-    Exchange(newPartitioning, child, None: Option[ExchangeCoordinator])
+    Exchange(newPartitioning, child, coordinator = None: Option[ExchangeCoordinator])
   }
 }
 
@@ -315,7 +315,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
             child.outputPartitioning match {
               case hash: HashPartitioning => true
               case collection: PartitioningCollection =>
-                collection.partitionings.exists(_.isInstanceOf[HashPartitioning])
+                collection.partitionings.forall(_.isInstanceOf[HashPartitioning])
               case _ => false
             }
         }
@@ -416,28 +416,48 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       // First check if the existing partitions of the children all match. This means they are
       // partitioned by the same partitioning into the same number of partitions. In that case,
       // don't try to make them match `defaultPartitions`, just use the existing partitioning.
-      // TODO: this should be a cost based decision. For example, a big relation should probably
-      // maintain its existing number of partitions and smaller partitions should be shuffled.
-      // defaultPartitions is arbitrary.
-      val numPartitions = children.head.outputPartitioning.numPartitions
+      val maxChildrenNumPartitions = children.map(_.outputPartitioning.numPartitions).max
       val useExistingPartitioning = children.zip(requiredChildDistributions).forall {
         case (child, distribution) => {
           child.outputPartitioning.guarantees(
-            createPartitioning(distribution, numPartitions))
+            createPartitioning(distribution, maxChildrenNumPartitions))
         }
       }
 
       children = if (useExistingPartitioning) {
+        // We do not need to shuffle any child's output.
         children
       } else {
+        // We need to shuffle at least one child's output.
+        // Now, we will determine the number of partitions that will be used by created
+        // partitioning schemes.
+        val numPartitions = {
+          // Let's see if we need to shuffle all child's outputs when we use
+          // maxChildrenNumPartitions.
+          val shufflesAllChildren = children.zip(requiredChildDistributions).forall {
+            case (child, distribution) => {
+              !child.outputPartitioning.guarantees(
+                createPartitioning(distribution, maxChildrenNumPartitions))
+            }
+          }
+          // If we need to shuffle all children, we use defaultNumPreShufflePartitions as the
+          // number of partitions. Otherwise, we use maxChildrenNumPartitions.
+          if (shufflesAllChildren) defaultNumPreShufflePartitions else maxChildrenNumPartitions
+        }
+
         children.zip(requiredChildDistributions).map {
           case (child, distribution) => {
             val targetPartitioning =
-              createPartitioning(distribution, defaultNumPreShufflePartitions)
+              createPartitioning(distribution, numPartitions)
             if (child.outputPartitioning.guarantees(targetPartitioning)) {
               child
             } else {
-              Exchange(targetPartitioning, child)
+              child match {
+                // If child is an exchange, we replace it with
+                // a new one having targetPartitioning.
+                case Exchange(_, c, _) => Exchange(targetPartitioning, c)
+                case _ => Exchange(targetPartitioning, child)
+              }
             }
           }
         }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExchangeCoordinator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExchangeCoordinator.scala
index 8dbd69e1f44b8..827fdd278460a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExchangeCoordinator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExchangeCoordinator.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.execution
 
 import java.util.{Map => JMap, HashMap => JHashMap}
+import javax.annotation.concurrent.GuardedBy
 
 import scala.collection.mutable.ArrayBuffer
 
@@ -97,6 +98,7 @@ private[sql] class ExchangeCoordinator(
    * Registers an [[Exchange]] operator to this coordinator. This method is only allowed to be
    * called in the `doPrepare` method of an [[Exchange]] operator.
    */
+  @GuardedBy("this")
   def registerExchange(exchange: Exchange): Unit = synchronized {
     exchanges += exchange
   }
@@ -109,7 +111,7 @@ private[sql] class ExchangeCoordinator(
    */
   private[sql] def estimatePartitionStartIndices(
       mapOutputStatistics: Array[MapOutputStatistics]): Array[Int] = {
-    // If we have mapOutputStatistics.length <= numExchange, it is because we do not submit
+    // If we have mapOutputStatistics.length < numExchange, it is because we do not submit
     // a stage when the number of partitions of this dependency is 0.
     assert(mapOutputStatistics.length <= numExchanges)
 
@@ -121,6 +123,8 @@ private[sql] class ExchangeCoordinator(
         val totalPostShuffleInputSize = mapOutputStatistics.map(_.bytesByPartitionId.sum).sum
         // The max at here is to make sure that when we have an empty table, we
         // only have a single post-shuffle partition.
+        // There is no particular reason that we pick 16. We just need a number to
+        // prevent maxPostShuffleInputSize from being set to 0.
         val maxPostShuffleInputSize =
           math.max(math.ceil(totalPostShuffleInputSize / numPartitions.toDouble).toLong, 16)
         math.min(maxPostShuffleInputSize, advisoryTargetPostShuffleInputSize)
@@ -135,6 +139,12 @@ private[sql] class ExchangeCoordinator(
     // Make sure we do get the same number of pre-shuffle partitions for those stages.
     val distinctNumPreShufflePartitions =
       mapOutputStatistics.map(stats => stats.bytesByPartitionId.length).distinct
+    // The reason that we are expecting a single value of the number of pre-shuffle partitions
+    // is that when we add Exchanges, we set the number of pre-shuffle partitions
+    // (i.e. map output partitions) using a static setting, which is the value of
+    // spark.sql.shuffle.partitions. Even if two input RDDs are having different
+    // number of partitions, they will have the same number of pre-shuffle partitions
+    // (i.e. map output partitions).
     assert(
       distinctNumPreShufflePartitions.length == 1,
       "There should be only one distinct value of the number pre-shuffle partitions " +
@@ -177,6 +187,7 @@ private[sql] class ExchangeCoordinator(
     partitionStartIndices.toArray
   }
 
+  @GuardedBy("this")
   private def doEstimationIfNecessary(): Unit = synchronized {
     // It is unlikely that this method will be called from multiple threads
     // (when multiple threads trigger the execution of THIS physical)
@@ -209,11 +220,11 @@ private[sql] class ExchangeCoordinator(
 
       // Wait for the finishes of those submitted map stages.
       val mapOutputStatistics = new Array[MapOutputStatistics](submittedStageFutures.length)
-      i = 0
-      while (i < submittedStageFutures.length) {
+      var j = 0
+      while (j < submittedStageFutures.length) {
         // This call is a blocking call. If the stage has not finished, we will wait at here.
-        mapOutputStatistics(i) = submittedStageFutures(i).get()
-        i += 1
+        mapOutputStatistics(j) = submittedStageFutures(j).get()
+        j += 1
       }
 
       // Now, we estimate partitionStartIndices. partitionStartIndices.length will be the
@@ -225,14 +236,14 @@ private[sql] class ExchangeCoordinator(
           Some(estimatePartitionStartIndices(mapOutputStatistics))
         }
 
-      i = 0
-      while (i < numExchanges) {
-        val exchange = exchanges(i)
+      var k = 0
+      while (k < numExchanges) {
+        val exchange = exchanges(k)
         val rdd =
-          exchange.preparePostShuffleRDD(shuffleDependencies(i), partitionStartIndices)
+          exchange.preparePostShuffleRDD(shuffleDependencies(k), partitionStartIndices)
         newPostShuffleRDDs.put(exchange, rdd)
 
-        i += 1
+        k += 1
       }
 
       // Finally, we set postShuffleRDDs and estimated.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
index dbcb011f603f7..bce94dafad755 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
@@ -29,12 +29,12 @@ import org.scalatest.concurrent.Eventually._
 import org.apache.spark.Accumulators
 import org.apache.spark.sql.columnar._
 import org.apache.spark.sql.functions._
-import org.apache.spark.sql.test.SharedSQLContext
+import org.apache.spark.sql.test.{SQLTestUtils, SharedSQLContext}
 import org.apache.spark.storage.{StorageLevel, RDDBlockId}
 
 private case class BigData(s: String)
 
-class CachedTableSuite extends QueryTest with SharedSQLContext {
+class CachedTableSuite extends QueryTest with SQLTestUtils with SharedSQLContext {
   import testImplicits._
 
   def rddIdOf(tableName: String): Int = {
@@ -375,53 +375,135 @@ class CachedTableSuite extends QueryTest with SharedSQLContext {
       sql("SELECT key, count(*) FROM orderedTable GROUP BY key ORDER BY key"),
       sql("SELECT key, count(*) FROM testData3x GROUP BY key ORDER BY key").collect())
     sqlContext.uncacheTable("orderedTable")
+    sqlContext.dropTempTable("orderedTable")
 
     // Set up two tables distributed in the same way. Try this with the data distributed into
     // different number of partitions.
     for (numPartitions <- 1 until 10 by 4) {
-      testData.repartition(numPartitions, $"key").registerTempTable("t1")
-      testData2.repartition(numPartitions, $"a").registerTempTable("t2")
+      withTempTable("t1", "t2") {
+        testData.repartition(numPartitions, $"key").registerTempTable("t1")
+        testData2.repartition(numPartitions, $"a").registerTempTable("t2")
+        sqlContext.cacheTable("t1")
+        sqlContext.cacheTable("t2")
+
+        // Joining them should result in no exchanges.
+        verifyNumExchanges(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"), 0)
+        checkAnswer(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"),
+          sql("SELECT * FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a"))
+
+        // Grouping on the partition key should result in no exchanges
+        verifyNumExchanges(sql("SELECT count(*) FROM t1 GROUP BY key"), 0)
+        checkAnswer(sql("SELECT count(*) FROM t1 GROUP BY key"),
+          sql("SELECT count(*) FROM testData GROUP BY key"))
+
+        sqlContext.uncacheTable("t1")
+        sqlContext.uncacheTable("t2")
+      }
+    }
+
+    // Distribute the tables into non-matching number of partitions. Need to shuffle one side.
+    withTempTable("t1", "t2") {
+      testData.repartition(6, $"key").registerTempTable("t1")
+      testData2.repartition(3, $"a").registerTempTable("t2")
       sqlContext.cacheTable("t1")
       sqlContext.cacheTable("t2")
 
-      // Joining them should result in no exchanges.
-      verifyNumExchanges(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"), 0)
-      checkAnswer(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"),
-        sql("SELECT * FROM testData t1 JOIN testData2 t2 ON t1.key = t2.a"))
+      val query = sql("SELECT key, value, a, b FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a")
+      verifyNumExchanges(query, 1)
+      assert(query.queryExecution.executedPlan.outputPartitioning.numPartitions === 6)
+      checkAnswer(
+        query,
+        testData.join(testData2, $"key" === $"a").select($"key", $"value", $"a", $"b"))
+      sqlContext.uncacheTable("t1")
+      sqlContext.uncacheTable("t2")
+    }
 
-      // Grouping on the partition key should result in no exchanges
-      verifyNumExchanges(sql("SELECT count(*) FROM t1 GROUP BY key"), 0)
-      checkAnswer(sql("SELECT count(*) FROM t1 GROUP BY key"),
-        sql("SELECT count(*) FROM testData GROUP BY key"))
+    // One side of join is not partitioned in the desired way. Need to shuffle one side.
+    withTempTable("t1", "t2") {
+      testData.repartition(6, $"value").registerTempTable("t1")
+      testData2.repartition(6, $"a").registerTempTable("t2")
+      sqlContext.cacheTable("t1")
+      sqlContext.cacheTable("t2")
 
+      val query = sql("SELECT key, value, a, b FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a")
+      verifyNumExchanges(query, 1)
+      assert(query.queryExecution.executedPlan.outputPartitioning.numPartitions === 6)
+      checkAnswer(
+        query,
+        testData.join(testData2, $"key" === $"a").select($"key", $"value", $"a", $"b"))
       sqlContext.uncacheTable("t1")
       sqlContext.uncacheTable("t2")
-      sqlContext.dropTempTable("t1")
-      sqlContext.dropTempTable("t2")
     }
 
-    // Distribute the tables into non-matching number of partitions. Need to shuffle.
-    testData.repartition(6, $"key").registerTempTable("t1")
-    testData2.repartition(3, $"a").registerTempTable("t2")
-    sqlContext.cacheTable("t1")
-    sqlContext.cacheTable("t2")
+    withTempTable("t1", "t2") {
+      testData.repartition(6, $"value").registerTempTable("t1")
+      testData2.repartition(12, $"a").registerTempTable("t2")
+      sqlContext.cacheTable("t1")
+      sqlContext.cacheTable("t2")
 
-    verifyNumExchanges(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"), 2)
-    sqlContext.uncacheTable("t1")
-    sqlContext.uncacheTable("t2")
-    sqlContext.dropTempTable("t1")
-    sqlContext.dropTempTable("t2")
+      val query = sql("SELECT key, value, a, b FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a")
+      verifyNumExchanges(query, 1)
+      assert(query.queryExecution.executedPlan.outputPartitioning.numPartitions === 12)
+      checkAnswer(
+        query,
+        testData.join(testData2, $"key" === $"a").select($"key", $"value", $"a", $"b"))
+      sqlContext.uncacheTable("t1")
+      sqlContext.uncacheTable("t2")
+    }
 
-    // One side of join is not partitioned in the desired way. Need to shuffle.
-    testData.repartition(6, $"value").registerTempTable("t1")
-    testData2.repartition(6, $"a").registerTempTable("t2")
-    sqlContext.cacheTable("t1")
-    sqlContext.cacheTable("t2")
+    // One side of join is not partitioned in the desired way. Since the number of partitions of
+    // the side that has already partitioned is smaller than the side that is not partitioned,
+    // we shuffle both side.
+    withTempTable("t1", "t2") {
+      testData.repartition(6, $"value").registerTempTable("t1")
+      testData2.repartition(3, $"a").registerTempTable("t2")
+      sqlContext.cacheTable("t1")
+      sqlContext.cacheTable("t2")
 
-    verifyNumExchanges(sql("SELECT * FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a"), 2)
-    sqlContext.uncacheTable("t1")
-    sqlContext.uncacheTable("t2")
-    sqlContext.dropTempTable("t1")
-    sqlContext.dropTempTable("t2")
+      val query = sql("SELECT key, value, a, b FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a")
+      verifyNumExchanges(query, 2)
+      checkAnswer(
+        query,
+        testData.join(testData2, $"key" === $"a").select($"key", $"value", $"a", $"b"))
+      sqlContext.uncacheTable("t1")
+      sqlContext.uncacheTable("t2")
+    }
+
+    // repartition's column ordering is different from group by column ordering.
+    // But they use the same set of columns.
+    withTempTable("t1") {
+      testData.repartition(6, $"value", $"key").registerTempTable("t1")
+      sqlContext.cacheTable("t1")
+
+      val query = sql("SELECT value, key from t1 group by key, value")
+      verifyNumExchanges(query, 0)
+      checkAnswer(
+        query,
+        testData.distinct().select($"value", $"key"))
+      sqlContext.uncacheTable("t1")
+    }
+
+    // repartition's column ordering is different from join condition's column ordering.
+    // We will still shuffle because hashcodes of a row depend on the column ordering.
+    // If we do not shuffle, we may actually partition two tables in totally two different way.
+    // See PartitioningSuite for more details.
+    withTempTable("t1", "t2") {
+      val df1 = testData
+      df1.repartition(6, $"value", $"key").registerTempTable("t1")
+      val df2 = testData2.select($"a", $"b".cast("string"))
+      df2.repartition(6, $"a", $"b").registerTempTable("t2")
+      sqlContext.cacheTable("t1")
+      sqlContext.cacheTable("t2")
+
+      val query =
+        sql("SELECT key, value, a, b FROM t1 t1 JOIN t2 t2 ON t1.key = t2.a and t1.value = t2.b")
+      verifyNumExchanges(query, 1)
+      assert(query.queryExecution.executedPlan.outputPartitioning.numPartitions === 6)
+      checkAnswer(
+        query,
+        df1.join(df2, $"key" === $"a" && $"value" === $"b").select($"key", $"value", $"a", $"b"))
+      sqlContext.uncacheTable("t1")
+      sqlContext.uncacheTable("t2")
+    }
   }
 }

From 62bb290773c9f9fa53cbe6d4eedc6e153761a763 Mon Sep 17 00:00:00 2001
From: Jacek Laskowski <jacek.laskowski@deepsense.io>
Date: Fri, 6 Nov 2015 20:05:18 +0000
Subject: [PATCH 097/518] Typo fixes + code readability improvements

Author: Jacek Laskowski <jacek.laskowski@deepsense.io>

Closes #9501 from jaceklaskowski/typos-with-style.
---
 .../scala/org/apache/spark/rdd/HadoopRDD.scala     | 14 ++++++--------
 .../org/apache/spark/scheduler/DAGScheduler.scala  | 12 +++++++++---
 .../apache/spark/scheduler/ShuffleMapTask.scala    | 10 +++++-----
 .../scala/org/apache/spark/scheduler/TaskSet.scala |  2 +-
 4 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala b/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
index d841f05ec52cf..0453614f6a1d3 100644
--- a/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
@@ -88,8 +88,8 @@ private[spark] class HadoopPartition(rddId: Int, idx: Int, s: InputSplit)
  *
  * @param sc The SparkContext to associate the RDD with.
  * @param broadcastedConf A general Hadoop Configuration, or a subclass of it. If the enclosed
- *     variabe references an instance of JobConf, then that JobConf will be used for the Hadoop job.
- *     Otherwise, a new JobConf will be created on each slave using the enclosed Configuration.
+ *   variable references an instance of JobConf, then that JobConf will be used for the Hadoop job.
+ *   Otherwise, a new JobConf will be created on each slave using the enclosed Configuration.
  * @param initLocalJobConfFuncOpt Optional closure used to initialize any JobConf that HadoopRDD
  *     creates.
  * @param inputFormatClass Storage format of the data to be read.
@@ -123,7 +123,7 @@ class HadoopRDD[K, V](
       sc,
       sc.broadcast(new SerializableConfiguration(conf))
         .asInstanceOf[Broadcast[SerializableConfiguration]],
-      None /* initLocalJobConfFuncOpt */,
+      initLocalJobConfFuncOpt = None,
       inputFormatClass,
       keyClass,
       valueClass,
@@ -184,8 +184,9 @@ class HadoopRDD[K, V](
   protected def getInputFormat(conf: JobConf): InputFormat[K, V] = {
     val newInputFormat = ReflectionUtils.newInstance(inputFormatClass.asInstanceOf[Class[_]], conf)
       .asInstanceOf[InputFormat[K, V]]
-    if (newInputFormat.isInstanceOf[Configurable]) {
-      newInputFormat.asInstanceOf[Configurable].setConf(conf)
+    newInputFormat match {
+      case c: Configurable => c.setConf(conf)
+      case _ =>
     }
     newInputFormat
   }
@@ -195,9 +196,6 @@ class HadoopRDD[K, V](
     // add the credentials here as this can be called before SparkContext initialized
     SparkHadoopUtil.get.addCredentials(jobConf)
     val inputFormat = getInputFormat(jobConf)
-    if (inputFormat.isInstanceOf[Configurable]) {
-      inputFormat.asInstanceOf[Configurable].setConf(jobConf)
-    }
     val inputSplits = inputFormat.getSplits(jobConf, minPartitions)
     val array = new Array[Partition](inputSplits.size)
     for (i <- 0 until inputSplits.size) {
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index a1f0fd05f661a..4a9518fff4e7b 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -541,8 +541,7 @@ class DAGScheduler(
   }
 
   /**
-   * Submit an action job to the scheduler and get a JobWaiter object back. The JobWaiter object
-   * can be used to block until the the job finishes executing or can be used to cancel the job.
+   * Submit an action job to the scheduler.
    *
    * @param rdd target RDD to run tasks on
    * @param func a function to run on each partition of the RDD
@@ -551,6 +550,11 @@ class DAGScheduler(
    * @param callSite where in the user program this job was called
    * @param resultHandler callback to pass each result to
    * @param properties scheduler properties to attach to this job, e.g. fair scheduler pool name
+   *
+   * @return a JobWaiter object that can be used to block until the job finishes executing
+   *         or can be used to cancel the job.
+   *
+   * @throws IllegalArgumentException when partitions ids are illegal
    */
   def submitJob[T, U](
       rdd: RDD[T],
@@ -584,7 +588,7 @@ class DAGScheduler(
 
   /**
    * Run an action job on the given RDD and pass all the results to the resultHandler function as
-   * they arrive. Throws an exception if the job fials, or returns normally if successful.
+   * they arrive.
    *
    * @param rdd target RDD to run tasks on
    * @param func a function to run on each partition of the RDD
@@ -593,6 +597,8 @@ class DAGScheduler(
    * @param callSite where in the user program this job was called
    * @param resultHandler callback to pass each result to
    * @param properties scheduler properties to attach to this job, e.g. fair scheduler pool name
+   *
+   * @throws Exception when the job fails
    */
   def runJob[T, U](
       rdd: RDD[T],
diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala
index f478f9982afef..ea97ef0e746d8 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala
@@ -27,11 +27,11 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.shuffle.ShuffleWriter
 
 /**
-* A ShuffleMapTask divides the elements of an RDD into multiple buckets (based on a partitioner
-* specified in the ShuffleDependency).
-*
-* See [[org.apache.spark.scheduler.Task]] for more information.
-*
+ * A ShuffleMapTask divides the elements of an RDD into multiple buckets (based on a partitioner
+ * specified in the ShuffleDependency).
+ *
+ * See [[org.apache.spark.scheduler.Task]] for more information.
+ *
  * @param stageId id of the stage this task belongs to
  * @param taskBinary broadcast version of the RDD and the ShuffleDependency. Once deserialized,
  *                   the type should be (RDD[_], ShuffleDependency[_, _, _]).
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala
index be8526ba9b94f..517c8991aed78 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala
@@ -29,7 +29,7 @@ private[spark] class TaskSet(
     val stageAttemptId: Int,
     val priority: Int,
     val properties: Properties) {
-    val id: String = stageId + "." + stageAttemptId
+  val id: String = stageId + "." + stageAttemptId
 
   override def toString: String = "TaskSet " + id
 }

From 49f1a820372d1cba41f3f00d07eb5728f2ed6705 Mon Sep 17 00:00:00 2001
From: Imran Rashid <irashid@cloudera.com>
Date: Fri, 6 Nov 2015 20:06:24 +0000
Subject: [PATCH 098/518] [SPARK-10116][CORE] XORShiftRandom.hashSeed is random
 in high bits

https://issues.apache.org/jira/browse/SPARK-10116

This is really trivial, just happened to notice it -- if `XORShiftRandom.hashSeed` is really supposed to have random bits throughout (as the comment implies), it needs to do something for the conversion to `long`.

mengxr mkolod

Author: Imran Rashid <irashid@cloudera.com>

Closes #8314 from squito/SPARK-10116.
---
 R/pkg/inst/tests/test_sparkSQL.R              |  8 +--
 .../spark/util/random/XORShiftRandom.scala    |  6 ++-
 .../java/org/apache/spark/JavaAPISuite.java   | 20 ++++---
 .../spark/rdd/PairRDDFunctionsSuite.scala     | 52 +++++++++++++------
 .../util/random/XORShiftRandomSuite.scala     | 15 ++++++
 .../MultilayerPerceptronClassifierSuite.scala |  5 +-
 .../spark/ml/feature/Word2VecSuite.scala      | 16 ++++--
 .../clustering/StreamingKMeansSuite.scala     | 13 +++--
 python/pyspark/ml/feature.py                  | 20 +++----
 python/pyspark/ml/recommendation.py           |  6 +--
 python/pyspark/mllib/recommendation.py        |  4 +-
 python/pyspark/sql/dataframe.py               |  6 +--
 .../catalyst/expressions/RandomSuite.scala    |  8 +--
 .../apache/spark/sql/JavaDataFrameSuite.java  |  6 ++-
 .../apache/spark/sql/DataFrameStatSuite.scala |  4 +-
 15 files changed, 128 insertions(+), 61 deletions(-)

diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 816315b1e4e13..92cff1fba7193 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -875,9 +875,9 @@ test_that("column binary mathfunctions", {
   expect_equal(collect(select(df, shiftRight(df$b, 1)))[4, 1], 4)
   expect_equal(collect(select(df, shiftRightUnsigned(df$b, 1)))[4, 1], 4)
   expect_equal(class(collect(select(df, rand()))[2, 1]), "numeric")
-  expect_equal(collect(select(df, rand(1)))[1, 1], 0.45, tolerance = 0.01)
+  expect_equal(collect(select(df, rand(1)))[1, 1], 0.134, tolerance = 0.01)
   expect_equal(class(collect(select(df, randn()))[2, 1]), "numeric")
-  expect_equal(collect(select(df, randn(1)))[1, 1], -0.0111, tolerance = 0.01)
+  expect_equal(collect(select(df, randn(1)))[1, 1], -1.03, tolerance = 0.01)
 })
 
 test_that("string operators", {
@@ -1458,8 +1458,8 @@ test_that("sampleBy() on a DataFrame", {
   fractions <- list("0" = 0.1, "1" = 0.2)
   sample <- sampleBy(df, "key", fractions, 0)
   result <- collect(orderBy(count(groupBy(sample, "key")), "key"))
-  expect_identical(as.list(result[1, ]), list(key = "0", count = 2))
-  expect_identical(as.list(result[2, ]), list(key = "1", count = 10))
+  expect_identical(as.list(result[1, ]), list(key = "0", count = 3))
+  expect_identical(as.list(result[2, ]), list(key = "1", count = 7))
 })
 
 test_that("SQL error message is returned from JVM", {
diff --git a/core/src/main/scala/org/apache/spark/util/random/XORShiftRandom.scala b/core/src/main/scala/org/apache/spark/util/random/XORShiftRandom.scala
index 85fb923cd9bc7..e8cdb6e98bf36 100644
--- a/core/src/main/scala/org/apache/spark/util/random/XORShiftRandom.scala
+++ b/core/src/main/scala/org/apache/spark/util/random/XORShiftRandom.scala
@@ -60,9 +60,11 @@ private[spark] class XORShiftRandom(init: Long) extends JavaRandom(init) {
 private[spark] object XORShiftRandom {
 
   /** Hash seeds to have 0/1 bits throughout. */
-  private def hashSeed(seed: Long): Long = {
+  private[random] def hashSeed(seed: Long): Long = {
     val bytes = ByteBuffer.allocate(java.lang.Long.SIZE).putLong(seed).array()
-    MurmurHash3.bytesHash(bytes)
+    val lowBits = MurmurHash3.bytesHash(bytes)
+    val highBits = MurmurHash3.bytesHash(bytes, lowBits)
+    (highBits.toLong << 32) | (lowBits.toLong & 0xFFFFFFFFL)
   }
 
   /**
diff --git a/core/src/test/java/org/apache/spark/JavaAPISuite.java b/core/src/test/java/org/apache/spark/JavaAPISuite.java
index fd8f7f39b7cc8..4d4e9820500e7 100644
--- a/core/src/test/java/org/apache/spark/JavaAPISuite.java
+++ b/core/src/test/java/org/apache/spark/JavaAPISuite.java
@@ -146,21 +146,29 @@ public void intersection() {
   public void sample() {
     List<Integer> ints = Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
     JavaRDD<Integer> rdd = sc.parallelize(ints);
-    JavaRDD<Integer> sample20 = rdd.sample(true, 0.2, 3);
+    // the seeds here are "magic" to make this work out nicely
+    JavaRDD<Integer> sample20 = rdd.sample(true, 0.2, 8);
     Assert.assertEquals(2, sample20.count());
-    JavaRDD<Integer> sample20WithoutReplacement = rdd.sample(false, 0.2, 5);
+    JavaRDD<Integer> sample20WithoutReplacement = rdd.sample(false, 0.2, 2);
     Assert.assertEquals(2, sample20WithoutReplacement.count());
   }
 
   @Test
   public void randomSplit() {
-    List<Integer> ints = Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
+    List<Integer> ints = new ArrayList<>(1000);
+    for (int i = 0; i < 1000; i++) {
+      ints.add(i);
+    }
     JavaRDD<Integer> rdd = sc.parallelize(ints);
     JavaRDD<Integer>[] splits = rdd.randomSplit(new double[] { 0.4, 0.6, 1.0 }, 31);
+    // the splits aren't perfect -- not enough data for them to be -- just check they're about right
     Assert.assertEquals(3, splits.length);
-    Assert.assertEquals(1, splits[0].count());
-    Assert.assertEquals(2, splits[1].count());
-    Assert.assertEquals(7, splits[2].count());
+    long s0 = splits[0].count();
+    long s1 = splits[1].count();
+    long s2 = splits[2].count();
+    Assert.assertTrue(s0 + " not within expected range", s0 > 150 && s0 < 250);
+    Assert.assertTrue(s1 + " not within expected range", s1 > 250 && s0 < 350);
+    Assert.assertTrue(s2 + " not within expected range", s2 > 430 && s2 < 570);
   }
 
   @Test
diff --git a/core/src/test/scala/org/apache/spark/rdd/PairRDDFunctionsSuite.scala b/core/src/test/scala/org/apache/spark/rdd/PairRDDFunctionsSuite.scala
index 1321ec84735b5..7d2cfcca9436a 100644
--- a/core/src/test/scala/org/apache/spark/rdd/PairRDDFunctionsSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rdd/PairRDDFunctionsSuite.scala
@@ -17,6 +17,7 @@
 
 package org.apache.spark.rdd
 
+import org.apache.commons.math3.distribution.{PoissonDistribution, BinomialDistribution}
 import org.apache.hadoop.fs.FileSystem
 import org.apache.hadoop.mapred._
 import org.apache.hadoop.util.Progressable
@@ -578,17 +579,36 @@ class PairRDDFunctionsSuite extends SparkFunSuite with SharedSparkContext {
       (x: Int) => if (x % 10 < (10 * fractionPositive).toInt) "1" else "0"
     }
 
-    def checkSize(exact: Boolean,
-        withReplacement: Boolean,
-        expected: Long,
-        actual: Long,
-        p: Double): Boolean = {
+    def assertBinomialSample(
+        exact: Boolean,
+        actual: Int,
+        trials: Int,
+        p: Double): Unit = {
+      if (exact) {
+        assert(actual == math.ceil(p * trials).toInt)
+      } else {
+        val dist = new BinomialDistribution(trials, p)
+        val q = dist.cumulativeProbability(actual)
+        withClue(s"p = $p: trials = $trials") {
+          assert(q >= 0.001 && q <= 0.999)
+        }
+      }
+    }
+
+    def assertPoissonSample(
+        exact: Boolean,
+        actual: Int,
+        trials: Int,
+        p: Double): Unit = {
       if (exact) {
-        return expected == actual
+        assert(actual == math.ceil(p * trials).toInt)
+      } else {
+        val dist = new PoissonDistribution(p * trials)
+        val q = dist.cumulativeProbability(actual)
+        withClue(s"p = $p: trials = $trials") {
+          assert(q >= 0.001 && q <= 0.999)
+        }
       }
-      val stdev = if (withReplacement) math.sqrt(expected) else math.sqrt(expected * p * (1 - p))
-      // Very forgiving margin since we're dealing with very small sample sizes most of the time
-      math.abs(actual - expected) <= 6 * stdev
     }
 
     def testSampleExact(stratifiedData: RDD[(String, Int)],
@@ -613,8 +633,7 @@ class PairRDDFunctionsSuite extends SparkFunSuite with SharedSparkContext {
         samplingRate: Double,
         seed: Long,
         n: Long): Unit = {
-      val expectedSampleSize = stratifiedData.countByKey()
-        .mapValues(count => math.ceil(count * samplingRate).toInt)
+      val trials = stratifiedData.countByKey()
       val fractions = Map("1" -> samplingRate, "0" -> samplingRate)
       val sample = if (exact) {
         stratifiedData.sampleByKeyExact(false, fractions, seed)
@@ -623,8 +642,10 @@ class PairRDDFunctionsSuite extends SparkFunSuite with SharedSparkContext {
       }
       val sampleCounts = sample.countByKey()
       val takeSample = sample.collect()
-      sampleCounts.foreach { case(k, v) =>
-        assert(checkSize(exact, false, expectedSampleSize(k), v, samplingRate)) }
+      sampleCounts.foreach { case (k, v) =>
+        assertBinomialSample(exact = exact, actual = v.toInt, trials = trials(k).toInt,
+          p = samplingRate)
+      }
       assert(takeSample.size === takeSample.toSet.size)
       takeSample.foreach { x => assert(1 <= x._2 && x._2 <= n, s"elements not in [1, $n]") }
     }
@@ -635,6 +656,7 @@ class PairRDDFunctionsSuite extends SparkFunSuite with SharedSparkContext {
         samplingRate: Double,
         seed: Long,
         n: Long): Unit = {
+      val trials = stratifiedData.countByKey()
       val expectedSampleSize = stratifiedData.countByKey().mapValues(count =>
         math.ceil(count * samplingRate).toInt)
       val fractions = Map("1" -> samplingRate, "0" -> samplingRate)
@@ -646,7 +668,7 @@ class PairRDDFunctionsSuite extends SparkFunSuite with SharedSparkContext {
       val sampleCounts = sample.countByKey()
       val takeSample = sample.collect()
       sampleCounts.foreach { case (k, v) =>
-        assert(checkSize(exact, true, expectedSampleSize(k), v, samplingRate))
+        assertPoissonSample(exact, actual = v.toInt, trials = trials(k).toInt, p = samplingRate)
       }
       val groupedByKey = takeSample.groupBy(_._1)
       for ((key, v) <- groupedByKey) {
@@ -657,7 +679,7 @@ class PairRDDFunctionsSuite extends SparkFunSuite with SharedSparkContext {
           if (exact) {
             assert(v.toSet.size <= expectedSampleSize(key))
           } else {
-            assert(checkSize(false, true, expectedSampleSize(key), v.toSet.size, samplingRate))
+            assertPoissonSample(false, actual = v.toSet.size, trials(key).toInt, p = samplingRate)
           }
         }
       }
diff --git a/core/src/test/scala/org/apache/spark/util/random/XORShiftRandomSuite.scala b/core/src/test/scala/org/apache/spark/util/random/XORShiftRandomSuite.scala
index d26667bf720cf..a5b50fce5c0a9 100644
--- a/core/src/test/scala/org/apache/spark/util/random/XORShiftRandomSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/random/XORShiftRandomSuite.scala
@@ -65,4 +65,19 @@ class XORShiftRandomSuite extends SparkFunSuite with Matchers {
     val random = new XORShiftRandom(0L)
     assert(random.nextInt() != 0)
   }
+
+  test ("hashSeed has random bits throughout") {
+    val totalBitCount = (0 until 10).map { seed =>
+      val hashed = XORShiftRandom.hashSeed(seed)
+      val bitCount = java.lang.Long.bitCount(hashed)
+      // make sure we have roughly equal numbers of 0s and 1s.  Mostly just check that we
+      // don't have all 0s or 1s in the high bits
+      bitCount should be > 20
+      bitCount should be < 44
+      bitCount
+    }.sum
+    // and over all the seeds, very close to equal numbers of 0s & 1s
+    totalBitCount should be > (32 * 10 - 30)
+    totalBitCount should be < (32 * 10 + 30)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifierSuite.scala
index 17db8c44777d4..a326432d017fc 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifierSuite.scala
@@ -61,8 +61,9 @@ class MultilayerPerceptronClassifierSuite extends SparkFunSuite with MLlibTestSp
 
     val xMean = Array(5.843, 3.057, 3.758, 1.199)
     val xVariance = Array(0.6856, 0.1899, 3.116, 0.581)
+    // the input seed is somewhat magic, to make this test pass
     val rdd = sc.parallelize(generateMultinomialLogisticInput(
-      coefficients, xMean, xVariance, true, nPoints, 42), 2)
+      coefficients, xMean, xVariance, true, nPoints, 1), 2)
     val dataFrame = sqlContext.createDataFrame(rdd).toDF("label", "features")
     val numClasses = 3
     val numIterations = 100
@@ -70,7 +71,7 @@ class MultilayerPerceptronClassifierSuite extends SparkFunSuite with MLlibTestSp
     val trainer = new MultilayerPerceptronClassifier()
       .setLayers(layers)
       .setBlockSize(1)
-      .setSeed(11L)
+      .setSeed(11L) // currently this seed is ignored
       .setMaxIter(numIterations)
     val model = trainer.fit(dataFrame)
     val numFeatures = dataFrame.select("features").first().getAs[Vector](0).size
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
index a2e46f2029956..23dfdaa9f8fc6 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
@@ -66,9 +66,12 @@ class Word2VecSuite extends SparkFunSuite with MLlibTestSparkContext {
     // copied model must have the same parent.
     MLTestingUtils.checkCopy(model)
 
+    // These expectations are just magic values, characterizing the current
+    // behavior.  The test needs to be updated to be more general, see SPARK-11502
+    val magicExp = Vectors.dense(0.30153007534417237, -0.6833061711354689, 0.5116530778733167)
     model.transform(docDF).select("result", "expected").collect().foreach {
       case Row(vector1: Vector, vector2: Vector) =>
-        assert(vector1 ~== vector2 absTol 1E-5, "Transformed vector is different with expected.")
+        assert(vector1 ~== magicExp absTol 1E-5, "Transformed vector is different with expected.")
     }
   }
 
@@ -99,8 +102,15 @@ class Word2VecSuite extends SparkFunSuite with MLlibTestSparkContext {
     val realVectors = model.getVectors.sort("word").select("vector").map {
       case Row(v: Vector) => v
     }.collect()
+    // These expectations are just magic values, characterizing the current
+    // behavior.  The test needs to be updated to be more general, see SPARK-11502
+    val magicExpected = Seq(
+      Vectors.dense(0.3326166272163391, -0.5603077411651611, -0.2309209555387497),
+      Vectors.dense(0.32463887333869934, -0.9306551218032837, 1.393115520477295),
+      Vectors.dense(-0.27150997519493103, 0.4372006058692932, -0.13465698063373566)
+    )
 
-    realVectors.zip(expectedVectors).foreach {
+    realVectors.zip(magicExpected).foreach {
       case (real, expected) =>
         assert(real ~== expected absTol 1E-5, "Actual vector is different from expected.")
     }
@@ -122,7 +132,7 @@ class Word2VecSuite extends SparkFunSuite with MLlibTestSparkContext {
       .setSeed(42L)
       .fit(docDF)
 
-    val expectedSimilarity = Array(0.2789285076917586, -0.6336972059851644)
+    val expectedSimilarity = Array(0.18032623242822343, -0.5717976464798823)
     val (synonyms, similarity) = model.findSynonyms("a", 2).map {
       case Row(w: String, sim: Double) => (w, sim)
     }.collect().unzip
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/clustering/StreamingKMeansSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/clustering/StreamingKMeansSuite.scala
index 3645d29dccdb2..65e37c64d404e 100644
--- a/mllib/src/test/scala/org/apache/spark/mllib/clustering/StreamingKMeansSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/mllib/clustering/StreamingKMeansSuite.scala
@@ -98,9 +98,16 @@ class StreamingKMeansSuite extends SparkFunSuite with TestSuiteBase {
     runStreams(ssc, numBatches, numBatches)
 
     // check that estimated centers are close to true centers
-    // NOTE exact assignment depends on the initialization!
-    assert(centers(0) ~== kMeans.latestModel().clusterCenters(0) absTol 1E-1)
-    assert(centers(1) ~== kMeans.latestModel().clusterCenters(1) absTol 1E-1)
+    // cluster ordering is arbitrary, so choose closest cluster
+    val d0 = Vectors.sqdist(kMeans.latestModel().clusterCenters(0), centers(0))
+    val d1 = Vectors.sqdist(kMeans.latestModel().clusterCenters(0), centers(1))
+    val (c0, c1) = if (d0 < d1) {
+      (centers(0), centers(1))
+    } else {
+      (centers(1), centers(0))
+    }
+    assert(c0 ~== kMeans.latestModel().clusterCenters(0) absTol 1E-1)
+    assert(c1 ~== kMeans.latestModel().clusterCenters(1) absTol 1E-1)
   }
 
   test("detecting dying clusters") {
diff --git a/python/pyspark/ml/feature.py b/python/pyspark/ml/feature.py
index c7b6dd926c3e8..b02d41b52ab25 100644
--- a/python/pyspark/ml/feature.py
+++ b/python/pyspark/ml/feature.py
@@ -1788,21 +1788,21 @@ class Word2Vec(JavaEstimator, HasStepSize, HasMaxIter, HasSeed, HasInputCol, Has
     +----+--------------------+
     |word|              vector|
     +----+--------------------+
-    |   a|[-0.3511952459812...|
-    |   b|[0.29077222943305...|
-    |   c|[0.02315592765808...|
+    |   a|[0.09461779892444...|
+    |   b|[1.15474212169647...|
+    |   c|[-0.3794820010662...|
     +----+--------------------+
     ...
     >>> model.findSynonyms("a", 2).show()
-    +----+-------------------+
-    |word|         similarity|
-    +----+-------------------+
-    |   b|0.29255685145799626|
-    |   c|-0.5414068302988307|
-    +----+-------------------+
+    +----+--------------------+
+    |word|          similarity|
+    +----+--------------------+
+    |   b| 0.16782984556103436|
+    |   c|-0.46761559092107646|
+    +----+--------------------+
     ...
     >>> model.transform(doc).head().model
-    DenseVector([-0.0422, -0.5138, -0.2546, 0.6885, 0.276])
+    DenseVector([0.5524, -0.4995, -0.3599, 0.0241, 0.3461])
 
     .. versionadded:: 1.4.0
     """
diff --git a/python/pyspark/ml/recommendation.py b/python/pyspark/ml/recommendation.py
index ec5748a1cfe94..b44c66f73cc49 100644
--- a/python/pyspark/ml/recommendation.py
+++ b/python/pyspark/ml/recommendation.py
@@ -76,11 +76,11 @@ class ALS(JavaEstimator, HasCheckpointInterval, HasMaxIter, HasPredictionCol, Ha
     >>> test = sqlContext.createDataFrame([(0, 2), (1, 0), (2, 0)], ["user", "item"])
     >>> predictions = sorted(model.transform(test).collect(), key=lambda r: r[0])
     >>> predictions[0]
-    Row(user=0, item=2, prediction=0.39...)
+    Row(user=0, item=2, prediction=-0.13807615637779236)
     >>> predictions[1]
-    Row(user=1, item=0, prediction=3.19...)
+    Row(user=1, item=0, prediction=2.6258413791656494)
     >>> predictions[2]
-    Row(user=2, item=0, prediction=-1.15...)
+    Row(user=2, item=0, prediction=-1.5018409490585327)
 
     .. versionadded:: 1.4.0
     """
diff --git a/python/pyspark/mllib/recommendation.py b/python/pyspark/mllib/recommendation.py
index b9442b0d16c0f..93e47a797f490 100644
--- a/python/pyspark/mllib/recommendation.py
+++ b/python/pyspark/mllib/recommendation.py
@@ -101,12 +101,12 @@ class MatrixFactorizationModel(JavaModelWrapper, JavaSaveable, JavaLoader):
 
     >>> model = ALS.train(ratings, 1, nonnegative=True, seed=10)
     >>> model.predict(2, 2)
-    3.8...
+    3.73...
 
     >>> df = sqlContext.createDataFrame([Rating(1, 1, 1.0), Rating(1, 2, 2.0), Rating(2, 1, 2.0)])
     >>> model = ALS.train(df, 1, nonnegative=True, seed=10)
     >>> model.predict(2, 2)
-    3.8...
+    3.73...
 
     >>> model = ALS.trainImplicit(ratings, 1, nonnegative=True, seed=10)
     >>> model.predict(2, 2)
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index 3baff8147753d..765a4511b64bc 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -436,7 +436,7 @@ def sample(self, withReplacement, fraction, seed=None):
         """Returns a sampled subset of this :class:`DataFrame`.
 
         >>> df.sample(False, 0.5, 42).count()
-        1
+        2
         """
         assert fraction >= 0.0, "Negative fraction value: %s" % fraction
         seed = seed if seed is not None else random.randint(0, sys.maxsize)
@@ -463,8 +463,8 @@ def sampleBy(self, col, fractions, seed=None):
         +---+-----+
         |key|count|
         +---+-----+
-        |  0|    3|
-        |  1|    8|
+        |  0|    5|
+        |  1|    9|
         +---+-----+
 
         """
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala
index 4a644d136f09c..b7a0d44fa7e57 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala
@@ -24,12 +24,12 @@ import org.apache.spark.SparkFunSuite
 class RandomSuite extends SparkFunSuite with ExpressionEvalHelper {
 
   test("random") {
-    checkDoubleEvaluation(Rand(30), 0.7363714192755834 +- 0.001)
-    checkDoubleEvaluation(Randn(30), 0.5181478766595276 +- 0.001)
+    checkDoubleEvaluation(Rand(30), 0.31429268272540556 +- 0.001)
+    checkDoubleEvaluation(Randn(30), -0.4798519469521663 +- 0.001)
   }
 
   test("SPARK-9127 codegen with long seed") {
-    checkDoubleEvaluation(Rand(5419823303878592871L), 0.4061913198963727 +- 0.001)
-    checkDoubleEvaluation(Randn(5419823303878592871L), -0.24417152005343168 +- 0.001)
+    checkDoubleEvaluation(Rand(5419823303878592871L), 0.2304755080444375 +- 0.001)
+    checkDoubleEvaluation(Randn(5419823303878592871L), -1.2824262718225607 +- 0.001)
   }
 }
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index 49f516e86d754..40bff57a17a03 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -257,7 +257,9 @@ public void testSampleBy() {
     DataFrame df = context.range(0, 100, 1, 2).select(col("id").mod(3).as("key"));
     DataFrame sampled = df.stat().<Integer>sampleBy("key", ImmutableMap.of(0, 0.1, 1, 0.2), 0L);
     Row[] actual = sampled.groupBy("key").count().orderBy("key").collect();
-    Row[] expected = {RowFactory.create(0, 5), RowFactory.create(1, 8)};
-    Assert.assertArrayEquals(expected, actual);
+    Assert.assertEquals(0, actual[0].getLong(0));
+    Assert.assertTrue(0 <= actual[0].getLong(1) && actual[0].getLong(1) <= 8);
+    Assert.assertEquals(1, actual[1].getLong(0));
+    Assert.assertTrue(2 <= actual[1].getLong(1) && actual[1].getLong(1) <= 13);
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala
index 6524abcf5e97f..b15af42caa3ab 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala
@@ -41,7 +41,7 @@ class DataFrameStatSuite extends QueryTest with SharedSQLContext {
     val data = sparkContext.parallelize(1 to n, 2).toDF("id")
     checkAnswer(
       data.sample(withReplacement = false, 0.05, seed = 13),
-      Seq(16, 23, 88, 100).map(Row(_))
+      Seq(3, 17, 27, 58, 62).map(Row(_))
     )
   }
 
@@ -186,6 +186,6 @@ class DataFrameStatSuite extends QueryTest with SharedSQLContext {
     val sampled = df.stat.sampleBy("key", Map(0 -> 0.1, 1 -> 0.2), 0L)
     checkAnswer(
       sampled.groupBy("key").count().orderBy("key"),
-      Seq(Row(0, 5), Row(1, 8)))
+      Seq(Row(0, 6), Row(1, 11)))
   }
 }

From f328fedafd7bd084470a5e402de0429b5b7f8cd7 Mon Sep 17 00:00:00 2001
From: Herman van Hovell <hvanhovell@questtec.nl>
Date: Fri, 6 Nov 2015 12:21:53 -0800
Subject: [PATCH 099/518] [SPARK-11450] [SQL] Add Unsafe Row processing to
 Expand

This PR enables the Expand operator to process and produce Unsafe Rows.

Author: Herman van Hovell <hvanhovell@questtec.nl>

Closes #9414 from hvanhovell/SPARK-11450.
---
 .../sql/catalyst/expressions/Projection.scala |  6 ++-
 .../apache/spark/sql/execution/Expand.scala   | 19 ++++---
 .../spark/sql/execution/basicOperators.scala  |  8 +--
 .../spark/sql/execution/ExpandSuite.scala     | 54 +++++++++++++++++++
 4 files changed, 73 insertions(+), 14 deletions(-)
 create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/ExpandSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
index a6fe730f6dad4..79dabe8e925ad 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
@@ -128,7 +128,11 @@ object UnsafeProjection {
    * Returns an UnsafeProjection for given sequence of Expressions (bounded).
    */
   def create(exprs: Seq[Expression]): UnsafeProjection = {
-    GenerateUnsafeProjection.generate(exprs)
+    val unsafeExprs = exprs.map(_ transform {
+      case CreateStruct(children) => CreateStructUnsafe(children)
+      case CreateNamedStruct(children) => CreateNamedStructUnsafe(children)
+    })
+    GenerateUnsafeProjection.generate(unsafeExprs)
   }
 
   def create(expr: Expression): UnsafeProjection = create(Seq(expr))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Expand.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Expand.scala
index a458881f40948..55e95769d3faa 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Expand.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Expand.scala
@@ -41,14 +41,21 @@ case class Expand(
   // as UNKNOWN partitioning
   override def outputPartitioning: Partitioning = UnknownPartitioning(0)
 
+  override def outputsUnsafeRows: Boolean = child.outputsUnsafeRows
+  override def canProcessUnsafeRows: Boolean = true
+  override def canProcessSafeRows: Boolean = true
+
+  private[this] val projection = {
+    if (outputsUnsafeRows) {
+      (exprs: Seq[Expression]) => UnsafeProjection.create(exprs, child.output)
+    } else {
+      (exprs: Seq[Expression]) => newMutableProjection(exprs, child.output)()
+    }
+  }
+
   protected override def doExecute(): RDD[InternalRow] = attachTree(this, "execute") {
     child.execute().mapPartitions { iter =>
-      // TODO Move out projection objects creation and transfer to
-      // workers via closure. However we can't assume the Projection
-      // is serializable because of the code gen, so we have to
-      // create the projections within each of the partition processing.
-      val groups = projections.map(ee => newProjection(ee, child.output)).toArray
-
+      val groups = projections.map(projection).toArray
       new Iterator[InternalRow] {
         private[this] var result: InternalRow = _
         private[this] var idx = -1  // -1 means the initial state
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index d5a803f8c4b24..799650a4f784f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -67,16 +67,10 @@ case class TungstenProject(projectList: Seq[NamedExpression], child: SparkPlan)
 
   override def output: Seq[Attribute] = projectList.map(_.toAttribute)
 
-  /** Rewrite the project list to use unsafe expressions as needed. */
-  protected val unsafeProjectList = projectList.map(_ transform {
-    case CreateStruct(children) => CreateStructUnsafe(children)
-    case CreateNamedStruct(children) => CreateNamedStructUnsafe(children)
-  })
-
   protected override def doExecute(): RDD[InternalRow] = {
     val numRows = longMetric("numRows")
     child.execute().mapPartitions { iter =>
-      val project = UnsafeProjection.create(unsafeProjectList, child.output)
+      val project = UnsafeProjection.create(projectList, child.output)
       iter.map { row =>
         numRows += 1
         project(row)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExpandSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExpandSuite.scala
new file mode 100644
index 0000000000000..faef76d52ae75
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExpandSuite.scala
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BoundReference, Alias, Literal}
+import org.apache.spark.sql.test.SharedSQLContext
+import org.apache.spark.sql.types.IntegerType
+
+class ExpandSuite extends SparkPlanTest with SharedSQLContext {
+  import testImplicits.localSeqToDataFrameHolder
+
+  private def testExpand(f: SparkPlan => SparkPlan): Unit = {
+    val input = (1 to 1000).map(Tuple1.apply)
+    val projections = Seq.tabulate(2) { i =>
+      Alias(BoundReference(0, IntegerType, false), "id")() :: Alias(Literal(i), "gid")() :: Nil
+    }
+    val attributes = projections.head.map(_.toAttribute)
+    checkAnswer(
+      input.toDF(),
+      plan => Expand(projections, attributes, f(plan)),
+      input.flatMap(i => Seq.tabulate(2)(j => Row(i._1, j)))
+    )
+  }
+
+  test("inheriting child row type") {
+    val exprs = AttributeReference("a", IntegerType, false)() :: Nil
+    val plan = Expand(Seq(exprs), exprs, ConvertToUnsafe(LocalTableScan(exprs, Seq.empty)))
+    assert(plan.outputsUnsafeRows, "Expand should inherits the created row type from its child.")
+  }
+
+  test("expanding UnsafeRows") {
+    testExpand(ConvertToUnsafe)
+  }
+
+  test("expanding SafeRows") {
+    testExpand(identity)
+  }
+}

From 3a652f691b220fada0286f8d0a562c5657973d4d Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Fri, 6 Nov 2015 14:47:41 -0800
Subject: [PATCH 100/518] [SPARK-11561][SQL] Rename text data source's column
 name to value.

Author: Reynold Xin <rxin@databricks.com>

Closes #9527 from rxin/SPARK-11561.
---
 .../sql/execution/datasources/text/DefaultSource.scala      | 6 ++----
 .../spark/sql/execution/datasources/text/TextSuite.scala    | 2 +-
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/DefaultSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/DefaultSource.scala
index 52c4421d7e87e..4b8b8e4e74dad 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/DefaultSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/DefaultSource.scala
@@ -30,14 +30,12 @@ import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.mapred.SparkHadoopMapRedUtil
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{UnsafeRow, GenericMutableRow}
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow
 import org.apache.spark.sql.catalyst.expressions.codegen.{UnsafeRowWriter, BufferHolder}
-import org.apache.spark.sql.columnar.MutableUnsafeRow
 import org.apache.spark.sql.{AnalysisException, Row, SQLContext}
 import org.apache.spark.sql.execution.datasources.PartitionSpec
 import org.apache.spark.sql.sources._
 import org.apache.spark.sql.types.{StringType, StructType}
-import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.SerializableConfiguration
 
 /**
@@ -78,7 +76,7 @@ private[sql] class TextRelation(
   extends HadoopFsRelation(maybePartitionSpec) {
 
   /** Data schema is always a single column, named "text". */
-  override def dataSchema: StructType = new StructType().add("text", StringType)
+  override def dataSchema: StructType = new StructType().add("value", StringType)
 
   /** This is an internal data source that outputs internal row format. */
   override val needConversion: Boolean = false
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/TextSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/TextSuite.scala
index 0a2306c06646c..914e516613f9e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/TextSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/TextSuite.scala
@@ -65,7 +65,7 @@ class TextSuite extends QueryTest with SharedSQLContext {
   /** Verifies data and schema. */
   private def verifyFrame(df: DataFrame): Unit = {
     // schema
-    assert(df.schema == new StructType().add("text", StringType))
+    assert(df.schema == new StructType().add("value", StringType))
 
     // verify content
     val data = df.collect()

From c447c9d54603890db7399fb80adc9fae40b71f64 Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Fri, 6 Nov 2015 14:51:03 -0800
Subject: [PATCH 101/518] [SPARK-11217][ML] save/load for non-meta estimators
 and transformers

This PR implements the default save/load for non-meta estimators and transformers using the JSON serialization of param values. The saved metadata includes:

* class name
* uid
* timestamp
* paramMap

The save/load interface is similar to DataFrames. We use the current active context by default, which should be sufficient for most use cases.

~~~scala
instance.save("path")
instance.write.context(sqlContext).overwrite().save("path")

Instance.load("path")
~~~

The param handling is different from the design doc. We didn't save default and user-set params separately, and when we load it back, all parameters are user-set. This does cause issues. But it also cause other issues if we modify the default params.

TODOs:

* [x] Java test
* [ ] a follow-up PR to implement default save/load for all non-meta estimators and transformers

cc jkbradley

Author: Xiangrui Meng <meng@databricks.com>

Closes #9454 from mengxr/SPARK-11217.
---
 .../apache/spark/ml/feature/Binarizer.scala   |  11 +-
 .../org/apache/spark/ml/param/params.scala    |   2 +-
 .../org/apache/spark/ml/util/ReadWrite.scala  | 220 ++++++++++++++++++
 .../ml/util/JavaDefaultReadWriteSuite.java    |  74 ++++++
 .../spark/ml/feature/BinarizerSuite.scala     |  11 +-
 .../spark/ml/util/DefaultReadWriteTest.scala  | 110 +++++++++
 .../apache/spark/ml/util/TempDirectory.scala  |  45 ++++
 7 files changed, 469 insertions(+), 4 deletions(-)
 create mode 100644 mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
 create mode 100644 mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
 create mode 100644 mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
 create mode 100644 mllib/src/test/scala/org/apache/spark/ml/util/TempDirectory.scala

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
index edad754436455..e5c25574d4b11 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
@@ -22,7 +22,7 @@ import org.apache.spark.ml.Transformer
 import org.apache.spark.ml.attribute.BinaryAttribute
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.{DoubleType, StructType}
@@ -33,7 +33,7 @@ import org.apache.spark.sql.types.{DoubleType, StructType}
  */
 @Experimental
 final class Binarizer(override val uid: String)
-  extends Transformer with HasInputCol with HasOutputCol {
+  extends Transformer with Writable with HasInputCol with HasOutputCol {
 
   def this() = this(Identifiable.randomUID("binarizer"))
 
@@ -86,4 +86,11 @@ final class Binarizer(override val uid: String)
   }
 
   override def copy(extra: ParamMap): Binarizer = defaultCopy(extra)
+
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+object Binarizer extends Readable[Binarizer] {
+
+  override def read: Reader[Binarizer] = new DefaultParamsReader[Binarizer]
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/param/params.scala b/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
index 8361406f87299..c9325709187c5 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
@@ -592,7 +592,7 @@ trait Params extends Identifiable with Serializable {
   /**
    * Sets a parameter in the embedded param map.
    */
-  protected final def set[T](param: Param[T], value: T): this.type = {
+  final def set[T](param: Param[T], value: T): this.type = {
     set(param -> value)
   }
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
new file mode 100644
index 0000000000000..ea790e0dddc7f
--- /dev/null
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -0,0 +1,220 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.util
+
+import java.io.IOException
+
+import org.apache.hadoop.fs.{FileSystem, Path}
+import org.json4s._
+import org.json4s.JsonDSL._
+import org.json4s.jackson.JsonMethods._
+
+import org.apache.spark.{Logging, SparkContext}
+import org.apache.spark.annotation.{Experimental, Since}
+import org.apache.spark.ml.param.{ParamPair, Params}
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.util.Utils
+
+/**
+ * Trait for [[Writer]] and [[Reader]].
+ */
+private[util] sealed trait BaseReadWrite {
+  private var optionSQLContext: Option[SQLContext] = None
+
+  /**
+   * Sets the SQL context to use for saving/loading.
+   */
+  @Since("1.6.0")
+  def context(sqlContext: SQLContext): this.type = {
+    optionSQLContext = Option(sqlContext)
+    this
+  }
+
+  /**
+   * Returns the user-specified SQL context or the default.
+   */
+  protected final def sqlContext: SQLContext = optionSQLContext.getOrElse {
+    SQLContext.getOrCreate(SparkContext.getOrCreate())
+  }
+}
+
+/**
+ * Abstract class for utility classes that can save ML instances.
+ */
+@Experimental
+@Since("1.6.0")
+abstract class Writer extends BaseReadWrite {
+
+  protected var shouldOverwrite: Boolean = false
+
+  /**
+   * Saves the ML instances to the input path.
+   */
+  @Since("1.6.0")
+  @throws[IOException]("If the input path already exists but overwrite is not enabled.")
+  def save(path: String): Unit
+
+  /**
+   * Overwrites if the output path already exists.
+   */
+  @Since("1.6.0")
+  def overwrite(): this.type = {
+    shouldOverwrite = true
+    this
+  }
+
+  // override for Java compatibility
+  override def context(sqlContext: SQLContext): this.type = super.context(sqlContext)
+}
+
+/**
+ * Trait for classes that provide [[Writer]].
+ */
+@Since("1.6.0")
+trait Writable {
+
+  /**
+   * Returns a [[Writer]] instance for this ML instance.
+   */
+  @Since("1.6.0")
+  def write: Writer
+
+  /**
+   * Saves this ML instance to the input path, a shortcut of `write.save(path)`.
+   */
+  @Since("1.6.0")
+  @throws[IOException]("If the input path already exists but overwrite is not enabled.")
+  def save(path: String): Unit = write.save(path)
+}
+
+/**
+ * Abstract class for utility classes that can load ML instances.
+ * @tparam T ML instance type
+ */
+@Experimental
+@Since("1.6.0")
+abstract class Reader[T] extends BaseReadWrite {
+
+  /**
+   * Loads the ML component from the input path.
+   */
+  @Since("1.6.0")
+  def load(path: String): T
+
+  // override for Java compatibility
+  override def context(sqlContext: SQLContext): this.type = super.context(sqlContext)
+}
+
+/**
+ * Trait for objects that provide [[Reader]].
+ * @tparam T ML instance type
+ */
+@Experimental
+@Since("1.6.0")
+trait Readable[T] {
+
+  /**
+   * Returns a [[Reader]] instance for this class.
+   */
+  @Since("1.6.0")
+  def read: Reader[T]
+
+  /**
+   * Reads an ML instance from the input path, a shortcut of `read.load(path)`.
+   */
+  @Since("1.6.0")
+  def load(path: String): T = read.load(path)
+}
+
+/**
+ * Default [[Writer]] implementation for transformers and estimators that contain basic
+ * (json4s-serializable) params and no data. This will not handle more complex params or types with
+ * data (e.g., models with coefficients).
+ * @param instance object to save
+ */
+private[ml] class DefaultParamsWriter(instance: Params) extends Writer with Logging {
+
+  /**
+   * Saves the ML component to the input path.
+   */
+  override def save(path: String): Unit = {
+    val sc = sqlContext.sparkContext
+
+    val hadoopConf = sc.hadoopConfiguration
+    val fs = FileSystem.get(hadoopConf)
+    val p = new Path(path)
+    if (fs.exists(p)) {
+      if (shouldOverwrite) {
+        logInfo(s"Path $path already exists. It will be overwritten.")
+        // TODO: Revert back to the original content if save is not successful.
+        fs.delete(p, true)
+      } else {
+        throw new IOException(
+          s"Path $path already exists. Please use write.overwrite().save(path) to overwrite it.")
+      }
+    }
+
+    val uid = instance.uid
+    val cls = instance.getClass.getName
+    val params = instance.extractParamMap().toSeq.asInstanceOf[Seq[ParamPair[Any]]]
+    val jsonParams = params.map { case ParamPair(p, v) =>
+      p.name -> parse(p.jsonEncode(v))
+    }.toList
+    val metadata = ("class" -> cls) ~
+      ("timestamp" -> System.currentTimeMillis()) ~
+      ("uid" -> uid) ~
+      ("paramMap" -> jsonParams)
+    val metadataPath = new Path(path, "metadata").toString
+    val metadataJson = compact(render(metadata))
+    sc.parallelize(Seq(metadataJson), 1).saveAsTextFile(metadataPath)
+  }
+}
+
+/**
+ * Default [[Reader]] implementation for transformers and estimators that contain basic
+ * (json4s-serializable) params and no data. This will not handle more complex params or types with
+ * data (e.g., models with coefficients).
+ * @tparam T ML instance type
+ */
+private[ml] class DefaultParamsReader[T] extends Reader[T] {
+
+  /**
+   * Loads the ML component from the input path.
+   */
+  override def load(path: String): T = {
+    implicit val format = DefaultFormats
+    val sc = sqlContext.sparkContext
+    val metadataPath = new Path(path, "metadata").toString
+    val metadataStr = sc.textFile(metadataPath, 1).first()
+    val metadata = parse(metadataStr)
+    val cls = Utils.classForName((metadata \ "class").extract[String])
+    val uid = (metadata \ "uid").extract[String]
+    val instance = cls.getConstructor(classOf[String]).newInstance(uid).asInstanceOf[Params]
+    (metadata \ "paramMap") match {
+      case JObject(pairs) =>
+        pairs.foreach { case (paramName, jsonValue) =>
+          val param = instance.getParam(paramName)
+          val value = param.jsonDecode(compact(render(jsonValue)))
+          instance.set(param, value)
+        }
+      case _ =>
+        throw new IllegalArgumentException(s"Cannot recognize JSON metadata: $metadataStr.")
+    }
+    instance.asInstanceOf[T]
+  }
+}
diff --git a/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java b/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
new file mode 100644
index 0000000000000..c39538014be81
--- /dev/null
+++ b/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
@@ -0,0 +1,74 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.util;
+
+import java.io.File;
+import java.io.IOException;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.util.Utils;
+
+public class JavaDefaultReadWriteSuite {
+
+  JavaSparkContext jsc = null;
+  File tempDir = null;
+
+  @Before
+  public void setUp() {
+    jsc = new JavaSparkContext("local[2]", "JavaDefaultReadWriteSuite");
+    tempDir = Utils.createTempDir(
+      System.getProperty("java.io.tmpdir"), "JavaDefaultReadWriteSuite");
+  }
+
+  @After
+  public void tearDown() {
+    if (jsc != null) {
+      jsc.stop();
+      jsc = null;
+    }
+    Utils.deleteRecursively(tempDir);
+  }
+
+  @Test
+  public void testDefaultReadWrite() throws IOException {
+    String uid = "my_params";
+    MyParams instance = new MyParams(uid);
+    instance.set(instance.intParam(), 2);
+    String outputPath = new File(tempDir, uid).getPath();
+    instance.save(outputPath);
+    try {
+      instance.save(outputPath);
+      Assert.fail(
+        "Write without overwrite enabled should fail if the output directory already exists.");
+    } catch (IOException e) {
+      // expected
+    }
+    SQLContext sqlContext = new SQLContext(jsc);
+    instance.write().context(sqlContext).overwrite().save(outputPath);
+    MyParams newInstance = MyParams.load(outputPath);
+    Assert.assertEquals("UID should match.", instance.uid(), newInstance.uid());
+    Assert.assertEquals("Params should be preserved.",
+      2, newInstance.getOrDefault(newInstance.intParam()));
+  }
+}
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
index 2086043983661..9dfa1439cc303 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
@@ -19,10 +19,11 @@ package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{DataFrame, Row}
 
-class BinarizerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class BinarizerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   @transient var data: Array[Double] = _
 
@@ -66,4 +67,12 @@ class BinarizerSuite extends SparkFunSuite with MLlibTestSparkContext {
         assert(x === y, "The feature value is not correct after binarization.")
     }
   }
+
+  test("read/write") {
+    val binarizer = new Binarizer()
+      .setInputCol("feature")
+      .setOutputCol("binarized_feature")
+      .setThreshold(0.1)
+    testDefaultReadWrite(binarizer)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
new file mode 100644
index 0000000000000..4545b0f281f5a
--- /dev/null
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
@@ -0,0 +1,110 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.util
+
+import java.io.{File, IOException}
+
+import org.scalatest.Suite
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.param._
+import org.apache.spark.mllib.util.MLlibTestSparkContext
+
+trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
+
+  /**
+   * Checks "overwrite" option and params.
+   * @param instance ML instance to test saving/loading
+   * @tparam T ML instance type
+   */
+  def testDefaultReadWrite[T <: Params with Writable](instance: T): Unit = {
+    val uid = instance.uid
+    val path = new File(tempDir, uid).getPath
+
+    instance.save(path)
+    intercept[IOException] {
+      instance.save(path)
+    }
+    instance.write.overwrite().save(path)
+    val loader = instance.getClass.getMethod("read").invoke(null).asInstanceOf[Reader[T]]
+    val newInstance = loader.load(path)
+
+    assert(newInstance.uid === instance.uid)
+    instance.params.foreach { p =>
+      if (instance.isDefined(p)) {
+        (instance.getOrDefault(p), newInstance.getOrDefault(p)) match {
+          case (Array(values), Array(newValues)) =>
+            assert(values === newValues, s"Values do not match on param ${p.name}.")
+          case (value, newValue) =>
+            assert(value === newValue, s"Values do not match on param ${p.name}.")
+        }
+      } else {
+        assert(!newInstance.isDefined(p), s"Param ${p.name} shouldn't be defined.")
+      }
+    }
+
+    val load = instance.getClass.getMethod("load", classOf[String])
+    val another = load.invoke(instance, path).asInstanceOf[T]
+    assert(another.uid === instance.uid)
+  }
+}
+
+class MyParams(override val uid: String) extends Params with Writable {
+
+  final val intParamWithDefault: IntParam = new IntParam(this, "intParamWithDefault", "doc")
+  final val intParam: IntParam = new IntParam(this, "intParam", "doc")
+  final val floatParam: FloatParam = new FloatParam(this, "floatParam", "doc")
+  final val doubleParam: DoubleParam = new DoubleParam(this, "doubleParam", "doc")
+  final val longParam: LongParam = new LongParam(this, "longParam", "doc")
+  final val stringParam: Param[String] = new Param[String](this, "stringParam", "doc")
+  final val intArrayParam: IntArrayParam = new IntArrayParam(this, "intArrayParam", "doc")
+  final val doubleArrayParam: DoubleArrayParam =
+    new DoubleArrayParam(this, "doubleArrayParam", "doc")
+  final val stringArrayParam: StringArrayParam =
+    new StringArrayParam(this, "stringArrayParam", "doc")
+
+  setDefault(intParamWithDefault -> 0)
+  set(intParam -> 1)
+  set(floatParam -> 2.0f)
+  set(doubleParam -> 3.0)
+  set(longParam -> 4L)
+  set(stringParam -> "5")
+  set(intArrayParam -> Array(6, 7))
+  set(doubleArrayParam -> Array(8.0, 9.0))
+  set(stringArrayParam -> Array("10", "11"))
+
+  override def copy(extra: ParamMap): Params = defaultCopy(extra)
+
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+object MyParams extends Readable[MyParams] {
+
+  override def read: Reader[MyParams] = new DefaultParamsReader[MyParams]
+
+  override def load(path: String): MyParams = read.load(path)
+}
+
+class DefaultReadWriteSuite extends SparkFunSuite with MLlibTestSparkContext
+  with DefaultReadWriteTest {
+
+  test("default read/write") {
+    val myParams = new MyParams("my_params")
+    testDefaultReadWrite(myParams)
+  }
+}
diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/TempDirectory.scala b/mllib/src/test/scala/org/apache/spark/ml/util/TempDirectory.scala
new file mode 100644
index 0000000000000..2742026a69c2e
--- /dev/null
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/TempDirectory.scala
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.util
+
+import java.io.File
+
+import org.scalatest.{BeforeAndAfterAll, Suite}
+
+import org.apache.spark.util.Utils
+
+/**
+ * Trait that creates a temporary directory before all tests and deletes it after all.
+ */
+trait TempDirectory extends BeforeAndAfterAll { self: Suite =>
+
+  private var _tempDir: File = _
+
+  /** Returns the temporary directory as a [[File]] instance. */
+  protected def tempDir: File = _tempDir
+
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+    _tempDir = Utils.createTempDir(this.getClass.getName)
+  }
+
+  override def afterAll(): Unit = {
+    Utils.deleteRecursively(_tempDir)
+    super.afterAll()
+  }
+}

From f6680cdc5d2912dea9768ef5c3e2cc101b06daf8 Mon Sep 17 00:00:00 2001
From: Thomas Graves <tgraves@staydecay.corp.gq1.yahoo.com>
Date: Fri, 6 Nov 2015 15:24:33 -0800
Subject: [PATCH 102/518] [SPARK-11555] spark on yarn spark-class --num-workers
 doesn't work

I tested the various options with both spark-submit and spark-class of specifying number of executors in both client and cluster mode where it applied.

--num-workers, --num-executors, spark.executor.instances, SPARK_EXECUTOR_INSTANCES, default nothing supplied

Author: Thomas Graves <tgraves@staydecay.corp.gq1.yahoo.com>

Closes #9523 from tgravescs/SPARK-11555.
---
 .../org/apache/spark/deploy/yarn/ClientArguments.scala     | 2 +-
 .../org/apache/spark/deploy/yarn/YarnSparkHadoopUtil.scala | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala
index 1165061db21e3..a9f4374357356 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala
@@ -81,7 +81,7 @@ private[spark] class ClientArguments(args: Array[String], sparkConf: SparkConf)
       .orNull
     // If dynamic allocation is enabled, start at the configured initial number of executors.
     // Default to minExecutors if no initialExecutors is set.
-    numExecutors = YarnSparkHadoopUtil.getInitialTargetExecutorNumber(sparkConf)
+    numExecutors = YarnSparkHadoopUtil.getInitialTargetExecutorNumber(sparkConf, numExecutors)
     principal = Option(principal)
       .orElse(sparkConf.getOption("spark.yarn.principal"))
       .orNull
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnSparkHadoopUtil.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnSparkHadoopUtil.scala
index 561ad79ee0228..a290ebeec9001 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnSparkHadoopUtil.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnSparkHadoopUtil.scala
@@ -392,8 +392,11 @@ object YarnSparkHadoopUtil {
   /**
    * Getting the initial target number of executors depends on whether dynamic allocation is
    * enabled.
+   * If not using dynamic allocation it gets the number of executors reqeusted by the user.
    */
-  def getInitialTargetExecutorNumber(conf: SparkConf): Int = {
+  def getInitialTargetExecutorNumber(
+      conf: SparkConf,
+      numExecutors: Int = DEFAULT_NUMBER_EXECUTORS): Int = {
     if (Utils.isDynamicAllocationEnabled(conf)) {
       val minNumExecutors = conf.getInt("spark.dynamicAllocation.minExecutors", 0)
       val initialNumExecutors =
@@ -406,7 +409,7 @@ object YarnSparkHadoopUtil {
       initialNumExecutors
     } else {
       val targetNumExecutors =
-        sys.env.get("SPARK_EXECUTOR_INSTANCES").map(_.toInt).getOrElse(DEFAULT_NUMBER_EXECUTORS)
+        sys.env.get("SPARK_EXECUTOR_INSTANCES").map(_.toInt).getOrElse(numExecutors)
       // System property can override environment variable.
       conf.getInt("spark.executor.instances", targetNumExecutors)
     }

From 7e9a9e603abce8689938bdd62d04b29299644aa4 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Fri, 6 Nov 2015 15:37:07 -0800
Subject: [PATCH 103/518] [SPARK-11269][SQL] Java API support & test cases for
 Dataset

This simply brings https://github.com/apache/spark/pull/9358 up-to-date.

Author: Wenchen Fan <wenchen@databricks.com>
Author: Reynold Xin <rxin@databricks.com>

Closes #9528 from rxin/dataset-java.
---
 .../spark/sql/catalyst/encoders/Encoder.scala | 123 +++++-
 .../sql/catalyst/expressions/objects.scala    |  21 ++
 .../scala/org/apache/spark/sql/Dataset.scala  | 126 ++++++-
 .../org/apache/spark/sql/DatasetHolder.scala  |   6 +-
 .../org/apache/spark/sql/GroupedDataset.scala |  17 +
 .../org/apache/spark/sql/SQLContext.scala     |   4 +
 .../apache/spark/sql/JavaDatasetSuite.java    | 357 ++++++++++++++++++
 .../spark/sql/DatasetPrimitiveSuite.scala     |   2 +-
 8 files changed, 644 insertions(+), 12 deletions(-)
 create mode 100644 sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala
index 329a132d3d8b2..f05e18288de2b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala
@@ -17,11 +17,11 @@
 
 package org.apache.spark.sql.catalyst.encoders
 
-
-
 import scala.reflect.ClassTag
 
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.Utils
+import org.apache.spark.sql.types.{DataType, ObjectType, StructField, StructType}
+import org.apache.spark.sql.catalyst.expressions._
 
 /**
  * Used to convert a JVM object of type `T` to and from the internal Spark SQL representation.
@@ -37,3 +37,120 @@ trait Encoder[T] extends Serializable {
   /** A ClassTag that can be used to construct and Array to contain a collection of `T`. */
   def clsTag: ClassTag[T]
 }
+
+object Encoder {
+  import scala.reflect.runtime.universe._
+
+  def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder(flat = true)
+  def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder(flat = true)
+  def SHORT: Encoder[java.lang.Short] = ExpressionEncoder(flat = true)
+  def INT: Encoder[java.lang.Integer] = ExpressionEncoder(flat = true)
+  def LONG: Encoder[java.lang.Long] = ExpressionEncoder(flat = true)
+  def FLOAT: Encoder[java.lang.Float] = ExpressionEncoder(flat = true)
+  def DOUBLE: Encoder[java.lang.Double] = ExpressionEncoder(flat = true)
+  def STRING: Encoder[java.lang.String] = ExpressionEncoder(flat = true)
+
+  def tuple[T1, T2](enc1: Encoder[T1], enc2: Encoder[T2]): Encoder[(T1, T2)] = {
+    tuple(Seq(enc1, enc2).map(_.asInstanceOf[ExpressionEncoder[_]]))
+      .asInstanceOf[ExpressionEncoder[(T1, T2)]]
+  }
+
+  def tuple[T1, T2, T3](
+      enc1: Encoder[T1],
+      enc2: Encoder[T2],
+      enc3: Encoder[T3]): Encoder[(T1, T2, T3)] = {
+    tuple(Seq(enc1, enc2, enc3).map(_.asInstanceOf[ExpressionEncoder[_]]))
+      .asInstanceOf[ExpressionEncoder[(T1, T2, T3)]]
+  }
+
+  def tuple[T1, T2, T3, T4](
+      enc1: Encoder[T1],
+      enc2: Encoder[T2],
+      enc3: Encoder[T3],
+      enc4: Encoder[T4]): Encoder[(T1, T2, T3, T4)] = {
+    tuple(Seq(enc1, enc2, enc3, enc4).map(_.asInstanceOf[ExpressionEncoder[_]]))
+      .asInstanceOf[ExpressionEncoder[(T1, T2, T3, T4)]]
+  }
+
+  def tuple[T1, T2, T3, T4, T5](
+      enc1: Encoder[T1],
+      enc2: Encoder[T2],
+      enc3: Encoder[T3],
+      enc4: Encoder[T4],
+      enc5: Encoder[T5]): Encoder[(T1, T2, T3, T4, T5)] = {
+    tuple(Seq(enc1, enc2, enc3, enc4, enc5).map(_.asInstanceOf[ExpressionEncoder[_]]))
+      .asInstanceOf[ExpressionEncoder[(T1, T2, T3, T4, T5)]]
+  }
+
+  private def tuple(encoders: Seq[ExpressionEncoder[_]]): ExpressionEncoder[_] = {
+    assert(encoders.length > 1)
+    // make sure all encoders are resolved, i.e. `Attribute` has been resolved to `BoundReference`.
+    assert(encoders.forall(_.constructExpression.find(_.isInstanceOf[Attribute]).isEmpty))
+
+    val schema = StructType(encoders.zipWithIndex.map {
+      case (e, i) => StructField(s"_${i + 1}", if (e.flat) e.schema.head.dataType else e.schema)
+    })
+
+    val cls = Utils.getContextOrSparkClassLoader.loadClass(s"scala.Tuple${encoders.size}")
+
+    val extractExpressions = encoders.map {
+      case e if e.flat => e.extractExpressions.head
+      case other => CreateStruct(other.extractExpressions)
+    }.zipWithIndex.map { case (expr, index) =>
+      expr.transformUp {
+        case BoundReference(0, t: ObjectType, _) =>
+          Invoke(
+            BoundReference(0, ObjectType(cls), true),
+            s"_${index + 1}",
+            t)
+      }
+    }
+
+    val constructExpressions = encoders.zipWithIndex.map { case (enc, index) =>
+      if (enc.flat) {
+        enc.constructExpression.transform {
+          case b: BoundReference => b.copy(ordinal = index)
+        }
+      } else {
+        enc.constructExpression.transformUp {
+          case BoundReference(ordinal, dt, _) =>
+            GetInternalRowField(BoundReference(index, enc.schema, true), ordinal, dt)
+        }
+      }
+    }
+
+    val constructExpression =
+      NewInstance(cls, constructExpressions, false, ObjectType(cls))
+
+    new ExpressionEncoder[Any](
+      schema,
+      false,
+      extractExpressions,
+      constructExpression,
+      ClassTag.apply(cls))
+  }
+
+
+  def typeTagOfTuple2[T1 : TypeTag, T2 : TypeTag]: TypeTag[(T1, T2)] = typeTag[(T1, T2)]
+
+  private def getTypeTag[T](c: Class[T]): TypeTag[T] = {
+    import scala.reflect.api
+
+    // val mirror = runtimeMirror(c.getClassLoader)
+    val mirror = rootMirror
+    val sym = mirror.staticClass(c.getName)
+    val tpe = sym.selfType
+    TypeTag(mirror, new api.TypeCreator {
+      def apply[U <: api.Universe with Singleton](m: api.Mirror[U]) =
+        if (m eq mirror) tpe.asInstanceOf[U # Type]
+        else throw new IllegalArgumentException(
+          s"Type tag defined in $mirror cannot be migrated to other mirrors.")
+    })
+  }
+
+  def forTuple2[T1, T2](c1: Class[T1], c2: Class[T2]): Encoder[(T1, T2)] = {
+    implicit val typeTag1 = getTypeTag(c1)
+    implicit val typeTag2 = getTypeTag(c2)
+    ExpressionEncoder[(T1, T2)]()
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index 81855289762c6..4f58464221b4b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -491,3 +491,24 @@ case class CreateExternalRow(children: Seq[Expression]) extends Expression {
       s"final ${classOf[Row].getName} ${ev.value} = new $rowClass($values);"
   }
 }
+
+case class GetInternalRowField(child: Expression, ordinal: Int, dataType: DataType)
+  extends UnaryExpression {
+
+  override def nullable: Boolean = true
+
+  override def eval(input: InternalRow): Any =
+    throw new UnsupportedOperationException("Only code-generated evaluation is supported")
+
+  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    val row = child.gen(ctx)
+    s"""
+      ${row.code}
+      final boolean ${ev.isNull} = ${row.isNull};
+      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+      if (!${ev.isNull}) {
+        ${ev.value} = ${ctx.getValue(row.value, dataType, ordinal.toString)};
+      }
+    """
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 4bca9c3b3fe54..fecbdac9a6004 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -17,9 +17,13 @@
 
 package org.apache.spark.sql
 
+import scala.collection.JavaConverters._
+
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAlias
+import org.apache.spark.api.java.function.{Function => JFunction, Function2 => JFunction2, _}
+
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.Inner
@@ -151,18 +155,37 @@ class Dataset[T] private[sql](
   def transform[U](t: Dataset[T] => Dataset[U]): Dataset[U] = t(this)
 
   /**
+   * (Scala-specific)
    * Returns a new [[Dataset]] that only contains elements where `func` returns `true`.
    * @since 1.6.0
    */
   def filter(func: T => Boolean): Dataset[T] = mapPartitions(_.filter(func))
 
   /**
+   * (Java-specific)
+   * Returns a new [[Dataset]] that only contains elements where `func` returns `true`.
+   * @since 1.6.0
+   */
+  def filter(func: JFunction[T, java.lang.Boolean]): Dataset[T] =
+    filter(t => func.call(t).booleanValue())
+
+  /**
+   * (Scala-specific)
    * Returns a new [[Dataset]] that contains the result of applying `func` to each element.
    * @since 1.6.0
    */
   def map[U : Encoder](func: T => U): Dataset[U] = mapPartitions(_.map(func))
 
   /**
+   * (Java-specific)
+   * Returns a new [[Dataset]] that contains the result of applying `func` to each element.
+   * @since 1.6.0
+   */
+  def map[U](func: JFunction[T, U], encoder: Encoder[U]): Dataset[U] =
+    map(t => func.call(t))(encoder)
+
+  /**
+   * (Scala-specific)
    * Returns a new [[Dataset]] that contains the result of applying `func` to each element.
    * @since 1.6.0
    */
@@ -177,30 +200,77 @@ class Dataset[T] private[sql](
         logicalPlan))
   }
 
+  /**
+   * (Java-specific)
+   * Returns a new [[Dataset]] that contains the result of applying `func` to each element.
+   * @since 1.6.0
+   */
+  def mapPartitions[U](
+      f: FlatMapFunction[java.util.Iterator[T], U],
+      encoder: Encoder[U]): Dataset[U] = {
+    val func: (Iterator[T]) => Iterator[U] = x => f.call(x.asJava).iterator().asScala
+    mapPartitions(func)(encoder)
+  }
+
+  /**
+   * (Scala-specific)
+   * Returns a new [[Dataset]] by first applying a function to all elements of this [[Dataset]],
+   * and then flattening the results.
+   * @since 1.6.0
+   */
   def flatMap[U : Encoder](func: T => TraversableOnce[U]): Dataset[U] =
     mapPartitions(_.flatMap(func))
 
+  /**
+   * (Java-specific)
+   * Returns a new [[Dataset]] by first applying a function to all elements of this [[Dataset]],
+   * and then flattening the results.
+   * @since 1.6.0
+   */
+  def flatMap[U](f: FlatMapFunction[T, U], encoder: Encoder[U]): Dataset[U] = {
+    val func: (T) => Iterable[U] = x => f.call(x).asScala
+    flatMap(func)(encoder)
+  }
+
   /* ************** *
    *  Side effects  *
    * ************** */
 
   /**
+   * (Scala-specific)
    * Runs `func` on each element of this Dataset.
    * @since 1.6.0
    */
   def foreach(func: T => Unit): Unit = rdd.foreach(func)
 
   /**
+   * (Java-specific)
+   * Runs `func` on each element of this Dataset.
+   * @since 1.6.0
+   */
+  def foreach(func: VoidFunction[T]): Unit = foreach(func.call(_))
+
+  /**
+   * (Scala-specific)
    * Runs `func` on each partition of this Dataset.
    * @since 1.6.0
    */
   def foreachPartition(func: Iterator[T] => Unit): Unit = rdd.foreachPartition(func)
 
+  /**
+   * (Java-specific)
+   * Runs `func` on each partition of this Dataset.
+   * @since 1.6.0
+   */
+  def foreachPartition(func: VoidFunction[java.util.Iterator[T]]): Unit =
+    foreachPartition(it => func.call(it.asJava))
+
   /* ************* *
    *  Aggregation  *
    * ************* */
 
   /**
+   * (Scala-specific)
    * Reduces the elements of this Dataset using the specified  binary function.  The given function
    * must be commutative and associative or the result may be non-deterministic.
    * @since 1.6.0
@@ -208,6 +278,15 @@ class Dataset[T] private[sql](
   def reduce(func: (T, T) => T): T = rdd.reduce(func)
 
   /**
+   * (Java-specific)
+   * Reduces the elements of this Dataset using the specified  binary function.  The given function
+   * must be commutative and associative or the result may be non-deterministic.
+   * @since 1.6.0
+   */
+  def reduce(func: JFunction2[T, T, T]): T = reduce(func.call(_, _))
+
+  /**
+   * (Scala-specific)
    * Aggregates the elements of each partition, and then the results for all the partitions, using a
    * given associative and commutative function and a neutral "zero value".
    *
@@ -221,6 +300,15 @@ class Dataset[T] private[sql](
   def fold(zeroValue: T)(op: (T, T) => T): T = rdd.fold(zeroValue)(op)
 
   /**
+   * (Java-specific)
+   * Aggregates the elements of each partition, and then the results for all the partitions, using a
+   * given associative and commutative function and a neutral "zero value".
+   * @since 1.6.0
+   */
+  def fold(zeroValue: T, func: JFunction2[T, T, T]): T = fold(zeroValue)(func.call(_, _))
+
+  /**
+   * (Scala-specific)
    * Returns a [[GroupedDataset]] where the data is grouped by the given key function.
    * @since 1.6.0
    */
@@ -258,6 +346,14 @@ class Dataset[T] private[sql](
       keyAttributes)
   }
 
+  /**
+   * (Java-specific)
+   * Returns a [[GroupedDataset]] where the data is grouped by the given key function.
+   * @since 1.6.0
+   */
+  def groupBy[K](f: JFunction[T, K], encoder: Encoder[K]): GroupedDataset[K, T] =
+    groupBy(f.call(_))(encoder)
+
   /* ****************** *
    *  Typed Relational  *
    * ****************** */
@@ -267,8 +363,7 @@ class Dataset[T] private[sql](
    * {{{
    *   df.select($"colA", $"colB" + 1)
    * }}}
-   * @group dfops
-   * @since 1.3.0
+   * @since 1.6.0
    */
   // Copied from Dataframe to make sure we don't have invalid overloads.
   @scala.annotation.varargs
@@ -279,7 +374,7 @@ class Dataset[T] private[sql](
    *
    * {{{
    *   val ds = Seq(1, 2, 3).toDS()
-   *   val newDS = ds.select(e[Int]("value + 1"))
+   *   val newDS = ds.select(expr("value + 1").as[Int])
    * }}}
    * @since 1.6.0
    */
@@ -405,6 +500,8 @@ class Dataset[T] private[sql](
    * This type of join can be useful both for preserving type-safety with the original object
    * types as well as working with relational data where either side of the join has column
    * names in common.
+   *
+   * @since 1.6.0
    */
   def joinWith[U](other: Dataset[U], condition: Column): Dataset[(T, U)] = {
     val left = this.logicalPlan
@@ -438,12 +535,31 @@ class Dataset[T] private[sql](
    *  Gather to Driver Actions  *
    * ************************** */
 
-  /** Returns the first element in this [[Dataset]]. */
+  /**
+   * Returns the first element in this [[Dataset]].
+   * @since 1.6.0
+   */
   def first(): T = rdd.first()
 
-  /** Collects the elements to an Array. */
+  /**
+   * Collects the elements to an Array.
+   * @since 1.6.0
+   */
   def collect(): Array[T] = rdd.collect()
 
+  /**
+   * (Java-specific)
+   * Collects the elements to a Java list.
+   *
+   * Due to the incompatibility problem between Scala and Java, the return type of [[collect()]] at
+   * Java side is `java.lang.Object`, which is not easy to use.  Java user can use this method
+   * instead and keep the generic type for result.
+   *
+   * @since 1.6.0
+   */
+  def collectAsList(): java.util.List[T] =
+    rdd.collect().toSeq.asJava
+
   /** Returns the first `num` elements of this [[Dataset]] as an Array. */
   def take(num: Int): Array[T] = rdd.take(num)
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala b/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala
index 45f0098b92887..08097e9f02084 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DatasetHolder.scala
@@ -27,9 +27,9 @@ package org.apache.spark.sql
  *
  * @since 1.6.0
  */
-case class DatasetHolder[T] private[sql](private val df: Dataset[T]) {
+case class DatasetHolder[T] private[sql](private val ds: Dataset[T]) {
 
   // This is declared with parentheses to prevent the Scala compiler from treating
-  // `rdd.toDF("1")` as invoking this toDF and then apply on the returned DataFrame.
-  def toDS(): Dataset[T] = df
+  // `rdd.toDS("1")` as invoking this toDS and then apply on the returned Dataset.
+  def toDS(): Dataset[T] = ds
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index b8fc373dffcf5..b2803d5a9a1e3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -17,7 +17,11 @@
 
 package org.apache.spark.sql
 
+import java.util.{Iterator => JIterator}
+import scala.collection.JavaConverters._
+
 import org.apache.spark.annotation.Experimental
+import org.apache.spark.api.java.function.{Function2 => JFunction2, Function3 => JFunction3, _}
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttribute}
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor, Encoder}
 import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression, Alias, Attribute}
@@ -104,6 +108,12 @@ class GroupedDataset[K, T] private[sql](
       MapGroups(f, groupingAttributes, logicalPlan))
   }
 
+  def mapGroups[U](
+      f: JFunction2[K, JIterator[T], JIterator[U]],
+      encoder: Encoder[U]): Dataset[U] = {
+    mapGroups((key, data) => f.call(key, data.asJava).asScala)(encoder)
+  }
+
   // To ensure valid overloading.
   protected def agg(expr: Column, exprs: Column*): DataFrame =
     groupedData.agg(expr, exprs: _*)
@@ -196,4 +206,11 @@ class GroupedDataset[K, T] private[sql](
         this.logicalPlan,
         other.logicalPlan))
   }
+
+  def cogroup[U, R](
+      other: GroupedDataset[K, U],
+      f: JFunction3[K, JIterator[T], JIterator[U], JIterator[R]],
+      encoder: Encoder[R]): Dataset[R] = {
+    cogroup(other)((key, left, right) => f.call(key, left.asJava, right.asJava).asScala)(encoder)
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index 5ad3871093fc8..5598731af5fcc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -508,6 +508,10 @@ class SQLContext private[sql](
     new Dataset[T](this, plan)
   }
 
+  def createDataset[T : Encoder](data: java.util.List[T]): Dataset[T] = {
+    createDataset(data.asScala)
+  }
+
   /**
    * Creates a DataFrame from an RDD[Row]. User can specify whether the input rows should be
    * converted to Catalyst rows.
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
new file mode 100644
index 0000000000000..a9493d576d179
--- /dev/null
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -0,0 +1,357 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package test.org.apache.spark.sql;
+
+import java.io.Serializable;
+import java.util.*;
+
+import scala.Tuple2;
+import scala.Tuple3;
+import scala.Tuple4;
+import scala.Tuple5;
+import org.junit.*;
+
+import org.apache.spark.Accumulator;
+import org.apache.spark.SparkContext;
+import org.apache.spark.api.java.function.*;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.sql.catalyst.encoders.Encoder;
+import org.apache.spark.sql.catalyst.encoders.Encoder$;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.GroupedDataset;
+import org.apache.spark.sql.test.TestSQLContext;
+
+import static org.apache.spark.sql.functions.*;
+
+public class JavaDatasetSuite implements Serializable {
+  private transient JavaSparkContext jsc;
+  private transient TestSQLContext context;
+  private transient Encoder$ e = Encoder$.MODULE$;
+
+  @Before
+  public void setUp() {
+    // Trigger static initializer of TestData
+    SparkContext sc = new SparkContext("local[*]", "testing");
+    jsc = new JavaSparkContext(sc);
+    context = new TestSQLContext(sc);
+    context.loadTestData();
+  }
+
+  @After
+  public void tearDown() {
+    context.sparkContext().stop();
+    context = null;
+    jsc = null;
+  }
+
+  private <T1, T2> Tuple2<T1, T2> tuple2(T1 t1, T2 t2) {
+    return new Tuple2<T1, T2>(t1, t2);
+  }
+
+  @Test
+  public void testCollect() {
+    List<String> data = Arrays.asList("hello", "world");
+    Dataset<String> ds = context.createDataset(data, e.STRING());
+    String[] collected = (String[]) ds.collect();
+    Assert.assertEquals(Arrays.asList("hello", "world"), Arrays.asList(collected));
+  }
+
+  @Test
+  public void testCommonOperation() {
+    List<String> data = Arrays.asList("hello", "world");
+    Dataset<String> ds = context.createDataset(data, e.STRING());
+    Assert.assertEquals("hello", ds.first());
+
+    Dataset<String> filtered = ds.filter(new Function<String, Boolean>() {
+      @Override
+      public Boolean call(String v) throws Exception {
+        return v.startsWith("h");
+      }
+    });
+    Assert.assertEquals(Arrays.asList("hello"), filtered.collectAsList());
+
+
+    Dataset<Integer> mapped = ds.map(new Function<String, Integer>() {
+      @Override
+      public Integer call(String v) throws Exception {
+        return v.length();
+      }
+    }, e.INT());
+    Assert.assertEquals(Arrays.asList(5, 5), mapped.collectAsList());
+
+    Dataset<String> parMapped = ds.mapPartitions(new FlatMapFunction<Iterator<String>, String>() {
+      @Override
+      public Iterable<String> call(Iterator<String> it) throws Exception {
+        List<String> ls = new LinkedList<String>();
+        while (it.hasNext()) {
+          ls.add(it.next().toUpperCase());
+        }
+        return ls;
+      }
+    }, e.STRING());
+    Assert.assertEquals(Arrays.asList("HELLO", "WORLD"), parMapped.collectAsList());
+
+    Dataset<String> flatMapped = ds.flatMap(new FlatMapFunction<String, String>() {
+      @Override
+      public Iterable<String> call(String s) throws Exception {
+        List<String> ls = new LinkedList<String>();
+        for (char c : s.toCharArray()) {
+          ls.add(String.valueOf(c));
+        }
+        return ls;
+      }
+    }, e.STRING());
+    Assert.assertEquals(
+      Arrays.asList("h", "e", "l", "l", "o", "w", "o", "r", "l", "d"),
+      flatMapped.collectAsList());
+  }
+
+  @Test
+  public void testForeach() {
+    final Accumulator<Integer> accum = jsc.accumulator(0);
+    List<String> data = Arrays.asList("a", "b", "c");
+    Dataset<String> ds = context.createDataset(data, e.STRING());
+
+    ds.foreach(new VoidFunction<String>() {
+      @Override
+      public void call(String s) throws Exception {
+        accum.add(1);
+      }
+    });
+    Assert.assertEquals(3, accum.value().intValue());
+  }
+
+  @Test
+  public void testReduce() {
+    List<Integer> data = Arrays.asList(1, 2, 3);
+    Dataset<Integer> ds = context.createDataset(data, e.INT());
+
+    int reduced = ds.reduce(new Function2<Integer, Integer, Integer>() {
+      @Override
+      public Integer call(Integer v1, Integer v2) throws Exception {
+        return v1 + v2;
+      }
+    });
+    Assert.assertEquals(6, reduced);
+
+    int folded = ds.fold(1, new Function2<Integer, Integer, Integer>() {
+      @Override
+      public Integer call(Integer v1, Integer v2) throws Exception {
+        return v1 * v2;
+      }
+    });
+    Assert.assertEquals(6, folded);
+  }
+
+  @Test
+  public void testGroupBy() {
+    List<String> data = Arrays.asList("a", "foo", "bar");
+    Dataset<String> ds = context.createDataset(data, e.STRING());
+    GroupedDataset<Integer, String> grouped = ds.groupBy(new Function<String, Integer>() {
+      @Override
+      public Integer call(String v) throws Exception {
+        return v.length();
+      }
+    }, e.INT());
+
+    Dataset<String> mapped = grouped.mapGroups(
+      new Function2<Integer, Iterator<String>, Iterator<String>>() {
+        @Override
+        public Iterator<String> call(Integer key, Iterator<String> data) throws Exception {
+          StringBuilder sb = new StringBuilder(key.toString());
+          while (data.hasNext()) {
+            sb.append(data.next());
+          }
+          return Collections.singletonList(sb.toString()).iterator();
+        }
+      },
+      e.STRING());
+
+    Assert.assertEquals(Arrays.asList("1a", "3foobar"), mapped.collectAsList());
+
+    List<Integer> data2 = Arrays.asList(2, 6, 10);
+    Dataset<Integer> ds2 = context.createDataset(data2, e.INT());
+    GroupedDataset<Integer, Integer> grouped2 = ds2.groupBy(new Function<Integer, Integer>() {
+      @Override
+      public Integer call(Integer v) throws Exception {
+        return v / 2;
+      }
+    }, e.INT());
+
+    Dataset<String> cogrouped = grouped.cogroup(
+      grouped2,
+      new Function3<Integer, Iterator<String>, Iterator<Integer>, Iterator<String>>() {
+        @Override
+        public Iterator<String> call(
+            Integer key,
+            Iterator<String> left,
+            Iterator<Integer> right) throws Exception {
+          StringBuilder sb = new StringBuilder(key.toString());
+          while (left.hasNext()) {
+            sb.append(left.next());
+          }
+          sb.append("#");
+          while (right.hasNext()) {
+            sb.append(right.next());
+          }
+          return Collections.singletonList(sb.toString()).iterator();
+        }
+      },
+      e.STRING());
+
+    Assert.assertEquals(Arrays.asList("1a#2", "3foobar#6", "5#10"), cogrouped.collectAsList());
+  }
+
+  @Test
+  public void testGroupByColumn() {
+    List<String> data = Arrays.asList("a", "foo", "bar");
+    Dataset<String> ds = context.createDataset(data, e.STRING());
+    GroupedDataset<Integer, String> grouped = ds.groupBy(length(col("value"))).asKey(e.INT());
+
+    Dataset<String> mapped = grouped.mapGroups(
+      new Function2<Integer, Iterator<String>, Iterator<String>>() {
+        @Override
+        public Iterator<String> call(Integer key, Iterator<String> data) throws Exception {
+          StringBuilder sb = new StringBuilder(key.toString());
+          while (data.hasNext()) {
+            sb.append(data.next());
+          }
+          return Collections.singletonList(sb.toString()).iterator();
+        }
+      },
+      e.STRING());
+
+    Assert.assertEquals(Arrays.asList("1a", "3foobar"), mapped.collectAsList());
+  }
+
+  @Test
+  public void testSelect() {
+    List<Integer> data = Arrays.asList(2, 6);
+    Dataset<Integer> ds = context.createDataset(data, e.INT());
+
+    Dataset<Tuple2<Integer, String>> selected = ds.select(
+      expr("value + 1").as(e.INT()),
+      col("value").cast("string").as(e.STRING()));
+
+    Assert.assertEquals(
+      Arrays.asList(tuple2(3, "2"), tuple2(7, "6")),
+      selected.collectAsList());
+  }
+
+  @Test
+  public void testSetOperation() {
+    List<String> data = Arrays.asList("abc", "abc", "xyz");
+    Dataset<String> ds = context.createDataset(data, e.STRING());
+
+    Assert.assertEquals(
+      Arrays.asList("abc", "xyz"),
+      sort(ds.distinct().collectAsList().toArray(new String[0])));
+
+    List<String> data2 = Arrays.asList("xyz", "foo", "foo");
+    Dataset<String> ds2 = context.createDataset(data2, e.STRING());
+
+    Dataset<String> intersected = ds.intersect(ds2);
+    Assert.assertEquals(Arrays.asList("xyz"), intersected.collectAsList());
+
+    Dataset<String> unioned = ds.union(ds2);
+    Assert.assertEquals(
+      Arrays.asList("abc", "abc", "foo", "foo", "xyz", "xyz"),
+      sort(unioned.collectAsList().toArray(new String[0])));
+
+    Dataset<String> subtracted = ds.subtract(ds2);
+    Assert.assertEquals(Arrays.asList("abc", "abc"), subtracted.collectAsList());
+  }
+
+  private <T extends Comparable<T>> List<T> sort(T[] data) {
+    Arrays.sort(data);
+    return Arrays.asList(data);
+  }
+
+  @Test
+  public void testJoin() {
+    List<Integer> data = Arrays.asList(1, 2, 3);
+    Dataset<Integer> ds = context.createDataset(data, e.INT()).as("a");
+    List<Integer> data2 = Arrays.asList(2, 3, 4);
+    Dataset<Integer> ds2 = context.createDataset(data2, e.INT()).as("b");
+
+    Dataset<Tuple2<Integer, Integer>> joined =
+      ds.joinWith(ds2, col("a.value").equalTo(col("b.value")));
+    Assert.assertEquals(
+      Arrays.asList(tuple2(2, 2), tuple2(3, 3)),
+      joined.collectAsList());
+  }
+
+  @Test
+  public void testTupleEncoder() {
+    Encoder<Tuple2<Integer, String>> encoder2 = e.tuple(e.INT(), e.STRING());
+    List<Tuple2<Integer, String>> data2 = Arrays.asList(tuple2(1, "a"), tuple2(2, "b"));
+    Dataset<Tuple2<Integer, String>> ds2 = context.createDataset(data2, encoder2);
+    Assert.assertEquals(data2, ds2.collectAsList());
+
+    Encoder<Tuple3<Integer, Long, String>> encoder3 = e.tuple(e.INT(), e.LONG(), e.STRING());
+    List<Tuple3<Integer, Long, String>> data3 =
+      Arrays.asList(new Tuple3<Integer, Long, String>(1, 2L, "a"));
+    Dataset<Tuple3<Integer, Long, String>> ds3 = context.createDataset(data3, encoder3);
+    Assert.assertEquals(data3, ds3.collectAsList());
+
+    Encoder<Tuple4<Integer, String, Long, String>> encoder4 =
+      e.tuple(e.INT(), e.STRING(), e.LONG(), e.STRING());
+    List<Tuple4<Integer, String, Long, String>> data4 =
+      Arrays.asList(new Tuple4<Integer, String, Long, String>(1, "b", 2L, "a"));
+    Dataset<Tuple4<Integer, String, Long, String>> ds4 = context.createDataset(data4, encoder4);
+    Assert.assertEquals(data4, ds4.collectAsList());
+
+    Encoder<Tuple5<Integer, String, Long, String, Boolean>> encoder5 =
+      e.tuple(e.INT(), e.STRING(), e.LONG(), e.STRING(), e.BOOLEAN());
+    List<Tuple5<Integer, String, Long, String, Boolean>> data5 =
+      Arrays.asList(new Tuple5<Integer, String, Long, String, Boolean>(1, "b", 2L, "a", true));
+    Dataset<Tuple5<Integer, String, Long, String, Boolean>> ds5 =
+      context.createDataset(data5, encoder5);
+    Assert.assertEquals(data5, ds5.collectAsList());
+  }
+
+  @Test
+  public void testNestedTupleEncoder() {
+    // test ((int, string), string)
+    Encoder<Tuple2<Tuple2<Integer, String>, String>> encoder =
+      e.tuple(e.tuple(e.INT(), e.STRING()), e.STRING());
+    List<Tuple2<Tuple2<Integer, String>, String>> data =
+      Arrays.asList(tuple2(tuple2(1, "a"), "a"), tuple2(tuple2(2, "b"), "b"));
+    Dataset<Tuple2<Tuple2<Integer, String>, String>> ds = context.createDataset(data, encoder);
+    Assert.assertEquals(data, ds.collectAsList());
+
+    // test (int, (string, string, long))
+    Encoder<Tuple2<Integer, Tuple3<String, String, Long>>> encoder2 =
+      e.tuple(e.INT(), e.tuple(e.STRING(), e.STRING(), e.LONG()));
+    List<Tuple2<Integer, Tuple3<String, String, Long>>> data2 =
+      Arrays.asList(tuple2(1, new Tuple3<String, String, Long>("a", "b", 3L)));
+    Dataset<Tuple2<Integer, Tuple3<String, String, Long>>> ds2 =
+      context.createDataset(data2, encoder2);
+    Assert.assertEquals(data2, ds2.collectAsList());
+
+    // test (int, ((string, long), string))
+    Encoder<Tuple2<Integer, Tuple2<Tuple2<String, Long>, String>>> encoder3 =
+      e.tuple(e.INT(), e.tuple(e.tuple(e.STRING(), e.LONG()), e.STRING()));
+    List<Tuple2<Integer, Tuple2<Tuple2<String, Long>, String>>> data3 =
+      Arrays.asList(tuple2(1, tuple2(tuple2("a", 2L), "b")));
+    Dataset<Tuple2<Integer, Tuple2<Tuple2<String, Long>, String>>> ds3 =
+      context.createDataset(data3, encoder3);
+    Assert.assertEquals(data3, ds3.collectAsList());
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
index 32443557fb8e0..e3b0346f857d3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
@@ -59,7 +59,7 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
   test("foreach") {
     val ds = Seq(1, 2, 3).toDS()
     val acc = sparkContext.accumulator(0)
-    ds.foreach(acc +=)
+    ds.foreach(acc += _)
     assert(acc.value == 6)
   }
 

From 1ab72b08601a1c8a674bdd3fab84d9804899b2c7 Mon Sep 17 00:00:00 2001
From: Nong Li <nong@databricks.com>
Date: Fri, 6 Nov 2015 15:48:20 -0800
Subject: [PATCH 104/518] =?UTF-8?q?[SPARK-11410]=20[PYSPARK]=20Add=20pytho?=
 =?UTF-8?q?n=20bindings=20for=20repartition=20and=20sortW=E2=80=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…ithinPartitions.

Author: Nong Li <nong@databricks.com>

Closes #9504 from nongli/spark-11410.
---
 python/pyspark/sql/dataframe.py | 117 +++++++++++++++++++++++++++-----
 1 file changed, 101 insertions(+), 16 deletions(-)

diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index 765a4511b64bc..b97c94dad834a 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -422,6 +422,67 @@ def repartition(self, numPartitions):
         """
         return DataFrame(self._jdf.repartition(numPartitions), self.sql_ctx)
 
+    @since(1.3)
+    def repartition(self, numPartitions, *cols):
+        """
+        Returns a new :class:`DataFrame` partitioned by the given partitioning expressions. The
+        resulting DataFrame is hash partitioned.
+
+        ``numPartitions`` can be an int to specify the target number of partitions or a Column.
+        If it is a Column, it will be used as the first partitioning column. If not specified,
+        the default number of partitions is used.
+
+        .. versionchanged:: 1.6
+           Added optional arguments to specify the partitioning columns. Also made numPartitions
+           optional if partitioning columns are specified.
+
+        >>> df.repartition(10).rdd.getNumPartitions()
+        10
+        >>> data = df.unionAll(df).repartition("age")
+        >>> data.show()
+        +---+-----+
+        |age| name|
+        +---+-----+
+        |  2|Alice|
+        |  2|Alice|
+        |  5|  Bob|
+        |  5|  Bob|
+        +---+-----+
+        >>> data = data.repartition(7, "age")
+        >>> data.show()
+        +---+-----+
+        |age| name|
+        +---+-----+
+        |  5|  Bob|
+        |  5|  Bob|
+        |  2|Alice|
+        |  2|Alice|
+        +---+-----+
+        >>> data.rdd.getNumPartitions()
+        7
+        >>> data = data.repartition("name", "age")
+        >>> data.show()
+        +---+-----+
+        |age| name|
+        +---+-----+
+        |  5|  Bob|
+        |  5|  Bob|
+        |  2|Alice|
+        |  2|Alice|
+        +---+-----+
+        """
+        if isinstance(numPartitions, int):
+            if len(cols) == 0:
+                return DataFrame(self._jdf.repartition(numPartitions), self.sql_ctx)
+            else:
+                return DataFrame(
+                    self._jdf.repartition(numPartitions, self._jcols(*cols)), self.sql_ctx)
+        elif isinstance(numPartitions, (basestring, Column)):
+            cols = (numPartitions, ) + cols
+            return DataFrame(self._jdf.repartition(self._jcols(*cols)), self.sql_ctx)
+        else:
+            raise TypeError("numPartitions should be an int or Column")
+
     @since(1.3)
     def distinct(self):
         """Returns a new :class:`DataFrame` containing the distinct rows in this :class:`DataFrame`.
@@ -589,6 +650,26 @@ def join(self, other, on=None, how=None):
                 jdf = self._jdf.join(other._jdf, on._jc, how)
         return DataFrame(jdf, self.sql_ctx)
 
+    @since(1.6)
+    def sortWithinPartitions(self, *cols, **kwargs):
+        """Returns a new :class:`DataFrame` with each partition sorted by the specified column(s).
+
+        :param cols: list of :class:`Column` or column names to sort by.
+        :param ascending: boolean or list of boolean (default True).
+            Sort ascending vs. descending. Specify list for multiple sort orders.
+            If a list is specified, length of the list must equal length of the `cols`.
+
+        >>> df.sortWithinPartitions("age", ascending=False).show()
+        +---+-----+
+        |age| name|
+        +---+-----+
+        |  2|Alice|
+        |  5|  Bob|
+        +---+-----+
+        """
+        jdf = self._jdf.sortWithinPartitions(self._sort_cols(cols, kwargs))
+        return DataFrame(jdf, self.sql_ctx)
+
     @ignore_unicode_prefix
     @since(1.3)
     def sort(self, *cols, **kwargs):
@@ -613,22 +694,7 @@ def sort(self, *cols, **kwargs):
         >>> df.orderBy(["age", "name"], ascending=[0, 1]).collect()
         [Row(age=5, name=u'Bob'), Row(age=2, name=u'Alice')]
         """
-        if not cols:
-            raise ValueError("should sort by at least one column")
-        if len(cols) == 1 and isinstance(cols[0], list):
-            cols = cols[0]
-        jcols = [_to_java_column(c) for c in cols]
-        ascending = kwargs.get('ascending', True)
-        if isinstance(ascending, (bool, int)):
-            if not ascending:
-                jcols = [jc.desc() for jc in jcols]
-        elif isinstance(ascending, list):
-            jcols = [jc if asc else jc.desc()
-                     for asc, jc in zip(ascending, jcols)]
-        else:
-            raise TypeError("ascending can only be boolean or list, but got %s" % type(ascending))
-
-        jdf = self._jdf.sort(self._jseq(jcols))
+        jdf = self._jdf.sort(self._sort_cols(cols, kwargs))
         return DataFrame(jdf, self.sql_ctx)
 
     orderBy = sort
@@ -650,6 +716,25 @@ def _jcols(self, *cols):
             cols = cols[0]
         return self._jseq(cols, _to_java_column)
 
+    def _sort_cols(self, cols, kwargs):
+        """ Return a JVM Seq of Columns that describes the sort order
+        """
+        if not cols:
+            raise ValueError("should sort by at least one column")
+        if len(cols) == 1 and isinstance(cols[0], list):
+            cols = cols[0]
+        jcols = [_to_java_column(c) for c in cols]
+        ascending = kwargs.get('ascending', True)
+        if isinstance(ascending, (bool, int)):
+            if not ascending:
+                jcols = [jc.desc() for jc in jcols]
+        elif isinstance(ascending, list):
+            jcols = [jc if asc else jc.desc()
+                     for asc, jc in zip(ascending, jcols)]
+        else:
+            raise TypeError("ascending can only be boolean or list, but got %s" % type(ascending))
+        return self._jseq(jcols)
+
     @since("1.3.1")
     def describe(self, *cols):
         """Computes statistics for numeric columns.

From 6d0ead322e72303c6444c6ac641378a4690cde96 Mon Sep 17 00:00:00 2001
From: Herman van Hovell <hvanhovell@questtec.nl>
Date: Fri, 6 Nov 2015 16:04:20 -0800
Subject: [PATCH 105/518] [SPARK-9241][SQL] Supporting multiple DISTINCT
 columns (2) - Rewriting Rule

The second PR for SPARK-9241, this adds support for multiple distinct columns to the new aggregation code path.

This PR solves the multiple DISTINCT column problem by rewriting these Aggregates into an Expand-Aggregate-Aggregate combination. See the [JIRA ticket](https://issues.apache.org/jira/browse/SPARK-9241) for some information on this. The advantages over the - competing - [first PR](https://github.com/apache/spark/pull/9280) are:
- This can use the faster TungstenAggregate code path.
- It is impossible to OOM due to an ```OpenHashSet``` allocating to much memory. However, this will multiply the number of input rows by the number of distinct clauses (plus one), and puts a lot more memory pressure on the aggregation code path itself.

The location of this Rule is a bit funny, and should probably change when the old aggregation path is changed.

cc yhuai - Could you also tell me where to add tests for this?

Author: Herman van Hovell <hvanhovell@questtec.nl>

Closes #9406 from hvanhovell/SPARK-9241-rewriter.
---
 .../expressions/aggregate/Count.scala         |   2 +
 .../expressions/aggregate/Utils.scala         | 186 +++++++++++++++++-
 .../expressions/aggregate/interfaces.scala    |   6 +
 .../sql/catalyst/optimizer/Optimizer.scala    |   6 +-
 .../plans/logical/basicOperators.scala        |  80 ++++----
 .../spark/sql/execution/SparkStrategies.scala |   2 +-
 6 files changed, 238 insertions(+), 44 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
index 54df96cd2446a..ec0c8b483a909 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
@@ -49,4 +49,6 @@ case class Count(child: Expression) extends DeclarativeAggregate {
   )
 
   override val evaluateExpression = Cast(count, LongType)
+
+  override def defaultResult: Option[Literal] = Option(Literal(0L))
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
index 644c6211d5f31..39010c3be6d4e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
@@ -20,8 +20,9 @@ package org.apache.spark.sql.catalyst.expressions.aggregate
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst._
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
-import org.apache.spark.sql.types.{StructType, MapType, ArrayType}
+import org.apache.spark.sql.catalyst.plans.logical.{Expand, Aggregate, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.types.{IntegerType, StructType, MapType, ArrayType}
 
 /**
  * Utility functions used by the query planner to convert our plan to new aggregation code path.
@@ -41,7 +42,7 @@ object Utils {
 
   private def doConvert(plan: LogicalPlan): Option[Aggregate] = plan match {
     case p: Aggregate if supportsGroupingKeySchema(p) =>
-      val converted = p.transformExpressionsDown {
+      val converted = MultipleDistinctRewriter.rewrite(p.transformExpressionsDown {
         case expressions.Average(child) =>
           aggregate.AggregateExpression2(
             aggregateFunction = aggregate.Average(child),
@@ -144,7 +145,8 @@ object Utils {
             aggregateFunction = aggregate.VarianceSamp(child),
             mode = aggregate.Complete,
             isDistinct = false)
-      }
+      })
+
       // Check if there is any expressions.AggregateExpression1 left.
       // If so, we cannot convert this plan.
       val hasAggregateExpression1 = converted.aggregateExpressions.exists { expr =>
@@ -156,6 +158,7 @@ object Utils {
       }
 
       // Check if there are multiple distinct columns.
+      // TODO remove this.
       val aggregateExpressions = converted.aggregateExpressions.flatMap { expr =>
         expr.collect {
           case agg: AggregateExpression2 => agg
@@ -213,3 +216,178 @@ object Utils {
     case other => None
   }
 }
+
+/**
+ * This rule rewrites an aggregate query with multiple distinct clauses into an expanded double
+ * aggregation in which the regular aggregation expressions and every distinct clause is aggregated
+ * in a separate group. The results are then combined in a second aggregate.
+ *
+ * TODO Expression cannocalization
+ * TODO Eliminate foldable expressions from distinct clauses.
+ * TODO This eliminates all distinct expressions. We could safely pass one to the aggregate
+ *      operator. Perhaps this is a good thing? It is much simpler to plan later on...
+ */
+object MultipleDistinctRewriter extends Rule[LogicalPlan] {
+
+  def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
+    case a: Aggregate => rewrite(a)
+    case p => p
+  }
+
+  def rewrite(a: Aggregate): Aggregate = {
+
+    // Collect all aggregate expressions.
+    val aggExpressions = a.aggregateExpressions.flatMap { e =>
+      e.collect {
+        case ae: AggregateExpression2 => ae
+      }
+    }
+
+    // Extract distinct aggregate expressions.
+    val distinctAggGroups = aggExpressions
+      .filter(_.isDistinct)
+      .groupBy(_.aggregateFunction.children.toSet)
+
+    // Only continue to rewrite if there is more than one distinct group.
+    if (distinctAggGroups.size > 1) {
+      // Create the attributes for the grouping id and the group by clause.
+      val gid = new AttributeReference("gid", IntegerType, false)()
+      val groupByMap = a.groupingExpressions.collect {
+        case ne: NamedExpression => ne -> ne.toAttribute
+        case e => e -> new AttributeReference(e.prettyName, e.dataType, e.nullable)()
+      }
+      val groupByAttrs = groupByMap.map(_._2)
+
+      // Functions used to modify aggregate functions and their inputs.
+      def evalWithinGroup(id: Literal, e: Expression) = If(EqualTo(gid, id), e, nullify(e))
+      def patchAggregateFunctionChildren(
+          af: AggregateFunction2,
+          id: Literal,
+          attrs: Map[Expression, Expression]): AggregateFunction2 = {
+        af.withNewChildren(af.children.map { case afc =>
+          evalWithinGroup(id, attrs(afc))
+        }).asInstanceOf[AggregateFunction2]
+      }
+
+      // Setup unique distinct aggregate children.
+      val distinctAggChildren = distinctAggGroups.keySet.flatten.toSeq
+      val distinctAggChildAttrMap = distinctAggChildren.map(expressionAttributePair).toMap
+      val distinctAggChildAttrs = distinctAggChildAttrMap.values.toSeq
+
+      // Setup expand & aggregate operators for distinct aggregate expressions.
+      val distinctAggOperatorMap = distinctAggGroups.toSeq.zipWithIndex.map {
+        case ((group, expressions), i) =>
+          val id = Literal(i + 1)
+
+          // Expand projection
+          val projection = distinctAggChildren.map {
+            case e if group.contains(e) => e
+            case e => nullify(e)
+          } :+ id
+
+          // Final aggregate
+          val operators = expressions.map { e =>
+            val af = e.aggregateFunction
+            val naf = patchAggregateFunctionChildren(af, id, distinctAggChildAttrMap)
+            (e, e.copy(aggregateFunction = naf, isDistinct = false))
+          }
+
+          (projection, operators)
+      }
+
+      // Setup expand for the 'regular' aggregate expressions.
+      val regularAggExprs = aggExpressions.filter(!_.isDistinct)
+      val regularAggChildren = regularAggExprs.flatMap(_.aggregateFunction.children).distinct
+      val regularAggChildAttrMap = regularAggChildren.map(expressionAttributePair).toMap
+
+      // Setup aggregates for 'regular' aggregate expressions.
+      val regularGroupId = Literal(0)
+      val regularAggOperatorMap = regularAggExprs.map { e =>
+        // Perform the actual aggregation in the initial aggregate.
+        val af = patchAggregateFunctionChildren(
+          e.aggregateFunction,
+          regularGroupId,
+          regularAggChildAttrMap)
+        val a = Alias(e.copy(aggregateFunction = af), e.toString)()
+
+        // Get the result of the first aggregate in the last aggregate.
+        val b = AggregateExpression2(
+          aggregate.First(evalWithinGroup(regularGroupId, a.toAttribute), Literal(true)),
+          mode = Complete,
+          isDistinct = false)
+
+        // Some aggregate functions (COUNT) have the special property that they can return a
+        // non-null result without any input. We need to make sure we return a result in this case.
+        val c = af.defaultResult match {
+          case Some(lit) => Coalesce(Seq(b, lit))
+          case None => b
+        }
+
+        (e, a, c)
+      }
+
+      // Construct the regular aggregate input projection only if we need one.
+      val regularAggProjection = if (regularAggExprs.nonEmpty) {
+        Seq(a.groupingExpressions ++
+          distinctAggChildren.map(nullify) ++
+          Seq(regularGroupId) ++
+          regularAggChildren)
+      } else {
+        Seq.empty[Seq[Expression]]
+      }
+
+      // Construct the distinct aggregate input projections.
+      val regularAggNulls = regularAggChildren.map(nullify)
+      val distinctAggProjections = distinctAggOperatorMap.map {
+        case (projection, _) =>
+          a.groupingExpressions ++
+            projection ++
+            regularAggNulls
+      }
+
+      // Construct the expand operator.
+      val expand = Expand(
+        regularAggProjection ++ distinctAggProjections,
+        groupByAttrs ++ distinctAggChildAttrs ++ Seq(gid) ++ regularAggChildAttrMap.values.toSeq,
+        a.child)
+
+      // Construct the first aggregate operator. This de-duplicates the all the children of
+      // distinct operators, and applies the regular aggregate operators.
+      val firstAggregateGroupBy = groupByAttrs ++ distinctAggChildAttrs :+ gid
+      val firstAggregate = Aggregate(
+        firstAggregateGroupBy,
+        firstAggregateGroupBy ++ regularAggOperatorMap.map(_._2),
+        expand)
+
+      // Construct the second aggregate
+      val transformations: Map[Expression, Expression] =
+        (distinctAggOperatorMap.flatMap(_._2) ++
+          regularAggOperatorMap.map(e => (e._1, e._3))).toMap
+
+      val patchedAggExpressions = a.aggregateExpressions.map { e =>
+        e.transformDown {
+          case e: Expression =>
+            // The same GROUP BY clauses can have different forms (different names for instance) in
+            // the groupBy and aggregate expressions of an aggregate. This makes a map lookup
+            // tricky. So we do a linear search for a semantically equal group by expression.
+            groupByMap
+              .find(ge => e.semanticEquals(ge._1))
+              .map(_._2)
+              .getOrElse(transformations.getOrElse(e, e))
+        }.asInstanceOf[NamedExpression]
+      }
+      Aggregate(groupByAttrs, patchedAggExpressions, firstAggregate)
+    } else {
+      a
+    }
+  }
+
+  private def nullify(e: Expression) = Literal.create(null, e.dataType)
+
+  private def expressionAttributePair(e: Expression) =
+    // We are creating a new reference here instead of reusing the attribute in case of a
+    // NamedExpression. This is done to prevent collisions between distinct and regular aggregate
+    // children, in this case attribute reuse causes the input of the regular aggregate to bound to
+    // the (nulled out) input of the distinct aggregate.
+    e -> new AttributeReference(e.prettyName, e.dataType, true)()
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala
index a2fab258fcac3..5c5b3d1ccd3cd 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala
@@ -133,6 +133,12 @@ sealed abstract class AggregateFunction2 extends Expression with ImplicitCastInp
    */
   def supportsPartial: Boolean = true
 
+  /**
+   * Result of the aggregate function when the input is empty. This is currently only used for the
+   * proper rewriting of distinct aggregate functions.
+   */
+  def defaultResult: Option[Literal] = None
+
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String =
     throw new UnsupportedOperationException(s"Cannot evaluate expression: $this")
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 338c5193cb7a2..d222dfa33ad8a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -200,9 +200,9 @@ object SetOperationPushDown extends Rule[LogicalPlan] with PredicateHelper {
  */
 object ColumnPruning extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
-    case a @ Aggregate(_, _, e @ Expand(_, groupByExprs, _, child))
-      if (child.outputSet -- AttributeSet(groupByExprs) -- a.references).nonEmpty =>
-      a.copy(child = e.copy(child = prunedChild(child, AttributeSet(groupByExprs) ++ a.references)))
+    case a @ Aggregate(_, _, e @ Expand(_, _, child))
+      if (child.outputSet -- AttributeSet(e.output) -- a.references).nonEmpty =>
+      a.copy(child = e.copy(child = prunedChild(child, AttributeSet(e.output) ++ a.references)))
 
     // Eliminate attributes that are not needed to calculate the specified aggregates.
     case a @ Aggregate(_, _, child) if (child.outputSet -- a.references).nonEmpty =>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index 4cb67aacf33ee..fb963e2f8f7e7 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -235,33 +235,17 @@ case class Window(
     projectList ++ windowExpressions.map(_.toAttribute)
 }
 
-/**
- * Apply the all of the GroupExpressions to every input row, hence we will get
- * multiple output rows for a input row.
- * @param bitmasks The bitmask set represents the grouping sets
- * @param groupByExprs The grouping by expressions
- * @param child       Child operator
- */
-case class Expand(
-    bitmasks: Seq[Int],
-    groupByExprs: Seq[Expression],
-    gid: Attribute,
-    child: LogicalPlan) extends UnaryNode {
-  override def statistics: Statistics = {
-    val sizeInBytes = child.statistics.sizeInBytes * projections.length
-    Statistics(sizeInBytes = sizeInBytes)
-  }
-
-  val projections: Seq[Seq[Expression]] = expand()
-
+private[sql] object Expand {
   /**
-   * Extract attribute set according to the grouping id
+   * Extract attribute set according to the grouping id.
+   *
    * @param bitmask bitmask to represent the selected of the attribute sequence
    * @param exprs the attributes in sequence
    * @return the attributes of non selected specified via bitmask (with the bit set to 1)
    */
-  private def buildNonSelectExprSet(bitmask: Int, exprs: Seq[Expression])
-  : OpenHashSet[Expression] = {
+  private def buildNonSelectExprSet(
+      bitmask: Int,
+      exprs: Seq[Expression]): OpenHashSet[Expression] = {
     val set = new OpenHashSet[Expression](2)
 
     var bit = exprs.length - 1
@@ -274,18 +258,28 @@ case class Expand(
   }
 
   /**
-   * Create an array of Projections for the child projection, and replace the projections'
-   * expressions which equal GroupBy expressions with Literal(null), if those expressions
-   * are not set for this grouping set (according to the bit mask).
+   * Apply the all of the GroupExpressions to every input row, hence we will get
+   * multiple output rows for a input row.
+   *
+   * @param bitmasks The bitmask set represents the grouping sets
+   * @param groupByExprs The grouping by expressions
+   * @param gid Attribute of the grouping id
+   * @param child Child operator
    */
-  private[this] def expand(): Seq[Seq[Expression]] = {
-    val result = new scala.collection.mutable.ArrayBuffer[Seq[Expression]]
-
-    bitmasks.foreach { bitmask =>
+  def apply(
+    bitmasks: Seq[Int],
+    groupByExprs: Seq[Expression],
+    gid: Attribute,
+    child: LogicalPlan): Expand = {
+    // Create an array of Projections for the child projection, and replace the projections'
+    // expressions which equal GroupBy expressions with Literal(null), if those expressions
+    // are not set for this grouping set (according to the bit mask).
+    val projections = bitmasks.map { bitmask =>
       // get the non selected grouping attributes according to the bit mask
       val nonSelectedGroupExprSet = buildNonSelectExprSet(bitmask, groupByExprs)
 
-      val substitution = (child.output :+ gid).map(expr => expr transformDown {
+      (child.output :+ gid).map(expr => expr transformDown {
+        // TODO this causes a problem when a column is used both for grouping and aggregation.
         case x: Expression if nonSelectedGroupExprSet.contains(x) =>
           // if the input attribute in the Invalid Grouping Expression set of for this group
           // replace it with constant null
@@ -294,15 +288,29 @@ case class Expand(
           // replace the groupingId with concrete value (the bit mask)
           Literal.create(bitmask, IntegerType)
       })
-
-      result += substitution
     }
-
-    result.toSeq
+    Expand(projections, child.output :+ gid, child)
   }
+}
 
-  override def output: Seq[Attribute] = {
-    child.output :+ gid
+/**
+ * Apply a number of projections to every input row, hence we will get multiple output rows for
+ * a input row.
+ *
+ * @param projections to apply
+ * @param output of all projections.
+ * @param child operator.
+ */
+case class Expand(
+    projections: Seq[Seq[Expression]],
+    output: Seq[Attribute],
+    child: LogicalPlan) extends UnaryNode {
+
+  override def statistics: Statistics = {
+    // TODO shouldn't we factor in the size of the projection versus the size of the backing child
+    //      row?
+    val sizeInBytes = child.statistics.sizeInBytes * projections.length
+    Statistics(sizeInBytes = sizeInBytes)
   }
 }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index f4464e0b916f8..dd3bb33c57287 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -420,7 +420,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         }
       case logical.Filter(condition, child) =>
         execution.Filter(condition, planLater(child)) :: Nil
-      case e @ logical.Expand(_, _, _, child) =>
+      case e @ logical.Expand(_, _, child) =>
         execution.Expand(e.projections, e.output, planLater(child)) :: Nil
       case a @ logical.Aggregate(group, agg, child) => {
         val useNewAggregation = sqlContext.conf.useSqlAggregate2 && sqlContext.conf.codegenEnabled

From 1c80d66e52c0bcc4e5adda78b3d8e5bf55e4f128 Mon Sep 17 00:00:00 2001
From: "navis.ryu" <navis@apache.org>
Date: Fri, 6 Nov 2015 17:13:46 -0800
Subject: [PATCH 106/518] [SPARK-11546] Thrift server makes too many logs about
 result schema

SparkExecuteStatementOperation logs result schema for each getNextRowSet() calls which is by default every 1000 rows, overwhelming whole log file.

Author: navis.ryu <navis@apache.org>

Closes #9514 from navis/SPARK-11546.
---
 .../SparkExecuteStatementOperation.scala      | 24 ++++++++++---------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala
index 719b03e1c7c71..82fef92dcb73b 100644
--- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala
+++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala
@@ -53,6 +53,18 @@ private[hive] class SparkExecuteStatementOperation(
   private var dataTypes: Array[DataType] = _
   private var statementId: String = _
 
+  private lazy val resultSchema: TableSchema = {
+    if (result == null || result.queryExecution.analyzed.output.size == 0) {
+      new TableSchema(Arrays.asList(new FieldSchema("Result", "string", "")))
+    } else {
+      logInfo(s"Result Schema: ${result.queryExecution.analyzed.output}")
+      val schema = result.queryExecution.analyzed.output.map { attr =>
+        new FieldSchema(attr.name, HiveMetastoreTypes.toMetastoreType(attr.dataType), "")
+      }
+      new TableSchema(schema.asJava)
+    }
+  }
+
   def close(): Unit = {
     // RDDs will be cleaned automatically upon garbage collection.
     hiveContext.sparkContext.clearJobGroup()
@@ -120,17 +132,7 @@ private[hive] class SparkExecuteStatementOperation(
     }
   }
 
-  def getResultSetSchema: TableSchema = {
-    if (result == null || result.queryExecution.analyzed.output.size == 0) {
-      new TableSchema(Arrays.asList(new FieldSchema("Result", "string", "")))
-    } else {
-      logInfo(s"Result Schema: ${result.queryExecution.analyzed.output}")
-      val schema = result.queryExecution.analyzed.output.map { attr =>
-        new FieldSchema(attr.name, HiveMetastoreTypes.toMetastoreType(attr.dataType), "")
-      }
-      new TableSchema(schema.asJava)
-    }
-  }
+  def getResultSetSchema: TableSchema = resultSchema
 
   override def run(): Unit = {
     setState(OperationState.PENDING)

From 105732dcc6b651b9779f4a5773a759c5b4fbd21d Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Fri, 6 Nov 2015 17:22:30 -0800
Subject: [PATCH 107/518] [HOTFIX] Fix python tests after #9527

#9527 missed updating the python tests.

Author: Michael Armbrust <michael@databricks.com>

Closes #9533 from marmbrus/hotfixTextValue.
---
 python/pyspark/sql/readwriter.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py
index 97bd90c4db829..927f4077424dc 100644
--- a/python/pyspark/sql/readwriter.py
+++ b/python/pyspark/sql/readwriter.py
@@ -203,7 +203,7 @@ def text(self, path):
 
         >>> df = sqlContext.read.text('python/test_support/sql/text-test.txt')
         >>> df.collect()
-        [Row(text=u'hello'), Row(text=u'this')]
+        [Row(value=u'hello'), Row(value=u'this')]
         """
         return self._df(self._jreader.text(path))
 

From 30b706b7b36482921ec04145a0121ca147984fa8 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Fri, 6 Nov 2015 18:17:34 -0800
Subject: [PATCH 108/518] [SPARK-11389][CORE] Add support for off-heap memory
 to MemoryManager

In order to lay the groundwork for proper off-heap memory support in SQL / Tungsten, we need to extend our MemoryManager to perform bookkeeping for off-heap memory.

## User-facing changes

This PR introduces a new configuration, `spark.memory.offHeapSize` (name subject to change), which specifies the absolute amount of off-heap memory that Spark and Spark SQL can use. If Tungsten is configured to use off-heap execution memory for allocating data pages, then all data page allocations must fit within this size limit.

## Internals changes

This PR contains a lot of internal refactoring of the MemoryManager. The key change at the heart of this patch is the introduction of a `MemoryPool` class (name subject to change) to manage the bookkeeping for a particular category of memory (storage, on-heap execution, and off-heap execution). These MemoryPools are not fixed-size; they can be dynamically grown and shrunk according to the MemoryManager's policies. In StaticMemoryManager, these pools have fixed sizes, proportional to the legacy `[storage|shuffle].memoryFraction`. In the new UnifiedMemoryManager, the sizes of these pools are dynamically adjusted according to its policies.

There are two subclasses of `MemoryPool`: `StorageMemoryPool` manages storage memory and `ExecutionMemoryPool` manages execution memory. The MemoryManager creates two execution pools, one for on-heap memory and one for off-heap. Instances of `ExecutionMemoryPool` manage the logic for fair sharing of their pooled memory across running tasks (in other words, the ShuffleMemoryManager-like logic has been moved out of MemoryManager and pushed into these ExecutionMemoryPool instances).

I think that this design is substantially easier to understand and reason about than the previous design, where most of these responsibilities were handled by MemoryManager and its subclasses. To see this, take at look at how simple the logic in `UnifiedMemoryManager` has become: it's now very easy to see when memory is dynamically shifted between storage and execution.

## TODOs

- [x] Fix handful of test failures in the MemoryManagerSuites.
- [x] Fix remaining TODO comments in code.
- [ ] Document new configuration.
- [x] Fix commented-out tests / asserts:
  - [x] UnifiedMemoryManagerSuite.
- [x] Write tests that exercise the new off-heap memory management policies.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9344 from JoshRosen/offheap-memory-accounting.
---
 .../apache/spark/memory/MemoryConsumer.java   |   7 +-
 .../org/apache/spark/memory/MemoryMode.java   |  26 ++
 .../spark/memory/TaskMemoryManager.java       |  72 +++--
 .../scala/org/apache/spark/SparkEnv.scala     |   2 +-
 .../spark/memory/ExecutionMemoryPool.scala    | 153 +++++++++++
 .../apache/spark/memory/MemoryManager.scala   | 246 ++++++------------
 .../org/apache/spark/memory/MemoryPool.scala  |  71 +++++
 .../spark/memory/StaticMemoryManager.scala    |  75 +-----
 .../spark/memory/StorageMemoryPool.scala      | 138 ++++++++++
 .../spark/memory/UnifiedMemoryManager.scala   | 138 +++++-----
 .../org/apache/spark/memory/package.scala     |  75 ++++++
 .../spark/util/collection/Spillable.scala     |   8 +-
 .../spark/memory/TaskMemoryManagerSuite.java  |   8 +-
 .../spark/memory/TestMemoryConsumer.java      |  10 +-
 .../sort/UnsafeShuffleWriterSuite.java        |   2 +-
 .../map/AbstractBytesToBytesMapSuite.java     |   4 +-
 .../spark/memory/MemoryManagerSuite.scala     | 104 +++++---
 .../memory/StaticMemoryManagerSuite.scala     |  39 +--
 .../spark/memory/TestMemoryManager.scala      |  20 +-
 .../memory/UnifiedMemoryManagerSuite.scala    |  93 +++----
 .../spark/storage/BlockManagerSuite.scala     |   2 +-
 21 files changed, 828 insertions(+), 465 deletions(-)
 create mode 100644 core/src/main/java/org/apache/spark/memory/MemoryMode.java
 create mode 100644 core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala
 create mode 100644 core/src/main/scala/org/apache/spark/memory/MemoryPool.scala
 create mode 100644 core/src/main/scala/org/apache/spark/memory/StorageMemoryPool.scala
 create mode 100644 core/src/main/scala/org/apache/spark/memory/package.scala

diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
index 8fbdb72832adf..36138cc9a297c 100644
--- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
+++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java
@@ -17,15 +17,15 @@
 
 package org.apache.spark.memory;
 
-
 import java.io.IOException;
 
 import org.apache.spark.unsafe.array.LongArray;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 
-
 /**
  * An memory consumer of TaskMemoryManager, which support spilling.
+ *
+ * Note: this only supports allocation / spilling of Tungsten memory.
  */
 public abstract class MemoryConsumer {
 
@@ -36,7 +36,6 @@ public abstract class MemoryConsumer {
   protected MemoryConsumer(TaskMemoryManager taskMemoryManager, long pageSize) {
     this.taskMemoryManager = taskMemoryManager;
     this.pageSize = pageSize;
-    this.used = 0;
   }
 
   protected MemoryConsumer(TaskMemoryManager taskMemoryManager) {
@@ -67,6 +66,8 @@ public void spill() throws IOException {
    *
    * Note: In order to avoid possible deadlock, should not call acquireMemory() from spill().
    *
+   * Note: today, this only frees Tungsten-managed pages.
+   *
    * @param size the amount of memory should be released
    * @param trigger the MemoryConsumer that trigger this spilling
    * @return the amount of released memory in bytes
diff --git a/core/src/main/java/org/apache/spark/memory/MemoryMode.java b/core/src/main/java/org/apache/spark/memory/MemoryMode.java
new file mode 100644
index 0000000000000..3a5e72d8aaec0
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/memory/MemoryMode.java
@@ -0,0 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.memory;
+
+import org.apache.spark.annotation.Private;
+
+@Private
+public enum MemoryMode {
+  ON_HEAP,
+  OFF_HEAP
+}
diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 6440f9c0f30de..5f743b28857b4 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -103,10 +103,10 @@ public class TaskMemoryManager {
    * without doing any masking or lookups. Since this branching should be well-predicted by the JIT,
    * this extra layer of indirection / abstraction hopefully shouldn't be too expensive.
    */
-  private final boolean inHeap;
+  final MemoryMode tungstenMemoryMode;
 
   /**
-   * The size of memory granted to each consumer.
+   * Tracks spillable memory consumers.
    */
   @GuardedBy("this")
   private final HashSet<MemoryConsumer> consumers;
@@ -115,7 +115,7 @@ public class TaskMemoryManager {
    * Construct a new TaskMemoryManager.
    */
   public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
-    this.inHeap = memoryManager.tungstenMemoryIsAllocatedInHeap();
+    this.tungstenMemoryMode = memoryManager.tungstenMemoryMode();
     this.memoryManager = memoryManager;
     this.taskAttemptId = taskAttemptId;
     this.consumers = new HashSet<>();
@@ -127,12 +127,19 @@ public TaskMemoryManager(MemoryManager memoryManager, long taskAttemptId) {
    *
    * @return number of bytes successfully granted (<= N).
    */
-  public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
+  public long acquireExecutionMemory(
+      long required,
+      MemoryMode mode,
+      MemoryConsumer consumer) {
     assert(required >= 0);
+    // If we are allocating Tungsten pages off-heap and receive a request to allocate on-heap
+    // memory here, then it may not make sense to spill since that would only end up freeing
+    // off-heap memory. This is subject to change, though, so it may be risky to make this
+    // optimization now in case we forget to undo it late when making changes.
     synchronized (this) {
-      long got = memoryManager.acquireExecutionMemory(required, taskAttemptId);
+      long got = memoryManager.acquireExecutionMemory(required, taskAttemptId, mode);
 
-      // try to release memory from other consumers first, then we can reduce the frequency of
+      // Try to release memory from other consumers first, then we can reduce the frequency of
       // spilling, avoid to have too many spilled files.
       if (got < required) {
         // Call spill() on other consumers to release memory
@@ -140,10 +147,10 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
           if (c != consumer && c.getUsed() > 0) {
             try {
               long released = c.spill(required - got, consumer);
-              if (released > 0) {
-                logger.info("Task {} released {} from {} for {}", taskAttemptId,
+              if (released > 0 && mode == tungstenMemoryMode) {
+                logger.debug("Task {} released {} from {} for {}", taskAttemptId,
                   Utils.bytesToString(released), c, consumer);
-                got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
+                got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId, mode);
                 if (got >= required) {
                   break;
                 }
@@ -161,10 +168,10 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
       if (got < required && consumer != null) {
         try {
           long released = consumer.spill(required - got, consumer);
-          if (released > 0) {
-            logger.info("Task {} released {} from itself ({})", taskAttemptId,
+          if (released > 0 && mode == tungstenMemoryMode) {
+            logger.debug("Task {} released {} from itself ({})", taskAttemptId,
               Utils.bytesToString(released), consumer);
-            got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId);
+            got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId, mode);
           }
         } catch (IOException e) {
           logger.error("error while calling spill() on " + consumer, e);
@@ -184,9 +191,9 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
   /**
    * Release N bytes of execution memory for a MemoryConsumer.
    */
-  public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
+  public void releaseExecutionMemory(long size, MemoryMode mode, MemoryConsumer consumer) {
     logger.debug("Task {} release {} from {}", taskAttemptId, Utils.bytesToString(size), consumer);
-    memoryManager.releaseExecutionMemory(size, taskAttemptId);
+    memoryManager.releaseExecutionMemory(size, taskAttemptId, mode);
   }
 
   /**
@@ -195,11 +202,19 @@ public void releaseExecutionMemory(long size, MemoryConsumer consumer) {
   public void showMemoryUsage() {
     logger.info("Memory used in task " + taskAttemptId);
     synchronized (this) {
+      long memoryAccountedForByConsumers = 0;
       for (MemoryConsumer c: consumers) {
-        if (c.getUsed() > 0) {
-          logger.info("Acquired by " + c + ": " + Utils.bytesToString(c.getUsed()));
+        long totalMemUsage = c.getUsed();
+        memoryAccountedForByConsumers += totalMemUsage;
+        if (totalMemUsage > 0) {
+          logger.info("Acquired by " + c + ": " + Utils.bytesToString(totalMemUsage));
         }
       }
+      long memoryNotAccountedFor =
+        memoryManager.getExecutionMemoryUsageForTask(taskAttemptId) - memoryAccountedForByConsumers;
+      logger.info(
+        "{} bytes of memory were used by task {} but are not associated with specific consumers",
+        memoryNotAccountedFor, taskAttemptId);
     }
   }
 
@@ -214,7 +229,8 @@ public long pageSizeBytes() {
    * Allocate a block of memory that will be tracked in the MemoryManager's page table; this is
    * intended for allocating large blocks of Tungsten memory that will be shared between operators.
    *
-   * Returns `null` if there was not enough memory to allocate the page.
+   * Returns `null` if there was not enough memory to allocate the page. May return a page that
+   * contains fewer bytes than requested, so callers should verify the size of returned pages.
    */
   public MemoryBlock allocatePage(long size, MemoryConsumer consumer) {
     if (size > MAXIMUM_PAGE_SIZE_BYTES) {
@@ -222,7 +238,7 @@ public MemoryBlock allocatePage(long size, MemoryConsumer consumer) {
         "Cannot allocate a page with more than " + MAXIMUM_PAGE_SIZE_BYTES + " bytes");
     }
 
-    long acquired = acquireExecutionMemory(size, consumer);
+    long acquired = acquireExecutionMemory(size, tungstenMemoryMode, consumer);
     if (acquired <= 0) {
       return null;
     }
@@ -231,7 +247,7 @@ public MemoryBlock allocatePage(long size, MemoryConsumer consumer) {
     synchronized (this) {
       pageNumber = allocatedPages.nextClearBit(0);
       if (pageNumber >= PAGE_TABLE_SIZE) {
-        releaseExecutionMemory(acquired, consumer);
+        releaseExecutionMemory(acquired, tungstenMemoryMode, consumer);
         throw new IllegalStateException(
           "Have already allocated a maximum of " + PAGE_TABLE_SIZE + " pages");
       }
@@ -262,7 +278,7 @@ public void freePage(MemoryBlock page, MemoryConsumer consumer) {
     }
     long pageSize = page.size();
     memoryManager.tungstenMemoryAllocator().free(page);
-    releaseExecutionMemory(pageSize, consumer);
+    releaseExecutionMemory(pageSize, tungstenMemoryMode, consumer);
   }
 
   /**
@@ -276,7 +292,7 @@ public void freePage(MemoryBlock page, MemoryConsumer consumer) {
    * @return an encoded page address.
    */
   public long encodePageNumberAndOffset(MemoryBlock page, long offsetInPage) {
-    if (!inHeap) {
+    if (tungstenMemoryMode == MemoryMode.OFF_HEAP) {
       // In off-heap mode, an offset is an absolute address that may require a full 64 bits to
       // encode. Due to our page size limitation, though, we can convert this into an offset that's
       // relative to the page's base offset; this relative offset will fit in 51 bits.
@@ -305,7 +321,7 @@ private static long decodeOffset(long pagePlusOffsetAddress) {
    * {@link TaskMemoryManager#encodePageNumberAndOffset(MemoryBlock, long)}
    */
   public Object getPage(long pagePlusOffsetAddress) {
-    if (inHeap) {
+    if (tungstenMemoryMode == MemoryMode.ON_HEAP) {
       final int pageNumber = decodePageNumber(pagePlusOffsetAddress);
       assert (pageNumber >= 0 && pageNumber < PAGE_TABLE_SIZE);
       final MemoryBlock page = pageTable[pageNumber];
@@ -323,7 +339,7 @@ public Object getPage(long pagePlusOffsetAddress) {
    */
   public long getOffsetInPage(long pagePlusOffsetAddress) {
     final long offsetInPage = decodeOffset(pagePlusOffsetAddress);
-    if (inHeap) {
+    if (tungstenMemoryMode == MemoryMode.ON_HEAP) {
       return offsetInPage;
     } else {
       // In off-heap mode, an offset is an absolute address. In encodePageNumberAndOffset, we
@@ -351,11 +367,19 @@ public long cleanUpAllAllocatedMemory() {
       }
       consumers.clear();
     }
+
+    for (MemoryBlock page : pageTable) {
+      if (page != null) {
+        memoryManager.tungstenMemoryAllocator().free(page);
+      }
+    }
+    Arrays.fill(pageTable, null);
+
     return memoryManager.releaseAllExecutionMemoryForTask(taskAttemptId);
   }
 
   /**
-   * Returns the memory consumption, in bytes, for the current task
+   * Returns the memory consumption, in bytes, for the current task.
    */
   public long getMemoryConsumptionForThisTask() {
     return memoryManager.getExecutionMemoryUsageForTask(taskAttemptId);
diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala b/core/src/main/scala/org/apache/spark/SparkEnv.scala
index 23ae9360f6a22..4474a83bedbdb 100644
--- a/core/src/main/scala/org/apache/spark/SparkEnv.scala
+++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala
@@ -341,7 +341,7 @@ object SparkEnv extends Logging {
       if (useLegacyMemoryManager) {
         new StaticMemoryManager(conf, numUsableCores)
       } else {
-        new UnifiedMemoryManager(conf, numUsableCores)
+        UnifiedMemoryManager(conf, numUsableCores)
       }
 
     val blockTransferService = new NettyBlockTransferService(conf, securityManager, numUsableCores)
diff --git a/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala b/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala
new file mode 100644
index 0000000000000..7825bae425877
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.memory
+
+import javax.annotation.concurrent.GuardedBy
+
+import scala.collection.mutable
+
+import org.apache.spark.Logging
+
+/**
+ * Implements policies and bookkeeping for sharing a adjustable-sized pool of memory between tasks.
+ *
+ * Tries to ensure that each task gets a reasonable share of memory, instead of some task ramping up
+ * to a large amount first and then causing others to spill to disk repeatedly.
+ *
+ * If there are N tasks, it ensures that each task can acquire at least 1 / 2N of the memory
+ * before it has to spill, and at most 1 / N. Because N varies dynamically, we keep track of the
+ * set of active tasks and redo the calculations of 1 / 2N and 1 / N in waiting tasks whenever this
+ * set changes. This is all done by synchronizing access to mutable state and using wait() and
+ * notifyAll() to signal changes to callers. Prior to Spark 1.6, this arbitration of memory across
+ * tasks was performed by the ShuffleMemoryManager.
+ *
+ * @param lock a [[MemoryManager]] instance to synchronize on
+ * @param poolName a human-readable name for this pool, for use in log messages
+ */
+class ExecutionMemoryPool(
+    lock: Object,
+    poolName: String
+  ) extends MemoryPool(lock) with Logging {
+
+  /**
+   * Map from taskAttemptId -> memory consumption in bytes
+   */
+  @GuardedBy("lock")
+  private val memoryForTask = new mutable.HashMap[Long, Long]()
+
+  override def memoryUsed: Long = lock.synchronized {
+    memoryForTask.values.sum
+  }
+
+  /**
+   * Returns the memory consumption, in bytes, for the given task.
+   */
+  def getMemoryUsageForTask(taskAttemptId: Long): Long = lock.synchronized {
+    memoryForTask.getOrElse(taskAttemptId, 0L)
+  }
+
+  /**
+   * Try to acquire up to `numBytes` of memory for the given task and return the number of bytes
+   * obtained, or 0 if none can be allocated.
+   *
+   * This call may block until there is enough free memory in some situations, to make sure each
+   * task has a chance to ramp up to at least 1 / 2N of the total memory pool (where N is the # of
+   * active tasks) before it is forced to spill. This can happen if the number of tasks increase
+   * but an older task had a lot of memory already.
+   *
+   * @return the number of bytes granted to the task.
+   */
+  def acquireMemory(numBytes: Long, taskAttemptId: Long): Long = lock.synchronized {
+    assert(numBytes > 0, s"invalid number of bytes requested: $numBytes")
+
+    // Add this task to the taskMemory map just so we can keep an accurate count of the number
+    // of active tasks, to let other tasks ramp down their memory in calls to `acquireMemory`
+    if (!memoryForTask.contains(taskAttemptId)) {
+      memoryForTask(taskAttemptId) = 0L
+      // This will later cause waiting tasks to wake up and check numTasks again
+      lock.notifyAll()
+    }
+
+    // Keep looping until we're either sure that we don't want to grant this request (because this
+    // task would have more than 1 / numActiveTasks of the memory) or we have enough free
+    // memory to give it (we always let each task get at least 1 / (2 * numActiveTasks)).
+    // TODO: simplify this to limit each task to its own slot
+    while (true) {
+      val numActiveTasks = memoryForTask.keys.size
+      val curMem = memoryForTask(taskAttemptId)
+
+      // How much we can grant this task; don't let it grow to more than 1 / numActiveTasks;
+      // don't let it be negative
+      val maxToGrant =
+        math.min(numBytes, math.max(0, (poolSize / numActiveTasks) - curMem))
+      // Only give it as much memory as is free, which might be none if it reached 1 / numTasks
+      val toGrant = math.min(maxToGrant, memoryFree)
+
+      if (curMem < poolSize / (2 * numActiveTasks)) {
+        // We want to let each task get at least 1 / (2 * numActiveTasks) before blocking;
+        // if we can't give it this much now, wait for other tasks to free up memory
+        // (this happens if older tasks allocated lots of memory before N grew)
+        if (memoryFree >= math.min(maxToGrant, poolSize / (2 * numActiveTasks) - curMem)) {
+          memoryForTask(taskAttemptId) += toGrant
+          return toGrant
+        } else {
+          logInfo(
+            s"TID $taskAttemptId waiting for at least 1/2N of $poolName pool to be free")
+          lock.wait()
+        }
+      } else {
+        memoryForTask(taskAttemptId) += toGrant
+        return toGrant
+      }
+    }
+    0L  // Never reached
+  }
+
+  /**
+   * Release `numBytes` of memory acquired by the given task.
+   */
+  def releaseMemory(numBytes: Long, taskAttemptId: Long): Unit = lock.synchronized {
+    val curMem = memoryForTask.getOrElse(taskAttemptId, 0L)
+    var memoryToFree = if (curMem < numBytes) {
+      logWarning(
+        s"Internal error: release called on $numBytes bytes but task only has $curMem bytes " +
+          s"of memory from the $poolName pool")
+      curMem
+    } else {
+      numBytes
+    }
+    if (memoryForTask.contains(taskAttemptId)) {
+      memoryForTask(taskAttemptId) -= memoryToFree
+      if (memoryForTask(taskAttemptId) <= 0) {
+        memoryForTask.remove(taskAttemptId)
+      }
+    }
+    lock.notifyAll() // Notify waiters in acquireMemory() that memory has been freed
+  }
+
+  /**
+   * Release all memory for the given task and mark it as inactive (e.g. when a task ends).
+   * @return the number of bytes freed.
+   */
+  def releaseAllMemoryForTask(taskAttemptId: Long): Long = lock.synchronized {
+    val numBytesToFree = getMemoryUsageForTask(taskAttemptId)
+    releaseMemory(numBytesToFree, taskAttemptId)
+    numBytesToFree
+  }
+
+}
diff --git a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
index b0cf2696a397f..ceb8ea434e1be 100644
--- a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
+++ b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala
@@ -20,12 +20,8 @@ package org.apache.spark.memory
 import javax.annotation.concurrent.GuardedBy
 
 import scala.collection.mutable
-import scala.collection.mutable.ArrayBuffer
 
-import com.google.common.annotations.VisibleForTesting
-
-import org.apache.spark.util.Utils
-import org.apache.spark.{SparkException, TaskContext, SparkConf, Logging}
+import org.apache.spark.{SparkConf, Logging}
 import org.apache.spark.storage.{BlockId, BlockStatus, MemoryStore}
 import org.apache.spark.unsafe.array.ByteArrayMethods
 import org.apache.spark.unsafe.memory.MemoryAllocator
@@ -36,53 +32,40 @@ import org.apache.spark.unsafe.memory.MemoryAllocator
  * In this context, execution memory refers to that used for computation in shuffles, joins,
  * sorts and aggregations, while storage memory refers to that used for caching and propagating
  * internal data across the cluster. There exists one MemoryManager per JVM.
- *
- * The MemoryManager abstract base class itself implements policies for sharing execution memory
- * between tasks; it tries to ensure that each task gets a reasonable share of memory, instead of
- * some task ramping up to a large amount first and then causing others to spill to disk repeatedly.
- * If there are N tasks, it ensures that each task can acquire at least 1 / 2N of the memory
- * before it has to spill, and at most 1 / N. Because N varies dynamically, we keep track of the
- * set of active tasks and redo the calculations of 1 / 2N and 1 / N in waiting tasks whenever
- * this set changes. This is all done by synchronizing access to mutable state and using wait() and
- * notifyAll() to signal changes to callers. Prior to Spark 1.6, this arbitration of memory across
- * tasks was performed by the ShuffleMemoryManager.
  */
-private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) extends Logging {
+private[spark] abstract class MemoryManager(
+    conf: SparkConf,
+    numCores: Int,
+    storageMemory: Long,
+    onHeapExecutionMemory: Long) extends Logging {
 
   // -- Methods related to memory allocation policies and bookkeeping ------------------------------
 
-  // The memory store used to evict cached blocks
-  private var _memoryStore: MemoryStore = _
-  protected def memoryStore: MemoryStore = {
-    if (_memoryStore == null) {
-      throw new IllegalArgumentException("memory store not initialized yet")
-    }
-    _memoryStore
-  }
+  @GuardedBy("this")
+  protected val storageMemoryPool = new StorageMemoryPool(this)
+  @GuardedBy("this")
+  protected val onHeapExecutionMemoryPool = new ExecutionMemoryPool(this, "on-heap execution")
+  @GuardedBy("this")
+  protected val offHeapExecutionMemoryPool = new ExecutionMemoryPool(this, "off-heap execution")
 
-  // Amount of execution/storage memory in use, accesses must be synchronized on `this`
-  @GuardedBy("this") protected var _executionMemoryUsed: Long = 0
-  @GuardedBy("this") protected var _storageMemoryUsed: Long = 0
-  // Map from taskAttemptId -> memory consumption in bytes
-  @GuardedBy("this") private val executionMemoryForTask = new mutable.HashMap[Long, Long]()
-
-  /**
-   * Set the [[MemoryStore]] used by this manager to evict cached blocks.
-   * This must be set after construction due to initialization ordering constraints.
-   */
-  final def setMemoryStore(store: MemoryStore): Unit = {
-    _memoryStore = store
-  }
+  storageMemoryPool.incrementPoolSize(storageMemory)
+  onHeapExecutionMemoryPool.incrementPoolSize(onHeapExecutionMemory)
+  offHeapExecutionMemoryPool.incrementPoolSize(conf.getSizeAsBytes("spark.memory.offHeapSize", 0))
 
   /**
-   * Total available memory for execution, in bytes.
+   * Total available memory for storage, in bytes. This amount can vary over time, depending on
+   * the MemoryManager implementation.
+   * In this model, this is equivalent to the amount of memory not occupied by execution.
    */
-  def maxExecutionMemory: Long
+  def maxStorageMemory: Long
 
   /**
-   * Total available memory for storage, in bytes.
+   * Set the [[MemoryStore]] used by this manager to evict cached blocks.
+   * This must be set after construction due to initialization ordering constraints.
    */
-  def maxStorageMemory: Long
+  final def setMemoryStore(store: MemoryStore): Unit = synchronized {
+    storageMemoryPool.setMemoryStore(store)
+  }
 
   // TODO: avoid passing evicted blocks around to simplify method signatures (SPARK-10985)
 
@@ -94,7 +77,9 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
   def acquireStorageMemory(
       blockId: BlockId,
       numBytes: Long,
-      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean
+      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
+    storageMemoryPool.acquireMemory(blockId, numBytes, evictedBlocks)
+  }
 
   /**
    * Acquire N bytes of memory to unroll the given block, evicting existing ones if necessary.
@@ -109,103 +94,25 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
   def acquireUnrollMemory(
       blockId: BlockId,
       numBytes: Long,
-      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
-    acquireStorageMemory(blockId, numBytes, evictedBlocks)
-  }
-
-  /**
-   * Acquire N bytes of memory for execution, evicting cached blocks if necessary.
-   * Blocks evicted in the process, if any, are added to `evictedBlocks`.
-   * @return number of bytes successfully granted (<= N).
-   */
-  @VisibleForTesting
-  private[memory] def doAcquireExecutionMemory(
-      numBytes: Long,
-      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Long
+      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean
 
   /**
-   * Try to acquire up to `numBytes` of execution memory for the current task and return the number
-   * of bytes obtained, or 0 if none can be allocated.
+   * Try to acquire up to `numBytes` of execution memory for the current task and return the
+   * number of bytes obtained, or 0 if none can be allocated.
    *
    * This call may block until there is enough free memory in some situations, to make sure each
    * task has a chance to ramp up to at least 1 / 2N of the total memory pool (where N is the # of
    * active tasks) before it is forced to spill. This can happen if the number of tasks increase
    * but an older task had a lot of memory already.
-   *
-   * Subclasses should override `doAcquireExecutionMemory` in order to customize the policies
-   * that control global sharing of memory between execution and storage.
    */
   private[memory]
-  final def acquireExecutionMemory(numBytes: Long, taskAttemptId: Long): Long = synchronized {
-    assert(numBytes > 0, "invalid number of bytes requested: " + numBytes)
-
-    // Add this task to the taskMemory map just so we can keep an accurate count of the number
-    // of active tasks, to let other tasks ramp down their memory in calls to tryToAcquire
-    if (!executionMemoryForTask.contains(taskAttemptId)) {
-      executionMemoryForTask(taskAttemptId) = 0L
-      // This will later cause waiting tasks to wake up and check numTasks again
-      notifyAll()
-    }
-
-    // Once the cross-task memory allocation policy has decided to grant more memory to a task,
-    // this method is called in order to actually obtain that execution memory, potentially
-    // triggering eviction of storage memory:
-    def acquire(toGrant: Long): Long = synchronized {
-      val evictedBlocks = new ArrayBuffer[(BlockId, BlockStatus)]
-      val acquired = doAcquireExecutionMemory(toGrant, evictedBlocks)
-      // Register evicted blocks, if any, with the active task metrics
-      Option(TaskContext.get()).foreach { tc =>
-        val metrics = tc.taskMetrics()
-        val lastUpdatedBlocks = metrics.updatedBlocks.getOrElse(Seq[(BlockId, BlockStatus)]())
-        metrics.updatedBlocks = Some(lastUpdatedBlocks ++ evictedBlocks.toSeq)
-      }
-      executionMemoryForTask(taskAttemptId) += acquired
-      acquired
-    }
-
-    // Keep looping until we're either sure that we don't want to grant this request (because this
-    // task would have more than 1 / numActiveTasks of the memory) or we have enough free
-    // memory to give it (we always let each task get at least 1 / (2 * numActiveTasks)).
-    // TODO: simplify this to limit each task to its own slot
-    while (true) {
-      val numActiveTasks = executionMemoryForTask.keys.size
-      val curMem = executionMemoryForTask(taskAttemptId)
-      val freeMemory = maxExecutionMemory - executionMemoryForTask.values.sum
-
-      // How much we can grant this task; don't let it grow to more than 1 / numActiveTasks;
-      // don't let it be negative
-      val maxToGrant =
-        math.min(numBytes, math.max(0, (maxExecutionMemory / numActiveTasks) - curMem))
-      // Only give it as much memory as is free, which might be none if it reached 1 / numTasks
-      val toGrant = math.min(maxToGrant, freeMemory)
-
-      if (curMem < maxExecutionMemory / (2 * numActiveTasks)) {
-        // We want to let each task get at least 1 / (2 * numActiveTasks) before blocking;
-        // if we can't give it this much now, wait for other tasks to free up memory
-        // (this happens if older tasks allocated lots of memory before N grew)
-        if (
-          freeMemory >= math.min(maxToGrant, maxExecutionMemory / (2 * numActiveTasks) - curMem)) {
-          return acquire(toGrant)
-        } else {
-          logInfo(
-            s"TID $taskAttemptId waiting for at least 1/2N of execution memory pool to be free")
-          wait()
-        }
-      } else {
-        return acquire(toGrant)
-      }
-    }
-    0L  // Never reached
-  }
-
-  @VisibleForTesting
-  private[memory] def releaseExecutionMemory(numBytes: Long): Unit = synchronized {
-    if (numBytes > _executionMemoryUsed) {
-      logWarning(s"Attempted to release $numBytes bytes of execution " +
-        s"memory when we only have ${_executionMemoryUsed} bytes")
-      _executionMemoryUsed = 0
-    } else {
-      _executionMemoryUsed -= numBytes
+  def acquireExecutionMemory(
+      numBytes: Long,
+      taskAttemptId: Long,
+      memoryMode: MemoryMode): Long = synchronized {
+    memoryMode match {
+      case MemoryMode.ON_HEAP => onHeapExecutionMemoryPool.acquireMemory(numBytes, taskAttemptId)
+      case MemoryMode.OFF_HEAP => offHeapExecutionMemoryPool.acquireMemory(numBytes, taskAttemptId)
     }
   }
 
@@ -213,24 +120,14 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
    * Release numBytes of execution memory belonging to the given task.
    */
   private[memory]
-  final def releaseExecutionMemory(numBytes: Long, taskAttemptId: Long): Unit = synchronized {
-    val curMem = executionMemoryForTask.getOrElse(taskAttemptId, 0L)
-    if (curMem < numBytes) {
-      if (Utils.isTesting) {
-        throw new SparkException(
-          s"Internal error: release called on $numBytes bytes but task only has $curMem")
-      } else {
-        logWarning(s"Internal error: release called on $numBytes bytes but task only has $curMem")
-      }
-    }
-    if (executionMemoryForTask.contains(taskAttemptId)) {
-      executionMemoryForTask(taskAttemptId) -= numBytes
-      if (executionMemoryForTask(taskAttemptId) <= 0) {
-        executionMemoryForTask.remove(taskAttemptId)
-      }
-      releaseExecutionMemory(numBytes)
+  def releaseExecutionMemory(
+      numBytes: Long,
+      taskAttemptId: Long,
+      memoryMode: MemoryMode): Unit = synchronized {
+    memoryMode match {
+      case MemoryMode.ON_HEAP => onHeapExecutionMemoryPool.releaseMemory(numBytes, taskAttemptId)
+      case MemoryMode.OFF_HEAP => offHeapExecutionMemoryPool.releaseMemory(numBytes, taskAttemptId)
     }
-    notifyAll() // Notify waiters in acquireExecutionMemory() that memory has been freed
   }
 
   /**
@@ -238,35 +135,28 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
    * @return the number of bytes freed.
    */
   private[memory] def releaseAllExecutionMemoryForTask(taskAttemptId: Long): Long = synchronized {
-    val numBytesToFree = getExecutionMemoryUsageForTask(taskAttemptId)
-    releaseExecutionMemory(numBytesToFree, taskAttemptId)
-    numBytesToFree
+    onHeapExecutionMemoryPool.releaseAllMemoryForTask(taskAttemptId) +
+      offHeapExecutionMemoryPool.releaseAllMemoryForTask(taskAttemptId)
   }
 
   /**
    * Release N bytes of storage memory.
    */
   def releaseStorageMemory(numBytes: Long): Unit = synchronized {
-    if (numBytes > _storageMemoryUsed) {
-      logWarning(s"Attempted to release $numBytes bytes of storage " +
-        s"memory when we only have ${_storageMemoryUsed} bytes")
-      _storageMemoryUsed = 0
-    } else {
-      _storageMemoryUsed -= numBytes
-    }
+    storageMemoryPool.releaseMemory(numBytes)
   }
 
   /**
    * Release all storage memory acquired.
    */
-  def releaseAllStorageMemory(): Unit = synchronized {
-    _storageMemoryUsed = 0
+  final def releaseAllStorageMemory(): Unit = synchronized {
+    storageMemoryPool.releaseAllMemory()
   }
 
   /**
    * Release N bytes of unroll memory.
    */
-  def releaseUnrollMemory(numBytes: Long): Unit = synchronized {
+  final def releaseUnrollMemory(numBytes: Long): Unit = synchronized {
     releaseStorageMemory(numBytes)
   }
 
@@ -274,25 +164,34 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
    * Execution memory currently in use, in bytes.
    */
   final def executionMemoryUsed: Long = synchronized {
-    _executionMemoryUsed
+    onHeapExecutionMemoryPool.memoryUsed + offHeapExecutionMemoryPool.memoryUsed
   }
 
   /**
    * Storage memory currently in use, in bytes.
    */
   final def storageMemoryUsed: Long = synchronized {
-    _storageMemoryUsed
+    storageMemoryPool.memoryUsed
   }
 
   /**
    * Returns the execution memory consumption, in bytes, for the given task.
    */
   private[memory] def getExecutionMemoryUsageForTask(taskAttemptId: Long): Long = synchronized {
-    executionMemoryForTask.getOrElse(taskAttemptId, 0L)
+    onHeapExecutionMemoryPool.getMemoryUsageForTask(taskAttemptId) +
+      offHeapExecutionMemoryPool.getMemoryUsageForTask(taskAttemptId)
   }
 
   // -- Fields related to Tungsten managed memory -------------------------------------------------
 
+  /**
+   * Tracks whether Tungsten memory will be allocated on the JVM heap or off-heap using
+   * sun.misc.Unsafe.
+   */
+  final val tungstenMemoryMode: MemoryMode = {
+    if (conf.getBoolean("spark.unsafe.offHeap", false)) MemoryMode.OFF_HEAP else MemoryMode.ON_HEAP
+  }
+
   /**
    * The default page size, in bytes.
    *
@@ -306,21 +205,22 @@ private[spark] abstract class MemoryManager(conf: SparkConf, numCores: Int) exte
     val cores = if (numCores > 0) numCores else Runtime.getRuntime.availableProcessors()
     // Because of rounding to next power of 2, we may have safetyFactor as 8 in worst case
     val safetyFactor = 16
-    val size = ByteArrayMethods.nextPowerOf2(maxExecutionMemory / cores / safetyFactor)
+    val maxTungstenMemory: Long = tungstenMemoryMode match {
+      case MemoryMode.ON_HEAP => onHeapExecutionMemoryPool.poolSize
+      case MemoryMode.OFF_HEAP => offHeapExecutionMemoryPool.poolSize
+    }
+    val size = ByteArrayMethods.nextPowerOf2(maxTungstenMemory / cores / safetyFactor)
     val default = math.min(maxPageSize, math.max(minPageSize, size))
     conf.getSizeAsBytes("spark.buffer.pageSize", default)
   }
 
-  /**
-   * Tracks whether Tungsten memory will be allocated on the JVM heap or off-heap using
-   * sun.misc.Unsafe.
-   */
-  final val tungstenMemoryIsAllocatedInHeap: Boolean =
-    !conf.getBoolean("spark.unsafe.offHeap", false)
-
   /**
    * Allocates memory for use by Unsafe/Tungsten code.
    */
-  private[memory] final val tungstenMemoryAllocator: MemoryAllocator =
-    if (tungstenMemoryIsAllocatedInHeap) MemoryAllocator.HEAP else MemoryAllocator.UNSAFE
+  private[memory] final val tungstenMemoryAllocator: MemoryAllocator = {
+    tungstenMemoryMode match {
+      case MemoryMode.ON_HEAP => MemoryAllocator.HEAP
+      case MemoryMode.OFF_HEAP => MemoryAllocator.UNSAFE
+    }
+  }
 }
diff --git a/core/src/main/scala/org/apache/spark/memory/MemoryPool.scala b/core/src/main/scala/org/apache/spark/memory/MemoryPool.scala
new file mode 100644
index 0000000000000..bfeec47e3892e
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/memory/MemoryPool.scala
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.memory
+
+import javax.annotation.concurrent.GuardedBy
+
+/**
+ * Manages bookkeeping for an adjustable-sized region of memory. This class is internal to
+ * the [[MemoryManager]]. See subclasses for more details.
+ *
+ * @param lock a [[MemoryManager]] instance, used for synchronization. We purposely erase the type
+ *             to `Object` to avoid programming errors, since this object should only be used for
+ *             synchronization purposes.
+ */
+abstract class MemoryPool(lock: Object) {
+
+  @GuardedBy("lock")
+  private[this] var _poolSize: Long = 0
+
+  /**
+   * Returns the current size of the pool, in bytes.
+   */
+  final def poolSize: Long = lock.synchronized {
+    _poolSize
+  }
+
+  /**
+   * Returns the amount of free memory in the pool, in bytes.
+   */
+  final def memoryFree: Long = lock.synchronized {
+    _poolSize - memoryUsed
+  }
+
+  /**
+   * Expands the pool by `delta` bytes.
+   */
+  final def incrementPoolSize(delta: Long): Unit = lock.synchronized {
+    require(delta >= 0)
+    _poolSize += delta
+  }
+
+  /**
+   * Shrinks the pool by `delta` bytes.
+   */
+  final def decrementPoolSize(delta: Long): Unit = lock.synchronized {
+    require(delta >= 0)
+    require(delta <= _poolSize)
+    require(_poolSize - delta >= memoryUsed)
+    _poolSize -= delta
+  }
+
+  /**
+   * Returns the amount of used memory in this pool (in bytes).
+   */
+  def memoryUsed: Long
+}
diff --git a/core/src/main/scala/org/apache/spark/memory/StaticMemoryManager.scala b/core/src/main/scala/org/apache/spark/memory/StaticMemoryManager.scala
index 9c2c2e90a2282..12a094306861f 100644
--- a/core/src/main/scala/org/apache/spark/memory/StaticMemoryManager.scala
+++ b/core/src/main/scala/org/apache/spark/memory/StaticMemoryManager.scala
@@ -22,7 +22,6 @@ import scala.collection.mutable
 import org.apache.spark.SparkConf
 import org.apache.spark.storage.{BlockId, BlockStatus}
 
-
 /**
  * A [[MemoryManager]] that statically partitions the heap space into disjoint regions.
  *
@@ -32,10 +31,14 @@ import org.apache.spark.storage.{BlockId, BlockStatus}
  */
 private[spark] class StaticMemoryManager(
     conf: SparkConf,
-    override val maxExecutionMemory: Long,
+    maxOnHeapExecutionMemory: Long,
     override val maxStorageMemory: Long,
     numCores: Int)
-  extends MemoryManager(conf, numCores) {
+  extends MemoryManager(
+    conf,
+    numCores,
+    maxStorageMemory,
+    maxOnHeapExecutionMemory) {
 
   def this(conf: SparkConf, numCores: Int) {
     this(
@@ -50,76 +53,15 @@ private[spark] class StaticMemoryManager(
     (maxStorageMemory * conf.getDouble("spark.storage.unrollFraction", 0.2)).toLong
   }
 
-  /**
-   * Acquire N bytes of memory for execution.
-   * @return number of bytes successfully granted (<= N).
-   */
-  override def doAcquireExecutionMemory(
-      numBytes: Long,
-      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Long = synchronized {
-    assert(numBytes >= 0)
-    assert(_executionMemoryUsed <= maxExecutionMemory)
-    val bytesToGrant = math.min(numBytes, maxExecutionMemory - _executionMemoryUsed)
-    _executionMemoryUsed += bytesToGrant
-    bytesToGrant
-  }
-
-  /**
-   * Acquire N bytes of memory to cache the given block, evicting existing ones if necessary.
-   * Blocks evicted in the process, if any, are added to `evictedBlocks`.
-   * @return whether all N bytes were successfully granted.
-   */
-  override def acquireStorageMemory(
-      blockId: BlockId,
-      numBytes: Long,
-      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
-    acquireStorageMemory(blockId, numBytes, numBytes, evictedBlocks)
-  }
-
-  /**
-   * Acquire N bytes of memory to unroll the given block, evicting existing ones if necessary.
-   *
-   * This evicts at most M bytes worth of existing blocks, where M is a fraction of the storage
-   * space specified by `spark.storage.unrollFraction`. Blocks evicted in the process, if any,
-   * are added to `evictedBlocks`.
-   *
-   * @return whether all N bytes were successfully granted.
-   */
   override def acquireUnrollMemory(
       blockId: BlockId,
       numBytes: Long,
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
-    val currentUnrollMemory = memoryStore.currentUnrollMemory
+    val currentUnrollMemory = storageMemoryPool.memoryStore.currentUnrollMemory
     val maxNumBytesToFree = math.max(0, maxMemoryToEvictForUnroll - currentUnrollMemory)
     val numBytesToFree = math.min(numBytes, maxNumBytesToFree)
-    acquireStorageMemory(blockId, numBytes, numBytesToFree, evictedBlocks)
+    storageMemoryPool.acquireMemory(blockId, numBytes, numBytesToFree, evictedBlocks)
   }
-
-  /**
-   * Acquire N bytes of storage memory for the given block, evicting existing ones if necessary.
-   *
-   * @param blockId the ID of the block we are acquiring storage memory for
-   * @param numBytesToAcquire the size of this block
-   * @param numBytesToFree the size of space to be freed through evicting blocks
-   * @param evictedBlocks a holder for blocks evicted in the process
-   * @return whether all N bytes were successfully granted.
-   */
-  private def acquireStorageMemory(
-      blockId: BlockId,
-      numBytesToAcquire: Long,
-      numBytesToFree: Long,
-      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
-    assert(numBytesToAcquire >= 0)
-    assert(numBytesToFree >= 0)
-    memoryStore.ensureFreeSpace(blockId, numBytesToFree, evictedBlocks)
-    assert(_storageMemoryUsed <= maxStorageMemory)
-    val enoughMemory = _storageMemoryUsed + numBytesToAcquire <= maxStorageMemory
-    if (enoughMemory) {
-      _storageMemoryUsed += numBytesToAcquire
-    }
-    enoughMemory
-  }
-
 }
 
 
@@ -135,7 +77,6 @@ private[spark] object StaticMemoryManager {
     (systemMaxMemory * memoryFraction * safetyFraction).toLong
   }
 
-
   /**
    * Return the total amount of memory available for the execution region, in bytes.
    */
diff --git a/core/src/main/scala/org/apache/spark/memory/StorageMemoryPool.scala b/core/src/main/scala/org/apache/spark/memory/StorageMemoryPool.scala
new file mode 100644
index 0000000000000..6a322eabf81ed
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/memory/StorageMemoryPool.scala
@@ -0,0 +1,138 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.memory
+
+import javax.annotation.concurrent.GuardedBy
+
+import scala.collection.mutable
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.spark.{TaskContext, Logging}
+import org.apache.spark.storage.{MemoryStore, BlockStatus, BlockId}
+
+/**
+ * Performs bookkeeping for managing an adjustable-size pool of memory that is used for storage
+ * (caching).
+ *
+ * @param lock a [[MemoryManager]] instance to synchronize on
+ */
+class StorageMemoryPool(lock: Object) extends MemoryPool(lock) with Logging {
+
+  @GuardedBy("lock")
+  private[this] var _memoryUsed: Long = 0L
+
+  override def memoryUsed: Long = lock.synchronized {
+    _memoryUsed
+  }
+
+  private var _memoryStore: MemoryStore = _
+  def memoryStore: MemoryStore = {
+    if (_memoryStore == null) {
+      throw new IllegalStateException("memory store not initialized yet")
+    }
+    _memoryStore
+  }
+
+  /**
+   * Set the [[MemoryStore]] used by this manager to evict cached blocks.
+   * This must be set after construction due to initialization ordering constraints.
+   */
+  final def setMemoryStore(store: MemoryStore): Unit = {
+    _memoryStore = store
+  }
+
+  /**
+   * Acquire N bytes of memory to cache the given block, evicting existing ones if necessary.
+   * Blocks evicted in the process, if any, are added to `evictedBlocks`.
+   * @return whether all N bytes were successfully granted.
+   */
+  def acquireMemory(
+      blockId: BlockId,
+      numBytes: Long,
+      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = lock.synchronized {
+    acquireMemory(blockId, numBytes, numBytes, evictedBlocks)
+  }
+
+  /**
+   * Acquire N bytes of storage memory for the given block, evicting existing ones if necessary.
+   *
+   * @param blockId the ID of the block we are acquiring storage memory for
+   * @param numBytesToAcquire the size of this block
+   * @param numBytesToFree the size of space to be freed through evicting blocks
+   * @return whether all N bytes were successfully granted.
+   */
+  def acquireMemory(
+      blockId: BlockId,
+      numBytesToAcquire: Long,
+      numBytesToFree: Long,
+      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = lock.synchronized {
+    assert(numBytesToAcquire >= 0)
+    assert(numBytesToFree >= 0)
+    assert(memoryUsed <= poolSize)
+    memoryStore.ensureFreeSpace(blockId, numBytesToFree, evictedBlocks)
+    // Register evicted blocks, if any, with the active task metrics
+    Option(TaskContext.get()).foreach { tc =>
+      val metrics = tc.taskMetrics()
+      val lastUpdatedBlocks = metrics.updatedBlocks.getOrElse(Seq[(BlockId, BlockStatus)]())
+      metrics.updatedBlocks = Some(lastUpdatedBlocks ++ evictedBlocks.toSeq)
+    }
+    // NOTE: If the memory store evicts blocks, then those evictions will synchronously call
+    // back into this StorageMemoryPool in order to free. Therefore, these variables should have
+    // been updated.
+    val enoughMemory = numBytesToAcquire <= memoryFree
+    if (enoughMemory) {
+      _memoryUsed += numBytesToAcquire
+    }
+    enoughMemory
+  }
+
+  def releaseMemory(size: Long): Unit = lock.synchronized {
+    if (size > _memoryUsed) {
+      logWarning(s"Attempted to release $size bytes of storage " +
+        s"memory when we only have ${_memoryUsed} bytes")
+      _memoryUsed = 0
+    } else {
+      _memoryUsed -= size
+    }
+  }
+
+  def releaseAllMemory(): Unit = lock.synchronized {
+    _memoryUsed = 0
+  }
+
+  /**
+   * Try to shrink the size of this storage memory pool by `spaceToFree` bytes. Return the number
+   * of bytes removed from the pool's capacity.
+   */
+  def shrinkPoolToFreeSpace(spaceToFree: Long): Long = lock.synchronized {
+    // First, shrink the pool by reclaiming free memory:
+    val spaceFreedByReleasingUnusedMemory = Math.min(spaceToFree, memoryFree)
+    decrementPoolSize(spaceFreedByReleasingUnusedMemory)
+    if (spaceFreedByReleasingUnusedMemory == spaceToFree) {
+      spaceFreedByReleasingUnusedMemory
+    } else {
+      // If reclaiming free memory did not adequately shrink the pool, begin evicting blocks:
+      val evictedBlocks = new ArrayBuffer[(BlockId, BlockStatus)]
+      memoryStore.ensureFreeSpace(spaceToFree - spaceFreedByReleasingUnusedMemory, evictedBlocks)
+      val spaceFreedByEviction = evictedBlocks.map(_._2.memSize).sum
+      _memoryUsed -= spaceFreedByEviction
+      decrementPoolSize(spaceFreedByEviction)
+      spaceFreedByReleasingUnusedMemory + spaceFreedByEviction
+    }
+  }
+}
diff --git a/core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala b/core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala
index a3093030a0f93..8be5b05419094 100644
--- a/core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala
+++ b/core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala
@@ -22,7 +22,6 @@ import scala.collection.mutable
 import org.apache.spark.SparkConf
 import org.apache.spark.storage.{BlockStatus, BlockId}
 
-
 /**
  * A [[MemoryManager]] that enforces a soft boundary between execution and storage such that
  * either side can borrow memory from the other.
@@ -41,98 +40,105 @@ import org.apache.spark.storage.{BlockStatus, BlockId}
  * The implication is that attempts to cache blocks may fail if execution has already eaten
  * up most of the storage space, in which case the new blocks will be evicted immediately
  * according to their respective storage levels.
+ *
+ * @param storageRegionSize Size of the storage region, in bytes.
+ *                          This region is not statically reserved; execution can borrow from
+ *                          it if necessary. Cached blocks can be evicted only if actual
+ *                          storage memory usage exceeds this region.
  */
-private[spark] class UnifiedMemoryManager(
+private[spark] class UnifiedMemoryManager private[memory] (
     conf: SparkConf,
     maxMemory: Long,
+    private val storageRegionSize: Long,
     numCores: Int)
-  extends MemoryManager(conf, numCores) {
-
-  def this(conf: SparkConf, numCores: Int) {
-    this(conf, UnifiedMemoryManager.getMaxMemory(conf), numCores)
-  }
-
-  /**
-   * Size of the storage region, in bytes.
-   *
-   * This region is not statically reserved; execution can borrow from it if necessary.
-   * Cached blocks can be evicted only if actual storage memory usage exceeds this region.
-   */
-  private val storageRegionSize: Long = {
-    (maxMemory * conf.getDouble("spark.memory.storageFraction", 0.5)).toLong
-  }
-
-  /**
-   * Total amount of memory, in bytes, not currently occupied by either execution or storage.
-   */
-  private def totalFreeMemory: Long = synchronized {
-    assert(_executionMemoryUsed <= maxMemory)
-    assert(_storageMemoryUsed <= maxMemory)
-    assert(_executionMemoryUsed + _storageMemoryUsed <= maxMemory)
-    maxMemory - _executionMemoryUsed - _storageMemoryUsed
-  }
+  extends MemoryManager(
+    conf,
+    numCores,
+    storageRegionSize,
+    maxMemory - storageRegionSize) {
 
-  /**
-   * Total available memory for execution, in bytes.
-   * In this model, this is equivalent to the amount of memory not occupied by storage.
-   */
-  override def maxExecutionMemory: Long = synchronized {
-    maxMemory - _storageMemoryUsed
-  }
+  // We always maintain this invariant:
+  assert(onHeapExecutionMemoryPool.poolSize + storageMemoryPool.poolSize == maxMemory)
 
-  /**
-   * Total available memory for storage, in bytes.
-   * In this model, this is equivalent to the amount of memory not occupied by execution.
-   */
   override def maxStorageMemory: Long = synchronized {
-    maxMemory - _executionMemoryUsed
+    maxMemory - onHeapExecutionMemoryPool.memoryUsed
   }
 
   /**
-   * Acquire N bytes of memory for execution, evicting cached blocks if necessary.
+   * Try to acquire up to `numBytes` of execution memory for the current task and return the
+   * number of bytes obtained, or 0 if none can be allocated.
    *
-   * This method evicts blocks only up to the amount of memory borrowed by storage.
-   * Blocks evicted in the process, if any, are added to `evictedBlocks`.
-   * @return number of bytes successfully granted (<= N).
+   * This call may block until there is enough free memory in some situations, to make sure each
+   * task has a chance to ramp up to at least 1 / 2N of the total memory pool (where N is the # of
+   * active tasks) before it is forced to spill. This can happen if the number of tasks increase
+   * but an older task had a lot of memory already.
    */
-  private[memory] override def doAcquireExecutionMemory(
+  override private[memory] def acquireExecutionMemory(
       numBytes: Long,
-      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Long = synchronized {
+      taskAttemptId: Long,
+      memoryMode: MemoryMode): Long = synchronized {
+    assert(onHeapExecutionMemoryPool.poolSize + storageMemoryPool.poolSize == maxMemory)
     assert(numBytes >= 0)
-    val memoryBorrowedByStorage = math.max(0, _storageMemoryUsed - storageRegionSize)
-    // If there is not enough free memory AND storage has borrowed some execution memory,
-    // then evict as much memory borrowed by storage as needed to grant this request
-    val shouldEvictStorage = totalFreeMemory < numBytes && memoryBorrowedByStorage > 0
-    if (shouldEvictStorage) {
-      val spaceToEnsure = math.min(numBytes, memoryBorrowedByStorage)
-      memoryStore.ensureFreeSpace(spaceToEnsure, evictedBlocks)
+    memoryMode match {
+      case MemoryMode.ON_HEAP =>
+        if (numBytes > onHeapExecutionMemoryPool.memoryFree) {
+          val extraMemoryNeeded = numBytes - onHeapExecutionMemoryPool.memoryFree
+          // There is not enough free memory in the execution pool, so try to reclaim memory from
+          // storage. We can reclaim any free memory from the storage pool. If the storage pool
+          // has grown to become larger than `storageRegionSize`, we can evict blocks and reclaim
+          // the memory that storage has borrowed from execution.
+          val memoryReclaimableFromStorage =
+            math.max(storageMemoryPool.memoryFree, storageMemoryPool.poolSize - storageRegionSize)
+          if (memoryReclaimableFromStorage > 0) {
+            // Only reclaim as much space as is necessary and available:
+            val spaceReclaimed = storageMemoryPool.shrinkPoolToFreeSpace(
+              math.min(extraMemoryNeeded, memoryReclaimableFromStorage))
+            onHeapExecutionMemoryPool.incrementPoolSize(spaceReclaimed)
+          }
+        }
+        onHeapExecutionMemoryPool.acquireMemory(numBytes, taskAttemptId)
+      case MemoryMode.OFF_HEAP =>
+        // For now, we only support on-heap caching of data, so we do not need to interact with
+        // the storage pool when allocating off-heap memory. This will change in the future, though.
+        super.acquireExecutionMemory(numBytes, taskAttemptId, memoryMode)
     }
-    val bytesToGrant = math.min(numBytes, totalFreeMemory)
-    _executionMemoryUsed += bytesToGrant
-    bytesToGrant
   }
 
-  /**
-   * Acquire N bytes of memory to cache the given block, evicting existing ones if necessary.
-   * Blocks evicted in the process, if any, are added to `evictedBlocks`.
-   * @return whether all N bytes were successfully granted.
-   */
   override def acquireStorageMemory(
       blockId: BlockId,
       numBytes: Long,
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
+    assert(onHeapExecutionMemoryPool.poolSize + storageMemoryPool.poolSize == maxMemory)
     assert(numBytes >= 0)
-    memoryStore.ensureFreeSpace(blockId, numBytes, evictedBlocks)
-    val enoughMemory = totalFreeMemory >= numBytes
-    if (enoughMemory) {
-      _storageMemoryUsed += numBytes
+    if (numBytes > storageMemoryPool.memoryFree) {
+      // There is not enough free memory in the storage pool, so try to borrow free memory from
+      // the execution pool.
+      val memoryBorrowedFromExecution = Math.min(onHeapExecutionMemoryPool.memoryFree, numBytes)
+      onHeapExecutionMemoryPool.decrementPoolSize(memoryBorrowedFromExecution)
+      storageMemoryPool.incrementPoolSize(memoryBorrowedFromExecution)
     }
-    enoughMemory
+    storageMemoryPool.acquireMemory(blockId, numBytes, evictedBlocks)
   }
 
+  override def acquireUnrollMemory(
+      blockId: BlockId,
+      numBytes: Long,
+      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = synchronized {
+    acquireStorageMemory(blockId, numBytes, evictedBlocks)
+  }
 }
 
-private object UnifiedMemoryManager {
+object UnifiedMemoryManager {
+
+  def apply(conf: SparkConf, numCores: Int): UnifiedMemoryManager = {
+    val maxMemory = getMaxMemory(conf)
+    new UnifiedMemoryManager(
+      conf,
+      maxMemory = maxMemory,
+      storageRegionSize =
+        (maxMemory * conf.getDouble("spark.memory.storageFraction", 0.5)).toLong,
+      numCores = numCores)
+  }
 
   /**
    * Return the total amount of memory shared between execution and storage, in bytes.
diff --git a/core/src/main/scala/org/apache/spark/memory/package.scala b/core/src/main/scala/org/apache/spark/memory/package.scala
new file mode 100644
index 0000000000000..564e30d2ffd66
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/memory/package.scala
@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark
+
+/**
+ * This package implements Spark's memory management system. This system consists of two main
+ * components, a JVM-wide memory manager and a per-task manager:
+ *
+ * - [[org.apache.spark.memory.MemoryManager]] manages Spark's overall memory usage within a JVM.
+ *   This component implements the policies for dividing the available memory across tasks and for
+ *   allocating memory between storage (memory used caching and data transfer) and execution (memory
+ *   used by computations, such as shuffles, joins, sorts, and aggregations).
+ * - [[org.apache.spark.memory.TaskMemoryManager]] manages the memory allocated by individual tasks.
+ *   Tasks interact with TaskMemoryManager and never directly interact with the JVM-wide
+ *   MemoryManager.
+ *
+ * Internally, each of these components have additional abstractions for memory bookkeeping:
+ *
+ *  - [[org.apache.spark.memory.MemoryConsumer]]s are clients of the TaskMemoryManager and
+ *    correspond to individual operators and data structures within a task. The TaskMemoryManager
+ *    receives memory allocation requests from MemoryConsumers and issues callbacks to consumers
+ *    in order to trigger spilling when running low on memory.
+ *  - [[org.apache.spark.memory.MemoryPool]]s are a bookkeeping abstraction used by the
+ *    MemoryManager to track the division of memory between storage and execution.
+ *
+ * Diagrammatically:
+ *
+ * {{{
+ *       +-------------+
+ *       | MemConsumer |----+                                   +------------------------+
+ *       +-------------+    |    +-------------------+          |     MemoryManager      |
+ *                          +--->| TaskMemoryManager |----+     |                        |
+ *       +-------------+    |    +-------------------+    |     |  +------------------+  |
+ *       | MemConsumer |----+                             |     |  |  StorageMemPool  |  |
+ *       +-------------+         +-------------------+    |     |  +------------------+  |
+ *                               | TaskMemoryManager |----+     |                        |
+ *                               +-------------------+    |     |  +------------------+  |
+ *                                                        +---->|  |OnHeapExecMemPool |  |
+ *                                        *               |     |  +------------------+  |
+ *                                        *               |     |                        |
+ *       +-------------+                  *               |     |  +------------------+  |
+ *       | MemConsumer |----+                             |     |  |OffHeapExecMemPool|  |
+ *       +-------------+    |    +-------------------+    |     |  +------------------+  |
+ *                          +--->| TaskMemoryManager |----+     |                        |
+ *                               +-------------------+          +------------------------+
+ * }}}
+ *
+ *
+ * There are two implementations of [[org.apache.spark.memory.MemoryManager]] which vary in how
+ * they handle the sizing of their memory pools:
+ *
+ *  - [[org.apache.spark.memory.UnifiedMemoryManager]], the default in Spark 1.6+, enforces soft
+ *    boundaries between storage and execution memory, allowing requests for memory in one region
+ *    to be fulfilled by borrowing memory from the other.
+ *  - [[org.apache.spark.memory.StaticMemoryManager]] enforces hard boundaries between storage
+ *    and execution memory by statically partitioning Spark's memory and preventing storage and
+ *    execution from borrowing memory from each other. This mode is retained only for legacy
+ *    compatibility purposes.
+ */
+package object memory
diff --git a/core/src/main/scala/org/apache/spark/util/collection/Spillable.scala b/core/src/main/scala/org/apache/spark/util/collection/Spillable.scala
index 9e002621a6909..3a48af82b1dae 100644
--- a/core/src/main/scala/org/apache/spark/util/collection/Spillable.scala
+++ b/core/src/main/scala/org/apache/spark/util/collection/Spillable.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.util.collection
 
-import org.apache.spark.memory.TaskMemoryManager
+import org.apache.spark.memory.{MemoryMode, TaskMemoryManager}
 import org.apache.spark.{Logging, SparkEnv}
 
 /**
@@ -78,7 +78,8 @@ private[spark] trait Spillable[C] extends Logging {
     if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
       // Claim up to double our current memory from the shuffle memory pool
       val amountToRequest = 2 * currentMemory - myMemoryThreshold
-      val granted = taskMemoryManager.acquireExecutionMemory(amountToRequest, null)
+      val granted =
+        taskMemoryManager.acquireExecutionMemory(amountToRequest, MemoryMode.ON_HEAP, null)
       myMemoryThreshold += granted
       // If we were granted too little memory to grow further (either tryToAcquire returned 0,
       // or we already had more memory than myMemoryThreshold), spill the current collection
@@ -107,7 +108,8 @@ private[spark] trait Spillable[C] extends Logging {
    */
   def releaseMemory(): Unit = {
     // The amount we requested does not include the initial memory tracking threshold
-    taskMemoryManager.releaseExecutionMemory(myMemoryThreshold - initialMemoryThreshold, null)
+    taskMemoryManager.releaseExecutionMemory(
+      myMemoryThreshold - initialMemoryThreshold, MemoryMode.ON_HEAP, null)
     myMemoryThreshold = initialMemoryThreshold
   }
 
diff --git a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
index c731317395612..711eed0193bc0 100644
--- a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
+++ b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java
@@ -28,8 +28,14 @@ public class TaskMemoryManagerSuite {
   @Test
   public void leakedPageMemoryIsDetected() {
     final TaskMemoryManager manager = new TaskMemoryManager(
-      new TestMemoryManager(new SparkConf().set("spark.unsafe.offHeap", "false")), 0);
+      new StaticMemoryManager(
+        new SparkConf().set("spark.unsafe.offHeap", "false"),
+        Long.MAX_VALUE,
+        Long.MAX_VALUE,
+        1),
+      0);
     manager.allocatePage(4096, null);  // leak memory
+    Assert.assertEquals(4096, manager.getMemoryConsumptionForThisTask());
     Assert.assertEquals(4096, manager.cleanUpAllAllocatedMemory());
   }
 
diff --git a/core/src/test/java/org/apache/spark/memory/TestMemoryConsumer.java b/core/src/test/java/org/apache/spark/memory/TestMemoryConsumer.java
index 8ae3642738509..e6e16fff80401 100644
--- a/core/src/test/java/org/apache/spark/memory/TestMemoryConsumer.java
+++ b/core/src/test/java/org/apache/spark/memory/TestMemoryConsumer.java
@@ -32,13 +32,19 @@ public long spill(long size, MemoryConsumer trigger) throws IOException {
   }
 
   void use(long size) {
-    long got = taskMemoryManager.acquireExecutionMemory(size, this);
+    long got = taskMemoryManager.acquireExecutionMemory(
+      size,
+      taskMemoryManager.tungstenMemoryMode,
+      this);
     used += got;
   }
 
   void free(long size) {
     used -= size;
-    taskMemoryManager.releaseExecutionMemory(size, this);
+    taskMemoryManager.releaseExecutionMemory(
+      size,
+      taskMemoryManager.tungstenMemoryMode,
+      this);
   }
 }
 
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
index 4763395d7d401..0e0eca515afc1 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
@@ -423,7 +423,7 @@ public void writeEnoughRecordsToTriggerSortBufferExpansionAndSpill() throws Exce
     memoryManager.limit(UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE * 16);
     final UnsafeShuffleWriter<Object, Object> writer = createWriter(false);
     final ArrayList<Product2<Object, Object>> dataToWrite = new ArrayList<>();
-    for (int i = 0; i < UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE; i++) {
+    for (int i = 0; i < UnsafeShuffleWriter.INITIAL_SORT_BUFFER_SIZE + 1; i++) {
       dataToWrite.add(new Tuple2<Object, Object>(i, i));
     }
     writer.write(dataToWrite.iterator());
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index 92bd45e5fa241..3bca790f30870 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -83,7 +83,9 @@ public OutputStream apply(OutputStream stream) {
   public void setup() {
     memoryManager =
       new TestMemoryManager(
-        new SparkConf().set("spark.unsafe.offHeap", "" + useOffHeapMemoryAllocator()));
+        new SparkConf()
+          .set("spark.unsafe.offHeap", "" + useOffHeapMemoryAllocator())
+          .set("spark.memory.offHeapSize", "256mb"));
     taskMemoryManager = new TaskMemoryManager(memoryManager, 0);
 
     tempDir = Utils.createTempDir(System.getProperty("java.io.tmpdir"), "unsafe-test");
diff --git a/core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala b/core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala
index 4a9479cf490fb..f55d435fa33a6 100644
--- a/core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala
@@ -19,6 +19,7 @@ package org.apache.spark.memory
 
 import java.util.concurrent.atomic.AtomicLong
 
+import scala.collection.mutable
 import scala.concurrent.duration.Duration
 import scala.concurrent.{Await, ExecutionContext, Future}
 
@@ -29,7 +30,7 @@ import org.mockito.stubbing.Answer
 import org.scalatest.time.SpanSugar._
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.storage.MemoryStore
+import org.apache.spark.storage.{BlockId, BlockStatus, MemoryStore, StorageLevel}
 
 
 /**
@@ -78,7 +79,12 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
         require(args(numBytesPos).isInstanceOf[Long], s"bad test: expected ensureFreeSpace " +
           s"argument at index $numBytesPos to be a Long: ${args.mkString(", ")}")
         val numBytes = args(numBytesPos).asInstanceOf[Long]
-        mockEnsureFreeSpace(mm, numBytes)
+        val success = mockEnsureFreeSpace(mm, numBytes)
+        if (success) {
+          args.last.asInstanceOf[mutable.Buffer[(BlockId, BlockStatus)]].append(
+            (null, BlockStatus(StorageLevel.MEMORY_ONLY, numBytes, 0L, 0L)))
+        }
+        success
       }
     }
   }
@@ -132,93 +138,95 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
   }
 
   /**
-   * Create a MemoryManager with the specified execution memory limit and no storage memory.
+   * Create a MemoryManager with the specified execution memory limits and no storage memory.
    */
-  protected def createMemoryManager(maxExecutionMemory: Long): MemoryManager
+  protected def createMemoryManager(
+     maxOnHeapExecutionMemory: Long,
+     maxOffHeapExecutionMemory: Long = 0L): MemoryManager
 
   // -- Tests of sharing of execution memory between tasks ----------------------------------------
   // Prior to Spark 1.6, these tests were part of ShuffleMemoryManagerSuite.
 
   implicit val ec = ExecutionContext.global
 
-  test("single task requesting execution memory") {
+  test("single task requesting on-heap execution memory") {
     val manager = createMemoryManager(1000L)
     val taskMemoryManager = new TaskMemoryManager(manager, 0)
 
-    assert(taskMemoryManager.acquireExecutionMemory(100L, null) === 100L)
-    assert(taskMemoryManager.acquireExecutionMemory(400L, null) === 400L)
-    assert(taskMemoryManager.acquireExecutionMemory(400L, null) === 400L)
-    assert(taskMemoryManager.acquireExecutionMemory(200L, null) === 100L)
-    assert(taskMemoryManager.acquireExecutionMemory(100L, null) === 0L)
-    assert(taskMemoryManager.acquireExecutionMemory(100L, null) === 0L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L, MemoryMode.ON_HEAP, null) === 100L)
+    assert(taskMemoryManager.acquireExecutionMemory(400L, MemoryMode.ON_HEAP, null) === 400L)
+    assert(taskMemoryManager.acquireExecutionMemory(400L, MemoryMode.ON_HEAP, null) === 400L)
+    assert(taskMemoryManager.acquireExecutionMemory(200L, MemoryMode.ON_HEAP, null) === 100L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L, MemoryMode.ON_HEAP, null) === 0L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L, MemoryMode.ON_HEAP, null) === 0L)
 
-    taskMemoryManager.releaseExecutionMemory(500L, null)
-    assert(taskMemoryManager.acquireExecutionMemory(300L, null) === 300L)
-    assert(taskMemoryManager.acquireExecutionMemory(300L, null) === 200L)
+    taskMemoryManager.releaseExecutionMemory(500L, MemoryMode.ON_HEAP, null)
+    assert(taskMemoryManager.acquireExecutionMemory(300L, MemoryMode.ON_HEAP, null) === 300L)
+    assert(taskMemoryManager.acquireExecutionMemory(300L, MemoryMode.ON_HEAP, null) === 200L)
 
     taskMemoryManager.cleanUpAllAllocatedMemory()
-    assert(taskMemoryManager.acquireExecutionMemory(1000L, null) === 1000L)
-    assert(taskMemoryManager.acquireExecutionMemory(100L, null) === 0L)
+    assert(taskMemoryManager.acquireExecutionMemory(1000L, MemoryMode.ON_HEAP, null) === 1000L)
+    assert(taskMemoryManager.acquireExecutionMemory(100L, MemoryMode.ON_HEAP, null) === 0L)
   }
 
-  test("two tasks requesting full execution memory") {
+  test("two tasks requesting full on-heap execution memory") {
     val memoryManager = createMemoryManager(1000L)
     val t1MemManager = new TaskMemoryManager(memoryManager, 1)
     val t2MemManager = new TaskMemoryManager(memoryManager, 2)
     val futureTimeout: Duration = 20.seconds
 
     // Have both tasks request 500 bytes, then wait until both requests have been granted:
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(500L, null) }
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t1Result1, futureTimeout) === 500L)
     assert(Await.result(t2Result1, futureTimeout) === 500L)
 
     // Have both tasks each request 500 bytes more; both should immediately return 0 as they are
     // both now at 1 / N
-    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L, null) }
-    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t1Result2, 200.millis) === 0L)
     assert(Await.result(t2Result2, 200.millis) === 0L)
   }
 
-  test("two tasks cannot grow past 1 / N of execution memory") {
+  test("two tasks cannot grow past 1 / N of on-heap execution memory") {
     val memoryManager = createMemoryManager(1000L)
     val t1MemManager = new TaskMemoryManager(memoryManager, 1)
     val t2MemManager = new TaskMemoryManager(memoryManager, 2)
     val futureTimeout: Duration = 20.seconds
 
     // Have both tasks request 250 bytes, then wait until both requests have been granted:
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(250L, null) }
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L, null) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(250L, MemoryMode.ON_HEAP, null) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t1Result1, futureTimeout) === 250L)
     assert(Await.result(t2Result1, futureTimeout) === 250L)
 
     // Have both tasks each request 500 bytes more.
     // We should only grant 250 bytes to each of them on this second request
-    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L, null) }
-    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t1Result2, futureTimeout) === 250L)
     assert(Await.result(t2Result2, futureTimeout) === 250L)
   }
 
-  test("tasks can block to get at least 1 / 2N of execution memory") {
+  test("tasks can block to get at least 1 / 2N of on-heap execution memory") {
     val memoryManager = createMemoryManager(1000L)
     val t1MemManager = new TaskMemoryManager(memoryManager, 1)
     val t2MemManager = new TaskMemoryManager(memoryManager, 2)
     val futureTimeout: Duration = 20.seconds
 
     // t1 grabs 1000 bytes and then waits until t2 is ready to make a request.
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L, null) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t1Result1, futureTimeout) === 1000L)
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L, null) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(250L, MemoryMode.ON_HEAP, null) }
     // Make sure that t2 didn't grab the memory right away. This is hacky but it would be difficult
     // to make sure the other thread blocks for some time otherwise.
     Thread.sleep(300)
-    t1MemManager.releaseExecutionMemory(250L, null)
+    t1MemManager.releaseExecutionMemory(250L, MemoryMode.ON_HEAP, null)
     // The memory freed from t1 should now be granted to t2.
     assert(Await.result(t2Result1, futureTimeout) === 250L)
     // Further requests by t2 should be denied immediately because it now has 1 / 2N of the memory.
-    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(100L, null) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(100L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t2Result2, 200.millis) === 0L)
   }
 
@@ -229,18 +237,18 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     val futureTimeout: Duration = 20.seconds
 
     // t1 grabs 1000 bytes and then waits until t2 is ready to make a request.
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L, null) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(1000L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t1Result1, futureTimeout) === 1000L)
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
     // Make sure that t2 didn't grab the memory right away. This is hacky but it would be difficult
     // to make sure the other thread blocks for some time otherwise.
     Thread.sleep(300)
     // t1 releases all of its memory, so t2 should be able to grab all of the memory
     t1MemManager.cleanUpAllAllocatedMemory()
     assert(Await.result(t2Result1, futureTimeout) === 500L)
-    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
+    val t2Result2 = Future { t2MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t2Result2, futureTimeout) === 500L)
-    val t2Result3 = Future { t2MemManager.acquireExecutionMemory(500L, null) }
+    val t2Result3 = Future { t2MemManager.acquireExecutionMemory(500L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t2Result3, 200.millis) === 0L)
   }
 
@@ -251,15 +259,35 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite {
     val t2MemManager = new TaskMemoryManager(memoryManager, 2)
     val futureTimeout: Duration = 20.seconds
 
-    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(700L, null) }
+    val t1Result1 = Future { t1MemManager.acquireExecutionMemory(700L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t1Result1, futureTimeout) === 700L)
 
-    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(300L, null) }
+    val t2Result1 = Future { t2MemManager.acquireExecutionMemory(300L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t2Result1, futureTimeout) === 300L)
 
-    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(300L, null) }
+    val t1Result2 = Future { t1MemManager.acquireExecutionMemory(300L, MemoryMode.ON_HEAP, null) }
     assert(Await.result(t1Result2, 200.millis) === 0L)
   }
+
+  test("off-heap execution allocations cannot exceed limit") {
+    val memoryManager = createMemoryManager(
+      maxOnHeapExecutionMemory = 0L,
+      maxOffHeapExecutionMemory = 1000L)
+
+    val tMemManager = new TaskMemoryManager(memoryManager, 1)
+    val result1 = Future { tMemManager.acquireExecutionMemory(1000L, MemoryMode.OFF_HEAP, null) }
+    assert(Await.result(result1, 200.millis) === 1000L)
+    assert(tMemManager.getMemoryConsumptionForThisTask === 1000L)
+
+    val result2 = Future { tMemManager.acquireExecutionMemory(300L, MemoryMode.OFF_HEAP, null) }
+    assert(Await.result(result2, 200.millis) === 0L)
+
+    assert(tMemManager.getMemoryConsumptionForThisTask === 1000L)
+    tMemManager.releaseExecutionMemory(500L, MemoryMode.OFF_HEAP, null)
+    assert(tMemManager.getMemoryConsumptionForThisTask === 500L)
+    tMemManager.releaseExecutionMemory(500L, MemoryMode.OFF_HEAP, null)
+    assert(tMemManager.getMemoryConsumptionForThisTask === 0L)
+  }
 }
 
 private object MemoryManagerSuite {
diff --git a/core/src/test/scala/org/apache/spark/memory/StaticMemoryManagerSuite.scala b/core/src/test/scala/org/apache/spark/memory/StaticMemoryManagerSuite.scala
index 885c450d6d4f5..54cb28c389c2f 100644
--- a/core/src/test/scala/org/apache/spark/memory/StaticMemoryManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/memory/StaticMemoryManagerSuite.scala
@@ -24,7 +24,6 @@ import org.mockito.Mockito.when
 import org.apache.spark.SparkConf
 import org.apache.spark.storage.{BlockId, BlockStatus, MemoryStore, TestBlockId}
 
-
 class StaticMemoryManagerSuite extends MemoryManagerSuite {
   private val conf = new SparkConf().set("spark.storage.unrollFraction", "0.4")
   private val evictedBlocks = new ArrayBuffer[(BlockId, BlockStatus)]
@@ -36,38 +35,47 @@ class StaticMemoryManagerSuite extends MemoryManagerSuite {
       maxExecutionMem: Long,
       maxStorageMem: Long): (StaticMemoryManager, MemoryStore) = {
     val mm = new StaticMemoryManager(
-      conf, maxExecutionMemory = maxExecutionMem, maxStorageMemory = maxStorageMem, numCores = 1)
+      conf,
+      maxOnHeapExecutionMemory = maxExecutionMem,
+      maxStorageMemory = maxStorageMem,
+      numCores = 1)
     val ms = makeMemoryStore(mm)
     (mm, ms)
   }
 
-  override protected def createMemoryManager(maxMemory: Long): MemoryManager = {
+  override protected def createMemoryManager(
+      maxOnHeapExecutionMemory: Long,
+      maxOffHeapExecutionMemory: Long): StaticMemoryManager = {
     new StaticMemoryManager(
-      conf,
-      maxExecutionMemory = maxMemory,
+      conf.clone
+        .set("spark.memory.fraction", "1")
+        .set("spark.testing.memory", maxOnHeapExecutionMemory.toString)
+        .set("spark.memory.offHeapSize", maxOffHeapExecutionMemory.toString),
+      maxOnHeapExecutionMemory = maxOnHeapExecutionMemory,
       maxStorageMemory = 0,
       numCores = 1)
   }
 
   test("basic execution memory") {
     val maxExecutionMem = 1000L
+    val taskAttemptId = 0L
     val (mm, _) = makeThings(maxExecutionMem, Long.MaxValue)
     assert(mm.executionMemoryUsed === 0L)
-    assert(mm.doAcquireExecutionMemory(10L, evictedBlocks) === 10L)
+    assert(mm.acquireExecutionMemory(10L, taskAttemptId, MemoryMode.ON_HEAP) === 10L)
     assert(mm.executionMemoryUsed === 10L)
-    assert(mm.doAcquireExecutionMemory(100L, evictedBlocks) === 100L)
+    assert(mm.acquireExecutionMemory(100L, taskAttemptId, MemoryMode.ON_HEAP) === 100L)
     // Acquire up to the max
-    assert(mm.doAcquireExecutionMemory(1000L, evictedBlocks) === 890L)
+    assert(mm.acquireExecutionMemory(1000L, taskAttemptId, MemoryMode.ON_HEAP) === 890L)
     assert(mm.executionMemoryUsed === maxExecutionMem)
-    assert(mm.doAcquireExecutionMemory(1L, evictedBlocks) === 0L)
+    assert(mm.acquireExecutionMemory(1L, taskAttemptId, MemoryMode.ON_HEAP) === 0L)
     assert(mm.executionMemoryUsed === maxExecutionMem)
-    mm.releaseExecutionMemory(800L)
+    mm.releaseExecutionMemory(800L, taskAttemptId, MemoryMode.ON_HEAP)
     assert(mm.executionMemoryUsed === 200L)
     // Acquire after release
-    assert(mm.doAcquireExecutionMemory(1L, evictedBlocks) === 1L)
+    assert(mm.acquireExecutionMemory(1L, taskAttemptId, MemoryMode.ON_HEAP) === 1L)
     assert(mm.executionMemoryUsed === 201L)
     // Release beyond what was acquired
-    mm.releaseExecutionMemory(maxExecutionMem)
+    mm.releaseExecutionMemory(maxExecutionMem, taskAttemptId, MemoryMode.ON_HEAP)
     assert(mm.executionMemoryUsed === 0L)
   }
 
@@ -113,13 +121,14 @@ class StaticMemoryManagerSuite extends MemoryManagerSuite {
   test("execution and storage isolation") {
     val maxExecutionMem = 200L
     val maxStorageMem = 1000L
+    val taskAttemptId = 0L
     val dummyBlock = TestBlockId("ain't nobody love like you do")
     val (mm, ms) = makeThings(maxExecutionMem, maxStorageMem)
     // Only execution memory should increase
-    assert(mm.doAcquireExecutionMemory(100L, evictedBlocks) === 100L)
+    assert(mm.acquireExecutionMemory(100L, taskAttemptId, MemoryMode.ON_HEAP) === 100L)
     assert(mm.storageMemoryUsed === 0L)
     assert(mm.executionMemoryUsed === 100L)
-    assert(mm.doAcquireExecutionMemory(1000L, evictedBlocks) === 100L)
+    assert(mm.acquireExecutionMemory(1000L, taskAttemptId, MemoryMode.ON_HEAP) === 100L)
     assert(mm.storageMemoryUsed === 0L)
     assert(mm.executionMemoryUsed === 200L)
     // Only storage memory should increase
@@ -128,7 +137,7 @@ class StaticMemoryManagerSuite extends MemoryManagerSuite {
     assert(mm.storageMemoryUsed === 50L)
     assert(mm.executionMemoryUsed === 200L)
     // Only execution memory should be released
-    mm.releaseExecutionMemory(133L)
+    mm.releaseExecutionMemory(133L, taskAttemptId, MemoryMode.ON_HEAP)
     assert(mm.storageMemoryUsed === 50L)
     assert(mm.executionMemoryUsed === 67L)
     // Only storage memory should be released
diff --git a/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala b/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
index 77e43554ee27c..0706a6e45de8f 100644
--- a/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
+++ b/core/src/test/scala/org/apache/spark/memory/TestMemoryManager.scala
@@ -22,19 +22,20 @@ import scala.collection.mutable
 import org.apache.spark.SparkConf
 import org.apache.spark.storage.{BlockStatus, BlockId}
 
-class TestMemoryManager(conf: SparkConf) extends MemoryManager(conf, numCores = 1) {
-  private[memory] override def doAcquireExecutionMemory(
+class TestMemoryManager(conf: SparkConf)
+  extends MemoryManager(conf, numCores = 1, Long.MaxValue, Long.MaxValue) {
+
+  override private[memory] def acquireExecutionMemory(
       numBytes: Long,
-      evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Long = synchronized {
+      taskAttemptId: Long,
+      memoryMode: MemoryMode): Long = {
     if (oomOnce) {
       oomOnce = false
       0
     } else if (available >= numBytes) {
-      _executionMemoryUsed += numBytes // To suppress warnings when freeing unallocated memory
       available -= numBytes
       numBytes
     } else {
-      _executionMemoryUsed += available
       val grant = available
       available = 0
       grant
@@ -48,12 +49,13 @@ class TestMemoryManager(conf: SparkConf) extends MemoryManager(conf, numCores =
       blockId: BlockId,
       numBytes: Long,
       evictedBlocks: mutable.Buffer[(BlockId, BlockStatus)]): Boolean = true
-  override def releaseExecutionMemory(numBytes: Long): Unit = {
+  override def releaseStorageMemory(numBytes: Long): Unit = {}
+  override private[memory] def releaseExecutionMemory(
+      numBytes: Long,
+      taskAttemptId: Long,
+      memoryMode: MemoryMode): Unit = {
     available += numBytes
-    _executionMemoryUsed -= numBytes
   }
-  override def releaseStorageMemory(numBytes: Long): Unit = {}
-  override def maxExecutionMemory: Long = Long.MaxValue
   override def maxStorageMemory: Long = Long.MaxValue
 
   private var oomOnce = false
diff --git a/core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala b/core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala
index 0c97f2bd89651..8cebe81c3bfff 100644
--- a/core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala
@@ -24,57 +24,52 @@ import org.scalatest.PrivateMethodTester
 import org.apache.spark.SparkConf
 import org.apache.spark.storage.{BlockId, BlockStatus, MemoryStore, TestBlockId}
 
-
 class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTester {
-  private val conf = new SparkConf().set("spark.memory.storageFraction", "0.5")
   private val dummyBlock = TestBlockId("--")
   private val evictedBlocks = new ArrayBuffer[(BlockId, BlockStatus)]
 
+  private val storageFraction: Double = 0.5
+
   /**
    * Make a [[UnifiedMemoryManager]] and a [[MemoryStore]] with limited class dependencies.
    */
   private def makeThings(maxMemory: Long): (UnifiedMemoryManager, MemoryStore) = {
-    val mm = new UnifiedMemoryManager(conf, maxMemory, numCores = 1)
+    val mm = createMemoryManager(maxMemory)
     val ms = makeMemoryStore(mm)
     (mm, ms)
   }
 
-  override protected def createMemoryManager(maxMemory: Long): MemoryManager = {
-    new UnifiedMemoryManager(conf, maxMemory, numCores = 1)
-  }
-
-  private def getStorageRegionSize(mm: UnifiedMemoryManager): Long = {
-    mm invokePrivate PrivateMethod[Long]('storageRegionSize)()
-  }
-
-  test("storage region size") {
-    val maxMemory = 1000L
-    val (mm, _) = makeThings(maxMemory)
-    val storageFraction = conf.get("spark.memory.storageFraction").toDouble
-    val expectedStorageRegionSize = maxMemory * storageFraction
-    val actualStorageRegionSize = getStorageRegionSize(mm)
-    assert(expectedStorageRegionSize === actualStorageRegionSize)
+  override protected def createMemoryManager(
+      maxOnHeapExecutionMemory: Long,
+      maxOffHeapExecutionMemory: Long): UnifiedMemoryManager = {
+    val conf = new SparkConf()
+      .set("spark.memory.fraction", "1")
+      .set("spark.testing.memory", maxOnHeapExecutionMemory.toString)
+      .set("spark.memory.offHeapSize", maxOffHeapExecutionMemory.toString)
+      .set("spark.memory.storageFraction", storageFraction.toString)
+    UnifiedMemoryManager(conf, numCores = 1)
   }
 
   test("basic execution memory") {
     val maxMemory = 1000L
+    val taskAttemptId = 0L
     val (mm, _) = makeThings(maxMemory)
     assert(mm.executionMemoryUsed === 0L)
-    assert(mm.doAcquireExecutionMemory(10L, evictedBlocks) === 10L)
+    assert(mm.acquireExecutionMemory(10L, taskAttemptId, MemoryMode.ON_HEAP) === 10L)
     assert(mm.executionMemoryUsed === 10L)
-    assert(mm.doAcquireExecutionMemory(100L, evictedBlocks) === 100L)
+    assert(mm.acquireExecutionMemory(100L, taskAttemptId, MemoryMode.ON_HEAP) === 100L)
     // Acquire up to the max
-    assert(mm.doAcquireExecutionMemory(1000L, evictedBlocks) === 890L)
+    assert(mm.acquireExecutionMemory(1000L, taskAttemptId, MemoryMode.ON_HEAP) === 890L)
     assert(mm.executionMemoryUsed === maxMemory)
-    assert(mm.doAcquireExecutionMemory(1L, evictedBlocks) === 0L)
+    assert(mm.acquireExecutionMemory(1L, taskAttemptId, MemoryMode.ON_HEAP) === 0L)
     assert(mm.executionMemoryUsed === maxMemory)
-    mm.releaseExecutionMemory(800L)
+    mm.releaseExecutionMemory(800L, taskAttemptId, MemoryMode.ON_HEAP)
     assert(mm.executionMemoryUsed === 200L)
     // Acquire after release
-    assert(mm.doAcquireExecutionMemory(1L, evictedBlocks) === 1L)
+    assert(mm.acquireExecutionMemory(1L, taskAttemptId, MemoryMode.ON_HEAP) === 1L)
     assert(mm.executionMemoryUsed === 201L)
     // Release beyond what was acquired
-    mm.releaseExecutionMemory(maxMemory)
+    mm.releaseExecutionMemory(maxMemory, taskAttemptId, MemoryMode.ON_HEAP)
     assert(mm.executionMemoryUsed === 0L)
   }
 
@@ -118,44 +113,34 @@ class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTes
 
   test("execution evicts storage") {
     val maxMemory = 1000L
+    val taskAttemptId = 0L
     val (mm, ms) = makeThings(maxMemory)
-    // First, ensure the test classes are set up as expected
-    val expectedStorageRegionSize = 500L
-    val expectedExecutionRegionSize = 500L
-    val storageRegionSize = getStorageRegionSize(mm)
-    val executionRegionSize = maxMemory - expectedStorageRegionSize
-    require(storageRegionSize === expectedStorageRegionSize,
-      "bad test: storage region size is unexpected")
-    require(executionRegionSize === expectedExecutionRegionSize,
-      "bad test: storage region size is unexpected")
     // Acquire enough storage memory to exceed the storage region
     assert(mm.acquireStorageMemory(dummyBlock, 750L, evictedBlocks))
     assertEnsureFreeSpaceCalled(ms, 750L)
     assert(mm.executionMemoryUsed === 0L)
     assert(mm.storageMemoryUsed === 750L)
-    require(mm.storageMemoryUsed > storageRegionSize,
-      s"bad test: storage memory used should exceed the storage region")
     // Execution needs to request 250 bytes to evict storage memory
-    assert(mm.doAcquireExecutionMemory(100L, evictedBlocks) === 100L)
+    assert(mm.acquireExecutionMemory(100L, taskAttemptId, MemoryMode.ON_HEAP) === 100L)
     assert(mm.executionMemoryUsed === 100L)
     assert(mm.storageMemoryUsed === 750L)
     assertEnsureFreeSpaceNotCalled(ms)
     // Execution wants 200 bytes but only 150 are free, so storage is evicted
-    assert(mm.doAcquireExecutionMemory(200L, evictedBlocks) === 200L)
-    assertEnsureFreeSpaceCalled(ms, 200L)
+    assert(mm.acquireExecutionMemory(200L, taskAttemptId, MemoryMode.ON_HEAP) === 200L)
+    assert(mm.executionMemoryUsed === 300L)
+    assertEnsureFreeSpaceCalled(ms, 50L)
     assert(mm.executionMemoryUsed === 300L)
     mm.releaseAllStorageMemory()
-    require(mm.executionMemoryUsed < executionRegionSize,
-      s"bad test: execution memory used should be within the execution region")
+    require(mm.executionMemoryUsed === 300L)
     require(mm.storageMemoryUsed === 0, "bad test: all storage memory should have been released")
     // Acquire some storage memory again, but this time keep it within the storage region
     assert(mm.acquireStorageMemory(dummyBlock, 400L, evictedBlocks))
     assertEnsureFreeSpaceCalled(ms, 400L)
-    require(mm.storageMemoryUsed < storageRegionSize,
-      s"bad test: storage memory used should be within the storage region")
+    assert(mm.storageMemoryUsed === 400L)
+    assert(mm.executionMemoryUsed === 300L)
     // Execution cannot evict storage because the latter is within the storage fraction,
     // so grant only what's remaining without evicting anything, i.e. 1000 - 300 - 400 = 300
-    assert(mm.doAcquireExecutionMemory(400L, evictedBlocks) === 300L)
+    assert(mm.acquireExecutionMemory(400L, taskAttemptId, MemoryMode.ON_HEAP) === 300L)
     assert(mm.executionMemoryUsed === 600L)
     assert(mm.storageMemoryUsed === 400L)
     assertEnsureFreeSpaceNotCalled(ms)
@@ -163,23 +148,13 @@ class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTes
 
   test("storage does not evict execution") {
     val maxMemory = 1000L
+    val taskAttemptId = 0L
     val (mm, ms) = makeThings(maxMemory)
-    // First, ensure the test classes are set up as expected
-    val expectedStorageRegionSize = 500L
-    val expectedExecutionRegionSize = 500L
-    val storageRegionSize = getStorageRegionSize(mm)
-    val executionRegionSize = maxMemory - expectedStorageRegionSize
-    require(storageRegionSize === expectedStorageRegionSize,
-      "bad test: storage region size is unexpected")
-    require(executionRegionSize === expectedExecutionRegionSize,
-      "bad test: storage region size is unexpected")
     // Acquire enough execution memory to exceed the execution region
-    assert(mm.doAcquireExecutionMemory(800L, evictedBlocks) === 800L)
+    assert(mm.acquireExecutionMemory(800L, taskAttemptId, MemoryMode.ON_HEAP) === 800L)
     assert(mm.executionMemoryUsed === 800L)
     assert(mm.storageMemoryUsed === 0L)
     assertEnsureFreeSpaceNotCalled(ms)
-    require(mm.executionMemoryUsed > executionRegionSize,
-      s"bad test: execution memory used should exceed the execution region")
     // Storage should not be able to evict execution
     assert(mm.acquireStorageMemory(dummyBlock, 100L, evictedBlocks))
     assert(mm.executionMemoryUsed === 800L)
@@ -189,15 +164,13 @@ class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTes
     assert(mm.executionMemoryUsed === 800L)
     assert(mm.storageMemoryUsed === 100L)
     assertEnsureFreeSpaceCalled(ms, 250L)
-    mm.releaseExecutionMemory(maxMemory)
+    mm.releaseExecutionMemory(maxMemory, taskAttemptId, MemoryMode.ON_HEAP)
     mm.releaseStorageMemory(maxMemory)
     // Acquire some execution memory again, but this time keep it within the execution region
-    assert(mm.doAcquireExecutionMemory(200L, evictedBlocks) === 200L)
+    assert(mm.acquireExecutionMemory(200L, taskAttemptId, MemoryMode.ON_HEAP) === 200L)
     assert(mm.executionMemoryUsed === 200L)
     assert(mm.storageMemoryUsed === 0L)
     assertEnsureFreeSpaceNotCalled(ms)
-    require(mm.executionMemoryUsed < executionRegionSize,
-      s"bad test: execution memory used should be within the execution region")
     // Storage should still not be able to evict execution
     assert(mm.acquireStorageMemory(dummyBlock, 750L, evictedBlocks))
     assert(mm.executionMemoryUsed === 200L)
diff --git a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala
index d49015afcd594..53991d8a1aede 100644
--- a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala
@@ -825,7 +825,7 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with BeforeAndAfterE
     val transfer = new NettyBlockTransferService(conf, securityMgr, numCores = 1)
     val memoryManager = new StaticMemoryManager(
       conf,
-      maxExecutionMemory = Long.MaxValue,
+      maxOnHeapExecutionMemory = Long.MaxValue,
       maxStorageMemory = 1200,
       numCores = 1)
     store = new BlockManager(SparkContext.DRIVER_IDENTIFIER, rpcEnv, master,

From 7f741905b06ed6d3dfbff6db41a3355dab71aa3c Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Sat, 7 Nov 2015 05:35:53 +0100
Subject: [PATCH 109/518] [SPARK-11112] DAG visualization: display RDD callsite

<img width="548" alt="screen shot 2015-11-01 at 9 42 33 am" src="https://cloud.githubusercontent.com/assets/2133137/10870343/2a8cd070-807d-11e5-857a-4ebcace77b5b.png">
mateiz sarutak

Author: Andrew Or <andrew@databricks.com>

Closes #9398 from andrewor14/rdd-callsite.
---
 .../apache/spark/ui/static/spark-dag-viz.css  |  4 ++
 .../org/apache/spark/storage/RDDInfo.scala    | 16 +++++++-
 .../spark/ui/scope/RDDOperationGraph.scala    | 10 +++--
 .../org/apache/spark/util/JsonProtocol.scala  | 17 ++++++++-
 .../scala/org/apache/spark/util/Utils.scala   |  1 +
 .../org/apache/spark/ui/UISeleniumSuite.scala | 14 +++----
 .../apache/spark/util/JsonProtocolSuite.scala | 37 ++++++++++++++++---
 7 files changed, 79 insertions(+), 20 deletions(-)

diff --git a/core/src/main/resources/org/apache/spark/ui/static/spark-dag-viz.css b/core/src/main/resources/org/apache/spark/ui/static/spark-dag-viz.css
index 3b4ae2ed354b8..9cc5c79f67346 100644
--- a/core/src/main/resources/org/apache/spark/ui/static/spark-dag-viz.css
+++ b/core/src/main/resources/org/apache/spark/ui/static/spark-dag-viz.css
@@ -122,3 +122,7 @@
   stroke: #52C366;
   stroke-width: 2px;
 }
+
+.tooltip-inner {
+  white-space: pre-wrap;
+}
diff --git a/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala b/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala
index 96062626b5045..3fa209b924170 100644
--- a/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala
+++ b/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala
@@ -19,7 +19,7 @@ package org.apache.spark.storage
 
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.rdd.{RDDOperationScope, RDD}
-import org.apache.spark.util.Utils
+import org.apache.spark.util.{CallSite, Utils}
 
 @DeveloperApi
 class RDDInfo(
@@ -28,9 +28,20 @@ class RDDInfo(
     val numPartitions: Int,
     var storageLevel: StorageLevel,
     val parentIds: Seq[Int],
+    val callSite: CallSite,
     val scope: Option[RDDOperationScope] = None)
   extends Ordered[RDDInfo] {
 
+  def this(
+      id: Int,
+      name: String,
+      numPartitions: Int,
+      storageLevel: StorageLevel,
+      parentIds: Seq[Int],
+      scope: Option[RDDOperationScope] = None) {
+    this(id, name, numPartitions, storageLevel, parentIds, CallSite.empty, scope)
+  }
+
   var numCachedPartitions = 0
   var memSize = 0L
   var diskSize = 0L
@@ -56,6 +67,7 @@ private[spark] object RDDInfo {
   def fromRdd(rdd: RDD[_]): RDDInfo = {
     val rddName = Option(rdd.name).getOrElse(Utils.getFormattedClassName(rdd))
     val parentIds = rdd.dependencies.map(_.rdd.id)
-    new RDDInfo(rdd.id, rddName, rdd.partitions.length, rdd.getStorageLevel, parentIds, rdd.scope)
+    new RDDInfo(rdd.id, rddName, rdd.partitions.length,
+      rdd.getStorageLevel, parentIds, rdd.creationSite, rdd.scope)
   }
 }
diff --git a/core/src/main/scala/org/apache/spark/ui/scope/RDDOperationGraph.scala b/core/src/main/scala/org/apache/spark/ui/scope/RDDOperationGraph.scala
index 81f168a447ead..24274562657b3 100644
--- a/core/src/main/scala/org/apache/spark/ui/scope/RDDOperationGraph.scala
+++ b/core/src/main/scala/org/apache/spark/ui/scope/RDDOperationGraph.scala
@@ -23,6 +23,7 @@ import scala.collection.mutable.{StringBuilder, ListBuffer}
 import org.apache.spark.Logging
 import org.apache.spark.scheduler.StageInfo
 import org.apache.spark.storage.StorageLevel
+import org.apache.spark.util.CallSite
 
 /**
  * A representation of a generic cluster graph used for storing information on RDD operations.
@@ -38,7 +39,7 @@ private[ui] case class RDDOperationGraph(
     rootCluster: RDDOperationCluster)
 
 /** A node in an RDDOperationGraph. This represents an RDD. */
-private[ui] case class RDDOperationNode(id: Int, name: String, cached: Boolean)
+private[ui] case class RDDOperationNode(id: Int, name: String, cached: Boolean, callsite: CallSite)
 
 /**
  * A directed edge connecting two nodes in an RDDOperationGraph.
@@ -104,8 +105,8 @@ private[ui] object RDDOperationGraph extends Logging {
       edges ++= rdd.parentIds.map { parentId => RDDOperationEdge(parentId, rdd.id) }
 
       // TODO: differentiate between the intention to cache an RDD and whether it's actually cached
-      val node = nodes.getOrElseUpdate(
-        rdd.id, RDDOperationNode(rdd.id, rdd.name, rdd.storageLevel != StorageLevel.NONE))
+      val node = nodes.getOrElseUpdate(rdd.id, RDDOperationNode(
+        rdd.id, rdd.name, rdd.storageLevel != StorageLevel.NONE, rdd.callSite))
 
       if (rdd.scope.isEmpty) {
         // This RDD has no encompassing scope, so we put it directly in the root cluster
@@ -177,7 +178,8 @@ private[ui] object RDDOperationGraph extends Logging {
 
   /** Return the dot representation of a node in an RDDOperationGraph. */
   private def makeDotNode(node: RDDOperationNode): String = {
-    s"""${node.id} [label="${node.name} [${node.id}]"]"""
+    val label = s"${node.name} [${node.id}]\n${node.callsite.shortForm}"
+    s"""${node.id} [label="$label"]"""
   }
 
   /** Update the dot representation of the RDDOperationGraph in cluster to subgraph. */
diff --git a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
index ee2eb58cf5e2a..c9beeb25e05af 100644
--- a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
+++ b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
@@ -398,6 +398,7 @@ private[spark] object JsonProtocol {
     ("RDD ID" -> rddInfo.id) ~
     ("Name" -> rddInfo.name) ~
     ("Scope" -> rddInfo.scope.map(_.toJson)) ~
+    ("Callsite" -> callsiteToJson(rddInfo.callSite)) ~
     ("Parent IDs" -> parentIds) ~
     ("Storage Level" -> storageLevel) ~
     ("Number of Partitions" -> rddInfo.numPartitions) ~
@@ -407,6 +408,11 @@ private[spark] object JsonProtocol {
     ("Disk Size" -> rddInfo.diskSize)
   }
 
+  def callsiteToJson(callsite: CallSite): JValue = {
+    ("Short Form" -> callsite.shortForm) ~
+    ("Long Form" -> callsite.longForm)
+  }
+
   def storageLevelToJson(storageLevel: StorageLevel): JValue = {
     ("Use Disk" -> storageLevel.useDisk) ~
     ("Use Memory" -> storageLevel.useMemory) ~
@@ -851,6 +857,9 @@ private[spark] object JsonProtocol {
     val scope = Utils.jsonOption(json \ "Scope")
       .map(_.extract[String])
       .map(RDDOperationScope.fromJson)
+    val callsite = Utils.jsonOption(json \ "Callsite")
+      .map(callsiteFromJson)
+      .getOrElse(CallSite.empty)
     val parentIds = Utils.jsonOption(json \ "Parent IDs")
       .map { l => l.extract[List[JValue]].map(_.extract[Int]) }
       .getOrElse(Seq.empty)
@@ -863,7 +872,7 @@ private[spark] object JsonProtocol {
       .getOrElse(json \ "Tachyon Size").extract[Long]
     val diskSize = (json \ "Disk Size").extract[Long]
 
-    val rddInfo = new RDDInfo(rddId, name, numPartitions, storageLevel, parentIds, scope)
+    val rddInfo = new RDDInfo(rddId, name, numPartitions, storageLevel, parentIds, callsite, scope)
     rddInfo.numCachedPartitions = numCachedPartitions
     rddInfo.memSize = memSize
     rddInfo.externalBlockStoreSize = externalBlockStoreSize
@@ -871,6 +880,12 @@ private[spark] object JsonProtocol {
     rddInfo
   }
 
+  def callsiteFromJson(json: JValue): CallSite = {
+    val shortForm = (json \ "Short Form").extract[String]
+    val longForm = (json \ "Long Form").extract[String]
+    CallSite(shortForm, longForm)
+  }
+
   def storageLevelFromJson(json: JValue): StorageLevel = {
     val useDisk = (json \ "Use Disk").extract[Boolean]
     val useMemory = (json \ "Use Memory").extract[Boolean]
diff --git a/core/src/main/scala/org/apache/spark/util/Utils.scala b/core/src/main/scala/org/apache/spark/util/Utils.scala
index 5a976ee839b1e..316c194ff3454 100644
--- a/core/src/main/scala/org/apache/spark/util/Utils.scala
+++ b/core/src/main/scala/org/apache/spark/util/Utils.scala
@@ -57,6 +57,7 @@ private[spark] case class CallSite(shortForm: String, longForm: String)
 private[spark] object CallSite {
   val SHORT_FORM = "callSite.short"
   val LONG_FORM = "callSite.long"
+  val empty = CallSite("", "")
 }
 
 /**
diff --git a/core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala b/core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala
index 18eec7da9763e..ceecfd665bf87 100644
--- a/core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala
@@ -615,29 +615,29 @@ class UISeleniumSuite extends SparkFunSuite with WebBrowser with Matchers with B
       assert(stage0.contains("digraph G {\n  subgraph clusterstage_0 {\n    " +
         "label=&quot;Stage 0&quot;;\n    subgraph "))
       assert(stage0.contains("{\n      label=&quot;parallelize&quot;;\n      " +
-        "0 [label=&quot;ParallelCollectionRDD [0]&quot;];\n    }"))
+        "0 [label=&quot;ParallelCollectionRDD [0]"))
       assert(stage0.contains("{\n      label=&quot;map&quot;;\n      " +
-        "1 [label=&quot;MapPartitionsRDD [1]&quot;];\n    }"))
+        "1 [label=&quot;MapPartitionsRDD [1]"))
       assert(stage0.contains("{\n      label=&quot;groupBy&quot;;\n      " +
-        "2 [label=&quot;MapPartitionsRDD [2]&quot;];\n    }"))
+        "2 [label=&quot;MapPartitionsRDD [2]"))
 
       val stage1 = Source.fromURL(sc.ui.get.appUIAddress +
         "/stages/stage/?id=1&attempt=0&expandDagViz=true").mkString
       assert(stage1.contains("digraph G {\n  subgraph clusterstage_1 {\n    " +
         "label=&quot;Stage 1&quot;;\n    subgraph "))
       assert(stage1.contains("{\n      label=&quot;groupBy&quot;;\n      " +
-        "3 [label=&quot;ShuffledRDD [3]&quot;];\n    }"))
+        "3 [label=&quot;ShuffledRDD [3]"))
       assert(stage1.contains("{\n      label=&quot;map&quot;;\n      " +
-        "4 [label=&quot;MapPartitionsRDD [4]&quot;];\n    }"))
+        "4 [label=&quot;MapPartitionsRDD [4]"))
       assert(stage1.contains("{\n      label=&quot;groupBy&quot;;\n      " +
-        "5 [label=&quot;MapPartitionsRDD [5]&quot;];\n    }"))
+        "5 [label=&quot;MapPartitionsRDD [5]"))
 
       val stage2 = Source.fromURL(sc.ui.get.appUIAddress +
         "/stages/stage/?id=2&attempt=0&expandDagViz=true").mkString
       assert(stage2.contains("digraph G {\n  subgraph clusterstage_2 {\n    " +
         "label=&quot;Stage 2&quot;;\n    subgraph "))
       assert(stage2.contains("{\n      label=&quot;groupBy&quot;;\n      " +
-        "6 [label=&quot;ShuffledRDD [6]&quot;];\n    }"))
+        "6 [label=&quot;ShuffledRDD [6]"))
     }
   }
 
diff --git a/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala b/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala
index 953456c2caa89..3f94ef7041914 100644
--- a/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala
@@ -111,6 +111,7 @@ class JsonProtocolSuite extends SparkFunSuite {
   test("Dependent Classes") {
     val logUrlMap = Map("stderr" -> "mystderr", "stdout" -> "mystdout").toMap
     testRDDInfo(makeRddInfo(2, 3, 4, 5L, 6L))
+    testCallsite(CallSite("happy", "birthday"))
     testStageInfo(makeStageInfo(10, 20, 30, 40L, 50L))
     testTaskInfo(makeTaskInfo(999L, 888, 55, 777L, false))
     testTaskMetrics(makeTaskMetrics(
@@ -163,6 +164,10 @@ class JsonProtocolSuite extends SparkFunSuite {
     testBlockId(StreamBlockId(1, 2L))
   }
 
+  /* ============================== *
+   |  Backward compatibility tests  |
+   * ============================== */
+
   test("ExceptionFailure backward compatibility") {
     val exceptionFailure = ExceptionFailure("To be", "or not to be", stackTrace, null,
       None, None)
@@ -334,14 +339,17 @@ class JsonProtocolSuite extends SparkFunSuite {
     assertEquals(expectedJobEnd, JsonProtocol.jobEndFromJson(oldEndEvent))
   }
 
-  test("RDDInfo backward compatibility (scope, parent IDs)") {
-    // Prior to Spark 1.4.0, RDDInfo did not have the "Scope" and "Parent IDs" properties
-    val rddInfo = new RDDInfo(
-      1, "one", 100, StorageLevel.NONE, Seq(1, 6, 8), Some(new RDDOperationScope("fable")))
+  test("RDDInfo backward compatibility (scope, parent IDs, callsite)") {
+    // "Scope" and "Parent IDs" were introduced in Spark 1.4.0
+    // "Callsite" was introduced in Spark 1.6.0
+    val rddInfo = new RDDInfo(1, "one", 100, StorageLevel.NONE, Seq(1, 6, 8),
+      CallSite("short", "long"), Some(new RDDOperationScope("fable")))
     val oldRddInfoJson = JsonProtocol.rddInfoToJson(rddInfo)
       .removeField({ _._1 == "Parent IDs"})
       .removeField({ _._1 == "Scope"})
-    val expectedRddInfo = new RDDInfo(1, "one", 100, StorageLevel.NONE, Seq.empty, scope = None)
+      .removeField({ _._1 == "Callsite"})
+    val expectedRddInfo = new RDDInfo(
+      1, "one", 100, StorageLevel.NONE, Seq.empty, CallSite.empty, scope = None)
     assertEquals(expectedRddInfo, JsonProtocol.rddInfoFromJson(oldRddInfoJson))
   }
 
@@ -389,6 +397,11 @@ class JsonProtocolSuite extends SparkFunSuite {
     assertEquals(info, newInfo)
   }
 
+  private def testCallsite(callsite: CallSite): Unit = {
+    val newCallsite = JsonProtocol.callsiteFromJson(JsonProtocol.callsiteToJson(callsite))
+    assert(callsite === newCallsite)
+  }
+
   private def testStageInfo(info: StageInfo) {
     val newInfo = JsonProtocol.stageInfoFromJson(JsonProtocol.stageInfoToJson(info))
     assertEquals(info, newInfo)
@@ -713,7 +726,8 @@ class JsonProtocolSuite extends SparkFunSuite {
   }
 
   private def makeRddInfo(a: Int, b: Int, c: Int, d: Long, e: Long) = {
-    val r = new RDDInfo(a, "mayor", b, StorageLevel.MEMORY_AND_DISK, Seq(1, 4, 7))
+    val r = new RDDInfo(a, "mayor", b, StorageLevel.MEMORY_AND_DISK,
+      Seq(1, 4, 7), CallSite(a.toString, b.toString))
     r.numCachedPartitions = c
     r.memSize = d
     r.diskSize = e
@@ -856,6 +870,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |      {
       |        "RDD ID": 101,
       |        "Name": "mayor",
+      |        "Callsite": {"Short Form": "101", "Long Form": "201"},
       |        "Parent IDs": [1, 4, 7],
       |        "Storage Level": {
       |          "Use Disk": true,
@@ -1258,6 +1273,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 1,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "1", "Long Form": "200"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1301,6 +1317,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 2,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "2", "Long Form": "400"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1318,6 +1335,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 3,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "3", "Long Form": "401"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1361,6 +1379,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 3,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "3", "Long Form": "600"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1378,6 +1397,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 4,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "4", "Long Form": "601"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1395,6 +1415,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 5,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "5", "Long Form": "602"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1438,6 +1459,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 4,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "4", "Long Form": "800"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1455,6 +1477,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 5,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "5", "Long Form": "801"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1472,6 +1495,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 6,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "6", "Long Form": "802"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,
@@ -1489,6 +1513,7 @@ class JsonProtocolSuite extends SparkFunSuite {
       |        {
       |          "RDD ID": 7,
       |          "Name": "mayor",
+      |          "Callsite": {"Short Form": "7", "Long Form": "803"},
       |          "Parent IDs": [1, 4, 7],
       |          "Storage Level": {
       |            "Use Disk": true,

From 2ff0e79a8647cca5c9c57f613a07e739ac4f677e Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Fri, 6 Nov 2015 22:56:29 -0800
Subject: [PATCH 110/518] [SPARK-8467] [MLLIB] [PYSPARK] Add
 LDAModel.describeTopics() in Python

Could jkbradley and davies review it?

- Create a wrapper class: `LDAModelWrapper` for `LDAModel`. Because we can't deal with the return value of`describeTopics` in Scala from pyspark directly. `Array[(Array[Int], Array[Double])]` is too complicated to convert it.
- Add `loadLDAModel` in `PythonMLlibAPI`. Since `LDAModel` in Scala is an abstract class and we need to call `load` of `DistributedLDAModel`.

[[SPARK-8467] Add LDAModel.describeTopics() in Python - ASF JIRA](https://issues.apache.org/jira/browse/SPARK-8467)

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>

Closes #8643 from yu-iskw/SPARK-8467-2.
---
 .../mllib/api/python/LDAModelWrapper.scala    | 46 +++++++++++++++++++
 .../mllib/api/python/PythonMLLibAPI.scala     | 13 +++++-
 python/pyspark/mllib/clustering.py            | 33 +++++++------
 3 files changed, 75 insertions(+), 17 deletions(-)
 create mode 100644 mllib/src/main/scala/org/apache/spark/mllib/api/python/LDAModelWrapper.scala

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/api/python/LDAModelWrapper.scala b/mllib/src/main/scala/org/apache/spark/mllib/api/python/LDAModelWrapper.scala
new file mode 100644
index 0000000000000..63282eee6e656
--- /dev/null
+++ b/mllib/src/main/scala/org/apache/spark/mllib/api/python/LDAModelWrapper.scala
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.mllib.api.python
+
+import scala.collection.JavaConverters
+
+import org.apache.spark.SparkContext
+import org.apache.spark.mllib.clustering.LDAModel
+import org.apache.spark.mllib.linalg.Matrix
+
+/**
+ * Wrapper around LDAModel to provide helper methods in Python
+ */
+private[python] class LDAModelWrapper(model: LDAModel) {
+
+  def topicsMatrix(): Matrix = model.topicsMatrix
+
+  def vocabSize(): Int = model.vocabSize
+
+  def describeTopics(): Array[Byte] = describeTopics(this.model.vocabSize)
+
+  def describeTopics(maxTermsPerTopic: Int): Array[Byte] = {
+    val topics = model.describeTopics(maxTermsPerTopic).map { case (terms, termWeights) =>
+      val jTerms = JavaConverters.seqAsJavaListConverter(terms).asJava
+      val jTermWeights = JavaConverters.seqAsJavaListConverter(termWeights).asJava
+      Array[Any](jTerms, jTermWeights)
+    }
+    SerDe.dumps(JavaConverters.seqAsJavaListConverter(topics).asJava)
+  }
+
+  def save(sc: SparkContext, path: String): Unit = model.save(sc, path)
+}
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
index 40c41806cdfea..54b03a9f90283 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
@@ -517,7 +517,7 @@ private[python] class PythonMLLibAPI extends Serializable {
       topicConcentration: Double,
       seed: java.lang.Long,
       checkpointInterval: Int,
-      optimizer: String): LDAModel = {
+      optimizer: String): LDAModelWrapper = {
     val algo = new LDA()
       .setK(k)
       .setMaxIterations(maxIterations)
@@ -535,7 +535,16 @@ private[python] class PythonMLLibAPI extends Serializable {
         case _ => throw new IllegalArgumentException("input values contains invalid type value.")
       }
     }
-    algo.run(documents)
+    val model = algo.run(documents)
+    new LDAModelWrapper(model)
+  }
+
+  /**
+   * Load a LDA model
+   */
+  def loadLDAModel(jsc: JavaSparkContext, path: String): LDAModelWrapper = {
+    val model = DistributedLDAModel.load(jsc.sc, path)
+    new LDAModelWrapper(model)
   }
 
 
diff --git a/python/pyspark/mllib/clustering.py b/python/pyspark/mllib/clustering.py
index 8629aa5a17164..12081f8c69075 100644
--- a/python/pyspark/mllib/clustering.py
+++ b/python/pyspark/mllib/clustering.py
@@ -671,7 +671,7 @@ def predictOnValues(self, dstream):
         return dstream.mapValues(lambda x: self._model.predict(x))
 
 
-class LDAModel(JavaModelWrapper):
+class LDAModel(JavaModelWrapper, JavaSaveable, Loader):
 
     """ A clustering model derived from the LDA method.
 
@@ -691,9 +691,14 @@ class LDAModel(JavaModelWrapper):
     ...     [2, SparseVector(2, {0: 1.0})],
     ... ]
     >>> rdd =  sc.parallelize(data)
-    >>> model = LDA.train(rdd, k=2)
+    >>> model = LDA.train(rdd, k=2, seed=1)
     >>> model.vocabSize()
     2
+    >>> model.describeTopics()
+    [([1, 0], [0.5..., 0.49...]), ([0, 1], [0.5..., 0.49...])]
+    >>> model.describeTopics(1)
+    [([1], [0.5...]), ([0], [0.5...])]
+
     >>> topics = model.topicsMatrix()
     >>> topics_expect = array([[0.5,  0.5], [0.5, 0.5]])
     >>> assert_almost_equal(topics, topics_expect, 1)
@@ -724,18 +729,17 @@ def vocabSize(self):
         """Vocabulary size (number of terms or terms in the vocabulary)"""
         return self.call("vocabSize")
 
-    @since('1.5.0')
-    def save(self, sc, path):
-        """Save the LDAModel on to disk.
+    @since('1.6.0')
+    def describeTopics(self, maxTermsPerTopic=None):
+        """Return the topics described by weighted terms.
 
-        :param sc: SparkContext
-        :param path: str, path to where the model needs to be stored.
+        WARNING: If vocabSize and k are large, this can return a large object!
         """
-        if not isinstance(sc, SparkContext):
-            raise TypeError("sc should be a SparkContext, got type %s" % type(sc))
-        if not isinstance(path, basestring):
-            raise TypeError("path should be a basestring, got type %s" % type(path))
-        self._java_model.save(sc._jsc.sc(), path)
+        if maxTermsPerTopic is None:
+            topics = self.call("describeTopics")
+        else:
+            topics = self.call("describeTopics", maxTermsPerTopic)
+        return topics
 
     @classmethod
     @since('1.5.0')
@@ -749,9 +753,8 @@ def load(cls, sc, path):
             raise TypeError("sc should be a SparkContext, got type %s" % type(sc))
         if not isinstance(path, basestring):
             raise TypeError("path should be a basestring, got type %s" % type(path))
-        java_model = sc._jvm.org.apache.spark.mllib.clustering.DistributedLDAModel.load(
-            sc._jsc.sc(), path)
-        return cls(java_model)
+        model = callMLlibFunc("loadLDAModel", sc, path)
+        return LDAModel(model)
 
 
 class LDA(object):

From ef362846eb448769bcf774fc9090a5013d459464 Mon Sep 17 00:00:00 2001
From: Herman van Hovell <hvanhovell@questtec.nl>
Date: Sat, 7 Nov 2015 13:37:37 -0800
Subject: [PATCH 111/518] [SPARK-9241][SQL] Supporting multiple DISTINCT
 columns - follow-up

This PR is a follow up for PR https://github.com/apache/spark/pull/9406. It adds more documentation to the rewriting rule, removes a redundant if expression in the non-distinct aggregation path and adds a multiple distinct test to the AggregationQuerySuite.

cc yhuai marmbrus

Author: Herman van Hovell <hvanhovell@questtec.nl>

Closes #9541 from hvanhovell/SPARK-9241-followup.
---
 .../expressions/aggregate/Utils.scala         | 114 ++++++++++++++----
 .../execution/AggregationQuerySuite.scala     |  17 +++
 2 files changed, 108 insertions(+), 23 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
index 39010c3be6d4e..ac23f727829b6 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
@@ -222,10 +222,76 @@ object Utils {
  * aggregation in which the regular aggregation expressions and every distinct clause is aggregated
  * in a separate group. The results are then combined in a second aggregate.
  *
- * TODO Expression cannocalization
- * TODO Eliminate foldable expressions from distinct clauses.
- * TODO This eliminates all distinct expressions. We could safely pass one to the aggregate
- *      operator. Perhaps this is a good thing? It is much simpler to plan later on...
+ * For example (in scala):
+ * {{{
+ *   val data = Seq(
+ *     ("a", "ca1", "cb1", 10),
+ *     ("a", "ca1", "cb2", 5),
+ *     ("b", "ca1", "cb1", 13))
+ *     .toDF("key", "cat1", "cat2", "value")
+ *   data.registerTempTable("data")
+ *
+ *   val agg = data.groupBy($"key")
+ *     .agg(
+ *       countDistinct($"cat1").as("cat1_cnt"),
+ *       countDistinct($"cat2").as("cat2_cnt"),
+ *       sum($"value").as("total"))
+ * }}}
+ *
+ * This translates to the following (pseudo) logical plan:
+ * {{{
+ * Aggregate(
+ *    key = ['key]
+ *    functions = [COUNT(DISTINCT 'cat1),
+ *                 COUNT(DISTINCT 'cat2),
+ *                 sum('value)]
+ *    output = ['key, 'cat1_cnt, 'cat2_cnt, 'total])
+ *   LocalTableScan [...]
+ * }}}
+ *
+ * This rule rewrites this logical plan to the following (pseudo) logical plan:
+ * {{{
+ * Aggregate(
+ *    key = ['key]
+ *    functions = [count(if (('gid = 1)) 'cat1 else null),
+ *                 count(if (('gid = 2)) 'cat2 else null),
+ *                 first(if (('gid = 0)) 'total else null) ignore nulls]
+ *    output = ['key, 'cat1_cnt, 'cat2_cnt, 'total])
+ *   Aggregate(
+ *      key = ['key, 'cat1, 'cat2, 'gid]
+ *      functions = [sum('value)]
+ *      output = ['key, 'cat1, 'cat2, 'gid, 'total])
+ *     Expand(
+ *        projections = [('key, null, null, 0, cast('value as bigint)),
+ *                       ('key, 'cat1, null, 1, null),
+ *                       ('key, null, 'cat2, 2, null)]
+ *        output = ['key, 'cat1, 'cat2, 'gid, 'value])
+ *       LocalTableScan [...]
+ * }}}
+ *
+ * The rule does the following things here:
+ * 1. Expand the data. There are three aggregation groups in this query:
+ *    i. the non-distinct group;
+ *    ii. the distinct 'cat1 group;
+ *    iii. the distinct 'cat2 group.
+ *    An expand operator is inserted to expand the child data for each group. The expand will null
+ *    out all unused columns for the given group; this must be done in order to ensure correctness
+ *    later on. Groups can by identified by a group id (gid) column added by the expand operator.
+ * 2. De-duplicate the distinct paths and aggregate the non-aggregate path. The group by clause of
+ *    this aggregate consists of the original group by clause, all the requested distinct columns
+ *    and the group id. Both de-duplication of distinct column and the aggregation of the
+ *    non-distinct group take advantage of the fact that we group by the group id (gid) and that we
+ *    have nulled out all non-relevant columns for the the given group.
+ * 3. Aggregating the distinct groups and combining this with the results of the non-distinct
+ *    aggregation. In this step we use the group id to filter the inputs for the aggregate
+ *    functions. The result of the non-distinct group are 'aggregated' by using the first operator,
+ *    it might be more elegant to use the native UDAF merge mechanism for this in the future.
+ *
+ * This rule duplicates the input data by two or more times (# distinct groups + an optional
+ * non-distinct group). This will put quite a bit of memory pressure of the used aggregate and
+ * exchange operators. Keeping the number of distinct groups as low a possible should be priority,
+ * we could improve this in the current rule by applying more advanced expression cannocalization
+ * techniques.
  */
 object MultipleDistinctRewriter extends Rule[LogicalPlan] {
 
@@ -261,11 +327,10 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
       // Functions used to modify aggregate functions and their inputs.
       def evalWithinGroup(id: Literal, e: Expression) = If(EqualTo(gid, id), e, nullify(e))
       def patchAggregateFunctionChildren(
-          af: AggregateFunction2,
-          id: Literal,
-          attrs: Map[Expression, Expression]): AggregateFunction2 = {
-        af.withNewChildren(af.children.map { case afc =>
-          evalWithinGroup(id, attrs(afc))
+          af: AggregateFunction2)(
+          attrs: Expression => Expression): AggregateFunction2 = {
+        af.withNewChildren(af.children.map {
+          case afc => attrs(afc)
         }).asInstanceOf[AggregateFunction2]
       }
 
@@ -288,7 +353,9 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
           // Final aggregate
           val operators = expressions.map { e =>
             val af = e.aggregateFunction
-            val naf = patchAggregateFunctionChildren(af, id, distinctAggChildAttrMap)
+            val naf = patchAggregateFunctionChildren(af) { x =>
+              evalWithinGroup(id, distinctAggChildAttrMap(x))
+            }
             (e, e.copy(aggregateFunction = naf, isDistinct = false))
           }
 
@@ -304,26 +371,27 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
       val regularGroupId = Literal(0)
       val regularAggOperatorMap = regularAggExprs.map { e =>
         // Perform the actual aggregation in the initial aggregate.
-        val af = patchAggregateFunctionChildren(
-          e.aggregateFunction,
-          regularGroupId,
-          regularAggChildAttrMap)
-        val a = Alias(e.copy(aggregateFunction = af), e.toString)()
-
-        // Get the result of the first aggregate in the last aggregate.
-        val b = AggregateExpression2(
-          aggregate.First(evalWithinGroup(regularGroupId, a.toAttribute), Literal(true)),
+        val af = patchAggregateFunctionChildren(e.aggregateFunction)(regularAggChildAttrMap)
+        val operator = Alias(e.copy(aggregateFunction = af), e.toString)()
+
+        // Select the result of the first aggregate in the last aggregate.
+        val result = AggregateExpression2(
+          aggregate.First(evalWithinGroup(regularGroupId, operator.toAttribute), Literal(true)),
           mode = Complete,
           isDistinct = false)
 
         // Some aggregate functions (COUNT) have the special property that they can return a
         // non-null result without any input. We need to make sure we return a result in this case.
-        val c = af.defaultResult match {
-          case Some(lit) => Coalesce(Seq(b, lit))
-          case None => b
+        val resultWithDefault = af.defaultResult match {
+          case Some(lit) => Coalesce(Seq(result, lit))
+          case None => result
         }
 
-        (e, a, c)
+        // Return a Tuple3 containing:
+        // i. The original aggregate expression (used for look ups).
+        // ii. The actual aggregation operator (used in the first aggregate).
+        // iii. The operator that selects and returns the result (used in the second aggregate).
+        (e, operator, resultWithDefault)
       }
 
       // Construct the regular aggregate input projection only if we need one.
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index ea80060e370e0..7f6fe339232ad 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -516,6 +516,23 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
         Row(3, 4, 4, 3, null) :: Nil)
   }
 
+  test("multiple distinct column sets") {
+    checkAnswer(
+      sqlContext.sql(
+        """
+          |SELECT
+          |  key,
+          |  count(distinct value1),
+          |  count(distinct value2)
+          |FROM agg2
+          |GROUP BY key
+        """.stripMargin),
+      Row(null, 3, 3) ::
+        Row(1, 2, 3) ::
+        Row(2, 2, 1) ::
+        Row(3, 0, 1) :: Nil)
+  }
+
   test("test count") {
     checkAnswer(
       sqlContext.sql(

From 4b69a42eda3aff08eb7437c353fe2cc87ed67181 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Sat, 7 Nov 2015 19:44:45 -0800
Subject: [PATCH 112/518] [SPARK-11362] [SQL] Use Spark BitSet in
 BroadcastNestedLoopJoin

JIRA: https://issues.apache.org/jira/browse/SPARK-11362

We use scala.collection.mutable.BitSet in BroadcastNestedLoopJoin now. We should use Spark's BitSet.

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #9316 from viirya/use-spark-bitset.
---
 .../joins/BroadcastNestedLoopJoin.scala        | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoin.scala
index 05d20f511aef8..aab177b2e8427 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoin.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.execution.{BinaryNode, SparkPlan}
 import org.apache.spark.sql.execution.metric.SQLMetrics
-import org.apache.spark.util.collection.CompactBuffer
+import org.apache.spark.util.collection.{BitSet, CompactBuffer}
 
 
 case class BroadcastNestedLoopJoin(
@@ -95,9 +95,7 @@ case class BroadcastNestedLoopJoin(
     /** All rows that either match both-way, or rows from streamed joined with nulls. */
     val matchesOrStreamedRowsWithNulls = streamed.execute().mapPartitions { streamedIter =>
       val matchedRows = new CompactBuffer[InternalRow]
-      // TODO: Use Spark's BitSet.
-      val includedBroadcastTuples =
-        new scala.collection.mutable.BitSet(broadcastedRelation.value.size)
+      val includedBroadcastTuples = new BitSet(broadcastedRelation.value.size)
       val joinedRow = new JoinedRow
 
       val leftNulls = new GenericMutableRow(left.output.size)
@@ -115,11 +113,11 @@ case class BroadcastNestedLoopJoin(
             case BuildRight if boundCondition(joinedRow(streamedRow, broadcastedRow)) =>
               matchedRows += resultProj(joinedRow(streamedRow, broadcastedRow)).copy()
               streamRowMatched = true
-              includedBroadcastTuples += i
+              includedBroadcastTuples.set(i)
             case BuildLeft if boundCondition(joinedRow(broadcastedRow, streamedRow)) =>
               matchedRows += resultProj(joinedRow(broadcastedRow, streamedRow)).copy()
               streamRowMatched = true
-              includedBroadcastTuples += i
+              includedBroadcastTuples.set(i)
             case _ =>
           }
           i += 1
@@ -138,8 +136,8 @@ case class BroadcastNestedLoopJoin(
 
     val includedBroadcastTuples = matchesOrStreamedRowsWithNulls.map(_._2)
     val allIncludedBroadcastTuples = includedBroadcastTuples.fold(
-      new scala.collection.mutable.BitSet(broadcastedRelation.value.size)
-    )(_ ++ _)
+      new BitSet(broadcastedRelation.value.size)
+    )(_ | _)
 
     val leftNulls = new GenericMutableRow(left.output.size)
     val rightNulls = new GenericMutableRow(right.output.size)
@@ -155,7 +153,7 @@ case class BroadcastNestedLoopJoin(
           val joinedRow = new JoinedRow
           joinedRow.withLeft(leftNulls)
           while (i < rel.length) {
-            if (!allIncludedBroadcastTuples.contains(i)) {
+            if (!allIncludedBroadcastTuples.get(i)) {
               buf += resultProj(joinedRow.withRight(rel(i))).copy()
             }
             i += 1
@@ -164,7 +162,7 @@ case class BroadcastNestedLoopJoin(
           val joinedRow = new JoinedRow
           joinedRow.withRight(rightNulls)
           while (i < rel.length) {
-            if (!allIncludedBroadcastTuples.contains(i)) {
+            if (!allIncludedBroadcastTuples.get(i)) {
               buf += resultProj(joinedRow.withLeft(rel(i))).copy()
             }
             i += 1

From d981902101767b32dc83a5a639311e197f5cbcc1 Mon Sep 17 00:00:00 2001
From: Sean Owen <sowen@cloudera.com>
Date: Sun, 8 Nov 2015 11:15:58 +0000
Subject: [PATCH 113/518] [SPARK-11476][DOCS] Incorrect function referred to in
 MLib Random data generation documentation

Fix Python example to use normalRDD as advertised

Author: Sean Owen <sowen@cloudera.com>

Closes #9529 from srowen/SPARK-11476.
---
 docs/mllib-statistics.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/mllib-statistics.md b/docs/mllib-statistics.md
index 2c7c9ed693fd4..ade5b0768aefe 100644
--- a/docs/mllib-statistics.md
+++ b/docs/mllib-statistics.md
@@ -594,7 +594,7 @@ sc = ... # SparkContext
 
 # Generate a random double RDD that contains 1 million i.i.d. values drawn from the
 # standard normal distribution `N(0, 1)`, evenly distributed in 10 partitions.
-u = RandomRDDs.uniformRDD(sc, 1000000L, 10)
+u = RandomRDDs.normalRDD(sc, 1000000L, 10)
 # Apply a transform to get a random double RDD following `N(1, 4)`.
 v = u.map(lambda x: 1.0 + 2.0 * x)
 {% endhighlight %}

From 5c4e6d7ec9157c02494a382dfb49e7fbde3be222 Mon Sep 17 00:00:00 2001
From: Rohit Agarwal <mindprince@gmail.com>
Date: Sun, 8 Nov 2015 14:24:26 +0000
Subject: [PATCH 114/518] [DOC][SQL] Remove redundant out-of-place python
 snippet

This snippet seems to be mistakenly introduced at two places in #5348.

Author: Rohit Agarwal <mindprince@gmail.com>

Closes #9540 from mindprince/patch-1.
---
 docs/sql-programming-guide.md | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/docs/sql-programming-guide.md b/docs/sql-programming-guide.md
index 2fe5c36338899..085874133d968 100644
--- a/docs/sql-programming-guide.md
+++ b/docs/sql-programming-guide.md
@@ -1089,15 +1089,6 @@ for (teenName in collect(teenNames)) {
 
 </div>
 
-<div data-lang="python"  markdown="1">
-
-{% highlight python %}
-# sqlContext is an existing HiveContext
-sqlContext.sql("REFRESH TABLE my_table")
-{% endhighlight %}
-
-</div>
-
 <div data-lang="sql"  markdown="1">
 
 {% highlight sql %}

From 30c8ba71a76788cbc6916bc1ba6bc8522925fc2b Mon Sep 17 00:00:00 2001
From: Herman van Hovell <hvanhovell@questtec.nl>
Date: Sun, 8 Nov 2015 11:06:10 -0800
Subject: [PATCH 115/518] [SPARK-11451][SQL] Support single distinct count on
 multiple columns.

This PR adds support for multiple column in a single count distinct aggregate to the new aggregation path.

cc yhuai

Author: Herman van Hovell <hvanhovell@questtec.nl>

Closes #9409 from hvanhovell/SPARK-11451.
---
 .../expressions/aggregate/Utils.scala         | 44 +++++++++++--------
 .../expressions/conditionalExpressions.scala  | 30 ++++++++++++-
 .../plans/logical/basicOperators.scala        |  3 ++
 .../ConditionalExpressionSuite.scala          | 14 ++++++
 .../spark/sql/DataFrameAggregateSuite.scala   | 25 +++++++++++
 .../execution/AggregationQuerySuite.scala     | 37 +++++++++++++---
 6 files changed, 127 insertions(+), 26 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
index ac23f727829b6..9b22ce2619731 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
@@ -22,26 +22,27 @@ import org.apache.spark.sql.catalyst._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.{Expand, Aggregate, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.types.{IntegerType, StructType, MapType, ArrayType}
+import org.apache.spark.sql.types._
 
 /**
  * Utility functions used by the query planner to convert our plan to new aggregation code path.
  */
 object Utils {
-  // Right now, we do not support complex types in the grouping key schema.
-  private def supportsGroupingKeySchema(aggregate: Aggregate): Boolean = {
-    val hasComplexTypes = aggregate.groupingExpressions.map(_.dataType).exists {
-      case array: ArrayType => true
-      case map: MapType => true
-      case struct: StructType => true
-      case _ => false
-    }
 
-    !hasComplexTypes
+  // Check if the DataType given cannot be part of a group by clause.
+  private def isUnGroupable(dt: DataType): Boolean = dt match {
+    case _: ArrayType | _: MapType => true
+    case s: StructType => s.fields.exists(f => isUnGroupable(f.dataType))
+    case _ => false
   }
 
+  // Right now, we do not support complex types in the grouping key schema.
+  private def supportsGroupingKeySchema(aggregate: Aggregate): Boolean =
+    !aggregate.groupingExpressions.exists(e => isUnGroupable(e.dataType))
+
   private def doConvert(plan: LogicalPlan): Option[Aggregate] = plan match {
     case p: Aggregate if supportsGroupingKeySchema(p) =>
+
       val converted = MultipleDistinctRewriter.rewrite(p.transformExpressionsDown {
         case expressions.Average(child) =>
           aggregate.AggregateExpression2(
@@ -55,10 +56,14 @@ object Utils {
             mode = aggregate.Complete,
             isDistinct = false)
 
-        // We do not support multiple COUNT DISTINCT columns for now.
-        case expressions.CountDistinct(children) if children.length == 1 =>
+        case expressions.CountDistinct(children) =>
+          val child = if (children.size > 1) {
+            DropAnyNull(CreateStruct(children))
+          } else {
+            children.head
+          }
           aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Count(children.head),
+            aggregateFunction = aggregate.Count(child),
             mode = aggregate.Complete,
             isDistinct = true)
 
@@ -320,7 +325,7 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
       val gid = new AttributeReference("gid", IntegerType, false)()
       val groupByMap = a.groupingExpressions.collect {
         case ne: NamedExpression => ne -> ne.toAttribute
-        case e => e -> new AttributeReference(e.prettyName, e.dataType, e.nullable)()
+        case e => e -> new AttributeReference(e.prettyString, e.dataType, e.nullable)()
       }
       val groupByAttrs = groupByMap.map(_._2)
 
@@ -365,14 +370,15 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
       // Setup expand for the 'regular' aggregate expressions.
       val regularAggExprs = aggExpressions.filter(!_.isDistinct)
       val regularAggChildren = regularAggExprs.flatMap(_.aggregateFunction.children).distinct
-      val regularAggChildAttrMap = regularAggChildren.map(expressionAttributePair).toMap
+      val regularAggChildAttrMap = regularAggChildren.map(expressionAttributePair)
 
       // Setup aggregates for 'regular' aggregate expressions.
       val regularGroupId = Literal(0)
+      val regularAggChildAttrLookup = regularAggChildAttrMap.toMap
       val regularAggOperatorMap = regularAggExprs.map { e =>
         // Perform the actual aggregation in the initial aggregate.
-        val af = patchAggregateFunctionChildren(e.aggregateFunction)(regularAggChildAttrMap)
-        val operator = Alias(e.copy(aggregateFunction = af), e.toString)()
+        val af = patchAggregateFunctionChildren(e.aggregateFunction)(regularAggChildAttrLookup)
+        val operator = Alias(e.copy(aggregateFunction = af), e.prettyString)()
 
         // Select the result of the first aggregate in the last aggregate.
         val result = AggregateExpression2(
@@ -416,7 +422,7 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
       // Construct the expand operator.
       val expand = Expand(
         regularAggProjection ++ distinctAggProjections,
-        groupByAttrs ++ distinctAggChildAttrs ++ Seq(gid) ++ regularAggChildAttrMap.values.toSeq,
+        groupByAttrs ++ distinctAggChildAttrs ++ Seq(gid) ++ regularAggChildAttrMap.map(_._2),
         a.child)
 
       // Construct the first aggregate operator. This de-duplicates the all the children of
@@ -457,5 +463,5 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
     // NamedExpression. This is done to prevent collisions between distinct and regular aggregate
     // children, in this case attribute reuse causes the input of the regular aggregate to bound to
     // the (nulled out) input of the distinct aggregate.
-    e -> new AttributeReference(e.prettyName, e.dataType, true)()
+    e -> new AttributeReference(e.prettyString, e.dataType, true)()
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
index d532629984bec..0d4af43978ea1 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
@@ -21,7 +21,7 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.util.TypeUtils
-import org.apache.spark.sql.types.{NullType, BooleanType, DataType}
+import org.apache.spark.sql.types._
 
 
 case class If(predicate: Expression, trueValue: Expression, falseValue: Expression)
@@ -419,3 +419,31 @@ case class Greatest(children: Seq[Expression]) extends Expression {
     """
   }
 }
+
+/** Operator that drops a row when it contains any nulls. */
+case class DropAnyNull(child: Expression) extends UnaryExpression with ExpectsInputTypes {
+  override def nullable: Boolean = true
+  override def dataType: DataType = child.dataType
+  override def inputTypes: Seq[AbstractDataType] = Seq(StructType)
+
+  protected override def nullSafeEval(input: Any): InternalRow = {
+    val row = input.asInstanceOf[InternalRow]
+    if (row.anyNull) {
+      null
+    } else {
+      row
+    }
+  }
+
+  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    nullSafeCodeGen(ctx, ev, eval => {
+      s"""
+        if ($eval.anyNull()) {
+          ${ev.isNull} = true;
+        } else {
+          ${ev.value} = $eval;
+        }
+      """
+    })
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index fb963e2f8f7e7..09aac00a455f9 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -306,6 +306,9 @@ case class Expand(
     output: Seq[Attribute],
     child: LogicalPlan) extends UnaryNode {
 
+  override def references: AttributeSet =
+    AttributeSet(projections.flatten.flatMap(_.references))
+
   override def statistics: Statistics = {
     // TODO shouldn't we factor in the size of the projection versus the size of the backing child
     //      row?
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConditionalExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConditionalExpressionSuite.scala
index 0df673bb9fa02..c1e3c17b87102 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConditionalExpressionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConditionalExpressionSuite.scala
@@ -231,4 +231,18 @@ class ConditionalExpressionSuite extends SparkFunSuite with ExpressionEvalHelper
       checkConsistencyBetweenInterpretedAndCodegen(Greatest, dt, 2)
     }
   }
+
+  test("function dropAnyNull") {
+    val drop = DropAnyNull(CreateStruct(Seq('a.string.at(0), 'b.string.at(1))))
+    val a = create_row("a", "q")
+    val nullStr: String = null
+    checkEvaluation(drop, a, a)
+    checkEvaluation(drop, null, create_row("b", nullStr))
+    checkEvaluation(drop, null, create_row(nullStr, nullStr))
+
+    val row = 'r.struct(
+      StructField("a", StringType, false),
+      StructField("b", StringType, true)).at(0)
+    checkEvaluation(DropAnyNull(row), null, create_row(null))
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index 2e679e7bc4e0a..eb1ee266c5d28 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -162,6 +162,31 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
     )
   }
 
+  test("multiple column distinct count") {
+    val df1 = Seq(
+      ("a", "b", "c"),
+      ("a", "b", "c"),
+      ("a", "b", "d"),
+      ("x", "y", "z"),
+      ("x", "q", null.asInstanceOf[String]))
+      .toDF("key1", "key2", "key3")
+
+    checkAnswer(
+      df1.agg(countDistinct('key1, 'key2)),
+      Row(3)
+    )
+
+    checkAnswer(
+      df1.agg(countDistinct('key1, 'key2, 'key3)),
+      Row(3)
+    )
+
+    checkAnswer(
+      df1.groupBy('key1).agg(countDistinct('key2, 'key3)),
+      Seq(Row("a", 2), Row("x", 1))
+    )
+  }
+
   test("zero count") {
     val emptyTableData = Seq.empty[(Int, Int)].toDF("a", "b")
     checkAnswer(
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 7f6fe339232ad..ea36c132bb190 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -516,21 +516,46 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
         Row(3, 4, 4, 3, null) :: Nil)
   }
 
-  test("multiple distinct column sets") {
+  test("single distinct multiple columns set") {
+    checkAnswer(
+      sqlContext.sql(
+        """
+          |SELECT
+          |  key,
+          |  count(distinct value1, value2)
+          |FROM agg2
+          |GROUP BY key
+        """.stripMargin),
+      Row(null, 3) ::
+        Row(1, 3) ::
+        Row(2, 1) ::
+        Row(3, 0) :: Nil)
+  }
+
+  test("multiple distinct multiple columns sets") {
     checkAnswer(
       sqlContext.sql(
         """
           |SELECT
           |  key,
           |  count(distinct value1),
-          |  count(distinct value2)
+          |  sum(distinct value1),
+          |  count(distinct value2),
+          |  sum(distinct value2),
+          |  count(distinct value1, value2),
+          |  count(value1),
+          |  sum(value1),
+          |  count(value2),
+          |  sum(value2),
+          |  count(*),
+          |  count(1)
           |FROM agg2
           |GROUP BY key
         """.stripMargin),
-      Row(null, 3, 3) ::
-        Row(1, 2, 3) ::
-        Row(2, 2, 1) ::
-        Row(3, 0, 1) :: Nil)
+      Row(null, 3, 30, 3, 60, 3, 3, 30, 3, 60, 4, 4) ::
+        Row(1, 2, 40, 3, -10, 3, 3, 70, 3, -10, 3, 3) ::
+        Row(2, 2, 0, 1, 1, 1, 3, 1, 3, 3, 4, 4) ::
+        Row(3, 0, null, 1, 3, 0, 0, null, 1, 3, 2, 2) :: Nil)
   }
 
   test("test count") {

From 26739059bc39cd7aa7e0b1c16089c1cf8d8e4d7d Mon Sep 17 00:00:00 2001
From: xin Wu <xinwu@us.ibm.com>
Date: Sun, 8 Nov 2015 12:28:19 -0800
Subject: [PATCH 116/518] =?UTF-8?q?[SPARK-10046][SQL]=20Hive=20warehouse?=
 =?UTF-8?q?=20dir=20not=20set=20in=20current=20directory=20when=20not=20?=
 =?UTF-8?q?=E2=80=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Doc change to align with HiveConf default in terms of where to create `warehouse` directory.

Author: xin Wu <xinwu@us.ibm.com>

Closes #9365 from xwu0226/spark-10046-commit.
---
 docs/sql-programming-guide.md | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/docs/sql-programming-guide.md b/docs/sql-programming-guide.md
index 085874133d968..52e03b951f966 100644
--- a/docs/sql-programming-guide.md
+++ b/docs/sql-programming-guide.md
@@ -1627,8 +1627,10 @@ YARN cluster. The convenient way to do this is adding them through the `--jars`
 When working with Hive one must construct a `HiveContext`, which inherits from `SQLContext`, and
 adds support for finding tables in the MetaStore and writing queries using HiveQL. Users who do
 not have an existing Hive deployment can still create a `HiveContext`.  When not configured by the
-hive-site.xml, the context automatically creates `metastore_db` and `warehouse` in the current
-directory.
+hive-site.xml, the context automatically creates `metastore_db` in the current directory and
+creates `warehouse` directory indicated by HiveConf, which defaults to `/user/hive/warehouse`.
+Note that you may need to grant write privilege on `/user/hive/warehouse` to the user who starts
+the spark application.
 
 {% highlight scala %}
 // sc is an existing SparkContext.

From b2d195e137fad88d567974659fa7023ff4da96cd Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Sun, 8 Nov 2015 12:59:35 -0800
Subject: [PATCH 117/518] [SPARK-11554][SQL] add map/flatMap to GroupedDataset

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9521 from cloud-fan/map.
---
 .../plans/logical/basicOperators.scala        |  4 +-
 .../org/apache/spark/sql/GroupedDataset.scala | 29 ++++++++++++--
 .../spark/sql/execution/basicOperators.scala  |  2 +-
 .../apache/spark/sql/JavaDatasetSuite.java    | 16 ++++----
 .../spark/sql/DatasetPrimitiveSuite.scala     | 16 ++++++--
 .../org/apache/spark/sql/DatasetSuite.scala   | 40 +++++++++----------
 6 files changed, 70 insertions(+), 37 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index 09aac00a455f9..e151ac04ede2a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -494,7 +494,7 @@ case class AppendColumn[T, U](
 /** Factory for constructing new `MapGroups` nodes. */
 object MapGroups {
   def apply[K : Encoder, T : Encoder, U : Encoder](
-      func: (K, Iterator[T]) => Iterator[U],
+      func: (K, Iterator[T]) => TraversableOnce[U],
       groupingAttributes: Seq[Attribute],
       child: LogicalPlan): MapGroups[K, T, U] = {
     new MapGroups(
@@ -514,7 +514,7 @@ object MapGroups {
  * object representation of all the rows with that key.
  */
 case class MapGroups[K, T, U](
-    func: (K, Iterator[T]) => Iterator[U],
+    func: (K, Iterator[T]) => TraversableOnce[U],
     kEncoder: ExpressionEncoder[K],
     tEncoder: ExpressionEncoder[T],
     uEncoder: ExpressionEncoder[U],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index b2803d5a9a1e3..5c3f626545875 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -102,16 +102,39 @@ class GroupedDataset[K, T] private[sql](
    * (for example, by calling `toList`) unless they are sure that this is possible given the memory
    * constraints of their cluster.
    */
-  def mapGroups[U : Encoder](f: (K, Iterator[T]) => Iterator[U]): Dataset[U] = {
+  def flatMap[U : Encoder](f: (K, Iterator[T]) => TraversableOnce[U]): Dataset[U] = {
     new Dataset[U](
       sqlContext,
       MapGroups(f, groupingAttributes, logicalPlan))
   }
 
-  def mapGroups[U](
+  def flatMap[U](
       f: JFunction2[K, JIterator[T], JIterator[U]],
       encoder: Encoder[U]): Dataset[U] = {
-    mapGroups((key, data) => f.call(key, data.asJava).asScala)(encoder)
+    flatMap((key, data) => f.call(key, data.asJava).asScala)(encoder)
+  }
+
+  /**
+   * Applies the given function to each group of data.  For each unique group, the function will
+   * be passed the group key and an iterator that contains all of the elements in the group. The
+   * function can return an element of arbitrary type which will be returned as a new [[Dataset]].
+   *
+   * Internally, the implementation will spill to disk if any given group is too large to fit into
+   * memory.  However, users must take care to avoid materializing the whole iterator for a group
+   * (for example, by calling `toList`) unless they are sure that this is possible given the memory
+   * constraints of their cluster.
+   */
+  def map[U : Encoder](f: (K, Iterator[T]) => U): Dataset[U] = {
+    val func = (key: K, it: Iterator[T]) => Iterator(f(key, it))
+    new Dataset[U](
+      sqlContext,
+      MapGroups(func, groupingAttributes, logicalPlan))
+  }
+
+  def map[U](
+      f: JFunction2[K, JIterator[T], U],
+      encoder: Encoder[U]): Dataset[U] = {
+    map((key, data) => f.call(key, data.asJava))(encoder)
   }
 
   // To ensure valid overloading.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index 799650a4f784f..2593b16b1c8d7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -356,7 +356,7 @@ case class AppendColumns[T, U](
  * being output.
  */
 case class MapGroups[K, T, U](
-    func: (K, Iterator[T]) => Iterator[U],
+    func: (K, Iterator[T]) => TraversableOnce[U],
     kEncoder: ExpressionEncoder[K],
     tEncoder: ExpressionEncoder[T],
     uEncoder: ExpressionEncoder[U],
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index a9493d576d179..0d3b1a5af52c4 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -170,15 +170,15 @@ public Integer call(String v) throws Exception {
       }
     }, e.INT());
 
-    Dataset<String> mapped = grouped.mapGroups(
-      new Function2<Integer, Iterator<String>, Iterator<String>>() {
+    Dataset<String> mapped = grouped.map(
+      new Function2<Integer, Iterator<String>, String>() {
         @Override
-        public Iterator<String> call(Integer key, Iterator<String> data) throws Exception {
+        public String call(Integer key, Iterator<String> data) throws Exception {
           StringBuilder sb = new StringBuilder(key.toString());
           while (data.hasNext()) {
             sb.append(data.next());
           }
-          return Collections.singletonList(sb.toString()).iterator();
+          return sb.toString();
         }
       },
       e.STRING());
@@ -224,15 +224,15 @@ public void testGroupByColumn() {
     Dataset<String> ds = context.createDataset(data, e.STRING());
     GroupedDataset<Integer, String> grouped = ds.groupBy(length(col("value"))).asKey(e.INT());
 
-    Dataset<String> mapped = grouped.mapGroups(
-      new Function2<Integer, Iterator<String>, Iterator<String>>() {
+    Dataset<String> mapped = grouped.map(
+      new Function2<Integer, Iterator<String>, String>() {
         @Override
-        public Iterator<String> call(Integer key, Iterator<String> data) throws Exception {
+        public String call(Integer key, Iterator<String> data) throws Exception {
           StringBuilder sb = new StringBuilder(key.toString());
           while (data.hasNext()) {
             sb.append(data.next());
           }
-          return Collections.singletonList(sb.toString()).iterator();
+          return sb.toString();
         }
       },
       e.STRING());
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
index e3b0346f857d3..fcf03f7180984 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
@@ -88,16 +88,26 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
       0, 1)
   }
 
-  test("groupBy function, mapGroups") {
+  test("groupBy function, map") {
     val ds = Seq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11).toDS()
     val grouped = ds.groupBy(_ % 2)
-    val agged = grouped.mapGroups { case (g, iter) =>
+    val agged = grouped.map { case (g, iter) =>
       val name = if (g == 0) "even" else "odd"
-      Iterator((name, iter.size))
+      (name, iter.size)
     }
 
     checkAnswer(
       agged,
       ("even", 5), ("odd", 6))
   }
+
+  test("groupBy function, flatMap") {
+    val ds = Seq("a", "b", "c", "xyz", "hello").toDS()
+    val grouped = ds.groupBy(_.length)
+    val agged = grouped.flatMap { case (g, iter) => Iterator(g.toString, iter.mkString) }
+
+    checkAnswer(
+      agged,
+      "1", "abc", "3", "xyz", "5", "hello")
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index d61e17edc64ed..6f1174e6577e3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -198,60 +198,60 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       (1, 1))
   }
 
-  test("groupBy function, mapGroups") {
+  test("groupBy function, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy(v => (v._1, "word"))
-    val agged = grouped.mapGroups { case (g, iter) =>
-      Iterator((g._1, iter.map(_._2).sum))
-    }
+    val agged = grouped.map { case (g, iter) => (g._1, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
       ("a", 30), ("b", 3), ("c", 1))
   }
 
-  test("groupBy columns, mapGroups") {
+  test("groupBy function, fatMap") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+    val grouped = ds.groupBy(v => (v._1, "word"))
+    val agged = grouped.flatMap { case (g, iter) => Iterator(g._1, iter.map(_._2).sum.toString) }
+
+    checkAnswer(
+      agged,
+      "a", "30", "b", "3", "c", "1")
+  }
+
+  test("groupBy columns, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1")
-    val agged = grouped.mapGroups { case (g, iter) =>
-      Iterator((g.getString(0), iter.map(_._2).sum))
-    }
+    val agged = grouped.map { case (g, iter) => (g.getString(0), iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
       ("a", 30), ("b", 3), ("c", 1))
   }
 
-  test("groupBy columns asKey, mapGroups") {
+  test("groupBy columns asKey, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1").asKey[String]
-    val agged = grouped.mapGroups { case (g, iter) =>
-      Iterator((g, iter.map(_._2).sum))
-    }
+    val agged = grouped.map { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
       ("a", 30), ("b", 3), ("c", 1))
   }
 
-  test("groupBy columns asKey tuple, mapGroups") {
+  test("groupBy columns asKey tuple, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1", lit(1)).asKey[(String, Int)]
-    val agged = grouped.mapGroups { case (g, iter) =>
-      Iterator((g, iter.map(_._2).sum))
-    }
+    val agged = grouped.map { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
       (("a", 1), 30), (("b", 1), 3), (("c", 1), 1))
   }
 
-  test("groupBy columns asKey class, mapGroups") {
+  test("groupBy columns asKey class, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1".as("a"), lit(1).as("b")).asKey[ClassData]
-    val agged = grouped.mapGroups { case (g, iter) =>
-      Iterator((g, iter.map(_._2).sum))
-    }
+    val agged = grouped.map { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,

From 97b7080cf2d2846c7257f8926f775f27d457fe7d Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Sun, 8 Nov 2015 20:57:09 -0800
Subject: [PATCH 118/518] [SPARK-11564][SQL] Dataset Java API audit

A few changes:

1. Removed fold, since it can be confusing for distributed collections.
2. Created specific interfaces for each Dataset function (e.g. MapFunction, ReduceFunction, MapPartitionsFunction)
3. Added more documentation and test cases.

The other thing I'm considering doing is to have a "collector" interface for FlatMapFunction and MapPartitionsFunction, similar to MapReduce's map function.

Author: Reynold Xin <rxin@databricks.com>

Closes #9531 from rxin/SPARK-11564.
---
 .../api/java/function/FilterFunction.java     |  29 +++++
 .../api/java/function/ForeachFunction.java    |  29 +++++
 .../function/ForeachPartitionFunction.java    |  28 +++++
 .../spark/api/java/function/Function0.java    |   2 +-
 .../spark/api/java/function/MapFunction.java  |  27 +++++
 .../java/function/MapPartitionsFunction.java  |  28 +++++
 .../api/java/function/ReduceFunction.java     |  27 +++++
 .../spark/sql/catalyst/encoders/Encoder.scala |  38 +++++--
 .../org/apache/spark/sql/DataFrame.scala      |  47 ++++++--
 .../scala/org/apache/spark/sql/Dataset.scala  | 100 +++++++++---------
 .../apache/spark/sql/JavaDataFrameSuite.java  |   7 ++
 .../apache/spark/sql/JavaDatasetSuite.java    |  36 +++----
 .../spark/sql/DatasetPrimitiveSuite.scala     |   5 -
 .../org/apache/spark/sql/DatasetSuite.scala   |  10 +-
 14 files changed, 316 insertions(+), 97 deletions(-)
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/FilterFunction.java
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/ForeachFunction.java
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/ForeachPartitionFunction.java
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/MapFunction.java
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/MapPartitionsFunction.java
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/ReduceFunction.java

diff --git a/core/src/main/java/org/apache/spark/api/java/function/FilterFunction.java b/core/src/main/java/org/apache/spark/api/java/function/FilterFunction.java
new file mode 100644
index 0000000000000..e8d999dd00135
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/FilterFunction.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+
+/**
+ * Base interface for a function used in Dataset's filter function.
+ *
+ * If the function returns true, the element is discarded in the returned Dataset.
+ */
+public interface FilterFunction<T> extends Serializable {
+  boolean call(T value) throws Exception;
+}
diff --git a/core/src/main/java/org/apache/spark/api/java/function/ForeachFunction.java b/core/src/main/java/org/apache/spark/api/java/function/ForeachFunction.java
new file mode 100644
index 0000000000000..07e54b28fa12c
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/ForeachFunction.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+
+/**
+ * Base interface for a function used in Dataset's foreach function.
+ *
+ * Spark will invoke the call function on each element in the input Dataset.
+ */
+public interface ForeachFunction<T> extends Serializable {
+  void call(T t) throws Exception;
+}
diff --git a/core/src/main/java/org/apache/spark/api/java/function/ForeachPartitionFunction.java b/core/src/main/java/org/apache/spark/api/java/function/ForeachPartitionFunction.java
new file mode 100644
index 0000000000000..4938a51bcd712
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/ForeachPartitionFunction.java
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+/**
+ * Base interface for a function used in Dataset's foreachPartition function.
+ */
+public interface ForeachPartitionFunction<T> extends Serializable {
+  void call(Iterator<T> t) throws Exception;
+}
diff --git a/core/src/main/java/org/apache/spark/api/java/function/Function0.java b/core/src/main/java/org/apache/spark/api/java/function/Function0.java
index 38e410c5debe6..c86928dd05408 100644
--- a/core/src/main/java/org/apache/spark/api/java/function/Function0.java
+++ b/core/src/main/java/org/apache/spark/api/java/function/Function0.java
@@ -23,5 +23,5 @@
  * A zero-argument function that returns an R.
  */
 public interface Function0<R> extends Serializable {
-  public R call() throws Exception;
+  R call() throws Exception;
 }
diff --git a/core/src/main/java/org/apache/spark/api/java/function/MapFunction.java b/core/src/main/java/org/apache/spark/api/java/function/MapFunction.java
new file mode 100644
index 0000000000000..3ae6ef44898e1
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/MapFunction.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+
+/**
+ * Base interface for a map function used in Dataset's map function.
+ */
+public interface MapFunction<T, U> extends Serializable {
+  U call(T value) throws Exception;
+}
diff --git a/core/src/main/java/org/apache/spark/api/java/function/MapPartitionsFunction.java b/core/src/main/java/org/apache/spark/api/java/function/MapPartitionsFunction.java
new file mode 100644
index 0000000000000..6cb569ce0cb6b
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/MapPartitionsFunction.java
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+/**
+ * Base interface for function used in Dataset's mapPartitions.
+ */
+public interface MapPartitionsFunction<T, U> extends Serializable {
+  Iterable<U> call(Iterator<T> input) throws Exception;
+}
diff --git a/core/src/main/java/org/apache/spark/api/java/function/ReduceFunction.java b/core/src/main/java/org/apache/spark/api/java/function/ReduceFunction.java
new file mode 100644
index 0000000000000..ee092d0058f44
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/ReduceFunction.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+
+/**
+ * Base interface for function used in Dataset's reduce.
+ */
+public interface ReduceFunction<T> extends Serializable {
+  T call(T v1, T v2) throws Exception;
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala
index f05e18288de2b..6569b900fed90 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala
@@ -20,7 +20,7 @@ package org.apache.spark.sql.catalyst.encoders
 import scala.reflect.ClassTag
 
 import org.apache.spark.util.Utils
-import org.apache.spark.sql.types.{DataType, ObjectType, StructField, StructType}
+import org.apache.spark.sql.types.{ObjectType, StructField, StructType}
 import org.apache.spark.sql.catalyst.expressions._
 
 /**
@@ -100,7 +100,7 @@ object Encoder {
       expr.transformUp {
         case BoundReference(0, t: ObjectType, _) =>
           Invoke(
-            BoundReference(0, ObjectType(cls), true),
+            BoundReference(0, ObjectType(cls), nullable = true),
             s"_${index + 1}",
             t)
       }
@@ -114,13 +114,13 @@ object Encoder {
       } else {
         enc.constructExpression.transformUp {
           case BoundReference(ordinal, dt, _) =>
-            GetInternalRowField(BoundReference(index, enc.schema, true), ordinal, dt)
+            GetInternalRowField(BoundReference(index, enc.schema, nullable = true), ordinal, dt)
         }
       }
     }
 
     val constructExpression =
-      NewInstance(cls, constructExpressions, false, ObjectType(cls))
+      NewInstance(cls, constructExpressions, propagateNull = false, ObjectType(cls))
 
     new ExpressionEncoder[Any](
       schema,
@@ -130,7 +130,6 @@ object Encoder {
       ClassTag.apply(cls))
   }
 
-
   def typeTagOfTuple2[T1 : TypeTag, T2 : TypeTag]: TypeTag[(T1, T2)] = typeTag[(T1, T2)]
 
   private def getTypeTag[T](c: Class[T]): TypeTag[T] = {
@@ -148,9 +147,36 @@ object Encoder {
     })
   }
 
-  def forTuple2[T1, T2](c1: Class[T1], c2: Class[T2]): Encoder[(T1, T2)] = {
+  def forTuple[T1, T2](c1: Class[T1], c2: Class[T2]): Encoder[(T1, T2)] = {
     implicit val typeTag1 = getTypeTag(c1)
     implicit val typeTag2 = getTypeTag(c2)
     ExpressionEncoder[(T1, T2)]()
   }
+
+  def forTuple[T1, T2, T3](c1: Class[T1], c2: Class[T2], c3: Class[T3]): Encoder[(T1, T2, T3)] = {
+    implicit val typeTag1 = getTypeTag(c1)
+    implicit val typeTag2 = getTypeTag(c2)
+    implicit val typeTag3 = getTypeTag(c3)
+    ExpressionEncoder[(T1, T2, T3)]()
+  }
+
+  def forTuple[T1, T2, T3, T4](
+      c1: Class[T1], c2: Class[T2], c3: Class[T3], c4: Class[T4]): Encoder[(T1, T2, T3, T4)] = {
+    implicit val typeTag1 = getTypeTag(c1)
+    implicit val typeTag2 = getTypeTag(c2)
+    implicit val typeTag3 = getTypeTag(c3)
+    implicit val typeTag4 = getTypeTag(c4)
+    ExpressionEncoder[(T1, T2, T3, T4)]()
+  }
+
+  def forTuple[T1, T2, T3, T4, T5](
+      c1: Class[T1], c2: Class[T2], c3: Class[T3], c4: Class[T4], c5: Class[T5])
+    : Encoder[(T1, T2, T3, T4, T5)] = {
+    implicit val typeTag1 = getTypeTag(c1)
+    implicit val typeTag2 = getTypeTag(c2)
+    implicit val typeTag3 = getTypeTag(c3)
+    implicit val typeTag4 = getTypeTag(c4)
+    implicit val typeTag5 = getTypeTag(c5)
+    ExpressionEncoder[(T1, T2, T3, T4, T5)]()
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index f2d4db5550273..8ab958adadcca 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -1478,18 +1478,54 @@ class DataFrame private[sql](
 
   /**
    * Returns the first `n` rows in the [[DataFrame]].
+   *
+   * Running take requires moving data into the application's driver process, and doing so on a
+   * very large dataset can crash the driver process with OutOfMemoryError.
+   *
    * @group action
    * @since 1.3.0
    */
   def take(n: Int): Array[Row] = head(n)
 
+  /**
+   * Returns the first `n` rows in the [[DataFrame]] as a list.
+   *
+   * Running take requires moving data into the application's driver process, and doing so with
+   * a very large `n` can crash the driver process with OutOfMemoryError.
+   *
+   * @group action
+   * @since 1.6.0
+   */
+  def takeAsList(n: Int): java.util.List[Row] = java.util.Arrays.asList(take(n) : _*)
+
   /**
    * Returns an array that contains all of [[Row]]s in this [[DataFrame]].
+   *
+   * Running take requires moving data into the application's driver process, and doing so with
+   * a very large `n` can crash the driver process with OutOfMemoryError.
+   *
+   * For Java API, use [[collectAsList]].
+   *
    * @group action
    * @since 1.3.0
    */
   def collect(): Array[Row] = collect(needCallback = true)
 
+  /**
+   * Returns a Java list that contains all of [[Row]]s in this [[DataFrame]].
+   *
+   * Running collect requires moving all the data into the application's driver process, and
+   * doing so on a very large dataset can crash the driver process with OutOfMemoryError.
+   *
+   * @group action
+   * @since 1.3.0
+   */
+  def collectAsList(): java.util.List[Row] = withCallback("collectAsList", this) { _ =>
+    withNewExecutionId {
+      java.util.Arrays.asList(rdd.collect() : _*)
+    }
+  }
+
   private def collect(needCallback: Boolean): Array[Row] = {
     def execute(): Array[Row] = withNewExecutionId {
       queryExecution.executedPlan.executeCollectPublic()
@@ -1502,17 +1538,6 @@ class DataFrame private[sql](
     }
   }
 
-  /**
-   * Returns a Java list that contains all of [[Row]]s in this [[DataFrame]].
-   * @group action
-   * @since 1.3.0
-   */
-  def collectAsList(): java.util.List[Row] = withCallback("collectAsList", this) { _ =>
-    withNewExecutionId {
-      java.util.Arrays.asList(rdd.collect() : _*)
-    }
-  }
-
   /**
    * Returns the number of rows in the [[DataFrame]].
    * @group action
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index fecbdac9a6004..959e0f5ba03e6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -22,7 +22,7 @@ import scala.collection.JavaConverters._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAlias
-import org.apache.spark.api.java.function.{Function => JFunction, Function2 => JFunction2, _}
+import org.apache.spark.api.java.function._
 
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
@@ -75,7 +75,11 @@ class Dataset[T] private[sql](
   private[sql] def this(sqlContext: SQLContext, plan: LogicalPlan)(implicit encoder: Encoder[T]) =
     this(sqlContext, new QueryExecution(sqlContext, plan), encoder)
 
-  /** Returns the schema of the encoded form of the objects in this [[Dataset]]. */
+  /**
+   * Returns the schema of the encoded form of the objects in this [[Dataset]].
+   *
+   * @since 1.6.0
+   */
   def schema: StructType = encoder.schema
 
   /* ************* *
@@ -103,6 +107,7 @@ class Dataset[T] private[sql](
   /**
    * Applies a logical alias to this [[Dataset]] that can be used to disambiguate columns that have
    * the same name after two Datasets have been joined.
+   * @since 1.6.0
    */
   def as(alias: String): Dataset[T] = withPlan(Subquery(alias, _))
 
@@ -166,8 +171,7 @@ class Dataset[T] private[sql](
    * Returns a new [[Dataset]] that only contains elements where `func` returns `true`.
    * @since 1.6.0
    */
-  def filter(func: JFunction[T, java.lang.Boolean]): Dataset[T] =
-    filter(t => func.call(t).booleanValue())
+  def filter(func: FilterFunction[T]): Dataset[T] = filter(t => func.call(t))
 
   /**
    * (Scala-specific)
@@ -181,7 +185,7 @@ class Dataset[T] private[sql](
    * Returns a new [[Dataset]] that contains the result of applying `func` to each element.
    * @since 1.6.0
    */
-  def map[U](func: JFunction[T, U], encoder: Encoder[U]): Dataset[U] =
+  def map[U](func: MapFunction[T, U], encoder: Encoder[U]): Dataset[U] =
     map(t => func.call(t))(encoder)
 
   /**
@@ -205,10 +209,8 @@ class Dataset[T] private[sql](
    * Returns a new [[Dataset]] that contains the result of applying `func` to each element.
    * @since 1.6.0
    */
-  def mapPartitions[U](
-      f: FlatMapFunction[java.util.Iterator[T], U],
-      encoder: Encoder[U]): Dataset[U] = {
-    val func: (Iterator[T]) => Iterator[U] = x => f.call(x.asJava).iterator().asScala
+  def mapPartitions[U](f: MapPartitionsFunction[T, U], encoder: Encoder[U]): Dataset[U] = {
+    val func: (Iterator[T]) => Iterator[U] = x => f.call(x.asJava).iterator.asScala
     mapPartitions(func)(encoder)
   }
 
@@ -248,7 +250,7 @@ class Dataset[T] private[sql](
    * Runs `func` on each element of this Dataset.
    * @since 1.6.0
    */
-  def foreach(func: VoidFunction[T]): Unit = foreach(func.call(_))
+  def foreach(func: ForeachFunction[T]): Unit = foreach(func.call(_))
 
   /**
    * (Scala-specific)
@@ -262,7 +264,7 @@ class Dataset[T] private[sql](
    * Runs `func` on each partition of this Dataset.
    * @since 1.6.0
    */
-  def foreachPartition(func: VoidFunction[java.util.Iterator[T]]): Unit =
+  def foreachPartition(func: ForeachPartitionFunction[T]): Unit =
     foreachPartition(it => func.call(it.asJava))
 
   /* ************* *
@@ -271,7 +273,7 @@ class Dataset[T] private[sql](
 
   /**
    * (Scala-specific)
-   * Reduces the elements of this Dataset using the specified  binary function.  The given function
+   * Reduces the elements of this Dataset using the specified binary function.  The given function
    * must be commutative and associative or the result may be non-deterministic.
    * @since 1.6.0
    */
@@ -279,33 +281,11 @@ class Dataset[T] private[sql](
 
   /**
    * (Java-specific)
-   * Reduces the elements of this Dataset using the specified  binary function.  The given function
+   * Reduces the elements of this Dataset using the specified binary function.  The given function
    * must be commutative and associative or the result may be non-deterministic.
    * @since 1.6.0
    */
-  def reduce(func: JFunction2[T, T, T]): T = reduce(func.call(_, _))
-
-  /**
-   * (Scala-specific)
-   * Aggregates the elements of each partition, and then the results for all the partitions, using a
-   * given associative and commutative function and a neutral "zero value".
-   *
-   * This behaves somewhat differently than the fold operations implemented for non-distributed
-   * collections in functional languages like Scala. This fold operation may be applied to
-   * partitions individually, and then those results will be folded into the final result.
-   * If op is not commutative, then the result may differ from that of a fold applied to a
-   * non-distributed collection.
-   * @since 1.6.0
-   */
-  def fold(zeroValue: T)(op: (T, T) => T): T = rdd.fold(zeroValue)(op)
-
-  /**
-   * (Java-specific)
-   * Aggregates the elements of each partition, and then the results for all the partitions, using a
-   * given associative and commutative function and a neutral "zero value".
-   * @since 1.6.0
-   */
-  def fold(zeroValue: T, func: JFunction2[T, T, T]): T = fold(zeroValue)(func.call(_, _))
+  def reduce(func: ReduceFunction[T]): T = reduce(func.call(_, _))
 
   /**
    * (Scala-specific)
@@ -351,7 +331,7 @@ class Dataset[T] private[sql](
    * Returns a [[GroupedDataset]] where the data is grouped by the given key function.
    * @since 1.6.0
    */
-  def groupBy[K](f: JFunction[T, K], encoder: Encoder[K]): GroupedDataset[K, T] =
+  def groupBy[K](f: MapFunction[T, K], encoder: Encoder[K]): GroupedDataset[K, T] =
     groupBy(f.call(_))(encoder)
 
   /* ****************** *
@@ -367,7 +347,7 @@ class Dataset[T] private[sql](
    */
   // Copied from Dataframe to make sure we don't have invalid overloads.
   @scala.annotation.varargs
-  def select(cols: Column*): DataFrame = toDF().select(cols: _*)
+  protected def select(cols: Column*): DataFrame = toDF().select(cols: _*)
 
   /**
    * Returns a new [[Dataset]] by computing the given [[Column]] expression for each element.
@@ -462,8 +442,7 @@ class Dataset[T] private[sql](
    * and thus is not affected by a custom `equals` function defined on `T`.
    * @since 1.6.0
    */
-  def intersect(other: Dataset[T]): Dataset[T] =
-    withPlan[T](other)(Intersect)
+  def intersect(other: Dataset[T]): Dataset[T] = withPlan[T](other)(Intersect)
 
   /**
    * Returns a new [[Dataset]] that contains the elements of both this and the `other` [[Dataset]]
@@ -473,8 +452,7 @@ class Dataset[T] private[sql](
    * duplicate items.  As such, it is analagous to `UNION ALL` in SQL.
    * @since 1.6.0
    */
-  def union(other: Dataset[T]): Dataset[T] =
-    withPlan[T](other)(Union)
+  def union(other: Dataset[T]): Dataset[T] = withPlan[T](other)(Union)
 
   /**
    * Returns a new [[Dataset]] where any elements present in `other` have been removed.
@@ -542,27 +520,47 @@ class Dataset[T] private[sql](
   def first(): T = rdd.first()
 
   /**
-   * Collects the elements to an Array.
+   * Returns an array that contains all the elements in this [[Dataset]].
+   *
+   * Running collect requires moving all the data into the application's driver process, and
+   * doing so on a very large dataset can crash the driver process with OutOfMemoryError.
+   *
+   * For Java API, use [[collectAsList]].
    * @since 1.6.0
    */
   def collect(): Array[T] = rdd.collect()
 
   /**
-   * (Java-specific)
-   * Collects the elements to a Java list.
+   * Returns an array that contains all the elements in this [[Dataset]].
    *
-   * Due to the incompatibility problem between Scala and Java, the return type of [[collect()]] at
-   * Java side is `java.lang.Object`, which is not easy to use.  Java user can use this method
-   * instead and keep the generic type for result.
+   * Running collect requires moving all the data into the application's driver process, and
+   * doing so on a very large dataset can crash the driver process with OutOfMemoryError.
    *
+   * For Java API, use [[collectAsList]].
    * @since 1.6.0
    */
-  def collectAsList(): java.util.List[T] =
-    rdd.collect().toSeq.asJava
+  def collectAsList(): java.util.List[T] = rdd.collect().toSeq.asJava
 
-  /** Returns the first `num` elements of this [[Dataset]] as an Array. */
+  /**
+   * Returns the first `num` elements of this [[Dataset]] as an array.
+   *
+   * Running take requires moving data into the application's driver process, and doing so with
+   * a very large `n` can crash the driver process with OutOfMemoryError.
+   *
+   * @since 1.6.0
+   */
   def take(num: Int): Array[T] = rdd.take(num)
 
+  /**
+   * Returns the first `num` elements of this [[Dataset]] as an array.
+   *
+   * Running take requires moving data into the application's driver process, and doing so with
+   * a very large `n` can crash the driver process with OutOfMemoryError.
+   *
+   * @since 1.6.0
+   */
+  def takeAsList(num: Int): java.util.List[T] = java.util.Arrays.asList(take(num) : _*)
+
   /* ******************** *
    *  Internal Functions  *
    * ******************** */
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index 40bff57a17a03..d191b50fa802e 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -65,6 +65,13 @@ public void testExecution() {
     Assert.assertEquals(1, df.select("key").collect()[0].get(0));
   }
 
+  @Test
+  public void testCollectAndTake() {
+    DataFrame df = context.table("testData").filter("key = 1 or key = 2 or key = 3");
+    Assert.assertEquals(3, df.select("key").collectAsList().size());
+    Assert.assertEquals(2, df.select("key").takeAsList(2).size());
+  }
+
   /**
    * See SPARK-5904. Abstract vararg methods defined in Scala do not work in Java.
    */
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index 0d3b1a5af52c4..0f90de774dd3e 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -68,8 +68,16 @@ private <T1, T2> Tuple2<T1, T2> tuple2(T1 t1, T2 t2) {
   public void testCollect() {
     List<String> data = Arrays.asList("hello", "world");
     Dataset<String> ds = context.createDataset(data, e.STRING());
-    String[] collected = (String[]) ds.collect();
-    Assert.assertEquals(Arrays.asList("hello", "world"), Arrays.asList(collected));
+    List<String> collected = ds.collectAsList();
+    Assert.assertEquals(Arrays.asList("hello", "world"), collected);
+  }
+
+  @Test
+  public void testTake() {
+    List<String> data = Arrays.asList("hello", "world");
+    Dataset<String> ds = context.createDataset(data, e.STRING());
+    List<String> collected = ds.takeAsList(1);
+    Assert.assertEquals(Arrays.asList("hello"), collected);
   }
 
   @Test
@@ -78,16 +86,16 @@ public void testCommonOperation() {
     Dataset<String> ds = context.createDataset(data, e.STRING());
     Assert.assertEquals("hello", ds.first());
 
-    Dataset<String> filtered = ds.filter(new Function<String, Boolean>() {
+    Dataset<String> filtered = ds.filter(new FilterFunction<String>() {
       @Override
-      public Boolean call(String v) throws Exception {
+      public boolean call(String v) throws Exception {
         return v.startsWith("h");
       }
     });
     Assert.assertEquals(Arrays.asList("hello"), filtered.collectAsList());
 
 
-    Dataset<Integer> mapped = ds.map(new Function<String, Integer>() {
+    Dataset<Integer> mapped = ds.map(new MapFunction<String, Integer>() {
       @Override
       public Integer call(String v) throws Exception {
         return v.length();
@@ -95,7 +103,7 @@ public Integer call(String v) throws Exception {
     }, e.INT());
     Assert.assertEquals(Arrays.asList(5, 5), mapped.collectAsList());
 
-    Dataset<String> parMapped = ds.mapPartitions(new FlatMapFunction<Iterator<String>, String>() {
+    Dataset<String> parMapped = ds.mapPartitions(new MapPartitionsFunction<String, String>() {
       @Override
       public Iterable<String> call(Iterator<String> it) throws Exception {
         List<String> ls = new LinkedList<String>();
@@ -128,7 +136,7 @@ public void testForeach() {
     List<String> data = Arrays.asList("a", "b", "c");
     Dataset<String> ds = context.createDataset(data, e.STRING());
 
-    ds.foreach(new VoidFunction<String>() {
+    ds.foreach(new ForeachFunction<String>() {
       @Override
       public void call(String s) throws Exception {
         accum.add(1);
@@ -142,28 +150,20 @@ public void testReduce() {
     List<Integer> data = Arrays.asList(1, 2, 3);
     Dataset<Integer> ds = context.createDataset(data, e.INT());
 
-    int reduced = ds.reduce(new Function2<Integer, Integer, Integer>() {
+    int reduced = ds.reduce(new ReduceFunction<Integer>() {
       @Override
       public Integer call(Integer v1, Integer v2) throws Exception {
         return v1 + v2;
       }
     });
     Assert.assertEquals(6, reduced);
-
-    int folded = ds.fold(1, new Function2<Integer, Integer, Integer>() {
-      @Override
-      public Integer call(Integer v1, Integer v2) throws Exception {
-        return v1 * v2;
-      }
-    });
-    Assert.assertEquals(6, folded);
   }
 
   @Test
   public void testGroupBy() {
     List<String> data = Arrays.asList("a", "foo", "bar");
     Dataset<String> ds = context.createDataset(data, e.STRING());
-    GroupedDataset<Integer, String> grouped = ds.groupBy(new Function<String, Integer>() {
+    GroupedDataset<Integer, String> grouped = ds.groupBy(new MapFunction<String, Integer>() {
       @Override
       public Integer call(String v) throws Exception {
         return v.length();
@@ -187,7 +187,7 @@ public String call(Integer key, Iterator<String> data) throws Exception {
 
     List<Integer> data2 = Arrays.asList(2, 6, 10);
     Dataset<Integer> ds2 = context.createDataset(data2, e.INT());
-    GroupedDataset<Integer, Integer> grouped2 = ds2.groupBy(new Function<Integer, Integer>() {
+    GroupedDataset<Integer, Integer> grouped2 = ds2.groupBy(new MapFunction<Integer, Integer>() {
       @Override
       public Integer call(Integer v) throws Exception {
         return v / 2;
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
index fcf03f7180984..63b00975e4eb1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
@@ -75,11 +75,6 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
     assert(ds.reduce(_ + _) == 6)
   }
 
-  test("fold") {
-    val ds = Seq(1, 2, 3).toDS()
-    assert(ds.fold(0)(_ + _) == 6)
-  }
-
   test("groupBy function, keys") {
     val ds = Seq(1, 2, 3, 4, 5).toDS()
     val grouped = ds.groupBy(_ % 2)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 6f1174e6577e3..aea5a700d0204 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -61,6 +61,11 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     assert(ds.collect() === Array(ClassData("a", 1), ClassData("b", 2), ClassData("c", 3)))
   }
 
+  test("as case class - take") {
+    val ds = Seq((1, "a"), (2, "b"), (3, "c")).toDF("b", "a").as[ClassData]
+    assert(ds.take(2) === Array(ClassData("a", 1), ClassData("b", 2)))
+  }
+
   test("map") {
     val ds = Seq(("a", 1) , ("b", 2), ("c", 3)).toDS()
     checkAnswer(
@@ -137,11 +142,6 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     assert(ds.reduce((a, b) => ("sum", a._2 + b._2)) == ("sum", 6))
   }
 
-  test("fold") {
-    val ds = Seq(("a", 1) , ("b", 2), ("c", 3)).toDS()
-    assert(ds.fold(("", 0))((a, b) => ("sum", a._2 + b._2)) == ("sum", 6))
-  }
-
   test("joinWith, flat schema") {
     val ds1 = Seq(1, 2, 3).toDS().as("a")
     val ds2 = Seq(1, 2).toDS().as("b")

From d8b50f70298dbf45e91074ee2d751fee7eecb119 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Sun, 8 Nov 2015 21:01:53 -0800
Subject: [PATCH 119/518] [SPARK-11453][SQL] append data to partitioned table
 will messes up the result

The reason is that:

1. For partitioned hive table, we will move the partitioned columns after data columns. (e.g. `<a: Int, b: Int>` partition by `a` will become `<b: Int, a: Int>`)
2. When append data to table, we use position to figure out how to match input columns to table's columns.

So when we append data to partitioned table, we will match wrong columns between input and table. A solution is reordering the input columns before match by position, like what we did for [`InsertIntoHadoopFsRelation`](https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelation.scala#L101-L105)

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9408 from cloud-fan/append.
---
 .../apache/spark/sql/DataFrameWriter.scala    | 29 ++++++++++++++++---
 .../sql/sources/PartitionedWriteSuite.scala   |  8 +++++
 .../sql/hive/execution/SQLQuerySuite.scala    | 20 +++++++++++++
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala
index 7887e559a3025..e63a4d5e8b10b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala
@@ -23,8 +23,8 @@ import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.sql.catalyst.{SqlParser, TableIdentifier}
-import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
-import org.apache.spark.sql.catalyst.plans.logical.InsertIntoTable
+import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedRelation}
+import org.apache.spark.sql.catalyst.plans.logical.{Project, InsertIntoTable}
 import org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils
 import org.apache.spark.sql.execution.datasources.{CreateTableUsingAsSelect, ResolvedDataSource}
 import org.apache.spark.sql.sources.HadoopFsRelation
@@ -167,17 +167,38 @@ final class DataFrameWriter private[sql](df: DataFrame) {
   }
 
   private def insertInto(tableIdent: TableIdentifier): Unit = {
-    val partitions = partitioningColumns.map(_.map(col => col -> (None: Option[String])).toMap)
+    val partitions = normalizedParCols.map(_.map(col => col -> (None: Option[String])).toMap)
     val overwrite = mode == SaveMode.Overwrite
+
+    // A partitioned relation's schema can be different from the input logicalPlan, since
+    // partition columns are all moved after data columns. We Project to adjust the ordering.
+    // TODO: this belongs to the analyzer.
+    val input = normalizedParCols.map { parCols =>
+      val (inputPartCols, inputDataCols) = df.logicalPlan.output.partition { attr =>
+        parCols.contains(attr.name)
+      }
+      Project(inputDataCols ++ inputPartCols, df.logicalPlan)
+    }.getOrElse(df.logicalPlan)
+
     df.sqlContext.executePlan(
       InsertIntoTable(
         UnresolvedRelation(tableIdent),
         partitions.getOrElse(Map.empty[String, Option[String]]),
-        df.logicalPlan,
+        input,
         overwrite,
         ifNotExists = false)).toRdd
   }
 
+  private def normalizedParCols: Option[Seq[String]] = partitioningColumns.map { parCols =>
+    parCols.map { col =>
+      df.logicalPlan.output
+        .map(_.name)
+        .find(df.sqlContext.analyzer.resolver(_, col))
+        .getOrElse(throw new AnalysisException(s"Partition column $col not found in existing " +
+          s"columns (${df.logicalPlan.output.map(_.name).mkString(", ")})"))
+    }
+  }
+
   /**
    * Saves the content of the [[DataFrame]] as the specified table.
    *
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/PartitionedWriteSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/PartitionedWriteSuite.scala
index c9791879ec74c..3eaa817f9c0b0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/PartitionedWriteSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/PartitionedWriteSuite.scala
@@ -53,4 +53,12 @@ class PartitionedWriteSuite extends QueryTest with SharedSQLContext {
 
     Utils.deleteRecursively(path)
   }
+
+  test("partitioned columns should appear at the end of schema") {
+    withTempPath { f =>
+      val path = f.getAbsolutePath
+      Seq(1 -> "a").toDF("i", "j").write.partitionBy("i").parquet(path)
+      assert(sqlContext.read.parquet(path).schema.map(_.name) == Seq("j", "i"))
+    }
+  }
 }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
index af48d478953b4..9a425d7f6b265 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
@@ -1428,4 +1428,24 @@ class SQLQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
       checkAnswer(sql("SELECT val FROM tbl10562 WHERE Year == 2012"), Row("a"))
     }
   }
+
+  test("SPARK-11453: append data to partitioned table") {
+    withTable("tbl11453") {
+      Seq("1" -> "10", "2" -> "20").toDF("i", "j")
+        .write.partitionBy("i").saveAsTable("tbl11453")
+
+      Seq("3" -> "30").toDF("i", "j")
+        .write.mode(SaveMode.Append).partitionBy("i").saveAsTable("tbl11453")
+      checkAnswer(
+        sqlContext.read.table("tbl11453").select("i", "j").orderBy("i"),
+        Row("1", "10") :: Row("2", "20") :: Row("3", "30") :: Nil)
+
+      // make sure case sensitivity is correct.
+      Seq("4" -> "40").toDF("i", "j")
+        .write.mode(SaveMode.Append).partitionBy("I").saveAsTable("tbl11453")
+      checkAnswer(
+        sqlContext.read.table("tbl11453").select("i", "j").orderBy("i"),
+        Row("1", "10") :: Row("2", "20") :: Row("3", "30") :: Row("4", "40") :: Nil)
+    }
+  }
 }

From 9e48cdfbdecc9554a425ba35c0252910fd1e8faa Mon Sep 17 00:00:00 2001
From: Charles Yeh <charlesyeh@dropbox.com>
Date: Mon, 9 Nov 2015 13:22:05 +0100
Subject: [PATCH 120/518] [SPARK-11218][CORE] show help messages for
 start-slave and start-master

Addressing https://issues.apache.org/jira/browse/SPARK-11218, mostly copied start-thriftserver.sh.
```
charlesyeh-mbp:spark charlesyeh$ ./sbin/start-master.sh --help
Usage: Master [options]

Options:
  -i HOST, --ip HOST     Hostname to listen on (deprecated, please use --host or -h)
  -h HOST, --host HOST   Hostname to listen on
  -p PORT, --port PORT   Port to listen on (default: 7077)
  --webui-port PORT      Port for web UI (default: 8080)
  --properties-file FILE Path to a custom Spark properties file.
                         Default is conf/spark-defaults.conf.
```
```
charlesyeh-mbp:spark charlesyeh$ ./sbin/start-slave.sh
Usage: Worker [options] <master>

Master must be a URL of the form spark://hostname:port

Options:
  -c CORES, --cores CORES  Number of cores to use
  -m MEM, --memory MEM     Amount of memory to use (e.g. 1000M, 2G)
  -d DIR, --work-dir DIR   Directory to run apps in (default: SPARK_HOME/work)
  -i HOST, --ip IP         Hostname to listen on (deprecated, please use --host or -h)
  -h HOST, --host HOST     Hostname to listen on
  -p PORT, --port PORT     Port to listen on (default: random)
  --webui-port PORT        Port for web UI (default: 8081)
  --properties-file FILE   Path to a custom Spark properties file.
                           Default is conf/spark-defaults.conf.
```

Author: Charles Yeh <charlesyeh@dropbox.com>

Closes #9432 from CharlesYeh/helpmsg.
---
 sbin/start-master.sh | 24 +++++++++++++++++++-----
 sbin/start-slave.sh  | 24 +++++++++++++++---------
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/sbin/start-master.sh b/sbin/start-master.sh
index c20e19a8412df..9f2e14dff609f 100755
--- a/sbin/start-master.sh
+++ b/sbin/start-master.sh
@@ -23,6 +23,20 @@ if [ -z "${SPARK_HOME}" ]; then
   export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
 fi
 
+# NOTE: This exact class name is matched downstream by SparkSubmit.
+# Any changes need to be reflected there.
+CLASS="org.apache.spark.deploy.master.Master"
+
+if [[ "$@" = *--help ]] || [[ "$@" = *-h ]]; then
+  echo "Usage: ./sbin/start-master.sh [options]"
+  pattern="Usage:"
+  pattern+="\|Using Spark's default log4j profile:"
+  pattern+="\|Registered signal handlers for"
+
+  "${SPARK_HOME}"/bin/spark-class $CLASS --help 2>&1 | grep -v "$pattern" 1>&2
+  exit 1
+fi
+
 ORIGINAL_ARGS="$@"
 
 START_TACHYON=false
@@ -30,7 +44,7 @@ START_TACHYON=false
 while (( "$#" )); do
 case $1 in
     --with-tachyon)
-      if [ ! -e "$sbin"/../tachyon/bin/tachyon ]; then
+      if [ ! -e "${SPARK_HOME}"/tachyon/bin/tachyon ]; then
         echo "Error: --with-tachyon specified, but tachyon not found."
         exit -1
       fi
@@ -56,12 +70,12 @@ if [ "$SPARK_MASTER_WEBUI_PORT" = "" ]; then
   SPARK_MASTER_WEBUI_PORT=8080
 fi
 
-"${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.master.Master 1 \
+"${SPARK_HOME}/sbin"/spark-daemon.sh start $CLASS 1 \
   --ip $SPARK_MASTER_IP --port $SPARK_MASTER_PORT --webui-port $SPARK_MASTER_WEBUI_PORT \
   $ORIGINAL_ARGS
 
 if [ "$START_TACHYON" == "true" ]; then
-  "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon bootstrap-conf $SPARK_MASTER_IP
-  "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon format -s
-  "${SPARK_HOME}/sbin"/../tachyon/bin/tachyon-start.sh master
+  "${SPARK_HOME}"/tachyon/bin/tachyon bootstrap-conf $SPARK_MASTER_IP
+  "${SPARK_HOME}"/tachyon/bin/tachyon format -s
+  "${SPARK_HOME}"/tachyon/bin/tachyon-start.sh master
 fi
diff --git a/sbin/start-slave.sh b/sbin/start-slave.sh
index 21455648d1c6d..8c268b8859155 100755
--- a/sbin/start-slave.sh
+++ b/sbin/start-slave.sh
@@ -31,18 +31,24 @@
 #                           worker.  Subsequent workers will increment this
 #                           number.  Default is 8081.
 
-usage="Usage: start-slave.sh <spark-master-URL> where <spark-master-URL> is like spark://localhost:7077"
-
-if [ $# -lt 1 ]; then
-  echo $usage
-  echo Called as start-slave.sh $*
-  exit 1
-fi
-
 if [ -z "${SPARK_HOME}" ]; then
   export SPARK_HOME="$(cd "`dirname "$0"`"/..; pwd)"
 fi
 
+# NOTE: This exact class name is matched downstream by SparkSubmit.
+# Any changes need to be reflected there.
+CLASS="org.apache.spark.deploy.worker.Worker"
+
+if [[ $# -lt 1 ]] || [[ "$@" = *--help ]] || [[ "$@" = *-h ]]; then
+  echo "Usage: ./sbin/start-slave.sh [options] <master>"
+  pattern="Usage:"
+  pattern+="\|Using Spark's default log4j profile:"
+  pattern+="\|Registered signal handlers for"
+
+  "${SPARK_HOME}"/bin/spark-class $CLASS --help 2>&1 | grep -v "$pattern" 1>&2
+  exit 1
+fi
+
 . "${SPARK_HOME}/sbin/spark-config.sh"
 
 . "${SPARK_HOME}/bin/load-spark-env.sh"
@@ -72,7 +78,7 @@ function start_instance {
   fi
   WEBUI_PORT=$(( $SPARK_WORKER_WEBUI_PORT + $WORKER_NUM - 1 ))
 
-  "${SPARK_HOME}/sbin"/spark-daemon.sh start org.apache.spark.deploy.worker.Worker $WORKER_NUM \
+  "${SPARK_HOME}/sbin"/spark-daemon.sh start $CLASS $WORKER_NUM \
      --webui-port "$WEBUI_PORT" $PORT_FLAG $PORT_NUM $MASTER "$@"
 }
 

From b541b31630b1b85b48d6096079d073ccf46a62e8 Mon Sep 17 00:00:00 2001
From: Rohit Agarwal <mindprince@gmail.com>
Date: Mon, 9 Nov 2015 13:28:00 +0100
Subject: [PATCH 121/518] [DOC][MINOR][SQL] Fix internal link

It doesn't show up as a hyperlink currently. It will show up as a hyperlink after this change.

Author: Rohit Agarwal <mindprince@gmail.com>

Closes #9544 from mindprince/patch-2.
---
 docs/sql-programming-guide.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/sql-programming-guide.md b/docs/sql-programming-guide.md
index 52e03b951f966..ccd26904329d3 100644
--- a/docs/sql-programming-guide.md
+++ b/docs/sql-programming-guide.md
@@ -2287,7 +2287,7 @@ Several caching related features are not supported yet:
 Spark SQL is designed to be compatible with the Hive Metastore, SerDes and UDFs.
 Currently Hive SerDes and UDFs are based on Hive 1.2.1,
 and Spark SQL can be connected to different versions of Hive Metastore
-(from 0.12.0 to 1.2.1. Also see http://spark.apache.org/docs/latest/sql-programming-guide.html#interacting-with-different-versions-of-hive-metastore).
+(from 0.12.0 to 1.2.1. Also see [Interacting with Different Versions of Hive Metastore] (#interacting-with-different-versions-of-hive-metastore)).
 
 #### Deploying in Existing Hive Warehouses
 

From 8c0e1b50e960d3e8e51d0618c462eed2bb4936f0 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Mon, 9 Nov 2015 08:56:22 -0800
Subject: [PATCH 122/518] [SPARK-11494][ML][R] Expose R-like summary statistics
 in SparkR::glm for linear regression

Expose R-like summary statistics in SparkR::glm for linear regression, the output of ```summary``` like
```Java
$DevianceResiduals
 Min        Max
 -0.9509607 0.7291832

$Coefficients
                   Estimate   Std. Error t value   Pr(>|t|)
(Intercept)        1.6765     0.2353597  7.123139  4.456124e-11
Sepal_Length       0.3498801  0.04630128 7.556598  4.187317e-12
Species_versicolor -0.9833885 0.07207471 -13.64402 0
Species_virginica  -1.00751   0.09330565 -10.79796 0
```

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9561 from yanboliang/spark-11494.
---
 R/pkg/R/mllib.R                               | 22 ++++++--
 R/pkg/inst/tests/test_mllib.R                 | 31 +++++++++---
 .../apache/spark/ml/r/SparkRWrappers.scala    | 50 +++++++++++++++++--
 3 files changed, 88 insertions(+), 15 deletions(-)

diff --git a/R/pkg/R/mllib.R b/R/pkg/R/mllib.R
index b0d73dd93a79d..7ff859741b4a0 100644
--- a/R/pkg/R/mllib.R
+++ b/R/pkg/R/mllib.R
@@ -91,12 +91,26 @@ setMethod("predict", signature(object = "PipelineModel"),
 #'}
 setMethod("summary", signature(x = "PipelineModel"),
           function(x, ...) {
+            modelName <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
+                                   "getModelName", x@model)
             features <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
                                    "getModelFeatures", x@model)
             coefficients <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
                                    "getModelCoefficients", x@model)
-            coefficients <- as.matrix(unlist(coefficients))
-            colnames(coefficients) <- c("Estimate")
-            rownames(coefficients) <- unlist(features)
-            return(list(coefficients = coefficients))
+            if (modelName == "LinearRegressionModel") {
+              devianceResiduals <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
+                                               "getModelDevianceResiduals", x@model)
+              devianceResiduals <- matrix(devianceResiduals, nrow = 1)
+              colnames(devianceResiduals) <- c("Min", "Max")
+              rownames(devianceResiduals) <- rep("", times = 1)
+              coefficients <- matrix(coefficients, ncol = 4)
+              colnames(coefficients) <- c("Estimate", "Std. Error", "t value", "Pr(>|t|)")
+              rownames(coefficients) <- unlist(features)
+              return(list(DevianceResiduals = devianceResiduals, Coefficients = coefficients))
+            } else {
+              coefficients <- as.matrix(unlist(coefficients))
+              colnames(coefficients) <- c("Estimate")
+              rownames(coefficients) <- unlist(features)
+              return(list(coefficients = coefficients))
+            }
           })
diff --git a/R/pkg/inst/tests/test_mllib.R b/R/pkg/inst/tests/test_mllib.R
index 4761e285a2479..2606407bdcb44 100644
--- a/R/pkg/inst/tests/test_mllib.R
+++ b/R/pkg/inst/tests/test_mllib.R
@@ -71,12 +71,23 @@ test_that("feature interaction vs native glm", {
 
 test_that("summary coefficients match with native glm", {
   training <- createDataFrame(sqlContext, iris)
-  stats <- summary(glm(Sepal_Width ~ Sepal_Length + Species, data = training, solver = "l-bfgs"))
-  coefs <- as.vector(stats$coefficients)
+  stats <- summary(glm(Sepal_Width ~ Sepal_Length + Species, data = training, solver = "normal"))
+  coefs <- unlist(stats$Coefficients)
+  devianceResiduals <- unlist(stats$DevianceResiduals)
+
   rCoefs <- as.vector(coef(glm(Sepal.Width ~ Sepal.Length + Species, data = iris)))
-  expect_true(all(abs(rCoefs - coefs) < 1e-6))
+  rStdError <- c(0.23536, 0.04630, 0.07207, 0.09331)
+  rTValue <- c(7.123, 7.557, -13.644, -10.798)
+  rPValue <- c(0.0, 0.0, 0.0, 0.0)
+  rDevianceResiduals <- c(-0.95096, 0.72918)
+
+  expect_true(all(abs(rCoefs - coefs[1:4]) < 1e-6))
+  expect_true(all(abs(rStdError - coefs[5:8]) < 1e-5))
+  expect_true(all(abs(rTValue - coefs[9:12]) < 1e-3))
+  expect_true(all(abs(rPValue - coefs[13:16]) < 1e-6))
+  expect_true(all(abs(rDevianceResiduals - devianceResiduals) < 1e-5))
   expect_true(all(
-    as.character(stats$features) ==
+    rownames(stats$Coefficients) ==
     c("(Intercept)", "Sepal_Length", "Species_versicolor", "Species_virginica")))
 })
 
@@ -85,14 +96,20 @@ test_that("summary coefficients match with native glm of family 'binomial'", {
   training <- filter(df, df$Species != "setosa")
   stats <- summary(glm(Species ~ Sepal_Length + Sepal_Width, data = training,
     family = "binomial"))
-  coefs <- as.vector(stats$coefficients)
+  coefs <- as.vector(stats$Coefficients)
 
   rTraining <- iris[iris$Species %in% c("versicolor","virginica"),]
   rCoefs <- as.vector(coef(glm(Species ~ Sepal.Length + Sepal.Width, data = rTraining,
     family = binomial(link = "logit"))))
+  rStdError <- c(3.0974, 0.5169, 0.8628)
+  rTValue <- c(-4.212, 3.680, 0.469)
+  rPValue <- c(0.000, 0.000, 0.639)
 
-  expect_true(all(abs(rCoefs - coefs) < 1e-4))
+  expect_true(all(abs(rCoefs - coefs[1:3]) < 1e-4))
+  expect_true(all(abs(rStdError - coefs[4:6]) < 1e-4))
+  expect_true(all(abs(rTValue - coefs[7:9]) < 1e-3))
+  expect_true(all(abs(rPValue - coefs[10:12]) < 1e-3))
   expect_true(all(
-    as.character(stats$features) ==
+    rownames(stats$Coefficients) ==
     c("(Intercept)", "Sepal_Length", "Sepal_Width")))
 })
diff --git a/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala b/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
index 5be2f86936211..4d82b90bfdf20 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/r/SparkRWrappers.scala
@@ -52,11 +52,36 @@ private[r] object SparkRWrappers {
   }
 
   def getModelCoefficients(model: PipelineModel): Array[Double] = {
+    model.stages.last match {
+      case m: LinearRegressionModel => {
+        val coefficientStandardErrorsR = Array(m.summary.coefficientStandardErrors.last) ++
+          m.summary.coefficientStandardErrors.dropRight(1)
+        val tValuesR = Array(m.summary.tValues.last) ++ m.summary.tValues.dropRight(1)
+        val pValuesR = Array(m.summary.pValues.last) ++ m.summary.pValues.dropRight(1)
+        if (m.getFitIntercept) {
+          Array(m.intercept) ++ m.coefficients.toArray ++ coefficientStandardErrorsR ++
+            tValuesR ++ pValuesR
+        } else {
+          m.coefficients.toArray ++ coefficientStandardErrorsR ++ tValuesR ++ pValuesR
+        }
+      }
+      case m: LogisticRegressionModel => {
+        if (m.getFitIntercept) {
+          Array(m.intercept) ++ m.coefficients.toArray
+        } else {
+          m.coefficients.toArray
+        }
+      }
+    }
+  }
+
+  def getModelDevianceResiduals(model: PipelineModel): Array[Double] = {
     model.stages.last match {
       case m: LinearRegressionModel =>
-        Array(m.intercept) ++ m.coefficients.toArray
+        m.summary.devianceResiduals
       case m: LogisticRegressionModel =>
-        Array(m.intercept) ++ m.coefficients.toArray
+        throw new UnsupportedOperationException(
+          "No deviance residuals available for LogisticRegressionModel")
     }
   }
 
@@ -65,11 +90,28 @@ private[r] object SparkRWrappers {
       case m: LinearRegressionModel =>
         val attrs = AttributeGroup.fromStructField(
           m.summary.predictions.schema(m.summary.featuresCol))
-        Array("(Intercept)") ++ attrs.attributes.get.map(_.name.get)
+        if (m.getFitIntercept) {
+          Array("(Intercept)") ++ attrs.attributes.get.map(_.name.get)
+        } else {
+          attrs.attributes.get.map(_.name.get)
+        }
       case m: LogisticRegressionModel =>
         val attrs = AttributeGroup.fromStructField(
           m.summary.predictions.schema(m.summary.featuresCol))
-        Array("(Intercept)") ++ attrs.attributes.get.map(_.name.get)
+        if (m.getFitIntercept) {
+          Array("(Intercept)") ++ attrs.attributes.get.map(_.name.get)
+        } else {
+          attrs.attributes.get.map(_.name.get)
+        }
+    }
+  }
+
+  def getModelName(model: PipelineModel): String = {
+    model.stages.last match {
+      case m: LinearRegressionModel =>
+        "LinearRegressionModel"
+      case m: LogisticRegressionModel =>
+        "LogisticRegressionModel"
     }
   }
 }

From d50a66cc04bfa1c483f04daffe465322316c745e Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Mon, 9 Nov 2015 08:57:29 -0800
Subject: [PATCH 123/518] [SPARK-10689][ML][DOC] User guide and example code
 for AFTSurvivalRegression

Add user guide and example code for ```AFTSurvivalRegression```.

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9491 from yanboliang/spark-10689.
---
 docs/ml-guide.md                              |  1 +
 docs/ml-survival-regression.md                | 96 +++++++++++++++++++
 .../ml/JavaAFTSurvivalRegressionExample.java  | 71 ++++++++++++++
 .../main/python/ml/aft_survival_regression.py | 51 ++++++++++
 .../ml/AFTSurvivalRegressionExample.scala     | 62 ++++++++++++
 5 files changed, 281 insertions(+)
 create mode 100644 docs/ml-survival-regression.md
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaAFTSurvivalRegressionExample.java
 create mode 100644 examples/src/main/python/ml/aft_survival_regression.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala

diff --git a/docs/ml-guide.md b/docs/ml-guide.md
index fd3a6167bc65e..c293e71d2870e 100644
--- a/docs/ml-guide.md
+++ b/docs/ml-guide.md
@@ -44,6 +44,7 @@ provide class probabilities, and linear models provide model summaries.
 * [Ensembles](ml-ensembles.html)
 * [Linear methods with elastic net regularization](ml-linear-methods.html)
 * [Multilayer perceptron classifier](ml-ann.html)
+* [Survival Regression](ml-survival-regression.html)
 
 
 # Main concepts in Pipelines
diff --git a/docs/ml-survival-regression.md b/docs/ml-survival-regression.md
new file mode 100644
index 0000000000000..ab275213b9a84
--- /dev/null
+++ b/docs/ml-survival-regression.md
@@ -0,0 +1,96 @@
+---
+layout: global
+title: Survival Regression - ML
+displayTitle: <a href="ml-guide.html">ML</a> - Survival Regression
+---
+
+
+`\[
+\newcommand{\R}{\mathbb{R}}
+\newcommand{\E}{\mathbb{E}}
+\newcommand{\x}{\mathbf{x}}
+\newcommand{\y}{\mathbf{y}}
+\newcommand{\wv}{\mathbf{w}}
+\newcommand{\av}{\mathbf{\alpha}}
+\newcommand{\bv}{\mathbf{b}}
+\newcommand{\N}{\mathbb{N}}
+\newcommand{\id}{\mathbf{I}}
+\newcommand{\ind}{\mathbf{1}}
+\newcommand{\0}{\mathbf{0}}
+\newcommand{\unit}{\mathbf{e}}
+\newcommand{\one}{\mathbf{1}}
+\newcommand{\zero}{\mathbf{0}}
+\]`
+
+
+In `spark.ml`, we implement the [Accelerated failure time (AFT)](https://en.wikipedia.org/wiki/Accelerated_failure_time_model) 
+model which is a parametric survival regression model for censored data. 
+It describes a model for the log of survival time, so it's often called 
+log-linear model for survival analysis. Different from 
+[Proportional hazards](https://en.wikipedia.org/wiki/Proportional_hazards_model) model
+designed for the same purpose, the AFT model is more easily to parallelize 
+because each instance contribute to the objective function independently.
+
+Given the values of the covariates $x^{'}$, for random lifetime $t_{i}$ of 
+subjects i = 1, ..., n, with possible right-censoring, 
+the likelihood function under the AFT model is given as:
+`\[
+L(\beta,\sigma)=\prod_{i=1}^n[\frac{1}{\sigma}f_{0}(\frac{\log{t_{i}}-x^{'}\beta}{\sigma})]^{\delta_{i}}S_{0}(\frac{\log{t_{i}}-x^{'}\beta}{\sigma})^{1-\delta_{i}}
+\]`
+Where $\delta_{i}$ is the indicator of the event has occurred i.e. uncensored or not.
+Using $\epsilon_{i}=\frac{\log{t_{i}}-x^{'}\beta}{\sigma}$, the log-likelihood function
+assumes the form:
+`\[
+\iota(\beta,\sigma)=\sum_{i=1}^{n}[-\delta_{i}\log\sigma+\delta_{i}\log{f_{0}}(\epsilon_{i})+(1-\delta_{i})\log{S_{0}(\epsilon_{i})}]
+\]`
+Where $S_{0}(\epsilon_{i})$ is the baseline survivor function,
+and $f_{0}(\epsilon_{i})$ is corresponding density function.
+
+The most commonly used AFT model is based on the Weibull distribution of the survival time. 
+The Weibull distribution for lifetime corresponding to extreme value distribution for 
+log of the lifetime, and the $S_{0}(\epsilon)$ function is:
+`\[   
+S_{0}(\epsilon_{i})=\exp(-e^{\epsilon_{i}})
+\]`
+the $f_{0}(\epsilon_{i})$ function is:
+`\[
+f_{0}(\epsilon_{i})=e^{\epsilon_{i}}\exp(-e^{\epsilon_{i}})
+\]`
+The log-likelihood function for AFT model with Weibull distribution of lifetime is:
+`\[
+\iota(\beta,\sigma)= -\sum_{i=1}^n[\delta_{i}\log\sigma-\delta_{i}\epsilon_{i}+e^{\epsilon_{i}}]
+\]`
+Due to minimizing the negative log-likelihood equivalent to maximum a posteriori probability,
+the loss function we use to optimize is $-\iota(\beta,\sigma)$.
+The gradient functions for $\beta$ and $\log\sigma$ respectively are:
+`\[   
+\frac{\partial (-\iota)}{\partial \beta}=\sum_{1=1}^{n}[\delta_{i}-e^{\epsilon_{i}}]\frac{x_{i}}{\sigma}
+\]`
+`\[ 
+\frac{\partial (-\iota)}{\partial (\log\sigma)}=\sum_{i=1}^{n}[\delta_{i}+(\delta_{i}-e^{\epsilon_{i}})\epsilon_{i}]
+\]`
+
+The AFT model can be formulated as a convex optimization problem, 
+i.e. the task of finding a minimizer of a convex function $-\iota(\beta,\sigma)$ 
+that depends coefficients vector $\beta$ and the log of scale parameter $\log\sigma$.
+The optimization algorithm underlying the implementation is L-BFGS.
+The implementation matches the result from R's survival function 
+[survreg](https://stat.ethz.ch/R-manual/R-devel/library/survival/html/survreg.html)
+
+## Example:
+
+<div class="codetabs">
+
+<div data-lang="scala" markdown="1">
+{% include_example scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala %}
+</div>
+
+<div data-lang="java" markdown="1">
+{% include_example java/org/apache/spark/examples/ml/JavaAFTSurvivalRegressionExample.java %}
+</div>
+
+<div data-lang="python" markdown="1">
+{% include_example python/ml/aft_survival_regression.py %}
+</div>
+
+</div>
\ No newline at end of file
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaAFTSurvivalRegressionExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaAFTSurvivalRegressionExample.java
new file mode 100644
index 0000000000000..69a174562fcf5
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaAFTSurvivalRegressionExample.java
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+// $example on$
+import java.util.Arrays;
+import java.util.List;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.ml.regression.AFTSurvivalRegression;
+import org.apache.spark.ml.regression.AFTSurvivalRegressionModel;
+import org.apache.spark.mllib.linalg.*;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.sql.types.*;
+// $example off$
+
+public class JavaAFTSurvivalRegressionExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaAFTSurvivalRegressionExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext jsql = new SQLContext(jsc);
+
+    // $example on$
+    List<Row> data = Arrays.asList(
+      RowFactory.create(1.218, 1.0, Vectors.dense(1.560, -0.605)),
+      RowFactory.create(2.949, 0.0, Vectors.dense(0.346, 2.158)),
+      RowFactory.create(3.627, 0.0, Vectors.dense(1.380, 0.231)),
+      RowFactory.create(0.273, 1.0, Vectors.dense(0.520, 1.151)),
+      RowFactory.create(4.199, 0.0, Vectors.dense(0.795, -0.226))
+    );
+    StructType schema = new StructType(new StructField[]{
+      new StructField("label", DataTypes.DoubleType, false, Metadata.empty()),
+      new StructField("censor", DataTypes.DoubleType, false, Metadata.empty()),
+      new StructField("features", new VectorUDT(), false, Metadata.empty())
+    });
+    DataFrame training = jsql.createDataFrame(data, schema);
+    double[] quantileProbabilities = new double[]{0.3, 0.6};
+    AFTSurvivalRegression aft = new AFTSurvivalRegression()
+      .setQuantileProbabilities(quantileProbabilities)
+      .setQuantilesCol("quantiles");
+
+    AFTSurvivalRegressionModel model = aft.fit(training);
+
+    // Print the coefficients, intercept and scale parameter for AFT survival regression
+    System.out.println("Coefficients: " + model.coefficients() + " Intercept: "
+      + model.intercept() + " Scale: " + model.scale());
+    model.transform(training).show(false);
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/python/ml/aft_survival_regression.py b/examples/src/main/python/ml/aft_survival_regression.py
new file mode 100644
index 0000000000000..0ee01fd8258df
--- /dev/null
+++ b/examples/src/main/python/ml/aft_survival_regression.py
@@ -0,0 +1,51 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import print_function
+
+from pyspark import SparkContext
+from pyspark.sql import SQLContext
+# $example on$
+from pyspark.ml.regression import AFTSurvivalRegression
+from pyspark.mllib.linalg import Vectors
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="AFTSurvivalRegressionExample")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    training = sqlContext.createDataFrame([
+        (1.218, 1.0, Vectors.dense(1.560, -0.605)),
+        (2.949, 0.0, Vectors.dense(0.346, 2.158)),
+        (3.627, 0.0, Vectors.dense(1.380, 0.231)),
+        (0.273, 1.0, Vectors.dense(0.520, 1.151)),
+        (4.199, 0.0, Vectors.dense(0.795, -0.226))], ["label", "censor", "features"])
+    quantileProbabilities = [0.3, 0.6]
+    aft = AFTSurvivalRegression(quantileProbabilities=quantileProbabilities,
+                                quantilesCol="quantiles")
+
+    model = aft.fit(training)
+
+    # Print the coefficients, intercept and scale parameter for AFT survival regression
+    print("Coefficients: " + str(model.coefficients))
+    print("Intercept: " + str(model.intercept))
+    print("Scale: " + str(model.scale))
+    model.transform(training).show(truncate=False)
+    # $example off$
+
+    sc.stop()
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala
new file mode 100644
index 0000000000000..5da285e83681f
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
+import org.apache.spark.ml.regression.AFTSurvivalRegression
+import org.apache.spark.mllib.linalg.Vectors
+// $example off$
+
+/**
+ * An example for AFTSurvivalRegression.
+ */
+object AFTSurvivalRegressionExample {
+
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("AFTSurvivalRegressionExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+
+    // $example on$
+    val training = sqlContext.createDataFrame(Seq(
+      (1.218, 1.0, Vectors.dense(1.560, -0.605)),
+      (2.949, 0.0, Vectors.dense(0.346, 2.158)),
+      (3.627, 0.0, Vectors.dense(1.380, 0.231)),
+      (0.273, 1.0, Vectors.dense(0.520, 1.151)),
+      (4.199, 0.0, Vectors.dense(0.795, -0.226))
+    )).toDF("label", "censor", "features")
+    val quantileProbabilities = Array(0.3, 0.6)
+    val aft = new AFTSurvivalRegression()
+      .setQuantileProbabilities(quantileProbabilities)
+      .setQuantilesCol("quantiles")
+
+    val model = aft.fit(training)
+
+    // Print the coefficients, intercept and scale parameter for AFT survival regression
+    println(s"Coefficients: ${model.coefficients} Intercept: " +
+      s"${model.intercept} Scale: ${model.scale}")
+    model.transform(training).show(false)
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:off println

From 9b88e1dcad6b5b14a22cf64a1055ad9870507b5a Mon Sep 17 00:00:00 2001
From: fazlan-nazeem <fazlann@wso2.com>
Date: Mon, 9 Nov 2015 08:58:55 -0800
Subject: [PATCH 124/518] [SPARK-11582][MLLIB] specifying pmml version
 attribute =4.2 in the root node of pmml model

The current pmml models generated do not specify the pmml version in its root node. This is a problem when using this pmml model in other tools because they expect the version attribute to be set explicitly. This fix adds the pmml version attribute to the generated pmml models and specifies its value as 4.2.

Author: fazlan-nazeem <fazlann@wso2.com>

Closes #9558 from fazlan-nazeem/master.
---
 .../org/apache/spark/mllib/pmml/export/PMMLModelExport.scala     | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/pmml/export/PMMLModelExport.scala b/mllib/src/main/scala/org/apache/spark/mllib/pmml/export/PMMLModelExport.scala
index c5fdecd3ca17f..9267e6dbdb857 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/pmml/export/PMMLModelExport.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/pmml/export/PMMLModelExport.scala
@@ -32,6 +32,7 @@ private[mllib] trait PMMLModelExport {
   @BeanProperty
   val pmml: PMML = new PMML
 
+  pmml.setVersion("4.2")
   setHeader(pmml)
 
   private def setHeader(pmml: PMML): Unit = {

From 08a7a836c393d6a62b9b216eeb01fad0b90b6c52 Mon Sep 17 00:00:00 2001
From: Charles Yeh <charlesyeh@dropbox.com>
Date: Mon, 9 Nov 2015 11:59:32 -0600
Subject: [PATCH 125/518] [SPARK-10565][CORE] add missing web UI stats to
 /api/v1/applications JSON

I looked at the other endpoints, and they don't seem to be missing any fields.
Added fields:
![image](https://cloud.githubusercontent.com/assets/613879/10948801/58159982-82e4-11e5-86dc-62da201af910.png)

Author: Charles Yeh <charlesyeh@dropbox.com>

Closes #9472 from CharlesYeh/api_vars.
---
 .../spark/deploy/master/ui/MasterWebUI.scala  |  7 +-
 .../api/v1/ApplicationListResource.scala      |  8 ++
 .../org/apache/spark/status/api/v1/api.scala  |  4 +
 .../scala/org/apache/spark/ui/SparkUI.scala   |  4 +
 .../deploy/master/ui/MasterWebUISuite.scala   | 90 +++++++++++++++++++
 project/MimaExcludes.scala                    |  3 +
 6 files changed, 114 insertions(+), 2 deletions(-)
 create mode 100644 core/src/test/scala/org/apache/spark/deploy/master/ui/MasterWebUISuite.scala

diff --git a/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterWebUI.scala b/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterWebUI.scala
index 6174fc11f83d8..e41554a5a6d26 100644
--- a/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterWebUI.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterWebUI.scala
@@ -28,14 +28,17 @@ import org.apache.spark.ui.JettyUtils._
  * Web UI server for the standalone master.
  */
 private[master]
-class MasterWebUI(val master: Master, requestedPort: Int)
+class MasterWebUI(
+    val master: Master,
+    requestedPort: Int,
+    customMasterPage: Option[MasterPage] = None)
   extends WebUI(master.securityMgr, requestedPort, master.conf, name = "MasterUI") with Logging
   with UIRoot {
 
   val masterEndpointRef = master.self
   val killEnabled = master.conf.getBoolean("spark.ui.killEnabled", true)
 
-  val masterPage = new MasterPage(this)
+  val masterPage = customMasterPage.getOrElse(new MasterPage(this))
 
   initialize()
 
diff --git a/core/src/main/scala/org/apache/spark/status/api/v1/ApplicationListResource.scala b/core/src/main/scala/org/apache/spark/status/api/v1/ApplicationListResource.scala
index 17b521f3e1d41..0fc0fb59d861f 100644
--- a/core/src/main/scala/org/apache/spark/status/api/v1/ApplicationListResource.scala
+++ b/core/src/main/scala/org/apache/spark/status/api/v1/ApplicationListResource.scala
@@ -62,6 +62,10 @@ private[spark] object ApplicationsListResource {
     new ApplicationInfo(
       id = app.id,
       name = app.name,
+      coresGranted = None,
+      maxCores = None,
+      coresPerExecutor = None,
+      memoryPerExecutorMB = None,
       attempts = app.attempts.map { internalAttemptInfo =>
         new ApplicationAttemptInfo(
           attemptId = internalAttemptInfo.attemptId,
@@ -81,6 +85,10 @@ private[spark] object ApplicationsListResource {
     new ApplicationInfo(
       id = internal.id,
       name = internal.desc.name,
+      coresGranted = Some(internal.coresGranted),
+      maxCores = internal.desc.maxCores,
+      coresPerExecutor = internal.desc.coresPerExecutor,
+      memoryPerExecutorMB = Some(internal.desc.memoryPerExecutorMB),
       attempts = Seq(new ApplicationAttemptInfo(
         attemptId = None,
         startTime = new Date(internal.startTime),
diff --git a/core/src/main/scala/org/apache/spark/status/api/v1/api.scala b/core/src/main/scala/org/apache/spark/status/api/v1/api.scala
index 2bec64f2ef02b..baddfc50c1a40 100644
--- a/core/src/main/scala/org/apache/spark/status/api/v1/api.scala
+++ b/core/src/main/scala/org/apache/spark/status/api/v1/api.scala
@@ -25,6 +25,10 @@ import org.apache.spark.JobExecutionStatus
 class ApplicationInfo private[spark](
     val id: String,
     val name: String,
+    val coresGranted: Option[Int],
+    val maxCores: Option[Int],
+    val coresPerExecutor: Option[Int],
+    val memoryPerExecutorMB: Option[Int],
     val attempts: Seq[ApplicationAttemptInfo])
 
 class ApplicationAttemptInfo private[spark](
diff --git a/core/src/main/scala/org/apache/spark/ui/SparkUI.scala b/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
index 99085ada9f0af..4608bce202ec8 100644
--- a/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
+++ b/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
@@ -102,6 +102,10 @@ private[spark] class SparkUI private (
     Iterator(new ApplicationInfo(
       id = appId,
       name = appName,
+      coresGranted = None,
+      maxCores = None,
+      coresPerExecutor = None,
+      memoryPerExecutorMB = None,
       attempts = Seq(new ApplicationAttemptInfo(
         attemptId = None,
         startTime = new Date(startTime),
diff --git a/core/src/test/scala/org/apache/spark/deploy/master/ui/MasterWebUISuite.scala b/core/src/test/scala/org/apache/spark/deploy/master/ui/MasterWebUISuite.scala
new file mode 100644
index 0000000000000..fba835f054f8a
--- /dev/null
+++ b/core/src/test/scala/org/apache/spark/deploy/master/ui/MasterWebUISuite.scala
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.deploy.master.ui
+
+import java.util.Date
+
+import scala.io.Source
+import scala.language.postfixOps
+
+import org.json4s.jackson.JsonMethods._
+import org.json4s.JsonAST.{JNothing, JString, JInt}
+import org.mockito.Mockito.{mock, when}
+import org.scalatest.BeforeAndAfter
+
+import org.apache.spark.{SparkConf, SecurityManager, SparkFunSuite}
+import org.apache.spark.deploy.DeployMessages.MasterStateResponse
+import org.apache.spark.deploy.DeployTestUtils._
+import org.apache.spark.deploy.master._
+import org.apache.spark.rpc.RpcEnv
+
+
+class MasterWebUISuite extends SparkFunSuite with BeforeAndAfter {
+
+  val masterPage = mock(classOf[MasterPage])
+  val master = {
+    val conf = new SparkConf
+    val securityMgr = new SecurityManager(conf)
+    val rpcEnv = RpcEnv.create(Master.SYSTEM_NAME, "localhost", 0, conf, securityMgr)
+    val master = new Master(rpcEnv, rpcEnv.address, 0, securityMgr, conf)
+    master
+  }
+  val masterWebUI = new MasterWebUI(master, 0, customMasterPage = Some(masterPage))
+
+  before {
+    masterWebUI.bind()
+  }
+
+  after {
+    masterWebUI.stop()
+  }
+
+  test("list applications") {
+    val worker = createWorkerInfo()
+    val appDesc = createAppDesc()
+    // use new start date so it isn't filtered by UI
+    val activeApp = new ApplicationInfo(
+      new Date().getTime, "id", appDesc, new Date(), null, Int.MaxValue)
+    activeApp.addExecutor(worker, 2)
+
+    val workers = Array[WorkerInfo](worker)
+    val activeApps = Array(activeApp)
+    val completedApps = Array[ApplicationInfo]()
+    val activeDrivers = Array[DriverInfo]()
+    val completedDrivers = Array[DriverInfo]()
+    val stateResponse = new MasterStateResponse(
+      "host", 8080, None, workers, activeApps, completedApps,
+      activeDrivers, completedDrivers, RecoveryState.ALIVE)
+
+    when(masterPage.getMasterState).thenReturn(stateResponse)
+
+    val resultJson = Source.fromURL(
+      s"http://localhost:${masterWebUI.boundPort}/api/v1/applications")
+      .mkString
+    val parsedJson = parse(resultJson)
+    val firstApp = parsedJson(0)
+
+    assert(firstApp \ "id" === JString(activeApp.id))
+    assert(firstApp \ "name" === JString(activeApp.desc.name))
+    assert(firstApp \ "coresGranted" === JInt(2))
+    assert(firstApp \ "maxCores" === JInt(4))
+    assert(firstApp \ "memoryPerExecutorMB" === JInt(1234))
+    assert(firstApp \ "coresPerExecutor" === JNothing)
+  }
+
+}
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index dacef911e397e..50220790d1f84 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -134,6 +134,9 @@ object MimaExcludes {
         ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.toString"),
         ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.jdbc.PostgresDialect.hashCode"),
         ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.sql.jdbc.NoopDialect$")
+      ) ++ Seq (
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.status.api.v1.ApplicationInfo.this")
       )
     case v if v.startsWith("1.5") =>
       Seq(

From 404a28f4edd09cf17361dcbd770e4cafde51bf6d Mon Sep 17 00:00:00 2001
From: tedyu <yuzhihong@gmail.com>
Date: Mon, 9 Nov 2015 10:07:58 -0800
Subject: [PATCH 126/518] [SPARK-11112] Fix Scala 2.11 compilation error in
 RDDInfo.scala

As shown in https://amplab.cs.berkeley.edu/jenkins/view/Spark-QA-Compile/job/Spark-Master-Scala211-Compile/1946/console , compilation fails with:
```
[error] /home/jenkins/workspace/Spark-Master-Scala211-Compile/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala:25: in class RDDInfo, multiple overloaded alternatives of constructor RDDInfo define default arguments.
[error] class RDDInfo(
[error]
```
This PR tries to fix the compilation error

Author: tedyu <yuzhihong@gmail.com>

Closes #9538 from tedyu/master.
---
 .../scala/org/apache/spark/storage/RDDInfo.scala     | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala b/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala
index 3fa209b924170..87c1b981e7e13 100644
--- a/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala
+++ b/core/src/main/scala/org/apache/spark/storage/RDDInfo.scala
@@ -28,20 +28,10 @@ class RDDInfo(
     val numPartitions: Int,
     var storageLevel: StorageLevel,
     val parentIds: Seq[Int],
-    val callSite: CallSite,
+    val callSite: CallSite = CallSite.empty,
     val scope: Option[RDDOperationScope] = None)
   extends Ordered[RDDInfo] {
 
-  def this(
-      id: Int,
-      name: String,
-      numPartitions: Int,
-      storageLevel: StorageLevel,
-      parentIds: Seq[Int],
-      scope: Option[RDDOperationScope] = None) {
-    this(id, name, numPartitions, storageLevel, parentIds, CallSite.empty, scope)
-  }
-
   var numCachedPartitions = 0
   var memSize = 0L
   var diskSize = 0L

From cd174882a5a211298d6e173fe989d567d08ebc0d Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Mon, 9 Nov 2015 10:26:09 -0800
Subject: [PATCH 127/518] [SPARK-9865][SPARKR] Flaky SparkR test:
 test_sparkSQL.R: sample on a DataFrame

Make sample test less flaky by setting the seed

Tested with
```
repeat {  if (count(sample(df, FALSE, 0.1)) == 3) { break } }
```

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9549 from felixcheung/rsample.
---
 R/pkg/inst/tests/test_sparkSQL.R | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 92cff1fba7193..fbdb9a8f1ef6b 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -647,11 +647,11 @@ test_that("sample on a DataFrame", {
   sampled <- sample(df, FALSE, 1.0)
   expect_equal(nrow(collect(sampled)), count(df))
   expect_is(sampled, "DataFrame")
-  sampled2 <- sample(df, FALSE, 0.1)
+  sampled2 <- sample(df, FALSE, 0.1, 0) # set seed for predictable result
   expect_true(count(sampled2) < 3)
 
   # Also test sample_frac
-  sampled3 <- sample_frac(df, FALSE, 0.1)
+  sampled3 <- sample_frac(df, FALSE, 0.1, 0) # set seed for predictable result
   expect_true(count(sampled3) < 3)
 })
 

From 874cd66d4b6d156d0ef112a3d0f3bc5683c6a0ec Mon Sep 17 00:00:00 2001
From: chriskang90 <jckang@uchicago.edu>
Date: Mon, 9 Nov 2015 19:39:22 +0100
Subject: [PATCH 128/518] [DOCS] Fix typo for Python section on unifying Kafka
 streams

1) kafkaStreams is a list.  The list should be unpacked when passing it into the streaming context union method, which accepts a variable number of streams.
2) print() should be pprint() for pyspark.

This contribution is my original work, and I license the work to the project under the project's open source license.

Author: chriskang90 <jckang@uchicago.edu>

Closes #9545 from c-kang/streaming_python_typo.
---
 docs/streaming-programming-guide.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/streaming-programming-guide.md b/docs/streaming-programming-guide.md
index c751dbb41785a..e9a27f446a898 100644
--- a/docs/streaming-programming-guide.md
+++ b/docs/streaming-programming-guide.md
@@ -1948,8 +1948,8 @@ unifiedStream.print();
 {% highlight python %}
 numStreams = 5
 kafkaStreams = [KafkaUtils.createStream(...) for _ in range (numStreams)]
-unifiedStream = streamingContext.union(kafkaStreams)
-unifiedStream.print()
+unifiedStream = streamingContext.union(*kafkaStreams)
+unifiedStream.pprint()
 {% endhighlight %}
 </div>
 </div>

From 860ea0d386b5fbbe26bf2954f402a9a73ad37edc Mon Sep 17 00:00:00 2001
From: Bharat Lal <bharat.iisc@gmail.com>
Date: Mon, 9 Nov 2015 11:33:01 -0800
Subject: [PATCH 129/518] [SPARK-11581][DOCS] Example mllib code in
 documentation incorrectly computes MSE

Author: Bharat Lal <bharat.iisc@gmail.com>

Closes #9560 from bharatl/SPARK-11581.
---
 docs/mllib-decision-tree.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/mllib-decision-tree.md b/docs/mllib-decision-tree.md
index f31c4f88936bd..b5b454bc69245 100644
--- a/docs/mllib-decision-tree.md
+++ b/docs/mllib-decision-tree.md
@@ -439,7 +439,7 @@ Double testMSE =
     public Double call(Double a, Double b) {
       return a + b;
     }
-  }) / data.count();
+  }) / testData.count();
 System.out.println("Test Mean Squared Error: " + testMSE);
 System.out.println("Learned regression tree model:\n" + model.toDebugString());
 

From 88a3fdcc783f880a8d01c7e194ec42fc114bdf8a Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Mon, 9 Nov 2015 13:16:04 -0800
Subject: [PATCH 130/518] [SPARK-10280][MLLIB][PYSPARK][DOCS] Add @since
 annotation to pyspark.ml.classification

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>

Closes #8690 from yu-iskw/SPARK-10280.
---
 python/pyspark/ml/classification.py | 56 +++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py
index 2e468f67b8987..603f2c7f798dc 100644
--- a/python/pyspark/ml/classification.py
+++ b/python/pyspark/ml/classification.py
@@ -67,6 +67,8 @@ class LogisticRegression(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPredicti
     Traceback (most recent call last):
         ...
     TypeError: Method setParams forces keyword arguments.
+
+    .. versionadded:: 1.3.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -99,6 +101,7 @@ def __init__(self, featuresCol="features", labelCol="label", predictionCol="pred
         self._checkThresholdConsistency()
 
     @keyword_only
+    @since("1.3.0")
     def setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                   maxIter=100, regParam=0.1, elasticNetParam=0.0, tol=1e-6, fitIntercept=True,
                   threshold=0.5, thresholds=None, probabilityCol="probability",
@@ -119,6 +122,7 @@ def setParams(self, featuresCol="features", labelCol="label", predictionCol="pre
     def _create_model(self, java_model):
         return LogisticRegressionModel(java_model)
 
+    @since("1.4.0")
     def setThreshold(self, value):
         """
         Sets the value of :py:attr:`threshold`.
@@ -129,6 +133,7 @@ def setThreshold(self, value):
             del self._paramMap[self.thresholds]
         return self
 
+    @since("1.4.0")
     def getThreshold(self):
         """
         Gets the value of threshold or its default value.
@@ -144,6 +149,7 @@ def getThreshold(self):
         else:
             return self.getOrDefault(self.threshold)
 
+    @since("1.5.0")
     def setThresholds(self, value):
         """
         Sets the value of :py:attr:`thresholds`.
@@ -154,6 +160,7 @@ def setThresholds(self, value):
             del self._paramMap[self.threshold]
         return self
 
+    @since("1.5.0")
     def getThresholds(self):
         """
         If :py:attr:`thresholds` is set, return its value.
@@ -185,9 +192,12 @@ def _checkThresholdConsistency(self):
 class LogisticRegressionModel(JavaModel):
     """
     Model fitted by LogisticRegression.
+
+    .. versionadded:: 1.3.0
     """
 
     @property
+    @since("1.4.0")
     def weights(self):
         """
         Model weights.
@@ -205,6 +215,7 @@ def coefficients(self):
         return self._call_java("coefficients")
 
     @property
+    @since("1.4.0")
     def intercept(self):
         """
         Model intercept.
@@ -215,6 +226,8 @@ def intercept(self):
 class TreeClassifierParams(object):
     """
     Private class to track supported impurity measures.
+
+    .. versionadded:: 1.4.0
     """
     supportedImpurities = ["entropy", "gini"]
 
@@ -231,6 +244,7 @@ def __init__(self):
                               "gain calculation (case-insensitive). Supported options: " +
                               ", ".join(self.supportedImpurities))
 
+    @since("1.6.0")
     def setImpurity(self, value):
         """
         Sets the value of :py:attr:`impurity`.
@@ -238,6 +252,7 @@ def setImpurity(self, value):
         self._paramMap[self.impurity] = value
         return self
 
+    @since("1.6.0")
     def getImpurity(self):
         """
         Gets the value of impurity or its default value.
@@ -248,6 +263,8 @@ def getImpurity(self):
 class GBTParams(TreeEnsembleParams):
     """
     Private class to track supported GBT params.
+
+    .. versionadded:: 1.4.0
     """
     supportedLossTypes = ["logistic"]
 
@@ -287,6 +304,8 @@ class DecisionTreeClassifier(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPred
     >>> test1 = sqlContext.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"])
     >>> model.transform(test1).head().prediction
     1.0
+
+    .. versionadded:: 1.4.0
     """
 
     @keyword_only
@@ -310,6 +329,7 @@ def __init__(self, featuresCol="features", labelCol="label", predictionCol="pred
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                   probabilityCol="probability", rawPredictionCol="rawPrediction",
                   maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0,
@@ -333,6 +353,8 @@ def _create_model(self, java_model):
 class DecisionTreeClassificationModel(DecisionTreeModel):
     """
     Model fitted by DecisionTreeClassifier.
+
+    .. versionadded:: 1.4.0
     """
 
 
@@ -371,6 +393,8 @@ class RandomForestClassifier(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPred
     >>> test1 = sqlContext.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"])
     >>> model.transform(test1).head().prediction
     1.0
+
+    .. versionadded:: 1.4.0
     """
 
     @keyword_only
@@ -396,6 +420,7 @@ def __init__(self, featuresCol="features", labelCol="label", predictionCol="pred
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                   probabilityCol="probability", rawPredictionCol="rawPrediction",
                   maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0,
@@ -419,6 +444,8 @@ def _create_model(self, java_model):
 class RandomForestClassificationModel(TreeEnsembleModels):
     """
     Model fitted by RandomForestClassifier.
+
+    .. versionadded:: 1.4.0
     """
 
 
@@ -450,6 +477,8 @@ class GBTClassifier(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPredictionCol
     >>> test1 = sqlContext.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"])
     >>> model.transform(test1).head().prediction
     1.0
+
+    .. versionadded:: 1.4.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -482,6 +511,7 @@ def __init__(self, featuresCol="features", labelCol="label", predictionCol="pred
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.4.0")
     def setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                   maxDepth=5, maxBins=32, minInstancesPerNode=1, minInfoGain=0.0,
                   maxMemoryInMB=256, cacheNodeIds=False, checkpointInterval=10,
@@ -499,6 +529,7 @@ def setParams(self, featuresCol="features", labelCol="label", predictionCol="pre
     def _create_model(self, java_model):
         return GBTClassificationModel(java_model)
 
+    @since("1.4.0")
     def setLossType(self, value):
         """
         Sets the value of :py:attr:`lossType`.
@@ -506,6 +537,7 @@ def setLossType(self, value):
         self._paramMap[self.lossType] = value
         return self
 
+    @since("1.4.0")
     def getLossType(self):
         """
         Gets the value of lossType or its default value.
@@ -516,6 +548,8 @@ def getLossType(self):
 class GBTClassificationModel(TreeEnsembleModels):
     """
     Model fitted by GBTClassifier.
+
+    .. versionadded:: 1.4.0
     """
 
 
@@ -555,6 +589,8 @@ class NaiveBayes(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPredictionCol, H
     >>> test1 = sc.parallelize([Row(features=Vectors.sparse(2, [0], [1.0]))]).toDF()
     >>> model.transform(test1).head().prediction
     1.0
+
+    .. versionadded:: 1.5.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -587,6 +623,7 @@ def __init__(self, featuresCol="features", labelCol="label", predictionCol="pred
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.5.0")
     def setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                   probabilityCol="probability", rawPredictionCol="rawPrediction", smoothing=1.0,
                   modelType="multinomial"):
@@ -602,6 +639,7 @@ def setParams(self, featuresCol="features", labelCol="label", predictionCol="pre
     def _create_model(self, java_model):
         return NaiveBayesModel(java_model)
 
+    @since("1.5.0")
     def setSmoothing(self, value):
         """
         Sets the value of :py:attr:`smoothing`.
@@ -609,12 +647,14 @@ def setSmoothing(self, value):
         self._paramMap[self.smoothing] = value
         return self
 
+    @since("1.5.0")
     def getSmoothing(self):
         """
         Gets the value of smoothing or its default value.
         """
         return self.getOrDefault(self.smoothing)
 
+    @since("1.5.0")
     def setModelType(self, value):
         """
         Sets the value of :py:attr:`modelType`.
@@ -622,6 +662,7 @@ def setModelType(self, value):
         self._paramMap[self.modelType] = value
         return self
 
+    @since("1.5.0")
     def getModelType(self):
         """
         Gets the value of modelType or its default value.
@@ -632,9 +673,12 @@ def getModelType(self):
 class NaiveBayesModel(JavaModel):
     """
     Model fitted by NaiveBayes.
+
+    .. versionadded:: 1.5.0
     """
 
     @property
+    @since("1.5.0")
     def pi(self):
         """
         log of class priors.
@@ -642,6 +686,7 @@ def pi(self):
         return self._call_java("pi")
 
     @property
+    @since("1.5.0")
     def theta(self):
         """
         log of class conditional probabilities.
@@ -681,6 +726,8 @@ class MultilayerPerceptronClassifier(JavaEstimator, HasFeaturesCol, HasLabelCol,
     |[0.0,0.0]|       0.0|
     +---------+----------+
     ...
+
+    .. versionadded:: 1.6.0
     """
 
     # a placeholder to make it appear in the generated doc
@@ -715,6 +762,7 @@ def __init__(self, featuresCol="features", labelCol="label", predictionCol="pred
         self.setParams(**kwargs)
 
     @keyword_only
+    @since("1.6.0")
     def setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                   maxIter=100, tol=1e-4, seed=None, layers=None, blockSize=128):
         """
@@ -731,6 +779,7 @@ def setParams(self, featuresCol="features", labelCol="label", predictionCol="pre
     def _create_model(self, java_model):
         return MultilayerPerceptronClassificationModel(java_model)
 
+    @since("1.6.0")
     def setLayers(self, value):
         """
         Sets the value of :py:attr:`layers`.
@@ -738,12 +787,14 @@ def setLayers(self, value):
         self._paramMap[self.layers] = value
         return self
 
+    @since("1.6.0")
     def getLayers(self):
         """
         Gets the value of layers or its default value.
         """
         return self.getOrDefault(self.layers)
 
+    @since("1.6.0")
     def setBlockSize(self, value):
         """
         Sets the value of :py:attr:`blockSize`.
@@ -751,6 +802,7 @@ def setBlockSize(self, value):
         self._paramMap[self.blockSize] = value
         return self
 
+    @since("1.6.0")
     def getBlockSize(self):
         """
         Gets the value of blockSize or its default value.
@@ -761,9 +813,12 @@ def getBlockSize(self):
 class MultilayerPerceptronClassificationModel(JavaModel):
     """
     Model fitted by MultilayerPerceptronClassifier.
+
+    .. versionadded:: 1.6.0
     """
 
     @property
+    @since("1.6.0")
     def layers(self):
         """
         array of layer sizes including input and output layers.
@@ -771,6 +826,7 @@ def layers(self):
         return self._call_java("javaLayers")
 
     @property
+    @since("1.6.0")
     def weights(self):
         """
         vector of initial weights for the model that consists of the weights of layers.

From 5039a49b636325f321daa089971107003fae9d4b Mon Sep 17 00:00:00 2001
From: Felix Bechstein <felix.bechstein@otto.de>
Date: Mon, 9 Nov 2015 13:36:14 -0800
Subject: [PATCH 131/518] [SPARK-10471][CORE][MESOS] prevent getting offers for
 unmet constraints

this change rejects offers for slaves with unmet constraints for 120s to mitigate offer starvation.
this prevents mesos to send us these offers again and again.
in return, we get more offers for slaves which might meet our constraints.
and it enables mesos to send the rejected offers to other frameworks.

Author: Felix Bechstein <felix.bechstein@otto.de>

Closes #8639 from felixb/decline_offers_constraint_mismatch.
---
 .../mesos/CoarseMesosSchedulerBackend.scala   | 92 +++++++++++--------
 .../cluster/mesos/MesosSchedulerBackend.scala | 48 +++++++---
 .../cluster/mesos/MesosSchedulerUtils.scala   |  4 +
 3 files changed, 91 insertions(+), 53 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
index d10a77f8e5c78..2de9b6a651692 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
@@ -101,6 +101,10 @@ private[spark] class CoarseMesosSchedulerBackend(
   private val slaveOfferConstraints =
     parseConstraintString(sc.conf.get("spark.mesos.constraints", ""))
 
+  // reject offers with mismatched constraints in seconds
+  private val rejectOfferDurationForUnmetConstraints =
+    getRejectOfferDurationForUnmetConstraints(sc)
+
   // A client for talking to the external shuffle service, if it is a
   private val mesosExternalShuffleClient: Option[MesosExternalShuffleClient] = {
     if (shuffleServiceEnabled) {
@@ -249,48 +253,56 @@ private[spark] class CoarseMesosSchedulerBackend(
         val mem = getResource(offer.getResourcesList, "mem")
         val cpus = getResource(offer.getResourcesList, "cpus").toInt
         val id = offer.getId.getValue
-        if (taskIdToSlaveId.size < executorLimit &&
-            totalCoresAcquired < maxCores &&
-            meetsConstraints &&
-            mem >= calculateTotalMemory(sc) &&
-            cpus >= 1 &&
-            failuresBySlaveId.getOrElse(slaveId, 0) < MAX_SLAVE_FAILURES &&
-            !slaveIdsWithExecutors.contains(slaveId)) {
-          // Launch an executor on the slave
-          val cpusToUse = math.min(cpus, maxCores - totalCoresAcquired)
-          totalCoresAcquired += cpusToUse
-          val taskId = newMesosTaskId()
-          taskIdToSlaveId.put(taskId, slaveId)
-          slaveIdsWithExecutors += slaveId
-          coresByTaskId(taskId) = cpusToUse
-          // Gather cpu resources from the available resources and use them in the task.
-          val (remainingResources, cpuResourcesToUse) =
-            partitionResources(offer.getResourcesList, "cpus", cpusToUse)
-          val (_, memResourcesToUse) =
-            partitionResources(remainingResources.asJava, "mem", calculateTotalMemory(sc))
-          val taskBuilder = MesosTaskInfo.newBuilder()
-            .setTaskId(TaskID.newBuilder().setValue(taskId.toString).build())
-            .setSlaveId(offer.getSlaveId)
-            .setCommand(createCommand(offer, cpusToUse + extraCoresPerSlave, taskId))
-            .setName("Task " + taskId)
-            .addAllResources(cpuResourcesToUse.asJava)
-            .addAllResources(memResourcesToUse.asJava)
-
-          sc.conf.getOption("spark.mesos.executor.docker.image").foreach { image =>
-            MesosSchedulerBackendUtil
-              .setupContainerBuilderDockerInfo(image, sc.conf, taskBuilder.getContainerBuilder())
+        if (meetsConstraints) {
+          if (taskIdToSlaveId.size < executorLimit &&
+              totalCoresAcquired < maxCores &&
+              mem >= calculateTotalMemory(sc) &&
+              cpus >= 1 &&
+              failuresBySlaveId.getOrElse(slaveId, 0) < MAX_SLAVE_FAILURES &&
+              !slaveIdsWithExecutors.contains(slaveId)) {
+            // Launch an executor on the slave
+            val cpusToUse = math.min(cpus, maxCores - totalCoresAcquired)
+            totalCoresAcquired += cpusToUse
+            val taskId = newMesosTaskId()
+            taskIdToSlaveId.put(taskId, slaveId)
+            slaveIdsWithExecutors += slaveId
+            coresByTaskId(taskId) = cpusToUse
+            // Gather cpu resources from the available resources and use them in the task.
+            val (remainingResources, cpuResourcesToUse) =
+              partitionResources(offer.getResourcesList, "cpus", cpusToUse)
+            val (_, memResourcesToUse) =
+              partitionResources(remainingResources.asJava, "mem", calculateTotalMemory(sc))
+            val taskBuilder = MesosTaskInfo.newBuilder()
+              .setTaskId(TaskID.newBuilder().setValue(taskId.toString).build())
+              .setSlaveId(offer.getSlaveId)
+              .setCommand(createCommand(offer, cpusToUse + extraCoresPerSlave, taskId))
+              .setName("Task " + taskId)
+              .addAllResources(cpuResourcesToUse.asJava)
+              .addAllResources(memResourcesToUse.asJava)
+
+            sc.conf.getOption("spark.mesos.executor.docker.image").foreach { image =>
+              MesosSchedulerBackendUtil
+                .setupContainerBuilderDockerInfo(image, sc.conf, taskBuilder.getContainerBuilder())
+            }
+
+            // Accept the offer and launch the task
+            logDebug(s"Accepting offer: $id with attributes: $offerAttributes mem: $mem cpu: $cpus")
+            slaveIdToHost(offer.getSlaveId.getValue) = offer.getHostname
+            d.launchTasks(
+              Collections.singleton(offer.getId),
+              Collections.singleton(taskBuilder.build()), filters)
+          } else {
+            // Decline the offer
+            logDebug(s"Declining offer: $id with attributes: $offerAttributes mem: $mem cpu: $cpus")
+            d.declineOffer(offer.getId)
           }
-
-          // accept the offer and launch the task
-          logDebug(s"Accepting offer: $id with attributes: $offerAttributes mem: $mem cpu: $cpus")
-          slaveIdToHost(offer.getSlaveId.getValue) = offer.getHostname
-          d.launchTasks(
-            Collections.singleton(offer.getId),
-            Collections.singleton(taskBuilder.build()), filters)
         } else {
-          // Decline the offer
-          logDebug(s"Declining offer: $id with attributes: $offerAttributes mem: $mem cpu: $cpus")
-          d.declineOffer(offer.getId)
+          // This offer does not meet constraints. We don't need to see it again.
+          // Decline the offer for a long period of time.
+          logDebug(s"Declining offer: $id with attributes: $offerAttributes mem: $mem cpu: $cpus"
+              + s" for $rejectOfferDurationForUnmetConstraints seconds")
+          d.declineOffer(offer.getId, Filters.newBuilder()
+            .setRefuseSeconds(rejectOfferDurationForUnmetConstraints).build())
         }
       }
     }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala
index aaffac604a885..281965a5981bb 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerBackend.scala
@@ -63,6 +63,10 @@ private[spark] class MesosSchedulerBackend(
   private[this] val slaveOfferConstraints =
     parseConstraintString(sc.conf.get("spark.mesos.constraints", ""))
 
+  // reject offers with mismatched constraints in seconds
+  private val rejectOfferDurationForUnmetConstraints =
+    getRejectOfferDurationForUnmetConstraints(sc)
+
   @volatile var appId: String = _
 
   override def start() {
@@ -212,29 +216,47 @@ private[spark] class MesosSchedulerBackend(
    */
   override def resourceOffers(d: SchedulerDriver, offers: JList[Offer]) {
     inClassLoader() {
-      // Fail-fast on offers we know will be rejected
-      val (usableOffers, unUsableOffers) = offers.asScala.partition { o =>
+      // Fail first on offers with unmet constraints
+      val (offersMatchingConstraints, offersNotMatchingConstraints) =
+        offers.asScala.partition { o =>
+          val offerAttributes = toAttributeMap(o.getAttributesList)
+          val meetsConstraints =
+            matchesAttributeRequirements(slaveOfferConstraints, offerAttributes)
+
+          // add some debug messaging
+          if (!meetsConstraints) {
+            val id = o.getId.getValue
+            logDebug(s"Declining offer: $id with attributes: $offerAttributes")
+          }
+
+          meetsConstraints
+        }
+
+      // These offers do not meet constraints. We don't need to see them again.
+      // Decline the offer for a long period of time.
+      offersNotMatchingConstraints.foreach { o =>
+        d.declineOffer(o.getId, Filters.newBuilder()
+          .setRefuseSeconds(rejectOfferDurationForUnmetConstraints).build())
+      }
+
+      // Of the matching constraints, see which ones give us enough memory and cores
+      val (usableOffers, unUsableOffers) = offersMatchingConstraints.partition { o =>
         val mem = getResource(o.getResourcesList, "mem")
         val cpus = getResource(o.getResourcesList, "cpus")
         val slaveId = o.getSlaveId.getValue
         val offerAttributes = toAttributeMap(o.getAttributesList)
 
-        // check if all constraints are satisfield
-        //  1. Attribute constraints
-        //  2. Memory requirements
-        //  3. CPU requirements - need at least 1 for executor, 1 for task
-        val meetsConstraints = matchesAttributeRequirements(slaveOfferConstraints, offerAttributes)
+        // check offers for
+        //  1. Memory requirements
+        //  2. CPU requirements - need at least 1 for executor, 1 for task
         val meetsMemoryRequirements = mem >= calculateTotalMemory(sc)
         val meetsCPURequirements = cpus >= (mesosExecutorCores + scheduler.CPUS_PER_TASK)
-
         val meetsRequirements =
-          (meetsConstraints && meetsMemoryRequirements && meetsCPURequirements) ||
+          (meetsMemoryRequirements && meetsCPURequirements) ||
           (slaveIdToExecutorInfo.contains(slaveId) && cpus >= scheduler.CPUS_PER_TASK)
-
-        // add some debug messaging
         val debugstr = if (meetsRequirements) "Accepting" else "Declining"
-        val id = o.getId.getValue
-        logDebug(s"$debugstr offer: $id with attributes: $offerAttributes mem: $mem cpu: $cpus")
+        logDebug(s"$debugstr offer: ${o.getId.getValue} with attributes: "
+          + s"$offerAttributes mem: $mem cpu: $cpus")
 
         meetsRequirements
       }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerUtils.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerUtils.scala
index 860c8e097b3b9..721861fbbc517 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerUtils.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/MesosSchedulerUtils.scala
@@ -336,4 +336,8 @@ private[mesos] trait MesosSchedulerUtils extends Logging {
     }
   }
 
+  protected def getRejectOfferDurationForUnmetConstraints(sc: SparkContext): Long = {
+    sc.conf.getTimeAsSeconds("spark.mesos.rejectOfferDurationForUnmetConstraints", "120s")
+  }
+
 }

From 51d41e4b1a3a25a3fde3a4345afcfe4766023d23 Mon Sep 17 00:00:00 2001
From: sachin aggarwal <different.sachin@gmail.com>
Date: Mon, 9 Nov 2015 14:25:42 -0800
Subject: [PATCH 132/518] [SPARK-11552][DOCS][Replaced example code in
 ml-decision-tree.md using include_example]

I have tested it on my local, it is working fine, please review

Author: sachin aggarwal <different.sachin@gmail.com>

Closes #9539 from agsachin/SPARK-11552-real.
---
 docs/ml-decision-tree.md                      | 338 +-----------------
 ...JavaDecisionTreeClassificationExample.java | 103 ++++++
 .../ml/JavaDecisionTreeRegressionExample.java |  90 +++++
 .../decision_tree_classification_example.py   |  77 ++++
 .../ml/decision_tree_regression_example.py    |  74 ++++
 .../DecisionTreeClassificationExample.scala   |  94 +++++
 .../ml/DecisionTreeRegressionExample.scala    |  81 +++++
 7 files changed, 527 insertions(+), 330 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java
 create mode 100644 examples/src/main/python/ml/decision_tree_classification_example.py
 create mode 100644 examples/src/main/python/ml/decision_tree_regression_example.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala

diff --git a/docs/ml-decision-tree.md b/docs/ml-decision-tree.md
index 542819e93e6dc..2bfac6f6c8378 100644
--- a/docs/ml-decision-tree.md
+++ b/docs/ml-decision-tree.md
@@ -118,196 +118,24 @@ We use two feature transformers to prepare the data; these help index categories
 
 More details on parameters can be found in the [Scala API documentation](api/scala/index.html#org.apache.spark.ml.classification.DecisionTreeClassifier).
 
-{% highlight scala %}
-import org.apache.spark.ml.Pipeline
-import org.apache.spark.ml.classification.DecisionTreeClassifier
-import org.apache.spark.ml.classification.DecisionTreeClassificationModel
-import org.apache.spark.ml.feature.{StringIndexer, IndexToString, VectorIndexer}
-import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
-import org.apache.spark.mllib.util.MLUtils
-
-// Load and parse the data file, converting it to a DataFrame.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
-
-// Index labels, adding metadata to the label column.
-// Fit on whole dataset to include all labels in index.
-val labelIndexer = new StringIndexer()
-  .setInputCol("label")
-  .setOutputCol("indexedLabel")
-  .fit(data)
-// Automatically identify categorical features, and index them.
-val featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4) // features with > 4 distinct values are treated as continuous
-  .fit(data)
-
-// Split the data into training and test sets (30% held out for testing)
-val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
-
-// Train a DecisionTree model.
-val dt = new DecisionTreeClassifier()
-  .setLabelCol("indexedLabel")
-  .setFeaturesCol("indexedFeatures")
-
-// Convert indexed labels back to original labels.
-val labelConverter = new IndexToString()
-  .setInputCol("prediction")
-  .setOutputCol("predictedLabel")
-  .setLabels(labelIndexer.labels)
-
-// Chain indexers and tree in a Pipeline
-val pipeline = new Pipeline()
-  .setStages(Array(labelIndexer, featureIndexer, dt, labelConverter))
-
-// Train model.  This also runs the indexers.
-val model = pipeline.fit(trainingData)
-
-// Make predictions.
-val predictions = model.transform(testData)
-
-// Select example rows to display.
-predictions.select("predictedLabel", "label", "features").show(5)
-
-// Select (prediction, true label) and compute test error
-val evaluator = new MulticlassClassificationEvaluator()
-  .setLabelCol("indexedLabel")
-  .setPredictionCol("prediction")
-  .setMetricName("precision")
-val accuracy = evaluator.evaluate(predictions)
-println("Test Error = " + (1.0 - accuracy))
-
-val treeModel = model.stages(2).asInstanceOf[DecisionTreeClassificationModel]
-println("Learned classification tree model:\n" + treeModel.toDebugString)
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala %}
+
 </div>
 
 <div data-lang="java" markdown="1">
 
 More details on parameters can be found in the [Java API documentation](api/java/org/apache/spark/ml/classification/DecisionTreeClassifier.html).
 
-{% highlight java %}
-import org.apache.spark.ml.Pipeline;
-import org.apache.spark.ml.PipelineModel;
-import org.apache.spark.ml.PipelineStage;
-import org.apache.spark.ml.classification.DecisionTreeClassifier;
-import org.apache.spark.ml.classification.DecisionTreeClassificationModel;
-import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
-import org.apache.spark.ml.feature.*;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.rdd.RDD;
-import org.apache.spark.sql.DataFrame;
-
-// Load and parse the data file, converting it to a DataFrame.
-RDD<LabeledPoint> rdd = MLUtils.loadLibSVMFile(sc.sc(), "data/mllib/sample_libsvm_data.txt");
-DataFrame data = jsql.createDataFrame(rdd, LabeledPoint.class);
-
-// Index labels, adding metadata to the label column.
-// Fit on whole dataset to include all labels in index.
-StringIndexerModel labelIndexer = new StringIndexer()
-  .setInputCol("label")
-  .setOutputCol("indexedLabel")
-  .fit(data);
-// Automatically identify categorical features, and index them.
-VectorIndexerModel featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4) // features with > 4 distinct values are treated as continuous
-  .fit(data);
-
-// Split the data into training and test sets (30% held out for testing)
-DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
-DataFrame trainingData = splits[0];
-DataFrame testData = splits[1];
-
-// Train a DecisionTree model.
-DecisionTreeClassifier dt = new DecisionTreeClassifier()
-  .setLabelCol("indexedLabel")
-  .setFeaturesCol("indexedFeatures");
-
-// Convert indexed labels back to original labels.
-IndexToString labelConverter = new IndexToString()
-  .setInputCol("prediction")
-  .setOutputCol("predictedLabel")
-  .setLabels(labelIndexer.labels());
-
-// Chain indexers and tree in a Pipeline
-Pipeline pipeline = new Pipeline()
-  .setStages(new PipelineStage[] {labelIndexer, featureIndexer, dt, labelConverter});
-
-// Train model.  This also runs the indexers.
-PipelineModel model = pipeline.fit(trainingData);
-
-// Make predictions.
-DataFrame predictions = model.transform(testData);
-
-// Select example rows to display.
-predictions.select("predictedLabel", "label", "features").show(5);
-
-// Select (prediction, true label) and compute test error
-MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
-  .setLabelCol("indexedLabel")
-  .setPredictionCol("prediction")
-  .setMetricName("precision");
-double accuracy = evaluator.evaluate(predictions);
-System.out.println("Test Error = " + (1.0 - accuracy));
-
-DecisionTreeClassificationModel treeModel =
-  (DecisionTreeClassificationModel)(model.stages()[2]);
-System.out.println("Learned classification tree model:\n" + treeModel.toDebugString());
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java %}
+
 </div>
 
 <div data-lang="python" markdown="1">
 
 More details on parameters can be found in the [Python API documentation](api/python/pyspark.ml.html#pyspark.ml.classification.DecisionTreeClassifier).
 
-{% highlight python %}
-from pyspark.ml import Pipeline
-from pyspark.ml.classification import DecisionTreeClassifier
-from pyspark.ml.feature import StringIndexer, VectorIndexer
-from pyspark.ml.evaluation import MulticlassClassificationEvaluator
-from pyspark.mllib.util import MLUtils
-
-# Load and parse the data file, converting it to a DataFrame.
-data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
-
-# Index labels, adding metadata to the label column.
-# Fit on whole dataset to include all labels in index.
-labelIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel").fit(data)
-# Automatically identify categorical features, and index them.
-# We specify maxCategories so features with > 4 distinct values are treated as continuous.
-featureIndexer =\
-    VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
-
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a DecisionTree model.
-dt = DecisionTreeClassifier(labelCol="indexedLabel", featuresCol="indexedFeatures")
-
-# Chain indexers and tree in a Pipeline
-pipeline = Pipeline(stages=[labelIndexer, featureIndexer, dt])
-
-# Train model.  This also runs the indexers.
-model = pipeline.fit(trainingData)
-
-# Make predictions.
-predictions = model.transform(testData)
-
-# Select example rows to display.
-predictions.select("prediction", "indexedLabel", "features").show(5)
-
-# Select (prediction, true label) and compute test error
-evaluator = MulticlassClassificationEvaluator(
-    labelCol="indexedLabel", predictionCol="prediction", metricName="precision")
-accuracy = evaluator.evaluate(predictions)
-print "Test Error = %g" % (1.0 - accuracy)
+{% include_example python/ml/decision_tree_classification_example.py %}
 
-treeModel = model.stages[2]
-print treeModel # summary only
-{% endhighlight %}
 </div>
 
 </div>
@@ -323,171 +151,21 @@ We use a feature transformer to index categorical features, adding metadata to t
 
 More details on parameters can be found in the [Scala API documentation](api/scala/index.html#org.apache.spark.ml.regression.DecisionTreeRegressor).
 
-{% highlight scala %}
-import org.apache.spark.ml.Pipeline
-import org.apache.spark.ml.regression.DecisionTreeRegressor
-import org.apache.spark.ml.regression.DecisionTreeRegressionModel
-import org.apache.spark.ml.feature.VectorIndexer
-import org.apache.spark.ml.evaluation.RegressionEvaluator
-import org.apache.spark.mllib.util.MLUtils
-
-// Load and parse the data file, converting it to a DataFrame.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
-
-// Automatically identify categorical features, and index them.
-// Here, we treat features with > 4 distinct values as continuous.
-val featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data)
-
-// Split the data into training and test sets (30% held out for testing)
-val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
-
-// Train a DecisionTree model.
-val dt = new DecisionTreeRegressor()
-  .setLabelCol("label")
-  .setFeaturesCol("indexedFeatures")
-
-// Chain indexer and tree in a Pipeline
-val pipeline = new Pipeline()
-  .setStages(Array(featureIndexer, dt))
-
-// Train model.  This also runs the indexer.
-val model = pipeline.fit(trainingData)
-
-// Make predictions.
-val predictions = model.transform(testData)
-
-// Select example rows to display.
-predictions.select("prediction", "label", "features").show(5)
-
-// Select (prediction, true label) and compute test error
-val evaluator = new RegressionEvaluator()
-  .setLabelCol("label")
-  .setPredictionCol("prediction")
-  .setMetricName("rmse")
-val rmse = evaluator.evaluate(predictions)
-println("Root Mean Squared Error (RMSE) on test data = " + rmse)
-
-val treeModel = model.stages(1).asInstanceOf[DecisionTreeRegressionModel]
-println("Learned regression tree model:\n" + treeModel.toDebugString)
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 
 More details on parameters can be found in the [Java API documentation](api/java/org/apache/spark/ml/regression/DecisionTreeRegressor.html).
 
-{% highlight java %}
-import org.apache.spark.ml.Pipeline;
-import org.apache.spark.ml.PipelineModel;
-import org.apache.spark.ml.PipelineStage;
-import org.apache.spark.ml.evaluation.RegressionEvaluator;
-import org.apache.spark.ml.feature.VectorIndexer;
-import org.apache.spark.ml.feature.VectorIndexerModel;
-import org.apache.spark.ml.regression.DecisionTreeRegressionModel;
-import org.apache.spark.ml.regression.DecisionTreeRegressor;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.rdd.RDD;
-import org.apache.spark.sql.DataFrame;
-
-// Load and parse the data file, converting it to a DataFrame.
-RDD<LabeledPoint> rdd = MLUtils.loadLibSVMFile(sc.sc(), "data/mllib/sample_libsvm_data.txt");
-DataFrame data = jsql.createDataFrame(rdd, LabeledPoint.class);
-
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-VectorIndexerModel featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data);
-
-// Split the data into training and test sets (30% held out for testing)
-DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
-DataFrame trainingData = splits[0];
-DataFrame testData = splits[1];
-
-// Train a DecisionTree model.
-DecisionTreeRegressor dt = new DecisionTreeRegressor()
-  .setFeaturesCol("indexedFeatures");
-
-// Chain indexer and tree in a Pipeline
-Pipeline pipeline = new Pipeline()
-  .setStages(new PipelineStage[] {featureIndexer, dt});
-
-// Train model.  This also runs the indexer.
-PipelineModel model = pipeline.fit(trainingData);
-
-// Make predictions.
-DataFrame predictions = model.transform(testData);
-
-// Select example rows to display.
-predictions.select("label", "features").show(5);
-
-// Select (prediction, true label) and compute test error
-RegressionEvaluator evaluator = new RegressionEvaluator()
-  .setLabelCol("label")
-  .setPredictionCol("prediction")
-  .setMetricName("rmse");
-double rmse = evaluator.evaluate(predictions);
-System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse);
-
-DecisionTreeRegressionModel treeModel =
-  (DecisionTreeRegressionModel)(model.stages()[1]);
-System.out.println("Learned regression tree model:\n" + treeModel.toDebugString());
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 
 More details on parameters can be found in the [Python API documentation](api/python/pyspark.ml.html#pyspark.ml.regression.DecisionTreeRegressor).
 
-{% highlight python %}
-from pyspark.ml import Pipeline
-from pyspark.ml.regression import DecisionTreeRegressor
-from pyspark.ml.feature import VectorIndexer
-from pyspark.ml.evaluation import RegressionEvaluator
-from pyspark.mllib.util import MLUtils
-
-# Load and parse the data file, converting it to a DataFrame.
-data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
-
-# Automatically identify categorical features, and index them.
-# We specify maxCategories so features with > 4 distinct values are treated as continuous.
-featureIndexer =\
-    VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
-
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a DecisionTree model.
-dt = DecisionTreeRegressor(featuresCol="indexedFeatures")
-
-# Chain indexer and tree in a Pipeline
-pipeline = Pipeline(stages=[featureIndexer, dt])
-
-# Train model.  This also runs the indexer.
-model = pipeline.fit(trainingData)
-
-# Make predictions.
-predictions = model.transform(testData)
-
-# Select example rows to display.
-predictions.select("prediction", "label", "features").show(5)
-
-# Select (prediction, true label) and compute test error
-evaluator = RegressionEvaluator(
-    labelCol="label", predictionCol="prediction", metricName="rmse")
-rmse = evaluator.evaluate(predictions)
-print "Root Mean Squared Error (RMSE) on test data = %g" % rmse
-
-treeModel = model.stages[1]
-print treeModel # summary only
-{% endhighlight %}
+{% include_example python/ml/decision_tree_regression_example.py %}
 </div>
 
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java
new file mode 100644
index 0000000000000..51c1730a8a085
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+// scalastyle:off println
+package org.apache.spark.examples.ml;
+// $example on$
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.ml.Pipeline;
+import org.apache.spark.ml.PipelineModel;
+import org.apache.spark.ml.PipelineStage;
+import org.apache.spark.ml.classification.DecisionTreeClassifier;
+import org.apache.spark.ml.classification.DecisionTreeClassificationModel;
+import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
+import org.apache.spark.ml.feature.*;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.util.MLUtils;
+import org.apache.spark.rdd.RDD;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+// $example off$
+
+public class JavaDecisionTreeClassificationExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaDecisionTreeClassificationExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    RDD<LabeledPoint> rdd = MLUtils.loadLibSVMFile(jsc.sc(), "data/mllib/sample_libsvm_data.txt");
+    DataFrame data = sqlContext.createDataFrame(rdd, LabeledPoint.class);
+
+    // Index labels, adding metadata to the label column.
+    // Fit on whole dataset to include all labels in index.
+    StringIndexerModel labelIndexer = new StringIndexer()
+      .setInputCol("label")
+      .setOutputCol("indexedLabel")
+      .fit(data);
+
+    // Automatically identify categorical features, and index them.
+    VectorIndexerModel featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4) // features with > 4 distinct values are treated as continuous
+      .fit(data);
+
+    // Split the data into training and test sets (30% held out for testing)
+    DataFrame[] splits = data.randomSplit(new double[]{0.7, 0.3});
+    DataFrame trainingData = splits[0];
+    DataFrame testData = splits[1];
+
+    // Train a DecisionTree model.
+    DecisionTreeClassifier dt = new DecisionTreeClassifier()
+      .setLabelCol("indexedLabel")
+      .setFeaturesCol("indexedFeatures");
+
+    // Convert indexed labels back to original labels.
+    IndexToString labelConverter = new IndexToString()
+      .setInputCol("prediction")
+      .setOutputCol("predictedLabel")
+      .setLabels(labelIndexer.labels());
+
+    // Chain indexers and tree in a Pipeline
+    Pipeline pipeline = new Pipeline()
+      .setStages(new PipelineStage[]{labelIndexer, featureIndexer, dt, labelConverter});
+
+    // Train model.  This also runs the indexers.
+    PipelineModel model = pipeline.fit(trainingData);
+
+    // Make predictions.
+    DataFrame predictions = model.transform(testData);
+
+    // Select example rows to display.
+    predictions.select("predictedLabel", "label", "features").show(5);
+
+    // Select (prediction, true label) and compute test error
+    MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
+      .setLabelCol("indexedLabel")
+      .setPredictionCol("prediction")
+      .setMetricName("precision");
+    double accuracy = evaluator.evaluate(predictions);
+    System.out.println("Test Error = " + (1.0 - accuracy));
+
+    DecisionTreeClassificationModel treeModel =
+      (DecisionTreeClassificationModel) (model.stages()[2]);
+    System.out.println("Learned classification tree model:\n" + treeModel.toDebugString());
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java
new file mode 100644
index 0000000000000..a4098a4233ec2
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+// scalastyle:off println
+package org.apache.spark.examples.ml;
+// $example on$
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.ml.Pipeline;
+import org.apache.spark.ml.PipelineModel;
+import org.apache.spark.ml.PipelineStage;
+import org.apache.spark.ml.evaluation.RegressionEvaluator;
+import org.apache.spark.ml.feature.VectorIndexer;
+import org.apache.spark.ml.feature.VectorIndexerModel;
+import org.apache.spark.ml.regression.DecisionTreeRegressionModel;
+import org.apache.spark.ml.regression.DecisionTreeRegressor;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.util.MLUtils;
+import org.apache.spark.rdd.RDD;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+// $example off$
+
+public class JavaDecisionTreeRegressionExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaDecisionTreeRegressionExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    RDD<LabeledPoint> rdd = MLUtils.loadLibSVMFile(jsc.sc(), "data/mllib/sample_libsvm_data.txt");
+    DataFrame data = sqlContext.createDataFrame(rdd, LabeledPoint.class);
+
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    VectorIndexerModel featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data);
+
+    // Split the data into training and test sets (30% held out for testing)
+    DataFrame[] splits = data.randomSplit(new double[]{0.7, 0.3});
+    DataFrame trainingData = splits[0];
+    DataFrame testData = splits[1];
+
+    // Train a DecisionTree model.
+    DecisionTreeRegressor dt = new DecisionTreeRegressor()
+      .setFeaturesCol("indexedFeatures");
+
+    // Chain indexer and tree in a Pipeline
+    Pipeline pipeline = new Pipeline()
+      .setStages(new PipelineStage[]{featureIndexer, dt});
+
+    // Train model.  This also runs the indexer.
+    PipelineModel model = pipeline.fit(trainingData);
+
+    // Make predictions.
+    DataFrame predictions = model.transform(testData);
+
+    // Select example rows to display.
+    predictions.select("label", "features").show(5);
+
+    // Select (prediction, true label) and compute test error
+    RegressionEvaluator evaluator = new RegressionEvaluator()
+      .setLabelCol("label")
+      .setPredictionCol("prediction")
+      .setMetricName("rmse");
+    double rmse = evaluator.evaluate(predictions);
+    System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse);
+
+    DecisionTreeRegressionModel treeModel =
+      (DecisionTreeRegressionModel) (model.stages()[1]);
+    System.out.println("Learned regression tree model:\n" + treeModel.toDebugString());
+    // $example off$
+  }
+}
diff --git a/examples/src/main/python/ml/decision_tree_classification_example.py b/examples/src/main/python/ml/decision_tree_classification_example.py
new file mode 100644
index 0000000000000..0af92050e3e3b
--- /dev/null
+++ b/examples/src/main/python/ml/decision_tree_classification_example.py
@@ -0,0 +1,77 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Decision Tree Classification Example.
+"""
+from __future__ import print_function
+
+import sys
+
+# $example on$
+from pyspark import SparkContext, SQLContext
+from pyspark.ml import Pipeline
+from pyspark.ml.classification import DecisionTreeClassifier
+from pyspark.ml.feature import StringIndexer, VectorIndexer
+from pyspark.ml.evaluation import MulticlassClassificationEvaluator
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="decision_tree_classification_example")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load and parse the data file, converting it to a DataFrame.
+    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+
+    # Index labels, adding metadata to the label column.
+    # Fit on whole dataset to include all labels in index.
+    labelIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel").fit(data)
+    # Automatically identify categorical features, and index them.
+    # We specify maxCategories so features with > 4 distinct values are treated as continuous.
+    featureIndexer =\
+        VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
+
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a DecisionTree model.
+    dt = DecisionTreeClassifier(labelCol="indexedLabel", featuresCol="indexedFeatures")
+
+    # Chain indexers and tree in a Pipeline
+    pipeline = Pipeline(stages=[labelIndexer, featureIndexer, dt])
+
+    # Train model.  This also runs the indexers.
+    model = pipeline.fit(trainingData)
+
+    # Make predictions.
+    predictions = model.transform(testData)
+
+    # Select example rows to display.
+    predictions.select("prediction", "indexedLabel", "features").show(5)
+
+    # Select (prediction, true label) and compute test error
+    evaluator = MulticlassClassificationEvaluator(
+        labelCol="indexedLabel", predictionCol="prediction", metricName="precision")
+    accuracy = evaluator.evaluate(predictions)
+    print("Test Error = %g " % (1.0 - accuracy))
+
+    treeModel = model.stages[2]
+    # summary only
+    print(treeModel)
+    # $example off$
diff --git a/examples/src/main/python/ml/decision_tree_regression_example.py b/examples/src/main/python/ml/decision_tree_regression_example.py
new file mode 100644
index 0000000000000..3857aed538da2
--- /dev/null
+++ b/examples/src/main/python/ml/decision_tree_regression_example.py
@@ -0,0 +1,74 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Decision Tree Regression Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext, SQLContext
+# $example on$
+from pyspark.ml import Pipeline
+from pyspark.ml.regression import DecisionTreeRegressor
+from pyspark.ml.feature import VectorIndexer
+from pyspark.ml.evaluation import RegressionEvaluator
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="decision_tree_classification_example")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load and parse the data file, converting it to a DataFrame.
+    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+
+    # Automatically identify categorical features, and index them.
+    # We specify maxCategories so features with > 4 distinct values are treated as continuous.
+    featureIndexer =\
+        VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
+
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a DecisionTree model.
+    dt = DecisionTreeRegressor(featuresCol="indexedFeatures")
+
+    # Chain indexer and tree in a Pipeline
+    pipeline = Pipeline(stages=[featureIndexer, dt])
+
+    # Train model.  This also runs the indexer.
+    model = pipeline.fit(trainingData)
+
+    # Make predictions.
+    predictions = model.transform(testData)
+
+    # Select example rows to display.
+    predictions.select("prediction", "label", "features").show(5)
+
+    # Select (prediction, true label) and compute test error
+    evaluator = RegressionEvaluator(
+        labelCol="label", predictionCol="prediction", metricName="rmse")
+    rmse = evaluator.evaluate(predictions)
+    print("Root Mean Squared Error (RMSE) on test data = %g" % rmse)
+
+    treeModel = model.stages[1]
+    # summary only
+    print(treeModel)
+    # $example off$
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
new file mode 100644
index 0000000000000..a24a344f1bcf4
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
+import org.apache.spark.ml.Pipeline
+import org.apache.spark.ml.classification.DecisionTreeClassifier
+import org.apache.spark.ml.classification.DecisionTreeClassificationModel
+import org.apache.spark.ml.feature.{StringIndexer, IndexToString, VectorIndexer}
+import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+
+object DecisionTreeClassificationExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("DecisionTreeClassificationExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+    import sqlContext.implicits._
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+
+    // Index labels, adding metadata to the label column.
+    // Fit on whole dataset to include all labels in index.
+    val labelIndexer = new StringIndexer()
+      .setInputCol("label")
+      .setOutputCol("indexedLabel")
+      .fit(data)
+    // Automatically identify categorical features, and index them.
+    val featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4) // features with > 4 distinct values are treated as continuous
+      .fit(data)
+
+    // Split the data into training and test sets (30% held out for testing)
+    val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
+
+    // Train a DecisionTree model.
+    val dt = new DecisionTreeClassifier()
+      .setLabelCol("indexedLabel")
+      .setFeaturesCol("indexedFeatures")
+
+    // Convert indexed labels back to original labels.
+    val labelConverter = new IndexToString()
+      .setInputCol("prediction")
+      .setOutputCol("predictedLabel")
+      .setLabels(labelIndexer.labels)
+
+    // Chain indexers and tree in a Pipeline
+    val pipeline = new Pipeline()
+      .setStages(Array(labelIndexer, featureIndexer, dt, labelConverter))
+
+    // Train model.  This also runs the indexers.
+    val model = pipeline.fit(trainingData)
+
+    // Make predictions.
+    val predictions = model.transform(testData)
+
+    // Select example rows to display.
+    predictions.select("predictedLabel", "label", "features").show(5)
+
+    // Select (prediction, true label) and compute test error
+    val evaluator = new MulticlassClassificationEvaluator()
+      .setLabelCol("indexedLabel")
+      .setPredictionCol("prediction")
+      .setMetricName("precision")
+    val accuracy = evaluator.evaluate(predictions)
+    println("Test Error = " + (1.0 - accuracy))
+
+    val treeModel = model.stages(2).asInstanceOf[DecisionTreeClassificationModel]
+    println("Learned classification tree model:\n" + treeModel.toDebugString)
+    // $example off$
+  }
+}
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala
new file mode 100644
index 0000000000000..64cd986129007
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
+import org.apache.spark.ml.Pipeline
+import org.apache.spark.ml.regression.DecisionTreeRegressor
+import org.apache.spark.ml.regression.DecisionTreeRegressionModel
+import org.apache.spark.ml.feature.VectorIndexer
+import org.apache.spark.ml.evaluation.RegressionEvaluator
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+object DecisionTreeRegressionExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("DecisionTreeRegressionExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+    import sqlContext.implicits._
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+
+    // Automatically identify categorical features, and index them.
+    // Here, we treat features with > 4 distinct values as continuous.
+    val featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data)
+
+    // Split the data into training and test sets (30% held out for testing)
+    val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
+
+    // Train a DecisionTree model.
+    val dt = new DecisionTreeRegressor()
+      .setLabelCol("label")
+      .setFeaturesCol("indexedFeatures")
+
+    // Chain indexer and tree in a Pipeline
+    val pipeline = new Pipeline()
+      .setStages(Array(featureIndexer, dt))
+
+    // Train model.  This also runs the indexer.
+    val model = pipeline.fit(trainingData)
+
+    // Make predictions.
+    val predictions = model.transform(testData)
+
+    // Select example rows to display.
+    predictions.select("prediction", "label", "features").show(5)
+
+    // Select (prediction, true label) and compute test error
+    val evaluator = new RegressionEvaluator()
+      .setLabelCol("label")
+      .setPredictionCol("prediction")
+      .setMetricName("rmse")
+    val rmse = evaluator.evaluate(predictions)
+    println("Root Mean Squared Error (RMSE) on test data = " + rmse)
+
+    val treeModel = model.stages(1).asInstanceOf[DecisionTreeRegressionModel]
+    println("Learned regression tree model:\n" + treeModel.toDebugString)
+    // $example off$
+  }
+}

From b7720fa45525cff6e812fa448d0841cb41f6c8a5 Mon Sep 17 00:00:00 2001
From: Rishabh Bhardwaj <rbnext29@gmail.com>
Date: Mon, 9 Nov 2015 14:27:36 -0800
Subject: [PATCH 133/518] [SPARK-11548][DOCS] Replaced example code in
 mllib-collaborative-filtering.md using include_example

Kindly review the changes.

Author: Rishabh Bhardwaj <rbnext29@gmail.com>

Closes #9519 from rishabhbhardwaj/SPARK-11337.
---
 docs/mllib-collaborative-filtering.md         | 138 +-----------------
 .../mllib/JavaRecommendationExample.java      |  97 ++++++++++++
 .../python/mllib/recommendation_example.py    |  54 +++++++
 .../mllib/RecommendationExample.scala         |  67 +++++++++
 4 files changed, 221 insertions(+), 135 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaRecommendationExample.java
 create mode 100644 examples/src/main/python/mllib/recommendation_example.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/RecommendationExample.scala

diff --git a/docs/mllib-collaborative-filtering.md b/docs/mllib-collaborative-filtering.md
index 1ad52123c74aa..7cd1b894e7cb5 100644
--- a/docs/mllib-collaborative-filtering.md
+++ b/docs/mllib-collaborative-filtering.md
@@ -66,43 +66,7 @@ recommendation model by measuring the Mean Squared Error of rating prediction.
 
 Refer to the [`ALS` Scala docs](api/scala/index.html#org.apache.spark.mllib.recommendation.ALS) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.recommendation.ALS
-import org.apache.spark.mllib.recommendation.MatrixFactorizationModel
-import org.apache.spark.mllib.recommendation.Rating
-
-// Load and parse the data
-val data = sc.textFile("data/mllib/als/test.data")
-val ratings = data.map(_.split(',') match { case Array(user, item, rate) =>
-    Rating(user.toInt, item.toInt, rate.toDouble)
-  })
-
-// Build the recommendation model using ALS
-val rank = 10
-val numIterations = 10
-val model = ALS.train(ratings, rank, numIterations, 0.01)
-
-// Evaluate the model on rating data
-val usersProducts = ratings.map { case Rating(user, product, rate) =>
-  (user, product)
-}
-val predictions = 
-  model.predict(usersProducts).map { case Rating(user, product, rate) => 
-    ((user, product), rate)
-  }
-val ratesAndPreds = ratings.map { case Rating(user, product, rate) => 
-  ((user, product), rate)
-}.join(predictions)
-val MSE = ratesAndPreds.map { case ((user, product), (r1, r2)) => 
-  val err = (r1 - r2)
-  err * err
-}.mean()
-println("Mean Squared Error = " + MSE)
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = MatrixFactorizationModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/RecommendationExample.scala %}
 
 If the rating matrix is derived from another source of information (e.g., it is inferred from
 other signals), you can use the `trainImplicit` method to get better results.
@@ -123,81 +87,7 @@ that is equivalent to the provided example in Scala is given below:
 
 Refer to the [`ALS` Java docs](api/java/org/apache/spark/mllib/recommendation/ALS.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-
-import org.apache.spark.api.java.*;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.mllib.recommendation.ALS;
-import org.apache.spark.mllib.recommendation.MatrixFactorizationModel;
-import org.apache.spark.mllib.recommendation.Rating;
-import org.apache.spark.SparkConf;
-
-public class CollaborativeFiltering {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf().setAppName("Collaborative Filtering Example");
-    JavaSparkContext sc = new JavaSparkContext(conf);
-
-    // Load and parse the data
-    String path = "data/mllib/als/test.data";
-    JavaRDD<String> data = sc.textFile(path);
-    JavaRDD<Rating> ratings = data.map(
-      new Function<String, Rating>() {
-        public Rating call(String s) {
-          String[] sarray = s.split(",");
-          return new Rating(Integer.parseInt(sarray[0]), Integer.parseInt(sarray[1]), 
-                            Double.parseDouble(sarray[2]));
-        }
-      }
-    );
-
-    // Build the recommendation model using ALS
-    int rank = 10;
-    int numIterations = 10;
-    MatrixFactorizationModel model = ALS.train(JavaRDD.toRDD(ratings), rank, numIterations, 0.01); 
-
-    // Evaluate the model on rating data
-    JavaRDD<Tuple2<Object, Object>> userProducts = ratings.map(
-      new Function<Rating, Tuple2<Object, Object>>() {
-        public Tuple2<Object, Object> call(Rating r) {
-          return new Tuple2<Object, Object>(r.user(), r.product());
-        }
-      }
-    );
-    JavaPairRDD<Tuple2<Integer, Integer>, Double> predictions = JavaPairRDD.fromJavaRDD(
-      model.predict(JavaRDD.toRDD(userProducts)).toJavaRDD().map(
-        new Function<Rating, Tuple2<Tuple2<Integer, Integer>, Double>>() {
-          public Tuple2<Tuple2<Integer, Integer>, Double> call(Rating r){
-            return new Tuple2<Tuple2<Integer, Integer>, Double>(
-              new Tuple2<Integer, Integer>(r.user(), r.product()), r.rating());
-          }
-        }
-    ));
-    JavaRDD<Tuple2<Double, Double>> ratesAndPreds = 
-      JavaPairRDD.fromJavaRDD(ratings.map(
-        new Function<Rating, Tuple2<Tuple2<Integer, Integer>, Double>>() {
-          public Tuple2<Tuple2<Integer, Integer>, Double> call(Rating r){
-            return new Tuple2<Tuple2<Integer, Integer>, Double>(
-              new Tuple2<Integer, Integer>(r.user(), r.product()), r.rating());
-          }
-        }
-    )).join(predictions).values();
-    double MSE = JavaDoubleRDD.fromRDD(ratesAndPreds.map(
-      new Function<Tuple2<Double, Double>, Object>() {
-        public Object call(Tuple2<Double, Double> pair) {
-          Double err = pair._1() - pair._2();
-          return err * err;
-        }
-      }
-    ).rdd()).mean();
-    System.out.println("Mean Squared Error = " + MSE);
-
-    // Save and load model
-    model.save(sc.sc(), "myModelPath");
-    MatrixFactorizationModel sameModel = MatrixFactorizationModel.load(sc.sc(), "myModelPath");
-  }
-}
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaRecommendationExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
@@ -207,29 +97,7 @@ recommendation by measuring the Mean Squared Error of rating prediction.
 
 Refer to the [`ALS` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.recommendation.ALS) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.recommendation import ALS, MatrixFactorizationModel, Rating
-
-# Load and parse the data
-data = sc.textFile("data/mllib/als/test.data")
-ratings = data.map(lambda l: l.split(',')).map(lambda l: Rating(int(l[0]), int(l[1]), float(l[2])))
-
-# Build the recommendation model using Alternating Least Squares
-rank = 10
-numIterations = 10
-model = ALS.train(ratings, rank, numIterations)
-
-# Evaluate the model on training data
-testdata = ratings.map(lambda p: (p[0], p[1]))
-predictions = model.predictAll(testdata).map(lambda r: ((r[0], r[1]), r[2]))
-ratesAndPreds = ratings.map(lambda r: ((r[0], r[1]), r[2])).join(predictions)
-MSE = ratesAndPreds.map(lambda r: (r[1][0] - r[1][1])**2).mean()
-print("Mean Squared Error = " + str(MSE))
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = MatrixFactorizationModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example python/mllib/recommendation_example.py %}
 
 If the rating matrix is derived from other source of information (i.e., it is inferred from other
 signals), you can use the trainImplicit method to get better results.
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaRecommendationExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRecommendationExample.java
new file mode 100644
index 0000000000000..1065fde953b96
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRecommendationExample.java
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import scala.Tuple2;
+
+import org.apache.spark.api.java.*;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.mllib.recommendation.ALS;
+import org.apache.spark.mllib.recommendation.MatrixFactorizationModel;
+import org.apache.spark.mllib.recommendation.Rating;
+import org.apache.spark.SparkConf;
+// $example off$
+
+public class JavaRecommendationExample {
+  public static void main(String args[]) {
+    // $example on$
+    SparkConf conf = new SparkConf().setAppName("Java Collaborative Filtering Example");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+
+    // Load and parse the data
+    String path = "data/mllib/als/test.data";
+    JavaRDD<String> data = jsc.textFile(path);
+    JavaRDD<Rating> ratings = data.map(
+      new Function<String, Rating>() {
+        public Rating call(String s) {
+          String[] sarray = s.split(",");
+          return new Rating(Integer.parseInt(sarray[0]), Integer.parseInt(sarray[1]),
+            Double.parseDouble(sarray[2]));
+        }
+      }
+    );
+
+    // Build the recommendation model using ALS
+    int rank = 10;
+    int numIterations = 10;
+    MatrixFactorizationModel model = ALS.train(JavaRDD.toRDD(ratings), rank, numIterations, 0.01);
+
+    // Evaluate the model on rating data
+    JavaRDD<Tuple2<Object, Object>> userProducts = ratings.map(
+      new Function<Rating, Tuple2<Object, Object>>() {
+        public Tuple2<Object, Object> call(Rating r) {
+          return new Tuple2<Object, Object>(r.user(), r.product());
+        }
+      }
+    );
+    JavaPairRDD<Tuple2<Integer, Integer>, Double> predictions = JavaPairRDD.fromJavaRDD(
+      model.predict(JavaRDD.toRDD(userProducts)).toJavaRDD().map(
+        new Function<Rating, Tuple2<Tuple2<Integer, Integer>, Double>>() {
+          public Tuple2<Tuple2<Integer, Integer>, Double> call(Rating r){
+            return new Tuple2<Tuple2<Integer, Integer>, Double>(
+              new Tuple2<Integer, Integer>(r.user(), r.product()), r.rating());
+          }
+        }
+      ));
+    JavaRDD<Tuple2<Double, Double>> ratesAndPreds =
+      JavaPairRDD.fromJavaRDD(ratings.map(
+        new Function<Rating, Tuple2<Tuple2<Integer, Integer>, Double>>() {
+          public Tuple2<Tuple2<Integer, Integer>, Double> call(Rating r){
+            return new Tuple2<Tuple2<Integer, Integer>, Double>(
+              new Tuple2<Integer, Integer>(r.user(), r.product()), r.rating());
+          }
+        }
+      )).join(predictions).values();
+    double MSE = JavaDoubleRDD.fromRDD(ratesAndPreds.map(
+      new Function<Tuple2<Double, Double>, Object>() {
+        public Object call(Tuple2<Double, Double> pair) {
+          Double err = pair._1() - pair._2();
+          return err * err;
+        }
+      }
+    ).rdd()).mean();
+    System.out.println("Mean Squared Error = " + MSE);
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myCollaborativeFilter");
+    MatrixFactorizationModel sameModel = MatrixFactorizationModel.load(jsc.sc(),
+      "target/tmp/myCollaborativeFilter");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/python/mllib/recommendation_example.py b/examples/src/main/python/mllib/recommendation_example.py
new file mode 100644
index 0000000000000..615db0749b182
--- /dev/null
+++ b/examples/src/main/python/mllib/recommendation_example.py
@@ -0,0 +1,54 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Collaborative Filtering Classification Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext
+
+# $example on$
+from pyspark.mllib.recommendation import ALS, MatrixFactorizationModel, Rating
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="PythonCollaborativeFilteringExample")
+    # $example on$
+    # Load and parse the data
+    data = sc.textFile("data/mllib/als/test.data")
+    ratings = data.map(lambda l: l.split(','))\
+        .map(lambda l: Rating(int(l[0]), int(l[1]), float(l[2])))
+
+    # Build the recommendation model using Alternating Least Squares
+    rank = 10
+    numIterations = 10
+    model = ALS.train(ratings, rank, numIterations)
+
+    # Evaluate the model on training data
+    testdata = ratings.map(lambda p: (p[0], p[1]))
+    predictions = model.predictAll(testdata).map(lambda r: ((r[0], r[1]), r[2]))
+    ratesAndPreds = ratings.map(lambda r: ((r[0], r[1]), r[2])).join(predictions)
+    MSE = ratesAndPreds.map(lambda r: (r[1][0] - r[1][1])**2).mean()
+    print("Mean Squared Error = " + str(MSE))
+
+    # Save and load model
+    model.save(sc, "target/tmp/myCollaborativeFilter")
+    sameModel = MatrixFactorizationModel.load(sc, "target/tmp/myCollaborativeFilter")
+    # $example off$
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/RecommendationExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/RecommendationExample.scala
new file mode 100644
index 0000000000000..64e4602465444
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/RecommendationExample.scala
@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
+import org.apache.spark.mllib.recommendation.ALS
+import org.apache.spark.mllib.recommendation.MatrixFactorizationModel
+import org.apache.spark.mllib.recommendation.Rating
+// $example off$
+
+object RecommendationExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("CollaborativeFilteringExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    // Load and parse the data
+    val data = sc.textFile("data/mllib/als/test.data")
+    val ratings = data.map(_.split(',') match { case Array(user, item, rate) =>
+      Rating(user.toInt, item.toInt, rate.toDouble)
+    })
+
+    // Build the recommendation model using ALS
+    val rank = 10
+    val numIterations = 10
+    val model = ALS.train(ratings, rank, numIterations, 0.01)
+
+    // Evaluate the model on rating data
+    val usersProducts = ratings.map { case Rating(user, product, rate) =>
+      (user, product)
+    }
+    val predictions =
+      model.predict(usersProducts).map { case Rating(user, product, rate) =>
+        ((user, product), rate)
+      }
+    val ratesAndPreds = ratings.map { case Rating(user, product, rate) =>
+      ((user, product), rate)
+    }.join(predictions)
+    val MSE = ratesAndPreds.map { case ((user, product), (r1, r2)) =>
+      val err = (r1 - r2)
+      err * err
+    }.mean()
+    println("Mean Squared Error = " + MSE)
+
+    // Save and load model
+    model.save(sc, "target/tmp/myCollaborativeFilter")
+    val sameModel = MatrixFactorizationModel.load(sc, "target/tmp/myCollaborativeFilter")
+    // $example off$
+  }
+}
+// scalastyle:on println

From f138cb873335654476d1cd1070900b552dd8b21a Mon Sep 17 00:00:00 2001
From: Nick Buroojy <nick.buroojy@civitaslearning.com>
Date: Mon, 9 Nov 2015 14:30:37 -0800
Subject: [PATCH 134/518] [SPARK-9301][SQL] Add collect_set and collect_list
 aggregate functions

For now they are thin wrappers around the corresponding Hive UDAFs.

One limitation with these in Hive 0.13.0 is they only support aggregating primitive types.

I chose snake_case here instead of camelCase because it seems to be used in the majority of the multi-word fns.

Do we also want to add these to `functions.py`?

This approach was recommended here: https://github.com/apache/spark/pull/8592#issuecomment-154247089

marmbrus rxin

Author: Nick Buroojy <nick.buroojy@civitaslearning.com>

Closes #9526 from nburoojy/nick/udaf-alias.

(cherry picked from commit a6ee4f989d020420dd08b97abb24802200ff23b2)
Signed-off-by: Michael Armbrust <michael@databricks.com>
---
 python/pyspark/sql/functions.py               | 25 +++++++++++--------
 python/pyspark/sql/tests.py                   | 17 +++++++++++++
 .../org/apache/spark/sql/functions.scala      | 20 +++++++++++++++
 .../hive/HiveDataFrameAnalyticsSuite.scala    | 15 +++++++++--
 4 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 2f7c2f4aacd47..962f676d406d8 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -124,17 +124,20 @@ def _():
 
 _functions_1_6 = {
     # unary math functions
-    "stddev": "Aggregate function: returns the unbiased sample standard deviation of" +
-              " the expression in a group.",
-    "stddev_samp": "Aggregate function: returns the unbiased sample standard deviation of" +
-              " the expression in a group.",
-    "stddev_pop": "Aggregate function: returns population standard deviation of" +
-              " the expression in a group.",
-    "variance": "Aggregate function: returns the population variance of the values in a group.",
-    "var_samp": "Aggregate function: returns the unbiased variance of the values in a group.",
-    "var_pop":  "Aggregate function: returns the population variance of the values in a group.",
-    "skewness": "Aggregate function: returns the skewness of the values in a group.",
-    "kurtosis": "Aggregate function: returns the kurtosis of the values in a group."
+    'stddev': 'Aggregate function: returns the unbiased sample standard deviation of' +
+              ' the expression in a group.',
+    'stddev_samp': 'Aggregate function: returns the unbiased sample standard deviation of' +
+                   ' the expression in a group.',
+    'stddev_pop': 'Aggregate function: returns population standard deviation of' +
+                  ' the expression in a group.',
+    'variance': 'Aggregate function: returns the population variance of the values in a group.',
+    'var_samp': 'Aggregate function: returns the unbiased variance of the values in a group.',
+    'var_pop':  'Aggregate function: returns the population variance of the values in a group.',
+    'skewness': 'Aggregate function: returns the skewness of the values in a group.',
+    'kurtosis': 'Aggregate function: returns the kurtosis of the values in a group.',
+    'collect_list': 'Aggregate function: returns a list of objects with duplicates.',
+    'collect_set': 'Aggregate function: returns a set of objects with duplicate elements' +
+                   ' eliminated.'
 }
 
 # math functions that take two arguments as input
diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
index 4c03a0d4ffe93..e224574bcb301 100644
--- a/python/pyspark/sql/tests.py
+++ b/python/pyspark/sql/tests.py
@@ -1230,6 +1230,23 @@ def test_window_functions_without_partitionBy(self):
         for r, ex in zip(rs, expected):
             self.assertEqual(tuple(r), ex[:len(r)])
 
+    def test_collect_functions(self):
+        df = self.sqlCtx.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"])
+        from pyspark.sql import functions
+
+        self.assertEqual(
+            sorted(df.select(functions.collect_set(df.key).alias('r')).collect()[0].r),
+            [1, 2])
+        self.assertEqual(
+            sorted(df.select(functions.collect_list(df.key).alias('r')).collect()[0].r),
+            [1, 1, 1, 2])
+        self.assertEqual(
+            sorted(df.select(functions.collect_set(df.value).alias('r')).collect()[0].r),
+            ["1", "2"])
+        self.assertEqual(
+            sorted(df.select(functions.collect_list(df.value).alias('r')).collect()[0].r),
+            ["1", "2", "2", "2"])
+
 
 if __name__ == "__main__":
     if xmlrunner:
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 04627589886a8..3f0b24b68b816 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -174,6 +174,26 @@ object functions {
    */
   def avg(columnName: String): Column = avg(Column(columnName))
 
+  /**
+   * Aggregate function: returns a list of objects with duplicates.
+   *
+   * For now this is an alias for the collect_list Hive UDAF.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def collect_list(e: Column): Column = callUDF("collect_list", e)
+
+  /**
+   * Aggregate function: returns a set of objects with duplicate elements eliminated.
+   *
+   * For now this is an alias for the collect_set Hive UDAF.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def collect_set(e: Column): Column = callUDF("collect_set", e)
+
   /**
    * Aggregate function: returns the Pearson Correlation Coefficient for two columns.
    *
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
index 2e5cae415e54b..9864acf765265 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.hive
 
-import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.{DataFrame, QueryTest, Row}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.hive.test.TestHiveSingleton
 import org.scalatest.BeforeAndAfterAll
@@ -32,7 +32,7 @@ class HiveDataFrameAnalyticsSuite extends QueryTest with TestHiveSingleton with
   private var testData: DataFrame = _
 
   override def beforeAll() {
-    testData = Seq((1, 2), (2, 4)).toDF("a", "b")
+    testData = Seq((1, 2), (2, 2), (3, 4)).toDF("a", "b")
     hiveContext.registerDataFrameAsTable(testData, "mytable")
   }
 
@@ -52,6 +52,17 @@ class HiveDataFrameAnalyticsSuite extends QueryTest with TestHiveSingleton with
     )
   }
 
+  test("collect functions") {
+    checkAnswer(
+      testData.select(collect_list($"a"), collect_list($"b")),
+      Seq(Row(Seq(1, 2, 3), Seq(2, 2, 4)))
+    )
+    checkAnswer(
+      testData.select(collect_set($"a"), collect_set($"b")),
+      Seq(Row(Seq(1, 2, 3), Seq(2, 4)))
+    )
+  }
+
   test("cube") {
     checkAnswer(
       testData.cube($"a" + $"b", $"b").agg(sum($"a" - $"b")),

From 150f6a89b79f0e5bc31aa83731429dc7ac5ea76b Mon Sep 17 00:00:00 2001
From: Cheng Lian <lian@databricks.com>
Date: Mon, 9 Nov 2015 14:32:52 -0800
Subject: [PATCH 135/518] [SPARK-11595] [SQL] Fixes ADD JAR when the input path
 contains URL scheme

Author: Cheng Lian <lian@databricks.com>

Closes #9569 from liancheng/spark-11595.fix-add-jar.
---
 .../hive/thriftserver/HiveThriftServer2Suites.scala   |  1 +
 .../apache/spark/sql/hive/client/ClientWrapper.scala  | 11 +++++++++--
 .../spark/sql/hive/client/IsolatedClientLoader.scala  |  9 +++------
 .../spark/sql/hive/execution/HiveQuerySuite.scala     |  8 +++++---
 4 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
index ff8ca0150649d..5903b9e71cdd2 100644
--- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
+++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
@@ -41,6 +41,7 @@ import org.apache.thrift.transport.TSocket
 import org.scalatest.BeforeAndAfterAll
 
 import org.apache.spark.sql.hive.HiveContext
+import org.apache.spark.sql.hive.test.TestHive
 import org.apache.spark.sql.test.ProcessTestUtils.ProcessOutputCapturer
 import org.apache.spark.util.Utils
 import org.apache.spark.{Logging, SparkFunSuite}
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/ClientWrapper.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/ClientWrapper.scala
index 3dce86c480747..f1c2489b38271 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/ClientWrapper.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/ClientWrapper.scala
@@ -19,7 +19,6 @@ package org.apache.spark.sql.hive.client
 
 import java.io.{File, PrintStream}
 import java.util.{Map => JMap}
-import javax.annotation.concurrent.GuardedBy
 
 import scala.collection.JavaConverters._
 import scala.language.reflectiveCalls
@@ -548,7 +547,15 @@ private[hive] class ClientWrapper(
   }
 
   def addJar(path: String): Unit = {
-    clientLoader.addJar(path)
+    val uri = new Path(path).toUri
+    val jarURL = if (uri.getScheme == null) {
+      // `path` is a local file path without a URL scheme
+      new File(path).toURI.toURL
+    } else {
+      // `path` is a URL with a scheme
+      uri.toURL
+    }
+    clientLoader.addJar(jarURL)
     runSqlHive(s"ADD JAR $path")
   }
 
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/IsolatedClientLoader.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/IsolatedClientLoader.scala
index f99c3ed2ae987..e041e0d8e5ae8 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/IsolatedClientLoader.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/IsolatedClientLoader.scala
@@ -22,7 +22,6 @@ import java.lang.reflect.InvocationTargetException
 import java.net.{URL, URLClassLoader}
 import java.util
 
-import scala.collection.mutable
 import scala.language.reflectiveCalls
 import scala.util.Try
 
@@ -30,10 +29,9 @@ import org.apache.commons.io.{FileUtils, IOUtils}
 
 import org.apache.spark.Logging
 import org.apache.spark.deploy.SparkSubmitUtils
-import org.apache.spark.util.{MutableURLClassLoader, Utils}
-
 import org.apache.spark.sql.catalyst.util.quietly
 import org.apache.spark.sql.hive.HiveContext
+import org.apache.spark.util.{MutableURLClassLoader, Utils}
 
 /** Factory for `IsolatedClientLoader` with specific versions of hive. */
 private[hive] object IsolatedClientLoader {
@@ -190,9 +188,8 @@ private[hive] class IsolatedClientLoader(
     new NonClosableMutableURLClassLoader(isolatedClassLoader)
   }
 
-  private[hive] def addJar(path: String): Unit = synchronized {
-    val jarURL = new java.io.File(path).toURI.toURL
-    classLoader.addURL(jarURL)
+  private[hive] def addJar(path: URL): Unit = synchronized {
+    classLoader.addURL(path)
   }
 
   /** The isolated client interface to Hive. */
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
index fc72e3c7dc6aa..78378c8b69c7a 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
@@ -927,7 +927,7 @@ class HiveQuerySuite extends HiveComparisonTest with BeforeAndAfter {
   test("SPARK-2263: Insert Map<K, V> values") {
     sql("CREATE TABLE m(value MAP<INT, STRING>)")
     sql("INSERT OVERWRITE TABLE m SELECT MAP(key, value) FROM src LIMIT 10")
-    sql("SELECT * FROM m").collect().zip(sql("SELECT * FROM src LIMIT 10").collect()).map {
+    sql("SELECT * FROM m").collect().zip(sql("SELECT * FROM src LIMIT 10").collect()).foreach {
       case (Row(map: Map[_, _]), Row(key: Int, value: String)) =>
         assert(map.size === 1)
         assert(map.head === (key, value))
@@ -961,10 +961,12 @@ class HiveQuerySuite extends HiveComparisonTest with BeforeAndAfter {
 
   test("CREATE TEMPORARY FUNCTION") {
     val funcJar = TestHive.getHiveFile("TestUDTF.jar").getCanonicalPath
-    sql(s"ADD JAR $funcJar")
+    val jarURL = s"file://$funcJar"
+    sql(s"ADD JAR $jarURL")
     sql(
       """CREATE TEMPORARY FUNCTION udtf_count2 AS
-        | 'org.apache.spark.sql.hive.execution.GenericUDTFCount2'""".stripMargin)
+        |'org.apache.spark.sql.hive.execution.GenericUDTFCount2'
+      """.stripMargin)
     assert(sql("DESCRIBE FUNCTION udtf_count2").count > 1)
     sql("DROP TEMPORARY FUNCTION udtf_count2")
   }

From a3a7c9103e136035d65a5564f9eb0fa04727c4f3 Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Mon, 9 Nov 2015 14:39:18 -0800
Subject: [PATCH 136/518] [SPARK-11359][STREAMING][KINESIS] Checkpoint to
 DynamoDB even when new data doesn't come in

Currently, the checkpoints to DynamoDB occur only when new data comes in, as we update the clock for the checkpointState. This PR makes the checkpoint a scheduled execution based on the `checkpointInterval`.

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9421 from brkyvz/kinesis-checkpoint.
---
 .../kinesis/KinesisCheckpointState.scala      |  54 -------
 .../kinesis/KinesisCheckpointer.scala         | 133 +++++++++++++++
 .../streaming/kinesis/KinesisReceiver.scala   |  38 ++++-
 .../kinesis/KinesisRecordProcessor.scala      |  59 ++-----
 .../kinesis/KinesisCheckpointerSuite.scala    | 152 ++++++++++++++++++
 .../kinesis/KinesisReceiverSuite.scala        |  96 +++--------
 6 files changed, 349 insertions(+), 183 deletions(-)
 delete mode 100644 extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointState.scala
 create mode 100644 extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointer.scala
 create mode 100644 extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointerSuite.scala

diff --git a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointState.scala b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointState.scala
deleted file mode 100644
index 83a4537559512..0000000000000
--- a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointState.scala
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.spark.streaming.kinesis
-
-import org.apache.spark.Logging
-import org.apache.spark.streaming.Duration
-import org.apache.spark.util.{Clock, ManualClock, SystemClock}
-
-/**
- * This is a helper class for managing checkpoint clocks.
- *
- * @param checkpointInterval
- * @param currentClock.  Default to current SystemClock if none is passed in (mocking purposes)
- */
-private[kinesis] class KinesisCheckpointState(
-    checkpointInterval: Duration,
-    currentClock: Clock = new SystemClock())
-  extends Logging {
-
-  /* Initialize the checkpoint clock using the given currentClock + checkpointInterval millis */
-  val checkpointClock = new ManualClock()
-  checkpointClock.setTime(currentClock.getTimeMillis() + checkpointInterval.milliseconds)
-
-  /**
-   * Check if it's time to checkpoint based on the current time and the derived time
-   *   for the next checkpoint
-   *
-   * @return true if it's time to checkpoint
-   */
-  def shouldCheckpoint(): Boolean = {
-    new SystemClock().getTimeMillis() > checkpointClock.getTimeMillis()
-  }
-
-  /**
-   * Advance the checkpoint clock by the checkpoint interval.
-   */
-  def advanceCheckpoint(): Unit = {
-    checkpointClock.advance(checkpointInterval.milliseconds)
-  }
-}
diff --git a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointer.scala b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointer.scala
new file mode 100644
index 0000000000000..1ca6d4302c2bb
--- /dev/null
+++ b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointer.scala
@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.streaming.kinesis
+
+import java.util.concurrent._
+
+import scala.util.control.NonFatal
+
+import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorCheckpointer
+import com.amazonaws.services.kinesis.clientlibrary.types.ShutdownReason
+
+import org.apache.spark.Logging
+import org.apache.spark.streaming.Duration
+import org.apache.spark.streaming.util.RecurringTimer
+import org.apache.spark.util.{Clock, SystemClock, ThreadUtils}
+
+/**
+ * This is a helper class for managing Kinesis checkpointing.
+ *
+ * @param receiver The receiver that keeps track of which sequence numbers we can checkpoint
+ * @param checkpointInterval How frequently we will checkpoint to DynamoDB
+ * @param workerId Worker Id of KCL worker for logging purposes
+ * @param clock In order to use ManualClocks for the purpose of testing
+ */
+private[kinesis] class KinesisCheckpointer(
+    receiver: KinesisReceiver[_],
+    checkpointInterval: Duration,
+    workerId: String,
+    clock: Clock = new SystemClock) extends Logging {
+
+  // a map from shardId's to checkpointers
+  private val checkpointers = new ConcurrentHashMap[String, IRecordProcessorCheckpointer]()
+
+  private val lastCheckpointedSeqNums = new ConcurrentHashMap[String, String]()
+
+  private val checkpointerThread: RecurringTimer = startCheckpointerThread()
+
+  /** Update the checkpointer instance to the most recent one for the given shardId. */
+  def setCheckpointer(shardId: String, checkpointer: IRecordProcessorCheckpointer): Unit = {
+    checkpointers.put(shardId, checkpointer)
+  }
+
+  /**
+   * Stop tracking the specified shardId.
+   *
+   * If a checkpointer is provided, e.g. on IRecordProcessor.shutdown [[ShutdownReason.TERMINATE]],
+   * we will use that to make the final checkpoint. If `null` is provided, we will not make the
+   * checkpoint, e.g. in case of [[ShutdownReason.ZOMBIE]].
+   */
+  def removeCheckpointer(shardId: String, checkpointer: IRecordProcessorCheckpointer): Unit = {
+    synchronized {
+      checkpointers.remove(shardId)
+      checkpoint(shardId, checkpointer)
+    }
+  }
+
+  /** Perform the checkpoint. */
+  private def checkpoint(shardId: String, checkpointer: IRecordProcessorCheckpointer): Unit = {
+    try {
+      if (checkpointer != null) {
+        receiver.getLatestSeqNumToCheckpoint(shardId).foreach { latestSeqNum =>
+          val lastSeqNum = lastCheckpointedSeqNums.get(shardId)
+          // Kinesis sequence numbers are monotonically increasing strings, therefore we can do
+          // safely do the string comparison
+          if (lastSeqNum == null || latestSeqNum > lastSeqNum) {
+            /* Perform the checkpoint */
+            KinesisRecordProcessor.retryRandom(checkpointer.checkpoint(latestSeqNum), 4, 100)
+            logDebug(s"Checkpoint:  WorkerId $workerId completed checkpoint at sequence number" +
+              s" $latestSeqNum for shardId $shardId")
+            lastCheckpointedSeqNums.put(shardId, latestSeqNum)
+          }
+        }
+      } else {
+        logDebug(s"Checkpointing skipped for shardId $shardId. Checkpointer not set.")
+      }
+    } catch {
+      case NonFatal(e) =>
+        logWarning(s"Failed to checkpoint shardId $shardId to DynamoDB.", e)
+    }
+  }
+
+  /** Checkpoint the latest saved sequence numbers for all active shardId's. */
+  private def checkpointAll(): Unit = synchronized {
+    // if this method throws an exception, then the scheduled task will not run again
+    try {
+      val shardIds = checkpointers.keys()
+      while (shardIds.hasMoreElements) {
+        val shardId = shardIds.nextElement()
+        checkpoint(shardId, checkpointers.get(shardId))
+      }
+    } catch {
+      case NonFatal(e) =>
+        logWarning("Failed to checkpoint to DynamoDB.", e)
+    }
+  }
+
+  /**
+   * Start the checkpointer thread with the given checkpoint duration.
+   */
+  private def startCheckpointerThread(): RecurringTimer = {
+    val period = checkpointInterval.milliseconds
+    val threadName = s"Kinesis Checkpointer - Worker $workerId"
+    val timer = new RecurringTimer(clock, period, _ => checkpointAll(), threadName)
+    timer.start()
+    logDebug(s"Started checkpointer thread: $threadName")
+    timer
+  }
+
+  /**
+   * Shutdown the checkpointer. Should be called on the onStop of the Receiver.
+   */
+  def shutdown(): Unit = {
+    // the recurring timer checkpoints for us one last time.
+    checkpointerThread.stop(interruptTimer = false)
+    checkpointers.clear()
+    lastCheckpointedSeqNums.clear()
+    logInfo("Successfully shutdown Kinesis Checkpointer.")
+  }
+}
diff --git a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisReceiver.scala b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisReceiver.scala
index 134d627cdaffa..50993f157cd95 100644
--- a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisReceiver.scala
+++ b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisReceiver.scala
@@ -23,7 +23,7 @@ import scala.collection.mutable
 import scala.util.control.NonFatal
 
 import com.amazonaws.auth.{AWSCredentials, AWSCredentialsProvider, DefaultAWSCredentialsProviderChain}
-import com.amazonaws.services.kinesis.clientlibrary.interfaces.{IRecordProcessor, IRecordProcessorFactory}
+import com.amazonaws.services.kinesis.clientlibrary.interfaces.{IRecordProcessorCheckpointer, IRecordProcessor, IRecordProcessorFactory}
 import com.amazonaws.services.kinesis.clientlibrary.lib.worker.{InitialPositionInStream, KinesisClientLibConfiguration, Worker}
 import com.amazonaws.services.kinesis.model.Record
 
@@ -31,8 +31,7 @@ import org.apache.spark.storage.{StorageLevel, StreamBlockId}
 import org.apache.spark.streaming.Duration
 import org.apache.spark.streaming.receiver.{BlockGenerator, BlockGeneratorListener, Receiver}
 import org.apache.spark.util.Utils
-import org.apache.spark.{Logging, SparkEnv}
-
+import org.apache.spark.Logging
 
 private[kinesis]
 case class SerializableAWSCredentials(accessKeyId: String, secretKey: String)
@@ -127,6 +126,11 @@ private[kinesis] class KinesisReceiver[T](
   private val blockIdToSeqNumRanges = new mutable.HashMap[StreamBlockId, SequenceNumberRanges]
     with mutable.SynchronizedMap[StreamBlockId, SequenceNumberRanges]
 
+  /**
+   * The centralized kinesisCheckpointer that checkpoints based on the given checkpointInterval.
+   */
+  @volatile private var kinesisCheckpointer: KinesisCheckpointer = null
+
   /**
    * Latest sequence number ranges that have been stored successfully.
    * This is used for checkpointing through KCL */
@@ -141,6 +145,7 @@ private[kinesis] class KinesisReceiver[T](
 
     workerId = Utils.localHostName() + ":" + UUID.randomUUID()
 
+    kinesisCheckpointer = new KinesisCheckpointer(receiver, checkpointInterval, workerId)
     // KCL config instance
     val awsCredProvider = resolveAWSCredentialsProvider()
     val kinesisClientLibConfiguration =
@@ -157,8 +162,8 @@ private[kinesis] class KinesisReceiver[T](
     *  We're using our custom KinesisRecordProcessor in this case.
     */
     val recordProcessorFactory = new IRecordProcessorFactory {
-      override def createProcessor: IRecordProcessor = new KinesisRecordProcessor(receiver,
-        workerId, new KinesisCheckpointState(checkpointInterval))
+      override def createProcessor: IRecordProcessor =
+        new KinesisRecordProcessor(receiver, workerId)
     }
 
     worker = new Worker(recordProcessorFactory, kinesisClientLibConfiguration)
@@ -198,6 +203,10 @@ private[kinesis] class KinesisReceiver[T](
       logInfo(s"Stopped receiver for workerId $workerId")
     }
     workerId = null
+    if (kinesisCheckpointer != null) {
+      kinesisCheckpointer.shutdown()
+      kinesisCheckpointer = null
+    }
   }
 
   /** Add records of the given shard to the current block being generated */
@@ -216,6 +225,25 @@ private[kinesis] class KinesisReceiver[T](
     shardIdToLatestStoredSeqNum.get(shardId)
   }
 
+  /**
+   * Set the checkpointer that will be used to checkpoint sequence numbers to DynamoDB for the
+   * given shardId.
+   */
+  def setCheckpointer(shardId: String, checkpointer: IRecordProcessorCheckpointer): Unit = {
+    assert(kinesisCheckpointer != null, "Kinesis Checkpointer not initialized!")
+    kinesisCheckpointer.setCheckpointer(shardId, checkpointer)
+  }
+
+  /**
+   * Remove the checkpointer for the given shardId. The provided checkpointer will be used to
+   * checkpoint one last time for the given shard. If `checkpointer` is `null`, then we will not
+   * checkpoint.
+   */
+  def removeCheckpointer(shardId: String, checkpointer: IRecordProcessorCheckpointer): Unit = {
+    assert(kinesisCheckpointer != null, "Kinesis Checkpointer not initialized!")
+    kinesisCheckpointer.removeCheckpointer(shardId, checkpointer)
+  }
+
   /**
    * Remember the range of sequence numbers that was added to the currently active block.
    * Internally, this is synchronized with `finalizeRangesForCurrentBlock()`.
diff --git a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisRecordProcessor.scala b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisRecordProcessor.scala
index 1d5178790ec4c..e381ffa0cbef4 100644
--- a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisRecordProcessor.scala
+++ b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisRecordProcessor.scala
@@ -27,26 +27,23 @@ import com.amazonaws.services.kinesis.clientlibrary.types.ShutdownReason
 import com.amazonaws.services.kinesis.model.Record
 
 import org.apache.spark.Logging
+import org.apache.spark.streaming.Duration
 
 /**
  * Kinesis-specific implementation of the Kinesis Client Library (KCL) IRecordProcessor.
  * This implementation operates on the Array[Byte] from the KinesisReceiver.
  * The Kinesis Worker creates an instance of this KinesisRecordProcessor for each
- *   shard in the Kinesis stream upon startup.  This is normally done in separate threads,
- *   but the KCLs within the KinesisReceivers will balance themselves out if you create
- *   multiple Receivers.
+ * shard in the Kinesis stream upon startup.  This is normally done in separate threads,
+ * but the KCLs within the KinesisReceivers will balance themselves out if you create
+ * multiple Receivers.
  *
  * @param receiver Kinesis receiver
  * @param workerId for logging purposes
- * @param checkpointState represents the checkpoint state including the next checkpoint time.
- *   It's injected here for mocking purposes.
  */
-private[kinesis] class KinesisRecordProcessor[T](
-    receiver: KinesisReceiver[T],
-    workerId: String,
-    checkpointState: KinesisCheckpointState) extends IRecordProcessor with Logging {
+private[kinesis] class KinesisRecordProcessor[T](receiver: KinesisReceiver[T], workerId: String)
+  extends IRecordProcessor with Logging {
 
-  // shardId to be populated during initialize()
+  // shardId populated during initialize()
   @volatile
   private var shardId: String = _
 
@@ -74,34 +71,7 @@ private[kinesis] class KinesisRecordProcessor[T](
       try {
         receiver.addRecords(shardId, batch)
         logDebug(s"Stored: Worker $workerId stored ${batch.size} records for shardId $shardId")
-
-        /*
-         *
-         * Checkpoint the sequence number of the last record successfully stored.
-         * Note that in this current implementation, the checkpointing occurs only when after
-         * checkpointIntervalMillis from the last checkpoint, AND when there is new record
-         * to process. This leads to the checkpointing lagging behind what records have been
-         * stored by the receiver. Ofcourse, this can lead records processed more than once,
-         * under failures and restarts.
-         *
-         * TODO: Instead of checkpointing here, run a separate timer task to perform
-         * checkpointing so that it checkpoints in a timely manner independent of whether
-         * new records are available or not.
-         */
-        if (checkpointState.shouldCheckpoint()) {
-          receiver.getLatestSeqNumToCheckpoint(shardId).foreach { latestSeqNum =>
-            /* Perform the checkpoint */
-            KinesisRecordProcessor.retryRandom(checkpointer.checkpoint(latestSeqNum), 4, 100)
-
-            /* Update the next checkpoint time */
-            checkpointState.advanceCheckpoint()
-
-            logDebug(s"Checkpoint:  WorkerId $workerId completed checkpoint of ${batch.size}" +
-              s" records for shardId $shardId")
-            logDebug(s"Checkpoint:  Next checkpoint is at " +
-              s" ${checkpointState.checkpointClock.getTimeMillis()} for shardId $shardId")
-          }
-        }
+        receiver.setCheckpointer(shardId, checkpointer)
       } catch {
         case NonFatal(e) => {
           /*
@@ -142,23 +112,18 @@ private[kinesis] class KinesisRecordProcessor[T](
        * It's now OK to read from the new shards that resulted from a resharding event.
        */
       case ShutdownReason.TERMINATE =>
-        val latestSeqNumToCheckpointOption = receiver.getLatestSeqNumToCheckpoint(shardId)
-        if (latestSeqNumToCheckpointOption.nonEmpty) {
-          KinesisRecordProcessor.retryRandom(
-            checkpointer.checkpoint(latestSeqNumToCheckpointOption.get), 4, 100)
-        }
+        receiver.removeCheckpointer(shardId, checkpointer)
 
       /*
-       * ZOMBIE Use Case.  NoOp.
+       * ZOMBIE Use Case or Unknown reason.  NoOp.
        * No checkpoint because other workers may have taken over and already started processing
        *    the same records.
        * This may lead to records being processed more than once.
        */
-      case ShutdownReason.ZOMBIE =>
-
-      /* Unknown reason.  NoOp */
       case _ =>
+        receiver.removeCheckpointer(shardId, null) // return null so that we don't checkpoint
     }
+
   }
 }
 
diff --git a/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointerSuite.scala b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointerSuite.scala
new file mode 100644
index 0000000000000..645e64a0bc3a0
--- /dev/null
+++ b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisCheckpointerSuite.scala
@@ -0,0 +1,152 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.kinesis
+
+import java.util.concurrent.{TimeoutException, ExecutorService}
+
+import scala.concurrent.{Await, ExecutionContext, Future}
+import scala.concurrent.duration._
+import scala.language.postfixOps
+
+import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorCheckpointer
+import org.mockito.Matchers._
+import org.mockito.Mockito._
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
+import org.scalatest.{PrivateMethodTester, BeforeAndAfterEach}
+import org.scalatest.concurrent.Eventually
+import org.scalatest.concurrent.Eventually._
+import org.scalatest.mock.MockitoSugar
+
+import org.apache.spark.streaming.{Duration, TestSuiteBase}
+import org.apache.spark.util.ManualClock
+
+class KinesisCheckpointerSuite extends TestSuiteBase
+  with MockitoSugar
+  with BeforeAndAfterEach
+  with PrivateMethodTester
+  with Eventually {
+
+  private val workerId = "dummyWorkerId"
+  private val shardId = "dummyShardId"
+  private val seqNum = "123"
+  private val otherSeqNum = "245"
+  private val checkpointInterval = Duration(10)
+  private val someSeqNum = Some(seqNum)
+  private val someOtherSeqNum = Some(otherSeqNum)
+
+  private var receiverMock: KinesisReceiver[Array[Byte]] = _
+  private var checkpointerMock: IRecordProcessorCheckpointer = _
+  private var kinesisCheckpointer: KinesisCheckpointer = _
+  private var clock: ManualClock = _
+
+  private val checkpoint = PrivateMethod[Unit]('checkpoint)
+
+  override def beforeEach(): Unit = {
+    receiverMock = mock[KinesisReceiver[Array[Byte]]]
+    checkpointerMock = mock[IRecordProcessorCheckpointer]
+    clock = new ManualClock()
+    kinesisCheckpointer = new KinesisCheckpointer(receiverMock, checkpointInterval, workerId, clock)
+  }
+
+  test("checkpoint is not called twice for the same sequence number") {
+    when(receiverMock.getLatestSeqNumToCheckpoint(shardId)).thenReturn(someSeqNum)
+    kinesisCheckpointer.invokePrivate(checkpoint(shardId, checkpointerMock))
+    kinesisCheckpointer.invokePrivate(checkpoint(shardId, checkpointerMock))
+
+    verify(checkpointerMock, times(1)).checkpoint(anyString())
+  }
+
+  test("checkpoint is called after sequence number increases") {
+    when(receiverMock.getLatestSeqNumToCheckpoint(shardId))
+      .thenReturn(someSeqNum).thenReturn(someOtherSeqNum)
+    kinesisCheckpointer.invokePrivate(checkpoint(shardId, checkpointerMock))
+    kinesisCheckpointer.invokePrivate(checkpoint(shardId, checkpointerMock))
+
+    verify(checkpointerMock, times(1)).checkpoint(seqNum)
+    verify(checkpointerMock, times(1)).checkpoint(otherSeqNum)
+  }
+
+  test("should checkpoint if we have exceeded the checkpoint interval") {
+    when(receiverMock.getLatestSeqNumToCheckpoint(shardId))
+      .thenReturn(someSeqNum).thenReturn(someOtherSeqNum)
+
+    kinesisCheckpointer.setCheckpointer(shardId, checkpointerMock)
+    clock.advance(5 * checkpointInterval.milliseconds)
+
+    eventually(timeout(1 second)) {
+      verify(checkpointerMock, times(1)).checkpoint(seqNum)
+      verify(checkpointerMock, times(1)).checkpoint(otherSeqNum)
+    }
+  }
+
+  test("shouldn't checkpoint if we have not exceeded the checkpoint interval") {
+    when(receiverMock.getLatestSeqNumToCheckpoint(shardId)).thenReturn(someSeqNum)
+
+    kinesisCheckpointer.setCheckpointer(shardId, checkpointerMock)
+    clock.advance(checkpointInterval.milliseconds / 2)
+
+    verify(checkpointerMock, never()).checkpoint(anyString())
+  }
+
+  test("should not checkpoint for the same sequence number") {
+    when(receiverMock.getLatestSeqNumToCheckpoint(shardId)).thenReturn(someSeqNum)
+
+    kinesisCheckpointer.setCheckpointer(shardId, checkpointerMock)
+
+    clock.advance(checkpointInterval.milliseconds * 5)
+    eventually(timeout(1 second)) {
+      verify(checkpointerMock, atMost(1)).checkpoint(anyString())
+    }
+  }
+
+  test("removing checkpointer checkpoints one last time") {
+    when(receiverMock.getLatestSeqNumToCheckpoint(shardId)).thenReturn(someSeqNum)
+
+    kinesisCheckpointer.removeCheckpointer(shardId, checkpointerMock)
+    verify(checkpointerMock, times(1)).checkpoint(anyString())
+  }
+
+  test("if checkpointing is going on, wait until finished before removing and checkpointing") {
+    when(receiverMock.getLatestSeqNumToCheckpoint(shardId))
+      .thenReturn(someSeqNum).thenReturn(someOtherSeqNum)
+    when(checkpointerMock.checkpoint(anyString)).thenAnswer(new Answer[Unit] {
+      override def answer(invocations: InvocationOnMock): Unit = {
+        clock.waitTillTime(clock.getTimeMillis() + checkpointInterval.milliseconds / 2)
+      }
+    })
+
+    kinesisCheckpointer.setCheckpointer(shardId, checkpointerMock)
+    clock.advance(checkpointInterval.milliseconds)
+    eventually(timeout(1 second)) {
+      verify(checkpointerMock, times(1)).checkpoint(anyString())
+    }
+    // don't block test thread
+    val f = Future(kinesisCheckpointer.removeCheckpointer(shardId, checkpointerMock))(
+      ExecutionContext.global)
+
+    intercept[TimeoutException] {
+      Await.ready(f, 50 millis)
+    }
+
+    clock.advance(checkpointInterval.milliseconds / 2)
+    eventually(timeout(1 second)) {
+      verify(checkpointerMock, times(2)).checkpoint(anyString())
+    }
+  }
+}
diff --git a/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisReceiverSuite.scala b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisReceiverSuite.scala
index 17ab444704f44..e5c70db554a27 100644
--- a/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisReceiverSuite.scala
+++ b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisReceiverSuite.scala
@@ -25,12 +25,13 @@ import com.amazonaws.services.kinesis.clientlibrary.interfaces.IRecordProcessorC
 import com.amazonaws.services.kinesis.clientlibrary.types.ShutdownReason
 import com.amazonaws.services.kinesis.model.Record
 import org.mockito.Matchers._
+import org.mockito.Matchers.{eq => meq}
 import org.mockito.Mockito._
 import org.scalatest.mock.MockitoSugar
 import org.scalatest.{BeforeAndAfter, Matchers}
 
-import org.apache.spark.streaming.{Milliseconds, TestSuiteBase}
-import org.apache.spark.util.{Clock, ManualClock, Utils}
+import org.apache.spark.streaming.{Duration, TestSuiteBase}
+import org.apache.spark.util.Utils
 
 /**
  * Suite of Kinesis streaming receiver tests focusing mostly on the KinesisRecordProcessor
@@ -44,6 +45,7 @@ class KinesisReceiverSuite extends TestSuiteBase with Matchers with BeforeAndAft
   val workerId = "dummyWorkerId"
   val shardId = "dummyShardId"
   val seqNum = "dummySeqNum"
+  val checkpointInterval = Duration(10)
   val someSeqNum = Some(seqNum)
 
   val record1 = new Record()
@@ -54,24 +56,10 @@ class KinesisReceiverSuite extends TestSuiteBase with Matchers with BeforeAndAft
 
   var receiverMock: KinesisReceiver[Array[Byte]] = _
   var checkpointerMock: IRecordProcessorCheckpointer = _
-  var checkpointClockMock: ManualClock = _
-  var checkpointStateMock: KinesisCheckpointState = _
-  var currentClockMock: Clock = _
 
   override def beforeFunction(): Unit = {
     receiverMock = mock[KinesisReceiver[Array[Byte]]]
     checkpointerMock = mock[IRecordProcessorCheckpointer]
-    checkpointClockMock = mock[ManualClock]
-    checkpointStateMock = mock[KinesisCheckpointState]
-    currentClockMock = mock[Clock]
-  }
-
-  override def afterFunction(): Unit = {
-    super.afterFunction()
-    // Since this suite was originally written using EasyMock, add this to preserve the old
-    // mocking semantics (see SPARK-5735 for more details)
-    verifyNoMoreInteractions(receiverMock, checkpointerMock, checkpointClockMock,
-      checkpointStateMock, currentClockMock)
   }
 
   test("check serializability of SerializableAWSCredentials") {
@@ -79,113 +67,67 @@ class KinesisReceiverSuite extends TestSuiteBase with Matchers with BeforeAndAft
       Utils.serialize(new SerializableAWSCredentials("x", "y")))
   }
 
-  test("process records including store and checkpoint") {
+  test("process records including store and set checkpointer") {
     when(receiverMock.isStopped()).thenReturn(false)
-    when(receiverMock.getLatestSeqNumToCheckpoint(shardId)).thenReturn(someSeqNum)
-    when(checkpointStateMock.shouldCheckpoint()).thenReturn(true)
 
-    val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId, checkpointStateMock)
+    val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId)
     recordProcessor.initialize(shardId)
     recordProcessor.processRecords(batch, checkpointerMock)
 
     verify(receiverMock, times(1)).isStopped()
     verify(receiverMock, times(1)).addRecords(shardId, batch)
-    verify(receiverMock, times(1)).getLatestSeqNumToCheckpoint(shardId)
-    verify(checkpointStateMock, times(1)).shouldCheckpoint()
-    verify(checkpointerMock, times(1)).checkpoint(anyString)
-    verify(checkpointStateMock, times(1)).advanceCheckpoint()
+    verify(receiverMock, times(1)).setCheckpointer(shardId, checkpointerMock)
   }
 
-  test("shouldn't store and checkpoint when receiver is stopped") {
+  test("shouldn't store and update checkpointer when receiver is stopped") {
     when(receiverMock.isStopped()).thenReturn(true)
 
-    val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId, checkpointStateMock)
+    val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId)
     recordProcessor.processRecords(batch, checkpointerMock)
 
     verify(receiverMock, times(1)).isStopped()
     verify(receiverMock, never).addRecords(anyString, anyListOf(classOf[Record]))
-    verify(checkpointerMock, never).checkpoint(anyString)
+    verify(receiverMock, never).setCheckpointer(anyString, meq(checkpointerMock))
   }
 
-  test("shouldn't checkpoint when exception occurs during store") {
+  test("shouldn't update checkpointer when exception occurs during store") {
     when(receiverMock.isStopped()).thenReturn(false)
     when(
       receiverMock.addRecords(anyString, anyListOf(classOf[Record]))
     ).thenThrow(new RuntimeException())
 
     intercept[RuntimeException] {
-      val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId, checkpointStateMock)
+      val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId)
       recordProcessor.initialize(shardId)
       recordProcessor.processRecords(batch, checkpointerMock)
     }
 
     verify(receiverMock, times(1)).isStopped()
     verify(receiverMock, times(1)).addRecords(shardId, batch)
-    verify(checkpointerMock, never).checkpoint(anyString)
-  }
-
-  test("should set checkpoint time to currentTime + checkpoint interval upon instantiation") {
-    when(currentClockMock.getTimeMillis()).thenReturn(0)
-
-    val checkpointIntervalMillis = 10
-    val checkpointState =
-      new KinesisCheckpointState(Milliseconds(checkpointIntervalMillis), currentClockMock)
-    assert(checkpointState.checkpointClock.getTimeMillis() == checkpointIntervalMillis)
-
-    verify(currentClockMock, times(1)).getTimeMillis()
-  }
-
-  test("should checkpoint if we have exceeded the checkpoint interval") {
-    when(currentClockMock.getTimeMillis()).thenReturn(0)
-
-    val checkpointState = new KinesisCheckpointState(Milliseconds(Long.MinValue), currentClockMock)
-    assert(checkpointState.shouldCheckpoint())
-
-    verify(currentClockMock, times(1)).getTimeMillis()
-  }
-
-  test("shouldn't checkpoint if we have not exceeded the checkpoint interval") {
-    when(currentClockMock.getTimeMillis()).thenReturn(0)
-
-    val checkpointState = new KinesisCheckpointState(Milliseconds(Long.MaxValue), currentClockMock)
-    assert(!checkpointState.shouldCheckpoint())
-
-    verify(currentClockMock, times(1)).getTimeMillis()
-  }
-
-  test("should add to time when advancing checkpoint") {
-    when(currentClockMock.getTimeMillis()).thenReturn(0)
-
-    val checkpointIntervalMillis = 10
-    val checkpointState =
-      new KinesisCheckpointState(Milliseconds(checkpointIntervalMillis), currentClockMock)
-    assert(checkpointState.checkpointClock.getTimeMillis() == checkpointIntervalMillis)
-    checkpointState.advanceCheckpoint()
-    assert(checkpointState.checkpointClock.getTimeMillis() == (2 * checkpointIntervalMillis))
-
-    verify(currentClockMock, times(1)).getTimeMillis()
+    verify(receiverMock, never).setCheckpointer(anyString, meq(checkpointerMock))
   }
 
   test("shutdown should checkpoint if the reason is TERMINATE") {
     when(receiverMock.getLatestSeqNumToCheckpoint(shardId)).thenReturn(someSeqNum)
 
-    val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId, checkpointStateMock)
+    val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId)
     recordProcessor.initialize(shardId)
     recordProcessor.shutdown(checkpointerMock, ShutdownReason.TERMINATE)
 
-    verify(receiverMock, times(1)).getLatestSeqNumToCheckpoint(shardId)
-    verify(checkpointerMock, times(1)).checkpoint(anyString)
+    verify(receiverMock, times(1)).removeCheckpointer(meq(shardId), meq(checkpointerMock))
   }
 
+
   test("shutdown should not checkpoint if the reason is something other than TERMINATE") {
     when(receiverMock.getLatestSeqNumToCheckpoint(shardId)).thenReturn(someSeqNum)
 
-    val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId, checkpointStateMock)
+    val recordProcessor = new KinesisRecordProcessor(receiverMock, workerId)
     recordProcessor.initialize(shardId)
     recordProcessor.shutdown(checkpointerMock, ShutdownReason.ZOMBIE)
     recordProcessor.shutdown(checkpointerMock, null)
 
-    verify(checkpointerMock, never).checkpoint(anyString)
+    verify(receiverMock, times(2)).removeCheckpointer(meq(shardId),
+      meq[IRecordProcessorCheckpointer](null))
   }
 
   test("retry success on first attempt") {

From 8a2336893a7ff610a6c4629dd567b85078730616 Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Mon, 9 Nov 2015 14:56:36 -0800
Subject: [PATCH 137/518] [SPARK-6517][MLLIB] Implement the Algorithm of
 Hierarchical Clustering

I implemented a hierarchical clustering algorithm again.  This PR doesn't include examples, documentation and spark.ml APIs. I am going to send another PRs later.
https://issues.apache.org/jira/browse/SPARK-6517

- This implementation based on a bi-sectiong K-means clustering.
    - It derives from the freeman-lab 's implementation
- The basic idea is not changed from the previous version. (#2906)
    - However, It is 1000x faster than the previous version through parallel processing.

Thank you for your great cooperation, RJ Nowling(rnowling), Jeremy Freeman(freeman-lab), Xiangrui Meng(mengxr) and Sean Owen(srowen).

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Author: Xiangrui Meng <meng@databricks.com>
Author: Yu ISHIKAWA <yu-iskw@users.noreply.github.com>

Closes #5267 from yu-iskw/new-hierarchical-clustering.
---
 .../mllib/clustering/BisectingKMeans.scala    | 491 ++++++++++++++++++
 .../clustering/BisectingKMeansModel.scala     |  95 ++++
 .../clustering/JavaBisectingKMeansSuite.java  |  73 +++
 .../clustering/BisectingKMeansSuite.scala     | 182 +++++++
 4 files changed, 841 insertions(+)
 create mode 100644 mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeans.scala
 create mode 100644 mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeansModel.scala
 create mode 100644 mllib/src/test/java/org/apache/spark/mllib/clustering/JavaBisectingKMeansSuite.java
 create mode 100644 mllib/src/test/scala/org/apache/spark/mllib/clustering/BisectingKMeansSuite.scala

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeans.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeans.scala
new file mode 100644
index 0000000000000..29a7aa0bb63f2
--- /dev/null
+++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeans.scala
@@ -0,0 +1,491 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.clustering
+
+import java.util.Random
+
+import scala.collection.mutable
+
+import org.apache.spark.Logging
+import org.apache.spark.annotation.{Experimental, Since}
+import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.mllib.linalg.{BLAS, Vector, Vectors}
+import org.apache.spark.mllib.util.MLUtils
+import org.apache.spark.rdd.RDD
+import org.apache.spark.storage.StorageLevel
+
+/**
+ * A bisecting k-means algorithm based on the paper "A comparison of document clustering techniques"
+ * by Steinbach, Karypis, and Kumar, with modification to fit Spark.
+ * The algorithm starts from a single cluster that contains all points.
+ * Iteratively it finds divisible clusters on the bottom level and bisects each of them using
+ * k-means, until there are `k` leaf clusters in total or no leaf clusters are divisible.
+ * The bisecting steps of clusters on the same level are grouped together to increase parallelism.
+ * If bisecting all divisible clusters on the bottom level would result more than `k` leaf clusters,
+ * larger clusters get higher priority.
+ *
+ * @param k the desired number of leaf clusters (default: 4). The actual number could be smaller if
+ *          there are no divisible leaf clusters.
+ * @param maxIterations the max number of k-means iterations to split clusters (default: 20)
+ * @param minDivisibleClusterSize the minimum number of points (if >= 1.0) or the minimum proportion
+ *                                of points (if < 1.0) of a divisible cluster (default: 1)
+ * @param seed a random seed (default: hash value of the class name)
+ *
+ * @see [[http://glaros.dtc.umn.edu/gkhome/fetch/papers/docclusterKDDTMW00.pdf
+ *     Steinbach, Karypis, and Kumar, A comparison of document clustering techniques,
+ *     KDD Workshop on Text Mining, 2000.]]
+ */
+@Since("1.6.0")
+@Experimental
+class BisectingKMeans private (
+    private var k: Int,
+    private var maxIterations: Int,
+    private var minDivisibleClusterSize: Double,
+    private var seed: Long) extends Logging {
+
+  import BisectingKMeans._
+
+  /**
+   * Constructs with the default configuration
+   */
+  @Since("1.6.0")
+  def this() = this(4, 20, 1.0, classOf[BisectingKMeans].getName.##)
+
+  /**
+   * Sets the desired number of leaf clusters (default: 4).
+   * The actual number could be smaller if there are no divisible leaf clusters.
+   */
+  @Since("1.6.0")
+  def setK(k: Int): this.type = {
+    require(k > 0, s"k must be positive but got $k.")
+    this.k = k
+    this
+  }
+
+  /**
+   * Gets the desired number of leaf clusters.
+   */
+  @Since("1.6.0")
+  def getK: Int = this.k
+
+  /**
+   * Sets the max number of k-means iterations to split clusters (default: 20).
+   */
+  @Since("1.6.0")
+  def setMaxIterations(maxIterations: Int): this.type = {
+    require(maxIterations > 0, s"maxIterations must be positive but got $maxIterations.")
+    this.maxIterations = maxIterations
+    this
+  }
+
+  /**
+   * Gets the max number of k-means iterations to split clusters.
+   */
+  @Since("1.6.0")
+  def getMaxIterations: Int = this.maxIterations
+
+  /**
+   * Sets the minimum number of points (if >= `1.0`) or the minimum proportion of points
+   * (if < `1.0`) of a divisible cluster (default: 1).
+   */
+  @Since("1.6.0")
+  def setMinDivisibleClusterSize(minDivisibleClusterSize: Double): this.type = {
+    require(minDivisibleClusterSize > 0.0,
+      s"minDivisibleClusterSize must be positive but got $minDivisibleClusterSize.")
+    this.minDivisibleClusterSize = minDivisibleClusterSize
+    this
+  }
+
+  /**
+   * Gets the minimum number of points (if >= `1.0`) or the minimum proportion of points
+   * (if < `1.0`) of a divisible cluster.
+   */
+  @Since("1.6.0")
+  def getMinDivisibleClusterSize: Double = minDivisibleClusterSize
+
+  /**
+   * Sets the random seed (default: hash value of the class name).
+   */
+  @Since("1.6.0")
+  def setSeed(seed: Long): this.type = {
+    this.seed = seed
+    this
+  }
+
+  /**
+   * Gets the random seed.
+   */
+  @Since("1.6.0")
+  def getSeed: Long = this.seed
+
+  /**
+   * Runs the bisecting k-means algorithm.
+   * @param input RDD of vectors
+   * @return model for the bisecting kmeans
+   */
+  @Since("1.6.0")
+  def run(input: RDD[Vector]): BisectingKMeansModel = {
+    if (input.getStorageLevel == StorageLevel.NONE) {
+      logWarning(s"The input RDD ${input.id} is not directly cached, which may hurt performance if"
+        + " its parent RDDs are also not cached.")
+    }
+    val d = input.map(_.size).first()
+    logInfo(s"Feature dimension: $d.")
+    // Compute and cache vector norms for fast distance computation.
+    val norms = input.map(v => Vectors.norm(v, 2.0)).persist(StorageLevel.MEMORY_AND_DISK)
+    val vectors = input.zip(norms).map { case (x, norm) => new VectorWithNorm(x, norm) }
+    var assignments = vectors.map(v => (ROOT_INDEX, v))
+    var activeClusters = summarize(d, assignments)
+    val rootSummary = activeClusters(ROOT_INDEX)
+    val n = rootSummary.size
+    logInfo(s"Number of points: $n.")
+    logInfo(s"Initial cost: ${rootSummary.cost}.")
+    val minSize = if (minDivisibleClusterSize >= 1.0) {
+      math.ceil(minDivisibleClusterSize).toLong
+    } else {
+      math.ceil(minDivisibleClusterSize * n).toLong
+    }
+    logInfo(s"The minimum number of points of a divisible cluster is $minSize.")
+    var inactiveClusters = mutable.Seq.empty[(Long, ClusterSummary)]
+    val random = new Random(seed)
+    var numLeafClustersNeeded = k - 1
+    var level = 1
+    while (activeClusters.nonEmpty && numLeafClustersNeeded > 0 && level < LEVEL_LIMIT) {
+      // Divisible clusters are sufficiently large and have non-trivial cost.
+      var divisibleClusters = activeClusters.filter { case (_, summary) =>
+        (summary.size >= minSize) && (summary.cost > MLUtils.EPSILON * summary.size)
+      }
+      // If we don't need all divisible clusters, take the larger ones.
+      if (divisibleClusters.size > numLeafClustersNeeded) {
+        divisibleClusters = divisibleClusters.toSeq.sortBy { case (_, summary) =>
+            -summary.size
+          }.take(numLeafClustersNeeded)
+          .toMap
+      }
+      if (divisibleClusters.nonEmpty) {
+        val divisibleIndices = divisibleClusters.keys.toSet
+        logInfo(s"Dividing ${divisibleIndices.size} clusters on level $level.")
+        var newClusterCenters = divisibleClusters.flatMap { case (index, summary) =>
+          val (left, right) = splitCenter(summary.center, random)
+          Iterator((leftChildIndex(index), left), (rightChildIndex(index), right))
+        }.map(identity) // workaround for a Scala bug (SI-7005) that produces a not serializable map
+        var newClusters: Map[Long, ClusterSummary] = null
+        var newAssignments: RDD[(Long, VectorWithNorm)] = null
+        for (iter <- 0 until maxIterations) {
+          newAssignments = updateAssignments(assignments, divisibleIndices, newClusterCenters)
+            .filter { case (index, _) =>
+            divisibleIndices.contains(parentIndex(index))
+          }
+          newClusters = summarize(d, newAssignments)
+          newClusterCenters = newClusters.mapValues(_.center).map(identity)
+        }
+        // TODO: Unpersist old indices.
+        val indices = updateAssignments(assignments, divisibleIndices, newClusterCenters).keys
+          .persist(StorageLevel.MEMORY_AND_DISK)
+        assignments = indices.zip(vectors)
+        inactiveClusters ++= activeClusters
+        activeClusters = newClusters
+        numLeafClustersNeeded -= divisibleClusters.size
+      } else {
+        logInfo(s"None active and divisible clusters left on level $level. Stop iterations.")
+        inactiveClusters ++= activeClusters
+        activeClusters = Map.empty
+      }
+      level += 1
+    }
+    val clusters = activeClusters ++ inactiveClusters
+    val root = buildTree(clusters)
+    new BisectingKMeansModel(root)
+  }
+
+  /**
+   * Java-friendly version of [[run(RDD[Vector])*]]
+   */
+  def run(data: JavaRDD[Vector]): BisectingKMeansModel = run(data.rdd)
+}
+
+private object BisectingKMeans extends Serializable {
+
+  /** The index of the root node of a tree. */
+  private val ROOT_INDEX: Long = 1
+
+  private val MAX_DIVISIBLE_CLUSTER_INDEX: Long = Long.MaxValue / 2
+
+  private val LEVEL_LIMIT = math.log10(Long.MaxValue) / math.log10(2)
+
+  /** Returns the left child index of the given node index. */
+  private def leftChildIndex(index: Long): Long = {
+    require(index <= MAX_DIVISIBLE_CLUSTER_INDEX, s"Child index out of bound: 2 * $index.")
+    2 * index
+  }
+
+  /** Returns the right child index of the given node index. */
+  private def rightChildIndex(index: Long): Long = {
+    require(index <= MAX_DIVISIBLE_CLUSTER_INDEX, s"Child index out of bound: 2 * $index + 1.")
+    2 * index + 1
+  }
+
+  /** Returns the parent index of the given node index, or 0 if the input is 1 (root). */
+  private def parentIndex(index: Long): Long = {
+    index / 2
+  }
+
+  /**
+   * Summarizes data by each cluster as Map.
+   * @param d feature dimension
+   * @param assignments pairs of point and its cluster index
+   * @return a map from cluster indices to corresponding cluster summaries
+   */
+  private def summarize(
+      d: Int,
+      assignments: RDD[(Long, VectorWithNorm)]): Map[Long, ClusterSummary] = {
+    assignments.aggregateByKey(new ClusterSummaryAggregator(d))(
+        seqOp = (agg, v) => agg.add(v),
+        combOp = (agg1, agg2) => agg1.merge(agg2)
+      ).mapValues(_.summary)
+      .collect().toMap
+  }
+
+  /**
+   * Cluster summary aggregator.
+   * @param d feature dimension
+   */
+  private class ClusterSummaryAggregator(val d: Int) extends Serializable {
+    private var n: Long = 0L
+    private val sum: Vector = Vectors.zeros(d)
+    private var sumSq: Double = 0.0
+
+    /** Adds a point. */
+    def add(v: VectorWithNorm): this.type = {
+      n += 1L
+      // TODO: use a numerically stable approach to estimate cost
+      sumSq += v.norm * v.norm
+      BLAS.axpy(1.0, v.vector, sum)
+      this
+    }
+
+    /** Merges another aggregator. */
+    def merge(other: ClusterSummaryAggregator): this.type = {
+      n += other.n
+      sumSq += other.sumSq
+      BLAS.axpy(1.0, other.sum, sum)
+      this
+    }
+
+    /** Returns the summary. */
+    def summary: ClusterSummary = {
+      val mean = sum.copy
+      if (n > 0L) {
+        BLAS.scal(1.0 / n, mean)
+      }
+      val center = new VectorWithNorm(mean)
+      val cost = math.max(sumSq - n * center.norm * center.norm, 0.0)
+      new ClusterSummary(n, center, cost)
+    }
+  }
+
+  /**
+   * Bisects a cluster center.
+   *
+   * @param center current cluster center
+   * @param random a random number generator
+   * @return initial centers
+   */
+  private def splitCenter(
+      center: VectorWithNorm,
+      random: Random): (VectorWithNorm, VectorWithNorm) = {
+    val d = center.vector.size
+    val norm = center.norm
+    val level = 1e-4 * norm
+    val noise = Vectors.dense(Array.fill(d)(random.nextDouble()))
+    val left = center.vector.copy
+    BLAS.axpy(-level, noise, left)
+    val right = center.vector.copy
+    BLAS.axpy(level, noise, right)
+    (new VectorWithNorm(left), new VectorWithNorm(right))
+  }
+
+  /**
+   * Updates assignments.
+   * @param assignments current assignments
+   * @param divisibleIndices divisible cluster indices
+   * @param newClusterCenters new cluster centers
+   * @return new assignments
+   */
+  private def updateAssignments(
+      assignments: RDD[(Long, VectorWithNorm)],
+      divisibleIndices: Set[Long],
+      newClusterCenters: Map[Long, VectorWithNorm]): RDD[(Long, VectorWithNorm)] = {
+    assignments.map { case (index, v) =>
+      if (divisibleIndices.contains(index)) {
+        val children = Seq(leftChildIndex(index), rightChildIndex(index))
+        val selected = children.minBy { child =>
+          KMeans.fastSquaredDistance(newClusterCenters(child), v)
+        }
+        (selected, v)
+      } else {
+        (index, v)
+      }
+    }
+  }
+
+  /**
+   * Builds a clustering tree by re-indexing internal and leaf clusters.
+   * @param clusters a map from cluster indices to corresponding cluster summaries
+   * @return the root node of the clustering tree
+   */
+  private def buildTree(clusters: Map[Long, ClusterSummary]): ClusteringTreeNode = {
+    var leafIndex = 0
+    var internalIndex = -1
+
+    /**
+     * Builds a subtree from this given node index.
+     */
+    def buildSubTree(rawIndex: Long): ClusteringTreeNode = {
+      val cluster = clusters(rawIndex)
+      val size = cluster.size
+      val center = cluster.center
+      val cost = cluster.cost
+      val isInternal = clusters.contains(leftChildIndex(rawIndex))
+      if (isInternal) {
+        val index = internalIndex
+        internalIndex -= 1
+        val leftIndex = leftChildIndex(rawIndex)
+        val rightIndex = rightChildIndex(rawIndex)
+        val height = math.sqrt(Seq(leftIndex, rightIndex).map { childIndex =>
+          KMeans.fastSquaredDistance(center, clusters(childIndex).center)
+        }.max)
+        val left = buildSubTree(leftIndex)
+        val right = buildSubTree(rightIndex)
+        new ClusteringTreeNode(index, size, center, cost, height, Array(left, right))
+      } else {
+        val index = leafIndex
+        leafIndex += 1
+        val height = 0.0
+        new ClusteringTreeNode(index, size, center, cost, height, Array.empty)
+      }
+    }
+
+    buildSubTree(ROOT_INDEX)
+  }
+
+  /**
+   * Summary of a cluster.
+   *
+   * @param size the number of points within this cluster
+   * @param center the center of the points within this cluster
+   * @param cost the sum of squared distances to the center
+   */
+  private case class ClusterSummary(size: Long, center: VectorWithNorm, cost: Double)
+}
+
+/**
+ * Represents a node in a clustering tree.
+ *
+ * @param index node index, negative for internal nodes and non-negative for leaf nodes
+ * @param size size of the cluster
+ * @param centerWithNorm cluster center with norm
+ * @param cost cost of the cluster, i.e., the sum of squared distances to the center
+ * @param height height of the node in the dendrogram. Currently this is defined as the max distance
+ *               from the center to the centers of the children's, but subject to change.
+ * @param children children nodes
+ */
+@Since("1.6.0")
+@Experimental
+class ClusteringTreeNode private[clustering] (
+    val index: Int,
+    val size: Long,
+    private val centerWithNorm: VectorWithNorm,
+    val cost: Double,
+    val height: Double,
+    val children: Array[ClusteringTreeNode]) extends Serializable {
+
+  /** Whether this is a leaf node. */
+  val isLeaf: Boolean = children.isEmpty
+
+  require((isLeaf && index >= 0) || (!isLeaf && index < 0))
+
+  /** Cluster center. */
+  def center: Vector = centerWithNorm.vector
+
+  /** Predicts the leaf cluster node index that the input point belongs to. */
+  def predict(point: Vector): Int = {
+    val (index, _) = predict(new VectorWithNorm(point))
+    index
+  }
+
+  /** Returns the full prediction path from root to leaf. */
+  def predictPath(point: Vector): Array[ClusteringTreeNode] = {
+    predictPath(new VectorWithNorm(point)).toArray
+  }
+
+  /** Returns the full prediction path from root to leaf. */
+  private def predictPath(pointWithNorm: VectorWithNorm): List[ClusteringTreeNode] = {
+    if (isLeaf) {
+      this :: Nil
+    } else {
+      val selected = children.minBy { child =>
+        KMeans.fastSquaredDistance(child.centerWithNorm, pointWithNorm)
+      }
+      selected :: selected.predictPath(pointWithNorm)
+    }
+  }
+
+  /**
+   * Computes the cost (squared distance to the predicted leaf cluster center) of the input point.
+   */
+  def computeCost(point: Vector): Double = {
+    val (_, cost) = predict(new VectorWithNorm(point))
+    cost
+  }
+
+  /**
+   * Predicts the cluster index and the cost of the input point.
+   */
+  private def predict(pointWithNorm: VectorWithNorm): (Int, Double) = {
+    predict(pointWithNorm, KMeans.fastSquaredDistance(centerWithNorm, pointWithNorm))
+  }
+
+  /**
+   * Predicts the cluster index and the cost of the input point.
+   * @param pointWithNorm input point
+   * @param cost the cost to the current center
+   * @return (predicted leaf cluster index, cost)
+   */
+  private def predict(pointWithNorm: VectorWithNorm, cost: Double): (Int, Double) = {
+    if (isLeaf) {
+      (index, cost)
+    } else {
+      val (selectedChild, minCost) = children.map { child =>
+        (child, KMeans.fastSquaredDistance(child.centerWithNorm, pointWithNorm))
+      }.minBy(_._2)
+      selectedChild.predict(pointWithNorm, minCost)
+    }
+  }
+
+  /**
+   * Returns all leaf nodes from this node.
+   */
+  def leafNodes: Array[ClusteringTreeNode] = {
+    if (isLeaf) {
+      Array(this)
+    } else {
+      children.flatMap(_.leafNodes)
+    }
+  }
+}
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeansModel.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeansModel.scala
new file mode 100644
index 0000000000000..5015f1540d920
--- /dev/null
+++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeansModel.scala
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.clustering
+
+import org.apache.spark.Logging
+import org.apache.spark.annotation.{Experimental, Since}
+import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.mllib.linalg.Vector
+import org.apache.spark.rdd.RDD
+
+/**
+ * Clustering model produced by [[BisectingKMeans]].
+ * The prediction is done level-by-level from the root node to a leaf node, and at each node among
+ * its children the closest to the input point is selected.
+ *
+ * @param root the root node of the clustering tree
+ */
+@Since("1.6.0")
+@Experimental
+class BisectingKMeansModel @Since("1.6.0") (
+    @Since("1.6.0") val root: ClusteringTreeNode
+  ) extends Serializable with Logging {
+
+  /**
+   * Leaf cluster centers.
+   */
+  @Since("1.6.0")
+  def clusterCenters: Array[Vector] = root.leafNodes.map(_.center)
+
+  /**
+   * Number of leaf clusters.
+   */
+  lazy val k: Int = clusterCenters.length
+
+  /**
+   * Predicts the index of the cluster that the input point belongs to.
+   */
+  @Since("1.6.0")
+  def predict(point: Vector): Int = {
+    root.predict(point)
+  }
+
+  /**
+   * Predicts the indices of the clusters that the input points belong to.
+   */
+  @Since("1.6.0")
+  def predict(points: RDD[Vector]): RDD[Int] = {
+    points.map { p => root.predict(p) }
+  }
+
+  /**
+   * Java-friendly version of [[predict(RDD[Vector])*]]
+   */
+  @Since("1.6.0")
+  def predict(points: JavaRDD[Vector]): JavaRDD[java.lang.Integer] =
+    predict(points.rdd).toJavaRDD().asInstanceOf[JavaRDD[java.lang.Integer]]
+
+  /**
+   * Computes the squared distance between the input point and the cluster center it belongs to.
+   */
+  @Since("1.6.0")
+  def computeCost(point: Vector): Double = {
+    root.computeCost(point)
+  }
+
+  /**
+   * Computes the sum of squared distances between the input points and their corresponding cluster
+   * centers.
+   */
+  @Since("1.6.0")
+  def computeCost(data: RDD[Vector]): Double = {
+    data.map(root.computeCost).sum()
+  }
+
+  /**
+   * Java-friendly version of [[computeCost(RDD[Vector])*]].
+   */
+  @Since("1.6.0")
+  def computeCost(data: JavaRDD[Vector]): Double = this.computeCost(data.rdd)
+}
diff --git a/mllib/src/test/java/org/apache/spark/mllib/clustering/JavaBisectingKMeansSuite.java b/mllib/src/test/java/org/apache/spark/mllib/clustering/JavaBisectingKMeansSuite.java
new file mode 100644
index 0000000000000..a714620ff7e4b
--- /dev/null
+++ b/mllib/src/test/java/org/apache/spark/mllib/clustering/JavaBisectingKMeansSuite.java
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.clustering;
+
+import java.io.Serializable;
+
+import com.google.common.collect.Lists;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.mllib.linalg.Vector;
+import org.apache.spark.mllib.linalg.Vectors;
+
+public class JavaBisectingKMeansSuite implements Serializable {
+  private transient JavaSparkContext sc;
+
+  @Before
+  public void setUp() {
+    sc = new JavaSparkContext("local", this.getClass().getSimpleName());
+  }
+
+  @After
+  public void tearDown() {
+    sc.stop();
+    sc = null;
+  }
+
+  @Test
+  public void twoDimensionalData() {
+    JavaRDD<Vector> points = sc.parallelize(Lists.newArrayList(
+      Vectors.dense(4, -1),
+      Vectors.dense(4, 1),
+      Vectors.sparse(2, new int[] {0}, new double[] {1.0})
+    ), 2);
+
+    BisectingKMeans bkm = new BisectingKMeans()
+      .setK(4)
+      .setMaxIterations(2)
+      .setSeed(1L);
+    BisectingKMeansModel model = bkm.run(points);
+    Assert.assertEquals(3, model.k());
+    Assert.assertArrayEquals(new double[] {3.0, 0.0}, model.root().center().toArray(), 1e-12);
+    for (ClusteringTreeNode child: model.root().children()) {
+      double[] center = child.center().toArray();
+      if (center[0] > 2) {
+        Assert.assertEquals(2, child.size());
+        Assert.assertArrayEquals(new double[] {4.0, 0.0}, center, 1e-12);
+      } else {
+        Assert.assertEquals(1, child.size());
+        Assert.assertArrayEquals(new double[] {1.0, 0.0}, center, 1e-12);
+      }
+    }
+  }
+}
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/clustering/BisectingKMeansSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/clustering/BisectingKMeansSuite.scala
new file mode 100644
index 0000000000000..41b9d5c0d93bb
--- /dev/null
+++ b/mllib/src/test/scala/org/apache/spark/mllib/clustering/BisectingKMeansSuite.scala
@@ -0,0 +1,182 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.clustering
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.mllib.linalg.Vectors
+import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.mllib.util.TestingUtils._
+
+class BisectingKMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
+
+  test("default values") {
+    val bkm0 = new BisectingKMeans()
+    assert(bkm0.getK === 4)
+    assert(bkm0.getMaxIterations === 20)
+    assert(bkm0.getMinDivisibleClusterSize === 1.0)
+    val bkm1 = new BisectingKMeans()
+    assert(bkm0.getSeed === bkm1.getSeed, "The default seed should be constant.")
+  }
+
+  test("setter/getter") {
+    val bkm = new BisectingKMeans()
+
+    val k = 10
+    assert(bkm.getK !== k)
+    assert(bkm.setK(k).getK === k)
+    val maxIter = 100
+    assert(bkm.getMaxIterations !== maxIter)
+    assert(bkm.setMaxIterations(maxIter).getMaxIterations === maxIter)
+    val minSize = 2.0
+    assert(bkm.getMinDivisibleClusterSize !== minSize)
+    assert(bkm.setMinDivisibleClusterSize(minSize).getMinDivisibleClusterSize === minSize)
+    val seed = 10L
+    assert(bkm.getSeed !== seed)
+    assert(bkm.setSeed(seed).getSeed === seed)
+
+    intercept[IllegalArgumentException] {
+      bkm.setK(0)
+    }
+    intercept[IllegalArgumentException] {
+      bkm.setMaxIterations(0)
+    }
+    intercept[IllegalArgumentException] {
+      bkm.setMinDivisibleClusterSize(0.0)
+    }
+  }
+
+  test("1D data") {
+    val points = Vectors.sparse(1, Array.empty, Array.empty) +:
+      (1 until 8).map(i => Vectors.dense(i))
+    val data = sc.parallelize(points, 2)
+    val bkm = new BisectingKMeans()
+      .setK(4)
+      .setMaxIterations(1)
+      .setSeed(1L)
+    // The clusters should be
+    // (0, 1, 2, 3, 4, 5, 6, 7)
+    //   - (0, 1, 2, 3)
+    //     - (0, 1)
+    //     - (2, 3)
+    //   - (4, 5, 6, 7)
+    //     - (4, 5)
+    //     - (6, 7)
+    val model = bkm.run(data)
+    assert(model.k === 4)
+    // The total cost should be 8 * 0.5 * 0.5 = 2.0.
+    assert(model.computeCost(data) ~== 2.0 relTol 1e-12)
+    val predictions = data.map(v => (v(0), model.predict(v))).collectAsMap()
+    Range(0, 8, 2).foreach { i =>
+      assert(predictions(i) === predictions(i + 1),
+        s"$i and ${i + 1} should belong to the same cluster.")
+    }
+    val root = model.root
+    assert(root.center(0) ~== 3.5 relTol 1e-12)
+    assert(root.height ~== 2.0 relTol 1e-12)
+    assert(root.children.length === 2)
+    assert(root.children(0).height ~== 1.0 relTol 1e-12)
+    assert(root.children(1).height ~== 1.0 relTol 1e-12)
+  }
+
+  test("points are the same") {
+    val data = sc.parallelize(Seq.fill(8)(Vectors.dense(1.0, 1.0)), 2)
+    val bkm = new BisectingKMeans()
+      .setK(2)
+      .setMaxIterations(1)
+      .setSeed(1L)
+    val model = bkm.run(data)
+    assert(model.k === 1)
+  }
+
+  test("more desired clusters than points") {
+    val data = sc.parallelize(Seq.tabulate(4)(i => Vectors.dense(i)), 2)
+    val bkm = new BisectingKMeans()
+      .setK(8)
+      .setMaxIterations(2)
+      .setSeed(1L)
+    val model = bkm.run(data)
+    assert(model.k === 4)
+  }
+
+  test("min divisible cluster") {
+    val data = sc.parallelize(
+      Seq.tabulate(16)(i => Vectors.dense(i)) ++ Seq.tabulate(4)(i => Vectors.dense(-100.0 - i)),
+      2)
+    val bkm = new BisectingKMeans()
+      .setK(4)
+      .setMinDivisibleClusterSize(10)
+      .setMaxIterations(1)
+      .setSeed(1L)
+    val model = bkm.run(data)
+    assert(model.k === 3)
+    assert(model.predict(Vectors.dense(-100)) === model.predict(Vectors.dense(-97)))
+    assert(model.predict(Vectors.dense(7)) !== model.predict(Vectors.dense(8)))
+
+    bkm.setMinDivisibleClusterSize(0.5)
+    val sameModel = bkm.run(data)
+    assert(sameModel.k === 3)
+  }
+
+  test("larger clusters get selected first") {
+    val data = sc.parallelize(
+      Seq.tabulate(16)(i => Vectors.dense(i)) ++ Seq.tabulate(4)(i => Vectors.dense(-100.0 - i)),
+      2)
+    val bkm = new BisectingKMeans()
+      .setK(3)
+      .setMaxIterations(1)
+      .setSeed(1L)
+    val model = bkm.run(data)
+    assert(model.k === 3)
+    assert(model.predict(Vectors.dense(-100)) === model.predict(Vectors.dense(-97)))
+    assert(model.predict(Vectors.dense(7)) !== model.predict(Vectors.dense(8)))
+  }
+
+  test("2D data") {
+    val points = Seq(
+      (11, 10), (9, 10), (10, 9), (10, 11),
+      (11, -10), (9, -10), (10, -9), (10, -11),
+      (0, 1), (0, -1)
+    ).map { case (x, y) =>
+      if (x == 0) {
+        Vectors.sparse(2, Array(1), Array(y))
+      } else {
+        Vectors.dense(x, y)
+      }
+    }
+    val data = sc.parallelize(points, 2)
+    val bkm = new BisectingKMeans()
+      .setK(3)
+      .setMaxIterations(4)
+      .setSeed(1L)
+    val model = bkm.run(data)
+    assert(model.k === 3)
+    assert(model.root.center ~== Vectors.dense(8, 0) relTol 1e-12)
+    model.root.leafNodes.foreach { node =>
+      if (node.center(0) < 5) {
+        assert(node.size === 2)
+        assert(node.center ~== Vectors.dense(0, 0) relTol 1e-12)
+      } else if (node.center(1) > 0) {
+        assert(node.size === 4)
+        assert(node.center ~== Vectors.dense(10, 10) relTol 1e-12)
+      } else {
+        assert(node.size === 4)
+        assert(node.center ~== Vectors.dense(10, -10) relTol 1e-12)
+      }
+    }
+  }
+}

From fcb57e9c7323e24b8563800deb035f94f616474e Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Mon, 9 Nov 2015 15:16:47 -0800
Subject: [PATCH 138/518] [SPARK-11564][SQL][FOLLOW-UP] improve java api for
 GroupedDataset

created `MapGroupFunction`, `FlatMapGroupFunction`, `CoGroupFunction`

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9564 from cloud-fan/map.
---
 .../api/java/function/CoGroupFunction.java    | 29 +++++++++++++++
 .../api/java/function/FlatMapFunction.java    |  2 +-
 .../api/java/function/FlatMapFunction2.java   |  2 +-
 .../java/function/FlatMapGroupFunction.java   | 28 +++++++++++++++
 .../api/java/function/MapGroupFunction.java   | 28 +++++++++++++++
 .../plans/logical/basicOperators.scala        |  4 +--
 .../org/apache/spark/sql/GroupedDataset.scala | 12 +++----
 .../spark/sql/execution/basicOperators.scala  |  2 +-
 .../apache/spark/sql/JavaDatasetSuite.java    | 36 ++++++++++++-------
 9 files changed, 118 insertions(+), 25 deletions(-)
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/CoGroupFunction.java
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupFunction.java
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java

diff --git a/core/src/main/java/org/apache/spark/api/java/function/CoGroupFunction.java b/core/src/main/java/org/apache/spark/api/java/function/CoGroupFunction.java
new file mode 100644
index 0000000000000..279639af5d430
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/CoGroupFunction.java
@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+/**
+ * A function that returns zero or more output records from each grouping key and its values from 2
+ * Datasets.
+ */
+public interface CoGroupFunction<K, V1, V2, R> extends Serializable {
+  Iterable<R> call(K key, Iterator<V1> left, Iterator<V2> right) throws Exception;
+}
diff --git a/core/src/main/java/org/apache/spark/api/java/function/FlatMapFunction.java b/core/src/main/java/org/apache/spark/api/java/function/FlatMapFunction.java
index 23f5fdd43631b..ef0d1824121ec 100644
--- a/core/src/main/java/org/apache/spark/api/java/function/FlatMapFunction.java
+++ b/core/src/main/java/org/apache/spark/api/java/function/FlatMapFunction.java
@@ -23,5 +23,5 @@
  * A function that returns zero or more output records from each input record.
  */
 public interface FlatMapFunction<T, R> extends Serializable {
-  public Iterable<R> call(T t) throws Exception;
+  Iterable<R> call(T t) throws Exception;
 }
diff --git a/core/src/main/java/org/apache/spark/api/java/function/FlatMapFunction2.java b/core/src/main/java/org/apache/spark/api/java/function/FlatMapFunction2.java
index c48e92f535ff5..14a98a38ef5ab 100644
--- a/core/src/main/java/org/apache/spark/api/java/function/FlatMapFunction2.java
+++ b/core/src/main/java/org/apache/spark/api/java/function/FlatMapFunction2.java
@@ -23,5 +23,5 @@
  * A function that takes two inputs and returns zero or more output records.
  */
 public interface FlatMapFunction2<T1, T2, R> extends Serializable {
-  public Iterable<R> call(T1 t1, T2 t2) throws Exception;
+  Iterable<R> call(T1 t1, T2 t2) throws Exception;
 }
diff --git a/core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupFunction.java b/core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupFunction.java
new file mode 100644
index 0000000000000..18a2d733ca70d
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupFunction.java
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+/**
+ * A function that returns zero or more output records from each grouping key and its values.
+ */
+public interface FlatMapGroupFunction<K, V, R> extends Serializable {
+  Iterable<R> call(K key, Iterator<V> values) throws Exception;
+}
diff --git a/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java b/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java
new file mode 100644
index 0000000000000..2935f9986a560
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java
@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+/**
+ * Base interface for a map function used in GroupedDataset's map function.
+ */
+public interface MapGroupFunction<K, V, R> extends Serializable {
+  R call(K key, Iterator<V> values) throws Exception;
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index e151ac04ede2a..d771088d69dea 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -527,7 +527,7 @@ case class MapGroups[K, T, U](
 /** Factory for constructing new `CoGroup` nodes. */
 object CoGroup {
   def apply[K : Encoder, Left : Encoder, Right : Encoder, R : Encoder](
-      func: (K, Iterator[Left], Iterator[Right]) => Iterator[R],
+      func: (K, Iterator[Left], Iterator[Right]) => TraversableOnce[R],
       leftGroup: Seq[Attribute],
       rightGroup: Seq[Attribute],
       left: LogicalPlan,
@@ -551,7 +551,7 @@ object CoGroup {
  * right children.
  */
 case class CoGroup[K, Left, Right, R](
-    func: (K, Iterator[Left], Iterator[Right]) => Iterator[R],
+    func: (K, Iterator[Left], Iterator[Right]) => TraversableOnce[R],
     kEncoder: ExpressionEncoder[K],
     leftEnc: ExpressionEncoder[Left],
     rightEnc: ExpressionEncoder[Right],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 5c3f626545875..850315e281dfe 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -108,9 +108,7 @@ class GroupedDataset[K, T] private[sql](
       MapGroups(f, groupingAttributes, logicalPlan))
   }
 
-  def flatMap[U](
-      f: JFunction2[K, JIterator[T], JIterator[U]],
-      encoder: Encoder[U]): Dataset[U] = {
+  def flatMap[U](f: FlatMapGroupFunction[K, T, U], encoder: Encoder[U]): Dataset[U] = {
     flatMap((key, data) => f.call(key, data.asJava).asScala)(encoder)
   }
 
@@ -131,9 +129,7 @@ class GroupedDataset[K, T] private[sql](
       MapGroups(func, groupingAttributes, logicalPlan))
   }
 
-  def map[U](
-      f: JFunction2[K, JIterator[T], U],
-      encoder: Encoder[U]): Dataset[U] = {
+  def map[U](f: MapGroupFunction[K, T, U], encoder: Encoder[U]): Dataset[U] = {
     map((key, data) => f.call(key, data.asJava))(encoder)
   }
 
@@ -218,7 +214,7 @@ class GroupedDataset[K, T] private[sql](
    */
   def cogroup[U, R : Encoder](
       other: GroupedDataset[K, U])(
-      f: (K, Iterator[T], Iterator[U]) => Iterator[R]): Dataset[R] = {
+      f: (K, Iterator[T], Iterator[U]) => TraversableOnce[R]): Dataset[R] = {
     implicit def uEnc: Encoder[U] = other.tEncoder
     new Dataset[R](
       sqlContext,
@@ -232,7 +228,7 @@ class GroupedDataset[K, T] private[sql](
 
   def cogroup[U, R](
       other: GroupedDataset[K, U],
-      f: JFunction3[K, JIterator[T], JIterator[U], JIterator[R]],
+      f: CoGroupFunction[K, T, U, R],
       encoder: Encoder[R]): Dataset[R] = {
     cogroup(other)((key, left, right) => f.call(key, left.asJava, right.asJava).asScala)(encoder)
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index 2593b16b1c8d7..145de0db9edaa 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -391,7 +391,7 @@ case class MapGroups[K, T, U](
  * The result of this function is encoded and flattened before being output.
  */
 case class CoGroup[K, Left, Right, R](
-    func: (K, Iterator[Left], Iterator[Right]) => Iterator[R],
+    func: (K, Iterator[Left], Iterator[Right]) => TraversableOnce[R],
     kEncoder: ExpressionEncoder[K],
     leftEnc: ExpressionEncoder[Left],
     rightEnc: ExpressionEncoder[Right],
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index 0f90de774dd3e..312cf33e4c2d4 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -29,7 +29,6 @@
 import org.apache.spark.Accumulator;
 import org.apache.spark.SparkContext;
 import org.apache.spark.api.java.function.*;
-import org.apache.spark.api.java.function.Function;
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.catalyst.encoders.Encoder;
 import org.apache.spark.sql.catalyst.encoders.Encoder$;
@@ -170,20 +169,33 @@ public Integer call(String v) throws Exception {
       }
     }, e.INT());
 
-    Dataset<String> mapped = grouped.map(
-      new Function2<Integer, Iterator<String>, String>() {
+    Dataset<String> mapped = grouped.map(new MapGroupFunction<Integer, String, String>() {
+      @Override
+      public String call(Integer key, Iterator<String> values) throws Exception {
+        StringBuilder sb = new StringBuilder(key.toString());
+        while (values.hasNext()) {
+          sb.append(values.next());
+        }
+        return sb.toString();
+      }
+    }, e.STRING());
+
+    Assert.assertEquals(Arrays.asList("1a", "3foobar"), mapped.collectAsList());
+
+    Dataset<String> flatMapped = grouped.flatMap(
+      new FlatMapGroupFunction<Integer, String, String>() {
         @Override
-        public String call(Integer key, Iterator<String> data) throws Exception {
+        public Iterable<String> call(Integer key, Iterator<String> values) throws Exception {
           StringBuilder sb = new StringBuilder(key.toString());
-          while (data.hasNext()) {
-            sb.append(data.next());
+          while (values.hasNext()) {
+            sb.append(values.next());
           }
-          return sb.toString();
+          return Collections.singletonList(sb.toString());
         }
       },
       e.STRING());
 
-    Assert.assertEquals(Arrays.asList("1a", "3foobar"), mapped.collectAsList());
+    Assert.assertEquals(Arrays.asList("1a", "3foobar"), flatMapped.collectAsList());
 
     List<Integer> data2 = Arrays.asList(2, 6, 10);
     Dataset<Integer> ds2 = context.createDataset(data2, e.INT());
@@ -196,9 +208,9 @@ public Integer call(Integer v) throws Exception {
 
     Dataset<String> cogrouped = grouped.cogroup(
       grouped2,
-      new Function3<Integer, Iterator<String>, Iterator<Integer>, Iterator<String>>() {
+      new CoGroupFunction<Integer, String, Integer, String>() {
         @Override
-        public Iterator<String> call(
+        public Iterable<String> call(
             Integer key,
             Iterator<String> left,
             Iterator<Integer> right) throws Exception {
@@ -210,7 +222,7 @@ public Iterator<String> call(
           while (right.hasNext()) {
             sb.append(right.next());
           }
-          return Collections.singletonList(sb.toString()).iterator();
+          return Collections.singletonList(sb.toString());
         }
       },
       e.STRING());
@@ -225,7 +237,7 @@ public void testGroupByColumn() {
     GroupedDataset<Integer, String> grouped = ds.groupBy(length(col("value"))).asKey(e.INT());
 
     Dataset<String> mapped = grouped.map(
-      new Function2<Integer, Iterator<String>, String>() {
+      new MapGroupFunction<Integer, String, String>() {
         @Override
         public String call(Integer key, Iterator<String> data) throws Exception {
           StringBuilder sb = new StringBuilder(key.toString());

From 9565c246eadecf4836d247d0067f2200f061d25f Mon Sep 17 00:00:00 2001
From: hyukjinkwon <gurwls223@gmail.com>
Date: Mon, 9 Nov 2015 15:20:50 -0800
Subject: [PATCH 139/518] [SPARK-9557][SQL] Refactor ParquetFilterSuite and
 remove old ParquetFilters code

Actually this was resolved by https://github.com/apache/spark/pull/8275.

But I found the JIRA issue for this is not marked as resolved since the PR above was made for another issue but the PR above resolved both.

I commented that this is resolved by the PR above; however, I opened this PR as I would like to just add
a little bit of corrections.

In the previous PR, I refactored the test by not reducing just collecting filters; however, this would not test  properly `And` filter (which is not given to the tests). I unintentionally changed this from the original way (before being refactored).

In this PR, I just followed the original way to collect filters by reducing.

I would like to close this if this PR is inappropriate and somebody would like this deal with it in the separate PR related with this.

Author: hyukjinkwon <gurwls223@gmail.com>

Closes #9554 from HyukjinKwon/SPARK-9557.
---
 .../datasources/parquet/ParquetFilterSuite.scala          | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
index c24c9f025dad7..579dabf73318b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
@@ -54,12 +54,12 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
         .select(output.map(e => Column(e)): _*)
         .where(Column(predicate))
 
-      val analyzedPredicate = query.queryExecution.optimizedPlan.collect {
+      val maybeAnalyzedPredicate = query.queryExecution.optimizedPlan.collect {
         case PhysicalOperation(_, filters, LogicalRelation(_: ParquetRelation, _)) => filters
-      }.flatten
-      assert(analyzedPredicate.nonEmpty)
+      }.flatten.reduceLeftOption(_ && _)
+      assert(maybeAnalyzedPredicate.isDefined)
 
-      val selectedFilters = analyzedPredicate.flatMap(DataSourceStrategy.translateFilter)
+      val selectedFilters = maybeAnalyzedPredicate.flatMap(DataSourceStrategy.translateFilter)
       assert(selectedFilters.nonEmpty)
 
       selectedFilters.foreach { pred =>

From 2f38378856fb56bdd9be7ccedf56427e81701f4e Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Mon, 9 Nov 2015 16:06:48 -0800
Subject: [PATCH 140/518] [SPARK-11360][DOC] Loss of nullability when writing
 parquet files

This fix is to add one line to explain the current behavior of Spark SQL when writing Parquet files. All columns are forced to be nullable for compatibility reasons.

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9314 from gatorsmile/lossNull.
---
 docs/sql-programming-guide.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/sql-programming-guide.md b/docs/sql-programming-guide.md
index ccd26904329d3..6e02d6564b002 100644
--- a/docs/sql-programming-guide.md
+++ b/docs/sql-programming-guide.md
@@ -982,7 +982,8 @@ when a table is dropped.
 
 [Parquet](http://parquet.io) is a columnar format that is supported by many other data processing systems.
 Spark SQL provides support for both reading and writing Parquet files that automatically preserves the schema
-of the original data.
+of the original data. When writing Parquet files, all columns are automatically converted to be nullable for 
+compatibility reasons.
 
 ### Loading Data Programmatically
 

From 9c740a9ddf6344a03b4b45380eaf0cfc6e2299b5 Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Mon, 9 Nov 2015 16:11:00 -0800
Subject: [PATCH 141/518] [SPARK-11578][SQL] User API for Typed Aggregation

This PR adds a new interface for user-defined aggregations, that can be used in `DataFrame` and `Dataset` operations to take all of the elements of a group and reduce them to a single value.

For example, the following aggregator extracts an `int` from a specific class and adds them up:

```scala
  case class Data(i: Int)

  val customSummer =  new Aggregator[Data, Int, Int] {
    def prepare(d: Data) = d.i
    def reduce(l: Int, r: Int) = l + r
    def present(r: Int) = r
  }.toColumn()

  val ds: Dataset[Data] = ...
  val aggregated = ds.select(customSummer)
```

By using helper functions, users can make a generic `Aggregator` that works on any input type:

```scala
/** An `Aggregator` that adds up any numeric type returned by the given function. */
class SumOf[I, N : Numeric](f: I => N) extends Aggregator[I, N, N] with Serializable {
  val numeric = implicitly[Numeric[N]]
  override def zero: N = numeric.zero
  override def reduce(b: N, a: I): N = numeric.plus(b, f(a))
  override def present(reduction: N): N = reduction
}

def sum[I, N : Numeric : Encoder](f: I => N): TypedColumn[I, N] = new SumOf(f).toColumn
```

These aggregators can then be used alongside other built-in SQL aggregations.

```scala
val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
ds
  .groupBy(_._1)
  .agg(
    sum(_._2),                // The aggregator defined above.
    expr("sum(_2)").as[Int],  // A built-in dynatically typed aggregation.
    count("*"))               // A built-in statically typed aggregation.
  .collect()

res0: ("a", 30, 30, 2L), ("b", 3, 3, 2L), ("c", 1, 1, 1L)
```

The current implementation focuses on integrating this into the typed API, but currently only supports running aggregations that return a single long value as explained in `TypedAggregateExpression`.  This will be improved in a followup PR.

Author: Michael Armbrust <michael@databricks.com>

Closes #9555 from marmbrus/dataset-useragg.
---
 .../scala/org/apache/spark/sql/Column.scala   |  11 +-
 .../scala/org/apache/spark/sql/Dataset.scala  |  30 ++--
 .../org/apache/spark/sql/GroupedDataset.scala |  51 ++++---
 .../org/apache/spark/sql/SQLContext.scala     |   1 -
 .../aggregate/TypedAggregateExpression.scala  | 129 ++++++++++++++++++
 .../spark/sql/expressions/Aggregator.scala    |  81 +++++++++++
 .../org/apache/spark/sql/functions.scala      |  30 +++-
 .../apache/spark/sql/JavaDatasetSuite.java    |   4 +-
 .../spark/sql/DatasetAggregatorSuite.scala    |  65 +++++++++
 9 files changed, 360 insertions(+), 42 deletions(-)
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
 create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index c32c93897ce0b..d26b6c3579205 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -23,7 +23,7 @@ import org.apache.spark.annotation.Experimental
 import org.apache.spark.Logging
 import org.apache.spark.sql.functions.lit
 import org.apache.spark.sql.catalyst.analysis._
-import org.apache.spark.sql.catalyst.encoders.Encoder
+import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.util.DataTypeParser
 import org.apache.spark.sql.types._
@@ -39,10 +39,13 @@ private[sql] object Column {
 }
 
 /**
- * A [[Column]] where an [[Encoder]] has been given for the expected return type.
+ * A [[Column]] where an [[Encoder]] has been given for the expected input and return type.
  * @since 1.6.0
+ * @tparam T The input type expected for this expression.  Can be `Any` if the expression is type
+ *           checked by the analyzer instead of the compiler (i.e. `expr("sum(...)")`).
+ * @tparam U The output type of this column.
  */
-class TypedColumn[T](expr: Expression)(implicit val encoder: Encoder[T]) extends Column(expr)
+class TypedColumn[-T, U](expr: Expression, val encoder: Encoder[U]) extends Column(expr)
 
 /**
  * :: Experimental ::
@@ -85,7 +88,7 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    * results into the correct JVM types.
    * @since 1.6.0
    */
-  def as[T : Encoder]: TypedColumn[T] = new TypedColumn[T](expr)
+  def as[U : Encoder]: TypedColumn[Any, U] = new TypedColumn[Any, U](expr, encoderFor[U])
 
   /**
    * Extracts a value or values from a complex type.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 959e0f5ba03e6..6d2968e2881f8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -358,7 +358,7 @@ class Dataset[T] private[sql](
    * }}}
    * @since 1.6.0
    */
-  def select[U1: Encoder](c1: TypedColumn[U1]): Dataset[U1] = {
+  def select[U1: Encoder](c1: TypedColumn[T, U1]): Dataset[U1] = {
     new Dataset[U1](sqlContext, Project(Alias(c1.expr, "_1")() :: Nil, logicalPlan))
   }
 
@@ -367,7 +367,7 @@ class Dataset[T] private[sql](
    * code reuse, we do this without the help of the type system and then use helper functions
    * that cast appropriately for the user facing interface.
    */
-  protected def selectUntyped(columns: TypedColumn[_]*): Dataset[_] = {
+  protected def selectUntyped(columns: TypedColumn[_, _]*): Dataset[_] = {
     val aliases = columns.zipWithIndex.map { case (c, i) => Alias(c.expr, s"_${i + 1}")() }
     val unresolvedPlan = Project(aliases, logicalPlan)
     val execution = new QueryExecution(sqlContext, unresolvedPlan)
@@ -385,7 +385,7 @@ class Dataset[T] private[sql](
    * Returns a new [[Dataset]] by computing the given [[Column]] expressions for each element.
    * @since 1.6.0
    */
-  def select[U1, U2](c1: TypedColumn[U1], c2: TypedColumn[U2]): Dataset[(U1, U2)] =
+  def select[U1, U2](c1: TypedColumn[T, U1], c2: TypedColumn[T, U2]): Dataset[(U1, U2)] =
     selectUntyped(c1, c2).asInstanceOf[Dataset[(U1, U2)]]
 
   /**
@@ -393,9 +393,9 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def select[U1, U2, U3](
-      c1: TypedColumn[U1],
-      c2: TypedColumn[U2],
-      c3: TypedColumn[U3]): Dataset[(U1, U2, U3)] =
+      c1: TypedColumn[T, U1],
+      c2: TypedColumn[T, U2],
+      c3: TypedColumn[T, U3]): Dataset[(U1, U2, U3)] =
     selectUntyped(c1, c2, c3).asInstanceOf[Dataset[(U1, U2, U3)]]
 
   /**
@@ -403,10 +403,10 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def select[U1, U2, U3, U4](
-      c1: TypedColumn[U1],
-      c2: TypedColumn[U2],
-      c3: TypedColumn[U3],
-      c4: TypedColumn[U4]): Dataset[(U1, U2, U3, U4)] =
+      c1: TypedColumn[T, U1],
+      c2: TypedColumn[T, U2],
+      c3: TypedColumn[T, U3],
+      c4: TypedColumn[T, U4]): Dataset[(U1, U2, U3, U4)] =
     selectUntyped(c1, c2, c3, c4).asInstanceOf[Dataset[(U1, U2, U3, U4)]]
 
   /**
@@ -414,11 +414,11 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def select[U1, U2, U3, U4, U5](
-      c1: TypedColumn[U1],
-      c2: TypedColumn[U2],
-      c3: TypedColumn[U3],
-      c4: TypedColumn[U4],
-      c5: TypedColumn[U5]): Dataset[(U1, U2, U3, U4, U5)] =
+      c1: TypedColumn[T, U1],
+      c2: TypedColumn[T, U2],
+      c3: TypedColumn[T, U3],
+      c4: TypedColumn[T, U4],
+      c5: TypedColumn[T, U5]): Dataset[(U1, U2, U3, U4, U5)] =
     selectUntyped(c1, c2, c3, c4, c5).asInstanceOf[Dataset[(U1, U2, U3, U4, U5)]]
 
   /* **************** *
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 850315e281dfe..db61499229284 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql
 
 import java.util.{Iterator => JIterator}
+
 import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
@@ -26,8 +27,10 @@ import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttrib
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor, Encoder}
 import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression, Alias, Attribute}
 import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
 import org.apache.spark.sql.execution.QueryExecution
 
+
 /**
  * :: Experimental ::
  * A [[Dataset]] has been logically grouped by a user specified grouping key.  Users should not
@@ -143,7 +146,7 @@ class GroupedDataset[K, T] private[sql](
    * that cast appropriately for the user facing interface.
    * TODO: does not handle aggrecations that return nonflat results,
    */
-  protected def aggUntyped(columns: TypedColumn[_]*): Dataset[_] = {
+  protected def aggUntyped(columns: TypedColumn[_, _]*): Dataset[_] = {
     val aliases = (groupingAttributes ++ columns.map(_.expr)).map {
       case u: UnresolvedAttribute => UnresolvedAlias(u)
       case expr: NamedExpression => expr
@@ -151,7 +154,15 @@ class GroupedDataset[K, T] private[sql](
     }
 
     val unresolvedPlan = Aggregate(groupingAttributes, aliases, logicalPlan)
-    val execution = new QueryExecution(sqlContext, unresolvedPlan)
+
+    // Fill in the input encoders for any aggregators in the plan.
+    val withEncoders = unresolvedPlan transformAllExpressions {
+      case ta: TypedAggregateExpression if ta.aEncoder.isEmpty =>
+        ta.copy(
+          aEncoder = Some(tEnc.asInstanceOf[ExpressionEncoder[Any]]),
+          children = dataAttributes)
+    }
+    val execution = new QueryExecution(sqlContext, withEncoders)
 
     val columnEncoders = columns.map(_.encoder.asInstanceOf[ExpressionEncoder[_]])
 
@@ -162,43 +173,47 @@ class GroupedDataset[K, T] private[sql](
       case (e, a) =>
         e.unbind(a :: Nil).resolve(execution.analyzed.output)
     }
-    new Dataset(sqlContext, execution, ExpressionEncoder.tuple(encoders))
+
+    new Dataset(
+      sqlContext,
+      execution,
+      ExpressionEncoder.tuple(encoders))
   }
 
   /**
    * Computes the given aggregation, returning a [[Dataset]] of tuples for each unique key
    * and the result of computing this aggregation over all elements in the group.
    */
-  def agg[A1](col1: TypedColumn[A1]): Dataset[(K, A1)] =
-    aggUntyped(col1).asInstanceOf[Dataset[(K, A1)]]
+  def agg[U1](col1: TypedColumn[T, U1]): Dataset[(K, U1)] =
+    aggUntyped(col1).asInstanceOf[Dataset[(K, U1)]]
 
   /**
    * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
    * and the result of computing these aggregations over all elements in the group.
    */
-  def agg[A1, A2](col1: TypedColumn[A1], col2: TypedColumn[A2]): Dataset[(K, A1, A2)] =
-    aggUntyped(col1, col2).asInstanceOf[Dataset[(K, A1, A2)]]
+  def agg[U1, U2](col1: TypedColumn[T, U1], col2: TypedColumn[T, U2]): Dataset[(K, U1, U2)] =
+    aggUntyped(col1, col2).asInstanceOf[Dataset[(K, U1, U2)]]
 
   /**
    * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
    * and the result of computing these aggregations over all elements in the group.
    */
-  def agg[A1, A2, A3](
-      col1: TypedColumn[A1],
-      col2: TypedColumn[A2],
-      col3: TypedColumn[A3]): Dataset[(K, A1, A2, A3)] =
-    aggUntyped(col1, col2, col3).asInstanceOf[Dataset[(K, A1, A2, A3)]]
+  def agg[U1, U2, U3](
+      col1: TypedColumn[T, U1],
+      col2: TypedColumn[T, U2],
+      col3: TypedColumn[T, U3]): Dataset[(K, U1, U2, U3)] =
+    aggUntyped(col1, col2, col3).asInstanceOf[Dataset[(K, U1, U2, U3)]]
 
   /**
    * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
    * and the result of computing these aggregations over all elements in the group.
    */
-  def agg[A1, A2, A3, A4](
-      col1: TypedColumn[A1],
-      col2: TypedColumn[A2],
-      col3: TypedColumn[A3],
-      col4: TypedColumn[A4]): Dataset[(K, A1, A2, A3, A4)] =
-    aggUntyped(col1, col2, col3, col4).asInstanceOf[Dataset[(K, A1, A2, A3, A4)]]
+  def agg[U1, U2, U3, U4](
+      col1: TypedColumn[T, U1],
+      col2: TypedColumn[T, U2],
+      col3: TypedColumn[T, U3],
+      col4: TypedColumn[T, U4]): Dataset[(K, U1, U2, U3, U4)] =
+    aggUntyped(col1, col2, col3, col4).asInstanceOf[Dataset[(K, U1, U2, U3, U4)]]
 
   /**
    * Returns a [[Dataset]] that contains a tuple with each key and the number of items present
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index 5598731af5fcc..1cf1e30f967cd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -21,7 +21,6 @@ import java.beans.{BeanInfo, Introspector}
 import java.util.Properties
 import java.util.concurrent.atomic.AtomicReference
 
-
 import scala.collection.JavaConverters._
 import scala.collection.immutable
 import scala.reflect.runtime.universe.TypeTag
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
new file mode 100644
index 0000000000000..24d8122b6222b
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.aggregate
+
+import scala.language.existentials
+
+import org.apache.spark.Logging
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
+import org.apache.spark.sql.catalyst.expressions.aggregate.ImperativeAggregate
+import org.apache.spark.sql.expressions.Aggregator
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types.{StructType, DataType}
+
+object TypedAggregateExpression {
+  def apply[A, B : Encoder, C : Encoder](
+      aggregator: Aggregator[A, B, C]): TypedAggregateExpression = {
+    new TypedAggregateExpression(
+      aggregator.asInstanceOf[Aggregator[Any, Any, Any]],
+      None,
+      encoderFor[B].asInstanceOf[ExpressionEncoder[Any]],
+      encoderFor[C].asInstanceOf[ExpressionEncoder[Any]],
+      Nil,
+      0,
+      0)
+  }
+}
+
+/**
+ * This class is a rough sketch of how to hook `Aggregator` into the Aggregation system.  It has
+ * the following limitations:
+ *  - It assumes the aggregator reduces and returns a single column of type `long`.
+ *  - It might only work when there is a single aggregator in the first column.
+ *  - It assumes the aggregator has a zero, `0`.
+ */
+case class TypedAggregateExpression(
+    aggregator: Aggregator[Any, Any, Any],
+    aEncoder: Option[ExpressionEncoder[Any]],
+    bEncoder: ExpressionEncoder[Any],
+    cEncoder: ExpressionEncoder[Any],
+    children: Seq[Expression],
+    mutableAggBufferOffset: Int,
+    inputAggBufferOffset: Int)
+  extends ImperativeAggregate with Logging {
+
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
+
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
+
+  override def nullable: Boolean = true
+
+  // TODO: this assumes flat results...
+  override def dataType: DataType = cEncoder.schema.head.dataType
+
+  override def deterministic: Boolean = true
+
+  override lazy val resolved: Boolean = aEncoder.isDefined
+
+  override lazy val inputTypes: Seq[DataType] =
+    aEncoder.map(_.schema.map(_.dataType)).getOrElse(Nil)
+
+  override val aggBufferSchema: StructType = bEncoder.schema
+
+  override val aggBufferAttributes: Seq[AttributeReference] = aggBufferSchema.toAttributes
+
+  // Note: although this simply copies aggBufferAttributes, this common code can not be placed
+  // in the superclass because that will lead to initialization ordering issues.
+  override val inputAggBufferAttributes: Seq[AttributeReference] =
+    aggBufferAttributes.map(_.newInstance())
+
+  lazy val inputAttributes = aEncoder.get.schema.toAttributes
+  lazy val inputMapping = AttributeMap(inputAttributes.zip(children))
+  lazy val boundA =
+    aEncoder.get.copy(constructExpression = aEncoder.get.constructExpression transform {
+      case a: AttributeReference => inputMapping(a)
+    })
+
+  // TODO: this probably only works when we are in the first column.
+  val bAttributes = bEncoder.schema.toAttributes
+  lazy val boundB = bEncoder.resolve(bAttributes).bind(bAttributes)
+
+  override def initialize(buffer: MutableRow): Unit = {
+    // TODO: We need to either force Aggregator to have a zero or we need to eliminate the need for
+    // this in execution.
+    buffer.setInt(mutableAggBufferOffset, aggregator.zero.asInstanceOf[Int])
+  }
+
+  override def update(buffer: MutableRow, input: InternalRow): Unit = {
+    val inputA = boundA.fromRow(input)
+    val currentB = boundB.fromRow(buffer)
+    val merged = aggregator.reduce(currentB, inputA)
+    val returned = boundB.toRow(merged)
+    buffer.setInt(mutableAggBufferOffset, returned.getInt(0))
+  }
+
+  override def merge(buffer1: MutableRow, buffer2: InternalRow): Unit = {
+    buffer1.setLong(
+      mutableAggBufferOffset,
+      buffer1.getLong(mutableAggBufferOffset) + buffer2.getLong(inputAggBufferOffset))
+  }
+
+  override def eval(buffer: InternalRow): Any = {
+    buffer.getInt(mutableAggBufferOffset)
+  }
+
+  override def toString: String = {
+    s"""${aggregator.getClass.getSimpleName}(${children.mkString(",")})"""
+  }
+
+  override def nodeName: String = aggregator.getClass.getSimpleName
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
new file mode 100644
index 0000000000000..0b3192a6da9d8
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.expressions
+
+import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, AggregateExpression2}
+import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
+import org.apache.spark.sql.{Dataset, DataFrame, TypedColumn}
+
+/**
+ * A base class for user-defined aggregations, which can be used in [[DataFrame]] and [[Dataset]]
+ * operations to take all of the elements of a group and reduce them to a single value.
+ *
+ * For example, the following aggregator extracts an `int` from a specific class and adds them up:
+ * {{{
+ *   case class Data(i: Int)
+ *
+ *   val customSummer =  new Aggregator[Data, Int, Int] {
+ *     def zero = 0
+ *     def reduce(b: Int, a: Data) = b + a.i
+ *     def present(r: Int) = r
+ *   }.toColumn()
+ *
+ *   val ds: Dataset[Data]
+ *   val aggregated = ds.select(customSummer)
+ * }}}
+ *
+ * Based loosely on Aggregator from Algebird: https://github.com/twitter/algebird
+ *
+ * @tparam A The input type for the aggregation.
+ * @tparam B The type of the intermediate value of the reduction.
+ * @tparam C The type of the final result.
+ */
+abstract class Aggregator[-A, B, C] {
+
+  /** A zero value for this aggregation. Should satisfy the property that any b + zero = b */
+  def zero: B
+
+  /**
+   * Combine two values to produce a new value.  For performance, the function may modify `b` and
+   * return it instead of constructing new object for b.
+   */
+  def reduce(b: B, a: A): B
+
+  /**
+   * Transform the output of the reduction.
+   */
+  def present(reduction: B): C
+
+  /**
+   * Returns this `Aggregator` as a [[TypedColumn]] that can be used in [[Dataset]] or [[DataFrame]]
+   * operations.
+   */
+  def toColumn(
+      implicit bEncoder: Encoder[B],
+      cEncoder: Encoder[C]): TypedColumn[A, C] = {
+    val expr =
+      new AggregateExpression2(
+        TypedAggregateExpression(this),
+        Complete,
+        false)
+
+    new TypedColumn[A, C](expr, encoderFor[C])
+  }
+}
+
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 3f0b24b68b816..6d56542ee0875 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql
 
+
+
 import scala.language.implicitConversions
 import scala.reflect.runtime.universe.{TypeTag, typeTag}
 import scala.util.Try
@@ -24,11 +26,32 @@ import scala.util.Try
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.sql.catalyst.{SqlParser, ScalaReflection}
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunction, Star}
+import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, Encoder}
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.BroadcastHint
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
 
+/**
+ * Ensures that java functions signatures for methods that now return a [[TypedColumn]] still have
+ * legacy equivalents in bytecode.  This compatibility is done by forcing the compiler to generate
+ * "bridge" methods due to the use of covariant return types.
+ *
+ * {{{
+ * In LegacyFunctions:
+ * public abstract org.apache.spark.sql.Column avg(java.lang.String);
+ *
+ * In functions:
+ * public static org.apache.spark.sql.TypedColumn<java.lang.Object, java.lang.Object> avg(...);
+ * }}}
+ *
+ * This allows us to use the same functions both in typed [[Dataset]] operations and untyped
+ * [[DataFrame]] operations when the return type for a given function is statically known.
+ */
+private[sql] abstract class LegacyFunctions {
+  def count(columnName: String): Column
+}
+
 /**
  * :: Experimental ::
  * Functions available for [[DataFrame]].
@@ -48,11 +71,14 @@ import org.apache.spark.util.Utils
  */
 @Experimental
 // scalastyle:off
-object functions {
+object functions extends LegacyFunctions {
 // scalastyle:on
 
   private def withExpr(expr: Expression): Column = Column(expr)
 
+  private implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder[Long](flat = true)
+
+
   /**
    * Returns a [[Column]] based on the given column name.
    *
@@ -234,7 +260,7 @@ object functions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def count(columnName: String): Column = count(Column(columnName))
+  def count(columnName: String): TypedColumn[Any, Long] = count(Column(columnName)).as[Long]
 
   /**
    * Aggregate function: returns the number of distinct items in a group.
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index 312cf33e4c2d4..2da63d1b96706 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -258,8 +258,8 @@ public void testSelect() {
     Dataset<Integer> ds = context.createDataset(data, e.INT());
 
     Dataset<Tuple2<Integer, String>> selected = ds.select(
-      expr("value + 1").as(e.INT()),
-      col("value").cast("string").as(e.STRING()));
+      expr("value + 1"),
+      col("value").cast("string")).as(e.tuple(e.INT(), e.STRING()));
 
     Assert.assertEquals(
       Arrays.asList(tuple2(3, "2"), tuple2(7, "6")),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
new file mode 100644
index 0000000000000..340470c096b87
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+import org.apache.spark.sql.catalyst.encoders.Encoder
+import org.apache.spark.sql.functions._
+
+import scala.language.postfixOps
+
+import org.apache.spark.sql.test.SharedSQLContext
+
+import org.apache.spark.sql.expressions.Aggregator
+
+/** An `Aggregator` that adds up any numeric type returned by the given function. */
+class SumOf[I, N : Numeric](f: I => N) extends Aggregator[I, N, N] with Serializable {
+  val numeric = implicitly[Numeric[N]]
+
+  override def zero: N = numeric.zero
+
+  override def reduce(b: N, a: I): N = numeric.plus(b, f(a))
+
+  override def present(reduction: N): N = reduction
+}
+
+class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
+
+  import testImplicits._
+
+  def sum[I, N : Numeric : Encoder](f: I => N): TypedColumn[I, N] =
+    new SumOf(f).toColumn
+
+  test("typed aggregation: TypedAggregator") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+
+    checkAnswer(
+      ds.groupBy(_._1).agg(sum(_._2)),
+      ("a", 30), ("b", 3), ("c", 1))
+  }
+
+  test("typed aggregation: TypedAggregator, expr, expr") {
+    val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
+
+    checkAnswer(
+      ds.groupBy(_._1).agg(
+        sum(_._2),
+        expr("sum(_2)").as[Int],
+        count("*")),
+      ("a", 30, 30, 2L), ("b", 3, 3, 2L), ("c", 1, 1, 1L))
+  }
+}

From 675c7e723cadff588405c23826a00686587728b8 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Mon, 9 Nov 2015 16:22:15 -0800
Subject: [PATCH 142/518] [SPARK-11564][SQL] Fix documentation for
 DataFrame.take/collect

Author: Reynold Xin <rxin@databricks.com>

Closes #9557 from rxin/SPARK-11564-1.
---
 .../src/main/scala/org/apache/spark/sql/DataFrame.scala   | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 8ab958adadcca..d25807cf8d09c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -1479,8 +1479,8 @@ class DataFrame private[sql](
   /**
    * Returns the first `n` rows in the [[DataFrame]].
    *
-   * Running take requires moving data into the application's driver process, and doing so on a
-   * very large dataset can crash the driver process with OutOfMemoryError.
+   * Running take requires moving data into the application's driver process, and doing so with
+   * a very large `n` can crash the driver process with OutOfMemoryError.
    *
    * @group action
    * @since 1.3.0
@@ -1501,8 +1501,8 @@ class DataFrame private[sql](
   /**
    * Returns an array that contains all of [[Row]]s in this [[DataFrame]].
    *
-   * Running take requires moving data into the application's driver process, and doing so with
-   * a very large `n` can crash the driver process with OutOfMemoryError.
+   * Running collect requires moving all the data into the application's driver process, and
+   * doing so on a very large dataset can crash the driver process with OutOfMemoryError.
    *
    * For Java API, use [[collectAsList]].
    *

From 7dc9d8dba6c4bc655896b137062d896dec4ef64a Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Mon, 9 Nov 2015 16:25:29 -0800
Subject: [PATCH 143/518] [SPARK-11610][MLLIB][PYTHON][DOCS] Make the docs of
 LDAModel.describeTopics in Python more specific

cc jkbradley

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>

Closes #9577 from yu-iskw/SPARK-11610.
---
 python/pyspark/mllib/clustering.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/python/pyspark/mllib/clustering.py b/python/pyspark/mllib/clustering.py
index 12081f8c69075..1fa061dc2da99 100644
--- a/python/pyspark/mllib/clustering.py
+++ b/python/pyspark/mllib/clustering.py
@@ -734,6 +734,12 @@ def describeTopics(self, maxTermsPerTopic=None):
         """Return the topics described by weighted terms.
 
         WARNING: If vocabSize and k are large, this can return a large object!
+
+        :param maxTermsPerTopic: Maximum number of terms to collect for each topic.
+            (default: vocabulary size)
+        :return: Array over topics. Each topic is represented as a pair of matching arrays:
+            (term indices, term weights in topic).
+            Each topic's terms are sorted in order of decreasing weight.
         """
         if maxTermsPerTopic is None:
             topics = self.call("describeTopics")

From 61f9c8711c79f35d67b0456155866da316b131d9 Mon Sep 17 00:00:00 2001
From: Yuhao Yang <hhbyyh@gmail.com>
Date: Mon, 9 Nov 2015 16:55:23 -0800
Subject: [PATCH 144/518] [SPARK-11069][ML] Add RegexTokenizer option to
 convert to lowercase

jira: https://issues.apache.org/jira/browse/SPARK-11069
quotes from jira:
Tokenizer converts strings to lowercase automatically, but RegexTokenizer does not. It would be nice to add an option to RegexTokenizer to convert to lowercase. Proposal:
call the Boolean Param "toLowercase"
set default to false (so behavior does not change)

Actually sklearn converts to lowercase before tokenizing too

Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #9092 from hhbyyh/tokenLower.
---
 .../apache/spark/ml/feature/Tokenizer.scala   | 19 ++++++++++++++--
 .../spark/ml/feature/JavaTokenizerSuite.java  |  1 +
 .../spark/ml/feature/TokenizerSuite.scala     | 22 ++++++++++++++-----
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
index 248288ca73e99..1b82b40caac18 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
@@ -100,10 +100,25 @@ class RegexTokenizer(override val uid: String)
   /** @group getParam */
   def getPattern: String = $(pattern)
 
-  setDefault(minTokenLength -> 1, gaps -> true, pattern -> "\\s+")
+  /**
+   * Indicates whether to convert all characters to lowercase before tokenizing.
+   * Default: true
+   * @group param
+   */
+  final val toLowercase: BooleanParam = new BooleanParam(this, "toLowercase",
+    "whether to convert all characters to lowercase before tokenizing.")
+
+  /** @group setParam */
+  def setToLowercase(value: Boolean): this.type = set(toLowercase, value)
+
+  /** @group getParam */
+  def getToLowercase: Boolean = $(toLowercase)
+
+  setDefault(minTokenLength -> 1, gaps -> true, pattern -> "\\s+", toLowercase -> true)
 
-  override protected def createTransformFunc: String => Seq[String] = { str =>
+  override protected def createTransformFunc: String => Seq[String] = { originStr =>
     val re = $(pattern).r
+    val str = if ($(toLowercase)) originStr.toLowerCase() else originStr
     val tokens = if ($(gaps)) re.split(str).toSeq else re.findAllIn(str).toSeq
     val minLength = $(minTokenLength)
     tokens.filter(_.length >= minLength)
diff --git a/mllib/src/test/java/org/apache/spark/ml/feature/JavaTokenizerSuite.java b/mllib/src/test/java/org/apache/spark/ml/feature/JavaTokenizerSuite.java
index 02309ce63219a..c407d98f1b795 100644
--- a/mllib/src/test/java/org/apache/spark/ml/feature/JavaTokenizerSuite.java
+++ b/mllib/src/test/java/org/apache/spark/ml/feature/JavaTokenizerSuite.java
@@ -53,6 +53,7 @@ public void regexTokenizer() {
       .setOutputCol("tokens")
       .setPattern("\\s")
       .setGaps(true)
+      .setToLowercase(false)
       .setMinTokenLength(3);
 
 
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/TokenizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/TokenizerSuite.scala
index e5fd21c3f6fca..a02992a2407b3 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/TokenizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/TokenizerSuite.scala
@@ -48,13 +48,13 @@ class RegexTokenizerSuite extends SparkFunSuite with MLlibTestSparkContext {
       .setInputCol("rawText")
       .setOutputCol("tokens")
     val dataset0 = sqlContext.createDataFrame(Seq(
-      TokenizerTestData("Test for tokenization.", Array("Test", "for", "tokenization", ".")),
-      TokenizerTestData("Te,st. punct", Array("Te", ",", "st", ".", "punct"))
+      TokenizerTestData("Test for tokenization.", Array("test", "for", "tokenization", ".")),
+      TokenizerTestData("Te,st. punct", Array("te", ",", "st", ".", "punct"))
     ))
     testRegexTokenizer(tokenizer0, dataset0)
 
     val dataset1 = sqlContext.createDataFrame(Seq(
-      TokenizerTestData("Test for tokenization.", Array("Test", "for", "tokenization")),
+      TokenizerTestData("Test for tokenization.", Array("test", "for", "tokenization")),
       TokenizerTestData("Te,st. punct", Array("punct"))
     ))
     tokenizer0.setMinTokenLength(3)
@@ -64,11 +64,23 @@ class RegexTokenizerSuite extends SparkFunSuite with MLlibTestSparkContext {
       .setInputCol("rawText")
       .setOutputCol("tokens")
     val dataset2 = sqlContext.createDataFrame(Seq(
-      TokenizerTestData("Test for tokenization.", Array("Test", "for", "tokenization.")),
-      TokenizerTestData("Te,st.  punct", Array("Te,st.", "punct"))
+      TokenizerTestData("Test for tokenization.", Array("test", "for", "tokenization.")),
+      TokenizerTestData("Te,st.  punct", Array("te,st.", "punct"))
     ))
     testRegexTokenizer(tokenizer2, dataset2)
   }
+
+  test("RegexTokenizer with toLowercase false") {
+    val tokenizer = new RegexTokenizer()
+      .setInputCol("rawText")
+      .setOutputCol("tokens")
+      .setToLowercase(false)
+    val dataset = sqlContext.createDataFrame(Seq(
+      TokenizerTestData("JAVA SCALA", Array("JAVA", "SCALA")),
+      TokenizerTestData("java scala", Array("java", "scala"))
+    ))
+    testRegexTokenizer(tokenizer, dataset)
+  }
 }
 
 object RegexTokenizerSuite extends SparkFunSuite {

From 26062d22607e1f9854bc2588ba22a4e0f8bba48c Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Mon, 9 Nov 2015 17:18:49 -0800
Subject: [PATCH 145/518] [SPARK-11198][STREAMING][KINESIS] Support
 de-aggregation of records during recovery

While the KCL handles de-aggregation during the regular operation, during recovery we use the lower level api, and therefore need to de-aggregate the records.

tdas Testing is an issue, we need protobuf magic to do the aggregated records. Maybe we could depend on KPL for tests?

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9403 from brkyvz/kinesis-deaggregation.
---
 extras/kinesis-asl/pom.xml                    |  6 ++
 .../kinesis/KinesisBackedBlockRDD.scala       |  6 +-
 .../streaming/kinesis/KinesisReceiver.scala   |  1 -
 .../kinesis/KinesisRecordProcessor.scala      |  2 +-
 .../kinesis/KinesisBackedBlockRDDSuite.scala  | 12 +++-
 .../kinesis/KinesisStreamSuite.scala          | 17 +++---
 .../streaming/kinesis/KinesisTestUtils.scala  | 55 +++++++++++++++----
 pom.xml                                       |  2 +
 8 files changed, 76 insertions(+), 25 deletions(-)
 rename extras/kinesis-asl/src/{main => test}/scala/org/apache/spark/streaming/kinesis/KinesisTestUtils.scala (80%)

diff --git a/extras/kinesis-asl/pom.xml b/extras/kinesis-asl/pom.xml
index ef72d97eae69d..519a920279c97 100644
--- a/extras/kinesis-asl/pom.xml
+++ b/extras/kinesis-asl/pom.xml
@@ -64,6 +64,12 @@
       <artifactId>aws-java-sdk</artifactId>
       <version>${aws.java.sdk.version}</version>
     </dependency>
+    <dependency>
+      <groupId>com.amazonaws</groupId>
+      <artifactId>amazon-kinesis-producer</artifactId>
+      <version>${aws.kinesis.producer.version}</version>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>
diff --git a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDD.scala b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDD.scala
index 000897a4e7290..691c1790b207f 100644
--- a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDD.scala
+++ b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDD.scala
@@ -23,6 +23,7 @@ import scala.util.control.NonFatal
 
 import com.amazonaws.auth.{AWSCredentials, DefaultAWSCredentialsProviderChain}
 import com.amazonaws.services.kinesis.AmazonKinesisClient
+import com.amazonaws.services.kinesis.clientlibrary.types.UserRecord
 import com.amazonaws.services.kinesis.model._
 
 import org.apache.spark._
@@ -210,7 +211,10 @@ class KinesisSequenceRangeIterator(
       s"getting records using shard iterator") {
         client.getRecords(getRecordsRequest)
       }
-    (getRecordsResult.getRecords.iterator().asScala, getRecordsResult.getNextShardIterator)
+    // De-aggregate records, if KPL was used in producing the records. The KCL automatically
+    // handles de-aggregation during regular operation. This code path is used during recovery
+    val recordIterator = UserRecord.deaggregate(getRecordsResult.getRecords)
+    (recordIterator.iterator().asScala, getRecordsResult.getNextShardIterator)
   }
 
   /**
diff --git a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisReceiver.scala b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisReceiver.scala
index 50993f157cd95..97dbb918573a3 100644
--- a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisReceiver.scala
+++ b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisReceiver.scala
@@ -216,7 +216,6 @@ private[kinesis] class KinesisReceiver[T](
       val metadata = SequenceNumberRange(streamName, shardId,
         records.get(0).getSequenceNumber(), records.get(records.size() - 1).getSequenceNumber())
       blockGenerator.addMultipleDataWithCallback(dataIterator, metadata)
-
     }
   }
 
diff --git a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisRecordProcessor.scala b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisRecordProcessor.scala
index e381ffa0cbef4..b5b76cb92d866 100644
--- a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisRecordProcessor.scala
+++ b/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisRecordProcessor.scala
@@ -80,7 +80,7 @@ private[kinesis] class KinesisRecordProcessor[T](receiver: KinesisReceiver[T], w
            *     more than once.
            */
           logError(s"Exception:  WorkerId $workerId encountered and exception while storing " +
-              " or checkpointing a batch for workerId $workerId and shardId $shardId.", e)
+              s" or checkpointing a batch for workerId $workerId and shardId $shardId.", e)
 
           /* Rethrow the exception to the Kinesis Worker that is managing this RecordProcessor. */
           throw e
diff --git a/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDDSuite.scala b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDDSuite.scala
index 9f9e146a08d46..52c61dfb1c023 100644
--- a/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDDSuite.scala
+++ b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisBackedBlockRDDSuite.scala
@@ -22,7 +22,8 @@ import org.scalatest.BeforeAndAfterAll
 import org.apache.spark.storage.{BlockId, BlockManager, StorageLevel, StreamBlockId}
 import org.apache.spark.{SparkConf, SparkContext, SparkException}
 
-class KinesisBackedBlockRDDSuite extends KinesisFunSuite with BeforeAndAfterAll {
+abstract class KinesisBackedBlockRDDTests(aggregateTestData: Boolean)
+  extends KinesisFunSuite with BeforeAndAfterAll {
 
   private val testData = 1 to 8
 
@@ -37,13 +38,12 @@ class KinesisBackedBlockRDDSuite extends KinesisFunSuite with BeforeAndAfterAll
   private var sc: SparkContext = null
   private var blockManager: BlockManager = null
 
-
   override def beforeAll(): Unit = {
     runIfTestsEnabled("Prepare KinesisTestUtils") {
       testUtils = new KinesisTestUtils()
       testUtils.createStream()
 
-      shardIdToDataAndSeqNumbers = testUtils.pushData(testData)
+      shardIdToDataAndSeqNumbers = testUtils.pushData(testData, aggregate = aggregateTestData)
       require(shardIdToDataAndSeqNumbers.size > 1, "Need data to be sent to multiple shards")
 
       shardIds = shardIdToDataAndSeqNumbers.keySet.toSeq
@@ -247,3 +247,9 @@ class KinesisBackedBlockRDDSuite extends KinesisFunSuite with BeforeAndAfterAll
     Array.tabulate(num) { i => new StreamBlockId(0, i) }
   }
 }
+
+class WithAggregationKinesisBackedBlockRDDSuite
+  extends KinesisBackedBlockRDDTests(aggregateTestData = true)
+
+class WithoutAggregationKinesisBackedBlockRDDSuite
+  extends KinesisBackedBlockRDDTests(aggregateTestData = false)
diff --git a/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisStreamSuite.scala b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisStreamSuite.scala
index ba84e557dfcc2..dee30444d8cc6 100644
--- a/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisStreamSuite.scala
+++ b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisStreamSuite.scala
@@ -39,7 +39,7 @@ import org.apache.spark.streaming.scheduler.ReceivedBlockInfo
 import org.apache.spark.util.Utils
 import org.apache.spark.{SparkConf, SparkContext}
 
-class KinesisStreamSuite extends KinesisFunSuite
+abstract class KinesisStreamTests(aggregateTestData: Boolean) extends KinesisFunSuite
   with Eventually with BeforeAndAfter with BeforeAndAfterAll {
 
   // This is the name that KCL will use to save metadata to DynamoDB
@@ -182,13 +182,13 @@ class KinesisStreamSuite extends KinesisFunSuite
     val collected = new mutable.HashSet[Int] with mutable.SynchronizedSet[Int]
     stream.map { bytes => new String(bytes).toInt }.foreachRDD { rdd =>
       collected ++= rdd.collect()
-      logInfo("Collected = " + rdd.collect().toSeq.mkString(", "))
+      logInfo("Collected = " + collected.mkString(", "))
     }
     ssc.start()
 
     val testData = 1 to 10
     eventually(timeout(120 seconds), interval(10 second)) {
-      testUtils.pushData(testData)
+      testUtils.pushData(testData, aggregateTestData)
       assert(collected === testData.toSet, "\nData received does not match data sent")
     }
     ssc.stop(stopSparkContext = false)
@@ -207,13 +207,13 @@ class KinesisStreamSuite extends KinesisFunSuite
     val collected = new mutable.HashSet[Int] with mutable.SynchronizedSet[Int]
     stream.foreachRDD { rdd =>
       collected ++= rdd.collect()
-      logInfo("Collected = " + rdd.collect().toSeq.mkString(", "))
+      logInfo("Collected = " + collected.mkString(", "))
     }
     ssc.start()
 
     val testData = 1 to 10
     eventually(timeout(120 seconds), interval(10 second)) {
-      testUtils.pushData(testData)
+      testUtils.pushData(testData, aggregateTestData)
       val modData = testData.map(_ + 5)
       assert(collected === modData.toSet, "\nData received does not match data sent")
     }
@@ -254,7 +254,7 @@ class KinesisStreamSuite extends KinesisFunSuite
     // If this times out because numBatchesWithData is empty, then its likely that foreachRDD
     // function failed with exceptions, and nothing got added to `collectedData`
     eventually(timeout(2 minutes), interval(1 seconds)) {
-      testUtils.pushData(1 to 5)
+      testUtils.pushData(1 to 5, aggregateTestData)
       assert(isCheckpointPresent && numBatchesWithData > 10)
     }
     ssc.stop(stopSparkContext = true)  // stop the SparkContext so that the blocks are not reused
@@ -285,5 +285,8 @@ class KinesisStreamSuite extends KinesisFunSuite
     }
     ssc.stop()
   }
-
 }
+
+class WithAggregationKinesisStreamSuite extends KinesisStreamTests(aggregateTestData = true)
+
+class WithoutAggregationKinesisStreamSuite extends KinesisStreamTests(aggregateTestData = false)
diff --git a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisTestUtils.scala b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisTestUtils.scala
similarity index 80%
rename from extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisTestUtils.scala
rename to extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisTestUtils.scala
index 634bf94521079..7487aa1c12639 100644
--- a/extras/kinesis-asl/src/main/scala/org/apache/spark/streaming/kinesis/KinesisTestUtils.scala
+++ b/extras/kinesis-asl/src/test/scala/org/apache/spark/streaming/kinesis/KinesisTestUtils.scala
@@ -31,6 +31,8 @@ import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient
 import com.amazonaws.services.dynamodbv2.document.DynamoDB
 import com.amazonaws.services.kinesis.AmazonKinesisClient
 import com.amazonaws.services.kinesis.model._
+import com.amazonaws.services.kinesis.producer.{KinesisProducer, KinesisProducerConfiguration, UserRecordResult}
+import com.google.common.util.concurrent.{FutureCallback, Futures}
 
 import org.apache.spark.Logging
 
@@ -64,6 +66,16 @@ private[kinesis] class KinesisTestUtils extends Logging {
     new DynamoDB(dynamoDBClient)
   }
 
+  private lazy val kinesisProducer: KinesisProducer = {
+    val conf = new KinesisProducerConfiguration()
+      .setRecordMaxBufferedTime(1000)
+      .setMaxConnections(1)
+      .setRegion(regionName)
+      .setMetricsLevel("none")
+
+    new KinesisProducer(conf)
+  }
+
   def streamName: String = {
     require(streamCreated, "Stream not yet created, call createStream() to create one")
     _streamName
@@ -90,22 +102,41 @@ private[kinesis] class KinesisTestUtils extends Logging {
    * Push data to Kinesis stream and return a map of
    * shardId -> seq of (data, seq number) pushed to corresponding shard
    */
-  def pushData(testData: Seq[Int]): Map[String, Seq[(Int, String)]] = {
+  def pushData(testData: Seq[Int], aggregate: Boolean): Map[String, Seq[(Int, String)]] = {
     require(streamCreated, "Stream not yet created, call createStream() to create one")
     val shardIdToSeqNumbers = new mutable.HashMap[String, ArrayBuffer[(Int, String)]]()
 
     testData.foreach { num =>
       val str = num.toString
-      val putRecordRequest = new PutRecordRequest().withStreamName(streamName)
-        .withData(ByteBuffer.wrap(str.getBytes()))
-        .withPartitionKey(str)
-
-      val putRecordResult = kinesisClient.putRecord(putRecordRequest)
-      val shardId = putRecordResult.getShardId
-      val seqNumber = putRecordResult.getSequenceNumber()
-      val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
-        new ArrayBuffer[(Int, String)]())
-      sentSeqNumbers += ((num, seqNumber))
+      val data = ByteBuffer.wrap(str.getBytes())
+      if (aggregate) {
+        val future = kinesisProducer.addUserRecord(streamName, str, data)
+        val kinesisCallBack = new FutureCallback[UserRecordResult]() {
+          override def onFailure(t: Throwable): Unit = {} // do nothing
+
+          override def onSuccess(result: UserRecordResult): Unit = {
+            val shardId = result.getShardId
+            val seqNumber = result.getSequenceNumber()
+            val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
+              new ArrayBuffer[(Int, String)]())
+            sentSeqNumbers += ((num, seqNumber))
+          }
+        }
+
+        Futures.addCallback(future, kinesisCallBack)
+        kinesisProducer.flushSync() // make sure we send all data before returning the map
+      } else {
+        val putRecordRequest = new PutRecordRequest().withStreamName(streamName)
+          .withData(data)
+          .withPartitionKey(str)
+
+        val putRecordResult = kinesisClient.putRecord(putRecordRequest)
+        val shardId = putRecordResult.getShardId
+        val seqNumber = putRecordResult.getSequenceNumber()
+        val sentSeqNumbers = shardIdToSeqNumbers.getOrElseUpdate(shardId,
+          new ArrayBuffer[(Int, String)]())
+        sentSeqNumbers += ((num, seqNumber))
+      }
     }
 
     logInfo(s"Pushed $testData:\n\t ${shardIdToSeqNumbers.mkString("\n\t")}")
@@ -116,7 +147,7 @@ private[kinesis] class KinesisTestUtils extends Logging {
    * Expose a Python friendly API.
    */
   def pushData(testData: java.util.List[Int]): Unit = {
-    pushData(testData.asScala)
+    pushData(testData.asScala, aggregate = false)
   }
 
   def deleteStream(): Unit = {
diff --git a/pom.xml b/pom.xml
index 4ed1c0c82dee6..fd8c773513881 100644
--- a/pom.xml
+++ b/pom.xml
@@ -154,6 +154,8 @@
     <jets3t.version>0.7.1</jets3t.version>
     <aws.java.sdk.version>1.9.40</aws.java.sdk.version>
     <aws.kinesis.client.version>1.4.0</aws.kinesis.client.version>
+    <!-- the producer is used in tests -->
+    <aws.kinesis.producer.version>0.10.1</aws.kinesis.producer.version>
     <!--  org.apache.httpcomponents/httpclient-->
     <commons.httpclient.version>4.3.2</commons.httpclient.version>
     <!--  commons-httpclient/commons-httpclient-->

From 0ce6f9b2d203ce67aeb4d3aedf19bbd997fe01b9 Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Mon, 9 Nov 2015 17:35:12 -0800
Subject: [PATCH 146/518] [SPARK-11141][STREAMING] Batch
 ReceivedBlockTrackerLogEvents for WAL writes

When using S3 as a directory for WALs, the writes take too long. The driver gets very easily bottlenecked when multiple receivers send AddBlock events to the ReceiverTracker. This PR adds batching of events in the ReceivedBlockTracker so that receivers don't get blocked by the driver for too long.

cc zsxwing tdas

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9143 from brkyvz/batch-wal-writes.
---
 .../scheduler/ReceivedBlockTracker.scala      |  62 ++-
 .../streaming/scheduler/ReceiverTracker.scala |  25 +-
 .../streaming/util/BatchedWriteAheadLog.scala | 223 ++++++++
 .../streaming/util/WriteAheadLogUtils.scala   |  21 +-
 .../streaming/util/WriteAheadLogSuite.scala   | 506 ++++++++++++------
 .../util/WriteAheadLogUtilsSuite.scala        | 122 +++++
 6 files changed, 767 insertions(+), 192 deletions(-)
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/util/BatchedWriteAheadLog.scala
 create mode 100644 streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogUtilsSuite.scala

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceivedBlockTracker.scala b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceivedBlockTracker.scala
index f2711d1355e60..500dc70c98506 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceivedBlockTracker.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceivedBlockTracker.scala
@@ -22,12 +22,13 @@ import java.nio.ByteBuffer
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.language.implicitConversions
+import scala.util.control.NonFatal
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.streaming.Time
-import org.apache.spark.streaming.util.{WriteAheadLog, WriteAheadLogUtils}
+import org.apache.spark.streaming.util.{BatchedWriteAheadLog, WriteAheadLog, WriteAheadLogUtils}
 import org.apache.spark.util.{Clock, Utils}
 import org.apache.spark.{Logging, SparkConf}
 
@@ -41,7 +42,6 @@ private[streaming] case class BatchAllocationEvent(time: Time, allocatedBlocks:
 private[streaming] case class BatchCleanupEvent(times: Seq[Time])
   extends ReceivedBlockTrackerLogEvent
 
-
 /** Class representing the blocks of all the streams allocated to a batch */
 private[streaming]
 case class AllocatedBlocks(streamIdToAllocatedBlocks: Map[Int, Seq[ReceivedBlockInfo]]) {
@@ -82,15 +82,22 @@ private[streaming] class ReceivedBlockTracker(
   }
 
   /** Add received block. This event will get written to the write ahead log (if enabled). */
-  def addBlock(receivedBlockInfo: ReceivedBlockInfo): Boolean = synchronized {
+  def addBlock(receivedBlockInfo: ReceivedBlockInfo): Boolean = {
     try {
-      writeToLog(BlockAdditionEvent(receivedBlockInfo))
-      getReceivedBlockQueue(receivedBlockInfo.streamId) += receivedBlockInfo
-      logDebug(s"Stream ${receivedBlockInfo.streamId} received " +
-        s"block ${receivedBlockInfo.blockStoreResult.blockId}")
-      true
+      val writeResult = writeToLog(BlockAdditionEvent(receivedBlockInfo))
+      if (writeResult) {
+        synchronized {
+          getReceivedBlockQueue(receivedBlockInfo.streamId) += receivedBlockInfo
+        }
+        logDebug(s"Stream ${receivedBlockInfo.streamId} received " +
+          s"block ${receivedBlockInfo.blockStoreResult.blockId}")
+      } else {
+        logDebug(s"Failed to acknowledge stream ${receivedBlockInfo.streamId} receiving " +
+          s"block ${receivedBlockInfo.blockStoreResult.blockId} in the Write Ahead Log.")
+      }
+      writeResult
     } catch {
-      case e: Exception =>
+      case NonFatal(e) =>
         logError(s"Error adding block $receivedBlockInfo", e)
         false
     }
@@ -106,10 +113,12 @@ private[streaming] class ReceivedBlockTracker(
           (streamId, getReceivedBlockQueue(streamId).dequeueAll(x => true))
       }.toMap
       val allocatedBlocks = AllocatedBlocks(streamIdToBlocks)
-      writeToLog(BatchAllocationEvent(batchTime, allocatedBlocks))
-      timeToAllocatedBlocks(batchTime) = allocatedBlocks
-      lastAllocatedBatchTime = batchTime
-      allocatedBlocks
+      if (writeToLog(BatchAllocationEvent(batchTime, allocatedBlocks))) {
+        timeToAllocatedBlocks.put(batchTime, allocatedBlocks)
+        lastAllocatedBatchTime = batchTime
+      } else {
+        logInfo(s"Possibly processed batch $batchTime need to be processed again in WAL recovery")
+      }
     } else {
       // This situation occurs when:
       // 1. WAL is ended with BatchAllocationEvent, but without BatchCleanupEvent,
@@ -157,9 +166,12 @@ private[streaming] class ReceivedBlockTracker(
     require(cleanupThreshTime.milliseconds < clock.getTimeMillis())
     val timesToCleanup = timeToAllocatedBlocks.keys.filter { _ < cleanupThreshTime }.toSeq
     logInfo("Deleting batches " + timesToCleanup)
-    writeToLog(BatchCleanupEvent(timesToCleanup))
-    timeToAllocatedBlocks --= timesToCleanup
-    writeAheadLogOption.foreach(_.clean(cleanupThreshTime.milliseconds, waitForCompletion))
+    if (writeToLog(BatchCleanupEvent(timesToCleanup))) {
+      timeToAllocatedBlocks --= timesToCleanup
+      writeAheadLogOption.foreach(_.clean(cleanupThreshTime.milliseconds, waitForCompletion))
+    } else {
+      logWarning("Failed to acknowledge batch clean up in the Write Ahead Log.")
+    }
   }
 
   /** Stop the block tracker. */
@@ -185,8 +197,8 @@ private[streaming] class ReceivedBlockTracker(
       logTrace(s"Recovery: Inserting allocated batch for time $batchTime to " +
         s"${allocatedBlocks.streamIdToAllocatedBlocks}")
       streamIdToUnallocatedBlockQueues.values.foreach { _.clear() }
-      lastAllocatedBatchTime = batchTime
       timeToAllocatedBlocks.put(batchTime, allocatedBlocks)
+      lastAllocatedBatchTime = batchTime
     }
 
     // Cleanup the batch allocations
@@ -213,12 +225,20 @@ private[streaming] class ReceivedBlockTracker(
   }
 
   /** Write an update to the tracker to the write ahead log */
-  private def writeToLog(record: ReceivedBlockTrackerLogEvent) {
+  private def writeToLog(record: ReceivedBlockTrackerLogEvent): Boolean = {
     if (isWriteAheadLogEnabled) {
-      logDebug(s"Writing to log $record")
-      writeAheadLogOption.foreach { logManager =>
-        logManager.write(ByteBuffer.wrap(Utils.serialize(record)), clock.getTimeMillis())
+      logTrace(s"Writing record: $record")
+      try {
+        writeAheadLogOption.get.write(ByteBuffer.wrap(Utils.serialize(record)),
+          clock.getTimeMillis())
+        true
+      } catch {
+        case NonFatal(e) =>
+          logWarning(s"Exception thrown while writing record: $record to the WriteAheadLog.", e)
+          false
       }
+    } else {
+      true
     }
   }
 
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala
index b183d856f50c3..ea5d12b50fcc5 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTracker.scala
@@ -20,7 +20,7 @@ package org.apache.spark.streaming.scheduler
 import java.util.concurrent.{CountDownLatch, TimeUnit}
 
 import scala.collection.mutable.HashMap
-import scala.concurrent.ExecutionContext
+import scala.concurrent.{Future, ExecutionContext}
 import scala.language.existentials
 import scala.util.{Failure, Success}
 
@@ -437,7 +437,12 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
     // TODO Remove this thread pool after https://github.com/apache/spark/issues/7385 is merged
     private val submitJobThreadPool = ExecutionContext.fromExecutorService(
-      ThreadUtils.newDaemonCachedThreadPool("submit-job-thead-pool"))
+      ThreadUtils.newDaemonCachedThreadPool("submit-job-thread-pool"))
+
+    private val walBatchingThreadPool = ExecutionContext.fromExecutorService(
+      ThreadUtils.newDaemonCachedThreadPool("wal-batching-thread-pool"))
+
+    @volatile private var active: Boolean = true
 
     override def receive: PartialFunction[Any, Unit] = {
       // Local messages
@@ -488,7 +493,19 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
           registerReceiver(streamId, typ, host, executorId, receiverEndpoint, context.senderAddress)
         context.reply(successful)
       case AddBlock(receivedBlockInfo) =>
-        context.reply(addBlock(receivedBlockInfo))
+        if (WriteAheadLogUtils.isBatchingEnabled(ssc.conf, isDriver = true)) {
+          walBatchingThreadPool.execute(new Runnable {
+            override def run(): Unit = Utils.tryLogNonFatalError {
+              if (active) {
+                context.reply(addBlock(receivedBlockInfo))
+              } else {
+                throw new IllegalStateException("ReceiverTracker RpcEndpoint shut down.")
+              }
+            }
+          })
+        } else {
+          context.reply(addBlock(receivedBlockInfo))
+        }
       case DeregisterReceiver(streamId, message, error) =>
         deregisterReceiver(streamId, message, error)
         context.reply(true)
@@ -599,6 +616,8 @@ class ReceiverTracker(ssc: StreamingContext, skipReceiverLaunch: Boolean = false
 
     override def onStop(): Unit = {
       submitJobThreadPool.shutdownNow()
+      active = false
+      walBatchingThreadPool.shutdown()
     }
 
     /**
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/BatchedWriteAheadLog.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/BatchedWriteAheadLog.scala
new file mode 100644
index 0000000000000..9727ed2ba1445
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/BatchedWriteAheadLog.scala
@@ -0,0 +1,223 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.util
+
+import java.nio.ByteBuffer
+import java.util.concurrent.LinkedBlockingQueue
+import java.util.{Iterator => JIterator}
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable.ArrayBuffer
+import scala.concurrent.{Await, Promise}
+import scala.concurrent.duration._
+import scala.util.control.NonFatal
+
+import org.apache.spark.{Logging, SparkConf}
+import org.apache.spark.util.Utils
+
+/**
+ * A wrapper for a WriteAheadLog that batches records before writing data. Handles aggregation
+ * during writes, and de-aggregation in the `readAll` method. The end consumer has to handle
+ * de-aggregation after the `read` method. In addition, the `WriteAheadLogRecordHandle` returned
+ * after the write will contain the batch of records rather than individual records.
+ *
+ * When writing a batch of records, the `time` passed to the `wrappedLog` will be the timestamp
+ * of the latest record in the batch. This is very important in achieving correctness. Consider the
+ * following example:
+ * We receive records with timestamps 1, 3, 5, 7. We use "log-1" as the filename. Once we receive
+ * a clean up request for timestamp 3, we would clean up the file "log-1", and lose data regarding
+ * 5 and 7.
+ *
+ * This means the caller can assume the same write semantics as any other WriteAheadLog
+ * implementation despite the batching in the background - when the write() returns, the data is
+ * written to the WAL and is durable. To take advantage of the batching, the caller can write from
+ * multiple threads, each of which will stay blocked until the corresponding data has been written.
+ *
+ * All other methods of the WriteAheadLog interface will be passed on to the wrapped WriteAheadLog.
+ */
+private[util] class BatchedWriteAheadLog(val wrappedLog: WriteAheadLog, conf: SparkConf)
+  extends WriteAheadLog with Logging {
+
+  import BatchedWriteAheadLog._
+
+  private val walWriteQueue = new LinkedBlockingQueue[Record]()
+
+  // Whether the writer thread is active
+  @volatile private var active: Boolean = true
+  private val buffer = new ArrayBuffer[Record]()
+
+  private val batchedWriterThread = startBatchedWriterThread()
+
+  /**
+   * Write a byte buffer to the log file. This method adds the byteBuffer to a queue and blocks
+   * until the record is properly written by the parent.
+   */
+  override def write(byteBuffer: ByteBuffer, time: Long): WriteAheadLogRecordHandle = {
+    val promise = Promise[WriteAheadLogRecordHandle]()
+    val putSuccessfully = synchronized {
+      if (active) {
+        walWriteQueue.offer(Record(byteBuffer, time, promise))
+        true
+      } else {
+        false
+      }
+    }
+    if (putSuccessfully) {
+      Await.result(promise.future, WriteAheadLogUtils.getBatchingTimeout(conf).milliseconds)
+    } else {
+      throw new IllegalStateException("close() was called on BatchedWriteAheadLog before " +
+        s"write request with time $time could be fulfilled.")
+    }
+  }
+
+  /**
+   * This method is not supported as the resulting ByteBuffer would actually require de-aggregation.
+   * This method is primarily used in testing, and to ensure that it is not used in production,
+   * we throw an UnsupportedOperationException.
+   */
+  override def read(segment: WriteAheadLogRecordHandle): ByteBuffer = {
+    throw new UnsupportedOperationException("read() is not supported for BatchedWriteAheadLog " +
+      "as the data may require de-aggregation.")
+  }
+
+  /**
+   * Read all the existing logs from the log directory. The output of the wrapped WriteAheadLog
+   * will be de-aggregated.
+   */
+  override def readAll(): JIterator[ByteBuffer] = {
+    wrappedLog.readAll().asScala.flatMap(deaggregate).asJava
+  }
+
+  /**
+   * Delete the log files that are older than the threshold time.
+   *
+   * This method is handled by the parent WriteAheadLog.
+   */
+  override def clean(threshTime: Long, waitForCompletion: Boolean): Unit = {
+    wrappedLog.clean(threshTime, waitForCompletion)
+  }
+
+
+  /**
+   * Stop the batched writer thread, fulfill promises with failures and close the wrapped WAL.
+   */
+  override def close(): Unit = {
+    logInfo(s"BatchedWriteAheadLog shutting down at time: ${System.currentTimeMillis()}.")
+    synchronized {
+      active = false
+    }
+    batchedWriterThread.interrupt()
+    batchedWriterThread.join()
+    while (!walWriteQueue.isEmpty) {
+      val Record(_, time, promise) = walWriteQueue.poll()
+      promise.failure(new IllegalStateException("close() was called on BatchedWriteAheadLog " +
+        s"before write request with time $time could be fulfilled."))
+    }
+    wrappedLog.close()
+  }
+
+  /** Start the actual log writer on a separate thread. */
+  private def startBatchedWriterThread(): Thread = {
+    val thread = new Thread(new Runnable {
+      override def run(): Unit = {
+        while (active) {
+          try {
+            flushRecords()
+          } catch {
+            case NonFatal(e) =>
+              logWarning("Encountered exception in Batched Writer Thread.", e)
+          }
+        }
+        logInfo("BatchedWriteAheadLog Writer thread exiting.")
+      }
+    }, "BatchedWriteAheadLog Writer")
+    thread.setDaemon(true)
+    thread.start()
+    thread
+  }
+
+  /** Write all the records in the buffer to the write ahead log. */
+  private def flushRecords(): Unit = {
+    try {
+      buffer.append(walWriteQueue.take())
+      val numBatched = walWriteQueue.drainTo(buffer.asJava) + 1
+      logDebug(s"Received $numBatched records from queue")
+    } catch {
+      case _: InterruptedException =>
+        logWarning("BatchedWriteAheadLog Writer queue interrupted.")
+    }
+    try {
+      var segment: WriteAheadLogRecordHandle = null
+      if (buffer.length > 0) {
+        logDebug(s"Batched ${buffer.length} records for Write Ahead Log write")
+        // We take the latest record for the timestamp. Please refer to the class Javadoc for
+        // detailed explanation
+        val time = buffer.last.time
+        segment = wrappedLog.write(aggregate(buffer), time)
+      }
+      buffer.foreach(_.promise.success(segment))
+    } catch {
+      case e: InterruptedException =>
+        logWarning("BatchedWriteAheadLog Writer queue interrupted.", e)
+        buffer.foreach(_.promise.failure(e))
+      case NonFatal(e) =>
+        logWarning(s"BatchedWriteAheadLog Writer failed to write $buffer", e)
+        buffer.foreach(_.promise.failure(e))
+    } finally {
+      buffer.clear()
+    }
+  }
+}
+
+/** Static methods for aggregating and de-aggregating records. */
+private[util] object BatchedWriteAheadLog {
+
+  /**
+   * Wrapper class for representing the records that we will write to the WriteAheadLog. Coupled
+   * with the timestamp for the write request of the record, and the promise that will block the
+   * write request, while a separate thread is actually performing the write.
+   */
+  case class Record(data: ByteBuffer, time: Long, promise: Promise[WriteAheadLogRecordHandle])
+
+  /** Copies the byte array of a ByteBuffer. */
+  private def getByteArray(buffer: ByteBuffer): Array[Byte] = {
+    val byteArray = new Array[Byte](buffer.remaining())
+    buffer.get(byteArray)
+    byteArray
+  }
+
+  /** Aggregate multiple serialized ReceivedBlockTrackerLogEvents in a single ByteBuffer. */
+  def aggregate(records: Seq[Record]): ByteBuffer = {
+    ByteBuffer.wrap(Utils.serialize[Array[Array[Byte]]](
+      records.map(record => getByteArray(record.data)).toArray))
+  }
+
+  /**
+   * De-aggregate serialized ReceivedBlockTrackerLogEvents in a single ByteBuffer.
+   * A stream may not have used batching initially, but started using it after a restart. This
+   * method therefore needs to be backwards compatible.
+   */
+  def deaggregate(buffer: ByteBuffer): Array[ByteBuffer] = {
+    try {
+      Utils.deserialize[Array[Array[Byte]]](getByteArray(buffer)).map(ByteBuffer.wrap)
+    } catch {
+      case _: ClassCastException => // users may restart a stream with batching enabled
+        Array(buffer)
+    }
+  }
+}
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/WriteAheadLogUtils.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/WriteAheadLogUtils.scala
index 0ea970e61b694..731a369fc92c0 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/WriteAheadLogUtils.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/WriteAheadLogUtils.scala
@@ -38,6 +38,8 @@ private[streaming] object WriteAheadLogUtils extends Logging {
   val DRIVER_WAL_ROLLING_INTERVAL_CONF_KEY =
     "spark.streaming.driver.writeAheadLog.rollingIntervalSecs"
   val DRIVER_WAL_MAX_FAILURES_CONF_KEY = "spark.streaming.driver.writeAheadLog.maxFailures"
+  val DRIVER_WAL_BATCHING_CONF_KEY = "spark.streaming.driver.writeAheadLog.allowBatching"
+  val DRIVER_WAL_BATCHING_TIMEOUT_CONF_KEY = "spark.streaming.driver.writeAheadLog.batchingTimeout"
   val DRIVER_WAL_CLOSE_AFTER_WRITE_CONF_KEY =
     "spark.streaming.driver.writeAheadLog.closeFileAfterWrite"
 
@@ -64,6 +66,18 @@ private[streaming] object WriteAheadLogUtils extends Logging {
     }
   }
 
+  def isBatchingEnabled(conf: SparkConf, isDriver: Boolean): Boolean = {
+    isDriver && conf.getBoolean(DRIVER_WAL_BATCHING_CONF_KEY, defaultValue = false)
+  }
+
+  /**
+   * How long we will wait for the wrappedLog in the BatchedWriteAheadLog to write the records
+   * before we fail the write attempt to unblock receivers.
+   */
+  def getBatchingTimeout(conf: SparkConf): Long = {
+    conf.getLong(DRIVER_WAL_BATCHING_TIMEOUT_CONF_KEY, defaultValue = 5000)
+  }
+
   def shouldCloseFileAfterWrite(conf: SparkConf, isDriver: Boolean): Boolean = {
     if (isDriver) {
       conf.getBoolean(DRIVER_WAL_CLOSE_AFTER_WRITE_CONF_KEY, defaultValue = false)
@@ -115,7 +129,7 @@ private[streaming] object WriteAheadLogUtils extends Logging {
     } else {
       sparkConf.getOption(RECEIVER_WAL_CLASS_CONF_KEY)
     }
-    classNameOption.map { className =>
+    val wal = classNameOption.map { className =>
       try {
         instantiateClass(
           Utils.classForName(className).asInstanceOf[Class[_ <: WriteAheadLog]], sparkConf)
@@ -128,6 +142,11 @@ private[streaming] object WriteAheadLogUtils extends Logging {
         getRollingIntervalSecs(sparkConf, isDriver), getMaxFailures(sparkConf, isDriver),
         shouldCloseFileAfterWrite(sparkConf, isDriver))
     }
+    if (isBatchingEnabled(sparkConf, isDriver)) {
+      new BatchedWriteAheadLog(wal, sparkConf)
+    } else {
+      wal
+    }
   }
 
   /** Instantiate the class, either using single arg constructor or zero arg constructor */
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
index 93ae41a3d2ecd..e96f4c2a29347 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
@@ -18,31 +18,47 @@ package org.apache.spark.streaming.util
 
 import java.io._
 import java.nio.ByteBuffer
-import java.util
+import java.util.concurrent.{ExecutionException, ThreadPoolExecutor}
+import java.util.concurrent.atomic.AtomicInteger
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
+import scala.concurrent._
 import scala.concurrent.duration._
 import scala.language.{implicitConversions, postfixOps}
-import scala.reflect.ClassTag
+import scala.util.{Failure, Success}
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
+import org.mockito.Matchers.{eq => meq}
+import org.mockito.Matchers._
+import org.mockito.Mockito._
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
+import org.scalatest.concurrent.Eventually
 import org.scalatest.concurrent.Eventually._
-import org.scalatest.BeforeAndAfter
+import org.scalatest.{BeforeAndAfterEach, BeforeAndAfter}
+import org.scalatest.mock.MockitoSugar
 
-import org.apache.spark.util.{ManualClock, Utils}
-import org.apache.spark.{SparkConf, SparkException, SparkFunSuite}
+import org.apache.spark.streaming.scheduler._
+import org.apache.spark.util.{ThreadUtils, ManualClock, Utils}
+import org.apache.spark.{SparkException, SparkConf, SparkFunSuite}
 
-class WriteAheadLogSuite extends SparkFunSuite with BeforeAndAfter {
+/** Common tests for WriteAheadLogs that we would like to test with different configurations. */
+abstract class CommonWriteAheadLogTests(
+    allowBatching: Boolean,
+    closeFileAfterWrite: Boolean,
+    testTag: String = "")
+  extends SparkFunSuite with BeforeAndAfter {
 
   import WriteAheadLogSuite._
 
-  val hadoopConf = new Configuration()
-  var tempDir: File = null
-  var testDir: String = null
-  var testFile: String = null
-  var writeAheadLog: FileBasedWriteAheadLog = null
+  protected val hadoopConf = new Configuration()
+  protected var tempDir: File = null
+  protected var testDir: String = null
+  protected var testFile: String = null
+  protected var writeAheadLog: WriteAheadLog = null
+  protected def testPrefix = if (testTag != "") testTag + " - " else testTag
 
   before {
     tempDir = Utils.createTempDir()
@@ -58,49 +74,130 @@ class WriteAheadLogSuite extends SparkFunSuite with BeforeAndAfter {
     Utils.deleteRecursively(tempDir)
   }
 
-  test("WriteAheadLogUtils - log selection and creation") {
-    val logDir = Utils.createTempDir().getAbsolutePath()
+  test(testPrefix + "read all logs") {
+    // Write data manually for testing reading through WriteAheadLog
+    val writtenData = (1 to 10).map { i =>
+      val data = generateRandomData()
+      val file = testDir + s"/log-$i-$i"
+      writeDataManually(data, file, allowBatching)
+      data
+    }.flatten
 
-    def assertDriverLogClass[T <: WriteAheadLog: ClassTag](conf: SparkConf): WriteAheadLog = {
-      val log = WriteAheadLogUtils.createLogForDriver(conf, logDir, hadoopConf)
-      assert(log.getClass === implicitly[ClassTag[T]].runtimeClass)
-      log
+    val logDirectoryPath = new Path(testDir)
+    val fileSystem = HdfsUtils.getFileSystemForPath(logDirectoryPath, hadoopConf)
+    assert(fileSystem.exists(logDirectoryPath) === true)
+
+    // Read data using manager and verify
+    val readData = readDataUsingWriteAheadLog(testDir, closeFileAfterWrite, allowBatching)
+    assert(readData === writtenData)
+  }
+
+  test(testPrefix + "write logs") {
+    // Write data with rotation using WriteAheadLog class
+    val dataToWrite = generateRandomData()
+    writeDataUsingWriteAheadLog(testDir, dataToWrite, closeFileAfterWrite = closeFileAfterWrite,
+      allowBatching = allowBatching)
+
+    // Read data manually to verify the written data
+    val logFiles = getLogFilesInDirectory(testDir)
+    assert(logFiles.size > 1)
+    val writtenData = readAndDeserializeDataManually(logFiles, allowBatching)
+    assert(writtenData === dataToWrite)
+  }
+
+  test(testPrefix + "read all logs after write") {
+    // Write data with manager, recover with new manager and verify
+    val dataToWrite = generateRandomData()
+    writeDataUsingWriteAheadLog(testDir, dataToWrite, closeFileAfterWrite, allowBatching)
+    val logFiles = getLogFilesInDirectory(testDir)
+    assert(logFiles.size > 1)
+    val readData = readDataUsingWriteAheadLog(testDir, closeFileAfterWrite, allowBatching)
+    assert(dataToWrite === readData)
+  }
+
+  test(testPrefix + "clean old logs") {
+    logCleanUpTest(waitForCompletion = false)
+  }
+
+  test(testPrefix + "clean old logs synchronously") {
+    logCleanUpTest(waitForCompletion = true)
+  }
+
+  private def logCleanUpTest(waitForCompletion: Boolean): Unit = {
+    // Write data with manager, recover with new manager and verify
+    val manualClock = new ManualClock
+    val dataToWrite = generateRandomData()
+    writeAheadLog = writeDataUsingWriteAheadLog(testDir, dataToWrite, closeFileAfterWrite,
+      allowBatching, manualClock, closeLog = false)
+    val logFiles = getLogFilesInDirectory(testDir)
+    assert(logFiles.size > 1)
+
+    writeAheadLog.clean(manualClock.getTimeMillis() / 2, waitForCompletion)
+
+    if (waitForCompletion) {
+      assert(getLogFilesInDirectory(testDir).size < logFiles.size)
+    } else {
+      eventually(Eventually.timeout(1 second), interval(10 milliseconds)) {
+        assert(getLogFilesInDirectory(testDir).size < logFiles.size)
+      }
     }
+  }
 
-    def assertReceiverLogClass[T: ClassTag](conf: SparkConf): WriteAheadLog = {
-      val log = WriteAheadLogUtils.createLogForReceiver(conf, logDir, hadoopConf)
-      assert(log.getClass === implicitly[ClassTag[T]].runtimeClass)
-      log
+  test(testPrefix + "handling file errors while reading rotating logs") {
+    // Generate a set of log files
+    val manualClock = new ManualClock
+    val dataToWrite1 = generateRandomData()
+    writeDataUsingWriteAheadLog(testDir, dataToWrite1, closeFileAfterWrite, allowBatching,
+      manualClock)
+    val logFiles1 = getLogFilesInDirectory(testDir)
+    assert(logFiles1.size > 1)
+
+
+    // Recover old files and generate a second set of log files
+    val dataToWrite2 = generateRandomData()
+    manualClock.advance(100000)
+    writeDataUsingWriteAheadLog(testDir, dataToWrite2, closeFileAfterWrite, allowBatching ,
+      manualClock)
+    val logFiles2 = getLogFilesInDirectory(testDir)
+    assert(logFiles2.size > logFiles1.size)
+
+    // Read the files and verify that all the written data can be read
+    val readData1 = readDataUsingWriteAheadLog(testDir, closeFileAfterWrite, allowBatching)
+    assert(readData1 === (dataToWrite1 ++ dataToWrite2))
+
+    // Corrupt the first set of files so that they are basically unreadable
+    logFiles1.foreach { f =>
+      val raf = new FileOutputStream(f, true).getChannel()
+      raf.truncate(1)
+      raf.close()
     }
 
-    val emptyConf = new SparkConf()  // no log configuration
-    assertDriverLogClass[FileBasedWriteAheadLog](emptyConf)
-    assertReceiverLogClass[FileBasedWriteAheadLog](emptyConf)
-
-    // Verify setting driver WAL class
-    val conf1 = new SparkConf().set("spark.streaming.driver.writeAheadLog.class",
-      classOf[MockWriteAheadLog0].getName())
-    assertDriverLogClass[MockWriteAheadLog0](conf1)
-    assertReceiverLogClass[FileBasedWriteAheadLog](conf1)
-
-    // Verify setting receiver WAL class
-    val receiverWALConf = new SparkConf().set("spark.streaming.receiver.writeAheadLog.class",
-      classOf[MockWriteAheadLog0].getName())
-    assertDriverLogClass[FileBasedWriteAheadLog](receiverWALConf)
-    assertReceiverLogClass[MockWriteAheadLog0](receiverWALConf)
-
-    // Verify setting receiver WAL class with 1-arg constructor
-    val receiverWALConf2 = new SparkConf().set("spark.streaming.receiver.writeAheadLog.class",
-      classOf[MockWriteAheadLog1].getName())
-    assertReceiverLogClass[MockWriteAheadLog1](receiverWALConf2)
-
-    // Verify failure setting receiver WAL class with 2-arg constructor
-    intercept[SparkException] {
-      val receiverWALConf3 = new SparkConf().set("spark.streaming.receiver.writeAheadLog.class",
-        classOf[MockWriteAheadLog2].getName())
-      assertReceiverLogClass[MockWriteAheadLog1](receiverWALConf3)
+    // Verify that the corrupted files do not prevent reading of the second set of data
+    val readData = readDataUsingWriteAheadLog(testDir, closeFileAfterWrite, allowBatching)
+    assert(readData === dataToWrite2)
+  }
+
+  test(testPrefix + "do not create directories or files unless write") {
+    val nonexistentTempPath = File.createTempFile("test", "")
+    nonexistentTempPath.delete()
+    assert(!nonexistentTempPath.exists())
+
+    val writtenSegment = writeDataManually(generateRandomData(), testFile, allowBatching)
+    val wal = createWriteAheadLog(testDir, closeFileAfterWrite, allowBatching)
+    assert(!nonexistentTempPath.exists(), "Directory created just by creating log object")
+    if (allowBatching) {
+      intercept[UnsupportedOperationException](wal.read(writtenSegment.head))
+    } else {
+      wal.read(writtenSegment.head)
     }
+    assert(!nonexistentTempPath.exists(), "Directory created just by attempting to read segment")
   }
+}
+
+class FileBasedWriteAheadLogSuite
+  extends CommonWriteAheadLogTests(false, false, "FileBasedWriteAheadLog") {
+
+  import WriteAheadLogSuite._
 
   test("FileBasedWriteAheadLogWriter - writing data") {
     val dataToWrite = generateRandomData()
@@ -122,7 +219,7 @@ class WriteAheadLogSuite extends SparkFunSuite with BeforeAndAfter {
 
   test("FileBasedWriteAheadLogReader - sequentially reading data") {
     val writtenData = generateRandomData()
-    writeDataManually(writtenData, testFile)
+    writeDataManually(writtenData, testFile, allowBatching = false)
     val reader = new FileBasedWriteAheadLogReader(testFile, hadoopConf)
     val readData = reader.toSeq.map(byteBufferToString)
     assert(readData === writtenData)
@@ -166,7 +263,7 @@ class WriteAheadLogSuite extends SparkFunSuite with BeforeAndAfter {
   test("FileBasedWriteAheadLogRandomReader - reading data using random reader") {
     // Write data manually for testing the random reader
     val writtenData = generateRandomData()
-    val segments = writeDataManually(writtenData, testFile)
+    val segments = writeDataManually(writtenData, testFile, allowBatching = false)
 
     // Get a random order of these segments and read them back
     val writtenDataAndSegments = writtenData.zip(segments).toSeq.permutations.take(10).flatten
@@ -190,163 +287,212 @@ class WriteAheadLogSuite extends SparkFunSuite with BeforeAndAfter {
     }
     reader.close()
   }
+}
 
-  test("FileBasedWriteAheadLog - write rotating logs") {
-    // Write data with rotation using WriteAheadLog class
-    val dataToWrite = generateRandomData()
-    writeDataUsingWriteAheadLog(testDir, dataToWrite)
-
-    // Read data manually to verify the written data
-    val logFiles = getLogFilesInDirectory(testDir)
-    assert(logFiles.size > 1)
-    val writtenData = logFiles.flatMap { file => readDataManually(file)}
-    assert(writtenData === dataToWrite)
-  }
+abstract class CloseFileAfterWriteTests(allowBatching: Boolean, testTag: String)
+  extends CommonWriteAheadLogTests(allowBatching, closeFileAfterWrite = true, testTag) {
 
-  test("FileBasedWriteAheadLog - close after write flag") {
+  import WriteAheadLogSuite._
+  test(testPrefix + "close after write flag") {
     // Write data with rotation using WriteAheadLog class
     val numFiles = 3
     val dataToWrite = Seq.tabulate(numFiles)(_.toString)
     // total advance time is less than 1000, therefore log shouldn't be rolled, but manually closed
     writeDataUsingWriteAheadLog(testDir, dataToWrite, closeLog = false, clockAdvanceTime = 100,
-      closeFileAfterWrite = true)
+      closeFileAfterWrite = true, allowBatching = allowBatching)
 
     // Read data manually to verify the written data
     val logFiles = getLogFilesInDirectory(testDir)
     assert(logFiles.size === numFiles)
-    val writtenData = logFiles.flatMap { file => readDataManually(file)}
+    val writtenData: Seq[String] = readAndDeserializeDataManually(logFiles, allowBatching)
     assert(writtenData === dataToWrite)
   }
+}
 
-  test("FileBasedWriteAheadLog - read rotating logs") {
-    // Write data manually for testing reading through WriteAheadLog
-    val writtenData = (1 to 10).map { i =>
-      val data = generateRandomData()
-      val file = testDir + s"/log-$i-$i"
-      writeDataManually(data, file)
-      data
-    }.flatten
+class FileBasedWriteAheadLogWithFileCloseAfterWriteSuite
+  extends CloseFileAfterWriteTests(allowBatching = false, "FileBasedWriteAheadLog")
 
-    val logDirectoryPath = new Path(testDir)
-    val fileSystem = HdfsUtils.getFileSystemForPath(logDirectoryPath, hadoopConf)
-    assert(fileSystem.exists(logDirectoryPath) === true)
+class BatchedWriteAheadLogSuite extends CommonWriteAheadLogTests(
+    allowBatching = true,
+    closeFileAfterWrite = false,
+    "BatchedWriteAheadLog") with MockitoSugar with BeforeAndAfterEach with Eventually {
 
-    // Read data using manager and verify
-    val readData = readDataUsingWriteAheadLog(testDir)
-    assert(readData === writtenData)
-  }
+  import BatchedWriteAheadLog._
+  import WriteAheadLogSuite._
 
-  test("FileBasedWriteAheadLog - recover past logs when creating new manager") {
-    // Write data with manager, recover with new manager and verify
-    val dataToWrite = generateRandomData()
-    writeDataUsingWriteAheadLog(testDir, dataToWrite)
-    val logFiles = getLogFilesInDirectory(testDir)
-    assert(logFiles.size > 1)
-    val readData = readDataUsingWriteAheadLog(testDir)
-    assert(dataToWrite === readData)
+  private var wal: WriteAheadLog = _
+  private var walHandle: WriteAheadLogRecordHandle = _
+  private var walBatchingThreadPool: ThreadPoolExecutor = _
+  private var walBatchingExecutionContext: ExecutionContextExecutorService = _
+  private val sparkConf = new SparkConf()
+
+  override def beforeEach(): Unit = {
+    wal = mock[WriteAheadLog]
+    walHandle = mock[WriteAheadLogRecordHandle]
+    walBatchingThreadPool = ThreadUtils.newDaemonFixedThreadPool(8, "wal-test-thread-pool")
+    walBatchingExecutionContext = ExecutionContext.fromExecutorService(walBatchingThreadPool)
   }
 
-  test("FileBasedWriteAheadLog - clean old logs") {
-    logCleanUpTest(waitForCompletion = false)
+  override def afterEach(): Unit = {
+    if (walBatchingExecutionContext != null) {
+      walBatchingExecutionContext.shutdownNow()
+    }
   }
 
-  test("FileBasedWriteAheadLog - clean old logs synchronously") {
-    logCleanUpTest(waitForCompletion = true)
-  }
+  test("BatchedWriteAheadLog - serializing and deserializing batched records") {
+    val events = Seq(
+      BlockAdditionEvent(ReceivedBlockInfo(0, None, None, null)),
+      BatchAllocationEvent(null, null),
+      BatchCleanupEvent(Nil)
+    )
 
-  private def logCleanUpTest(waitForCompletion: Boolean): Unit = {
-    // Write data with manager, recover with new manager and verify
-    val manualClock = new ManualClock
-    val dataToWrite = generateRandomData()
-    writeAheadLog = writeDataUsingWriteAheadLog(testDir, dataToWrite, manualClock, closeLog = false)
-    val logFiles = getLogFilesInDirectory(testDir)
-    assert(logFiles.size > 1)
+    val buffers = events.map(e => Record(ByteBuffer.wrap(Utils.serialize(e)), 0L, null))
+    val batched = BatchedWriteAheadLog.aggregate(buffers)
+    val deaggregate = BatchedWriteAheadLog.deaggregate(batched).map(buffer =>
+      Utils.deserialize[ReceivedBlockTrackerLogEvent](buffer.array()))
 
-    writeAheadLog.clean(manualClock.getTimeMillis() / 2, waitForCompletion)
+    assert(deaggregate.toSeq === events)
+  }
 
-    if (waitForCompletion) {
-      assert(getLogFilesInDirectory(testDir).size < logFiles.size)
-    } else {
-      eventually(timeout(1 second), interval(10 milliseconds)) {
-        assert(getLogFilesInDirectory(testDir).size < logFiles.size)
-      }
+  test("BatchedWriteAheadLog - failures in wrappedLog get bubbled up") {
+    when(wal.write(any[ByteBuffer], anyLong)).thenThrow(new RuntimeException("Hello!"))
+    // the BatchedWriteAheadLog should bubble up any exceptions that may have happened during writes
+    val batchedWal = new BatchedWriteAheadLog(wal, sparkConf)
+
+    intercept[RuntimeException] {
+      val buffer = mock[ByteBuffer]
+      batchedWal.write(buffer, 2L)
     }
   }
 
-  test("FileBasedWriteAheadLog - handling file errors while reading rotating logs") {
-    // Generate a set of log files
-    val manualClock = new ManualClock
-    val dataToWrite1 = generateRandomData()
-    writeDataUsingWriteAheadLog(testDir, dataToWrite1, manualClock)
-    val logFiles1 = getLogFilesInDirectory(testDir)
-    assert(logFiles1.size > 1)
+  // we make the write requests in separate threads so that we don't block the test thread
+  private def promiseWriteEvent(wal: WriteAheadLog, event: String, time: Long): Promise[Unit] = {
+    val p = Promise[Unit]()
+    p.completeWith(Future {
+      val v = wal.write(event, time)
+      assert(v === walHandle)
+    }(walBatchingExecutionContext))
+    p
+  }
 
+  /**
+   * In order to block the writes on the writer thread, we mock the write method, and block it
+   * for some time with a promise.
+   */
+  private def writeBlockingPromise(wal: WriteAheadLog): Promise[Any] = {
+    // we would like to block the write so that we can queue requests
+    val promise = Promise[Any]()
+    when(wal.write(any[ByteBuffer], any[Long])).thenAnswer(
+      new Answer[WriteAheadLogRecordHandle] {
+        override def answer(invocation: InvocationOnMock): WriteAheadLogRecordHandle = {
+          Await.ready(promise.future, 4.seconds)
+          walHandle
+        }
+      }
+    )
+    promise
+  }
 
-    // Recover old files and generate a second set of log files
-    val dataToWrite2 = generateRandomData()
-    manualClock.advance(100000)
-    writeDataUsingWriteAheadLog(testDir, dataToWrite2, manualClock)
-    val logFiles2 = getLogFilesInDirectory(testDir)
-    assert(logFiles2.size > logFiles1.size)
+  test("BatchedWriteAheadLog - name log with aggregated entries with the timestamp of last entry") {
+    val batchedWal = new BatchedWriteAheadLog(wal, sparkConf)
+    // block the write so that we can batch some records
+    val promise = writeBlockingPromise(wal)
+
+    val event1 = "hello"
+    val event2 = "world"
+    val event3 = "this"
+    val event4 = "is"
+    val event5 = "doge"
+
+    // The queue.take() immediately takes the 3, and there is nothing left in the queue at that
+    // moment. Then the promise blocks the writing of 3. The rest get queued.
+    promiseWriteEvent(batchedWal, event1, 3L)
+    // rest of the records will be batched while it takes 3 to get written
+    promiseWriteEvent(batchedWal, event2, 5L)
+    promiseWriteEvent(batchedWal, event3, 8L)
+    promiseWriteEvent(batchedWal, event4, 12L)
+    promiseWriteEvent(batchedWal, event5, 10L)
+    eventually(timeout(1 second)) {
+      assert(walBatchingThreadPool.getActiveCount === 5)
+    }
+    promise.success(true)
 
-    // Read the files and verify that all the written data can be read
-    val readData1 = readDataUsingWriteAheadLog(testDir)
-    assert(readData1 === (dataToWrite1 ++ dataToWrite2))
+    val buffer1 = wrapArrayArrayByte(Array(event1))
+    val buffer2 = wrapArrayArrayByte(Array(event2, event3, event4, event5))
 
-    // Corrupt the first set of files so that they are basically unreadable
-    logFiles1.foreach { f =>
-      val raf = new FileOutputStream(f, true).getChannel()
-      raf.truncate(1)
-      raf.close()
+    eventually(timeout(1 second)) {
+      verify(wal, times(1)).write(meq(buffer1), meq(3L))
+      // the file name should be the timestamp of the last record, as events should be naturally
+      // in order of timestamp, and we need the last element.
+      verify(wal, times(1)).write(meq(buffer2), meq(10L))
     }
-
-    // Verify that the corrupted files do not prevent reading of the second set of data
-    val readData = readDataUsingWriteAheadLog(testDir)
-    assert(readData === dataToWrite2)
   }
 
-  test("FileBasedWriteAheadLog - do not create directories or files unless write") {
-    val nonexistentTempPath = File.createTempFile("test", "")
-    nonexistentTempPath.delete()
-    assert(!nonexistentTempPath.exists())
+  test("BatchedWriteAheadLog - shutdown properly") {
+    val batchedWal = new BatchedWriteAheadLog(wal, sparkConf)
+    batchedWal.close()
+    verify(wal, times(1)).close()
 
-    val writtenSegment = writeDataManually(generateRandomData(), testFile)
-    val wal = new FileBasedWriteAheadLog(new SparkConf(), tempDir.getAbsolutePath,
-      new Configuration(), 1, 1, closeFileAfterWrite = false)
-    assert(!nonexistentTempPath.exists(), "Directory created just by creating log object")
-    wal.read(writtenSegment.head)
-    assert(!nonexistentTempPath.exists(), "Directory created just by attempting to read segment")
+    intercept[IllegalStateException](batchedWal.write(mock[ByteBuffer], 12L))
   }
-}
 
-object WriteAheadLogSuite {
+  test("BatchedWriteAheadLog - fail everything in queue during shutdown") {
+    val batchedWal = new BatchedWriteAheadLog(wal, sparkConf)
 
-  class MockWriteAheadLog0() extends WriteAheadLog {
-    override def write(record: ByteBuffer, time: Long): WriteAheadLogRecordHandle = { null }
-    override def read(handle: WriteAheadLogRecordHandle): ByteBuffer = { null }
-    override def readAll(): util.Iterator[ByteBuffer] = { null }
-    override def clean(threshTime: Long, waitForCompletion: Boolean): Unit = { }
-    override def close(): Unit = { }
-  }
+    // block the write so that we can batch some records
+    writeBlockingPromise(wal)
+
+    val event1 = ("hello", 3L)
+    val event2 = ("world", 5L)
+    val event3 = ("this", 8L)
+    val event4 = ("is", 9L)
+    val event5 = ("doge", 10L)
+
+    // The queue.take() immediately takes the 3, and there is nothing left in the queue at that
+    // moment. Then the promise blocks the writing of 3. The rest get queued.
+    val writePromises = Seq(event1, event2, event3, event4, event5).map { event =>
+      promiseWriteEvent(batchedWal, event._1, event._2)
+    }
 
-  class MockWriteAheadLog1(val conf: SparkConf) extends MockWriteAheadLog0()
+    eventually(timeout(1 second)) {
+      assert(walBatchingThreadPool.getActiveCount === 5)
+    }
+
+    batchedWal.close()
+    eventually(timeout(1 second)) {
+      assert(writePromises.forall(_.isCompleted))
+      assert(writePromises.forall(_.future.value.get.isFailure)) // all should have failed
+    }
+  }
+}
 
-  class MockWriteAheadLog2(val conf: SparkConf, x: Int) extends MockWriteAheadLog0()
+class BatchedWriteAheadLogWithCloseFileAfterWriteSuite
+  extends CloseFileAfterWriteTests(allowBatching = true, "BatchedWriteAheadLog")
 
+object WriteAheadLogSuite {
 
   private val hadoopConf = new Configuration()
 
   /** Write data to a file directly and return an array of the file segments written. */
-  def writeDataManually(data: Seq[String], file: String): Seq[FileBasedWriteAheadLogSegment] = {
+  def writeDataManually(
+      data: Seq[String],
+      file: String,
+      allowBatching: Boolean): Seq[FileBasedWriteAheadLogSegment] = {
     val segments = new ArrayBuffer[FileBasedWriteAheadLogSegment]()
     val writer = HdfsUtils.getOutputStream(file, hadoopConf)
-    data.foreach { item =>
+    def writeToStream(bytes: Array[Byte]): Unit = {
       val offset = writer.getPos
-      val bytes = Utils.serialize(item)
       writer.writeInt(bytes.size)
       writer.write(bytes)
       segments += FileBasedWriteAheadLogSegment(file, offset, bytes.size)
     }
+    if (allowBatching) {
+      writeToStream(wrapArrayArrayByte(data.toArray[String]).array())
+    } else {
+      data.foreach { item =>
+        writeToStream(Utils.serialize(item))
+      }
+    }
     writer.close()
     segments
   }
@@ -356,8 +502,7 @@ object WriteAheadLogSuite {
    */
   def writeDataUsingWriter(
       filePath: String,
-      data: Seq[String]
-    ): Seq[FileBasedWriteAheadLogSegment] = {
+      data: Seq[String]): Seq[FileBasedWriteAheadLogSegment] = {
     val writer = new FileBasedWriteAheadLogWriter(filePath, hadoopConf)
     val segments = data.map {
       item => writer.write(item)
@@ -370,13 +515,13 @@ object WriteAheadLogSuite {
   def writeDataUsingWriteAheadLog(
       logDirectory: String,
       data: Seq[String],
+      closeFileAfterWrite: Boolean,
+      allowBatching: Boolean,
       manualClock: ManualClock = new ManualClock,
       closeLog: Boolean = true,
-      clockAdvanceTime: Int = 500,
-      closeFileAfterWrite: Boolean = false): FileBasedWriteAheadLog = {
+      clockAdvanceTime: Int = 500): WriteAheadLog = {
     if (manualClock.getTimeMillis() < 100000) manualClock.setTime(10000)
-    val wal = new FileBasedWriteAheadLog(new SparkConf(), logDirectory, hadoopConf, 1, 1,
-      closeFileAfterWrite)
+    val wal = createWriteAheadLog(logDirectory, closeFileAfterWrite, allowBatching)
 
     // Ensure that 500 does not get sorted after 2000, so put a high base value.
     data.foreach { item =>
@@ -406,16 +551,16 @@ object WriteAheadLogSuite {
   }
 
   /** Read all the data from a log file directly and return the list of byte buffers. */
-  def readDataManually(file: String): Seq[String] = {
+  def readDataManually[T](file: String): Seq[T] = {
     val reader = HdfsUtils.getInputStream(file, hadoopConf)
-    val buffer = new ArrayBuffer[String]
+    val buffer = new ArrayBuffer[T]
     try {
       while (true) {
         // Read till EOF is thrown
         val length = reader.readInt()
         val bytes = new Array[Byte](length)
         reader.read(bytes)
-        buffer += Utils.deserialize[String](bytes)
+        buffer += Utils.deserialize[T](bytes)
       }
     } catch {
       case ex: EOFException =>
@@ -434,15 +579,17 @@ object WriteAheadLogSuite {
   }
 
   /** Read all the data in the log file in a directory using the WriteAheadLog class. */
-  def readDataUsingWriteAheadLog(logDirectory: String): Seq[String] = {
-    val wal = new FileBasedWriteAheadLog(new SparkConf(), logDirectory, hadoopConf, 1, 1,
-      closeFileAfterWrite = false)
+  def readDataUsingWriteAheadLog(
+      logDirectory: String,
+      closeFileAfterWrite: Boolean,
+      allowBatching: Boolean): Seq[String] = {
+    val wal = createWriteAheadLog(logDirectory, closeFileAfterWrite, allowBatching)
     val data = wal.readAll().asScala.map(byteBufferToString).toSeq
     wal.close()
     data
   }
 
-  /** Get the log files in a direction */
+  /** Get the log files in a directory. */
   def getLogFilesInDirectory(directory: String): Seq[String] = {
     val logDirectoryPath = new Path(directory)
     val fileSystem = HdfsUtils.getFileSystemForPath(logDirectoryPath, hadoopConf)
@@ -458,10 +605,31 @@ object WriteAheadLogSuite {
     }
   }
 
+  def createWriteAheadLog(
+      logDirectory: String,
+      closeFileAfterWrite: Boolean,
+      allowBatching: Boolean): WriteAheadLog = {
+    val sparkConf = new SparkConf
+    val wal = new FileBasedWriteAheadLog(sparkConf, logDirectory, hadoopConf, 1, 1,
+      closeFileAfterWrite)
+    if (allowBatching) new BatchedWriteAheadLog(wal, sparkConf) else wal
+  }
+
   def generateRandomData(): Seq[String] = {
     (1 to 100).map { _.toString }
   }
 
+  def readAndDeserializeDataManually(logFiles: Seq[String], allowBatching: Boolean): Seq[String] = {
+    if (allowBatching) {
+      logFiles.flatMap { file =>
+        val data = readDataManually[Array[Array[Byte]]](file)
+        data.flatMap(byteArray => byteArray.map(Utils.deserialize[String]))
+      }
+    } else {
+      logFiles.flatMap { file => readDataManually[String](file)}
+    }
+  }
+
   implicit def stringToByteBuffer(str: String): ByteBuffer = {
     ByteBuffer.wrap(Utils.serialize(str))
   }
@@ -469,4 +637,8 @@ object WriteAheadLogSuite {
   implicit def byteBufferToString(byteBuffer: ByteBuffer): String = {
     Utils.deserialize[String](byteBuffer.array)
   }
+
+  def wrapArrayArrayByte[T](records: Array[T]): ByteBuffer = {
+    ByteBuffer.wrap(Utils.serialize[Array[Array[Byte]]](records.map(Utils.serialize[T])))
+  }
 }
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogUtilsSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogUtilsSuite.scala
new file mode 100644
index 0000000000000..9152728191ea1
--- /dev/null
+++ b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogUtilsSuite.scala
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.util
+
+import java.nio.ByteBuffer
+import java.util
+
+import scala.reflect.ClassTag
+
+import org.apache.hadoop.conf.Configuration
+
+import org.apache.spark.{SparkException, SparkConf, SparkFunSuite}
+import org.apache.spark.util.Utils
+
+class WriteAheadLogUtilsSuite extends SparkFunSuite {
+  import WriteAheadLogUtilsSuite._
+
+  private val logDir = Utils.createTempDir().getAbsolutePath()
+  private val hadoopConf = new Configuration()
+
+  def assertDriverLogClass[T <: WriteAheadLog: ClassTag](
+      conf: SparkConf,
+      isBatched: Boolean = false): WriteAheadLog = {
+    val log = WriteAheadLogUtils.createLogForDriver(conf, logDir, hadoopConf)
+    if (isBatched) {
+      assert(log.isInstanceOf[BatchedWriteAheadLog])
+      val parentLog = log.asInstanceOf[BatchedWriteAheadLog].wrappedLog
+      assert(parentLog.getClass === implicitly[ClassTag[T]].runtimeClass)
+    } else {
+      assert(log.getClass === implicitly[ClassTag[T]].runtimeClass)
+    }
+    log
+  }
+
+  def assertReceiverLogClass[T <: WriteAheadLog: ClassTag](conf: SparkConf): WriteAheadLog = {
+    val log = WriteAheadLogUtils.createLogForReceiver(conf, logDir, hadoopConf)
+    assert(log.getClass === implicitly[ClassTag[T]].runtimeClass)
+    log
+  }
+
+  test("log selection and creation") {
+
+    val emptyConf = new SparkConf()  // no log configuration
+    assertDriverLogClass[FileBasedWriteAheadLog](emptyConf)
+    assertReceiverLogClass[FileBasedWriteAheadLog](emptyConf)
+
+    // Verify setting driver WAL class
+    val driverWALConf = new SparkConf().set("spark.streaming.driver.writeAheadLog.class",
+      classOf[MockWriteAheadLog0].getName())
+    assertDriverLogClass[MockWriteAheadLog0](driverWALConf)
+    assertReceiverLogClass[FileBasedWriteAheadLog](driverWALConf)
+
+    // Verify setting receiver WAL class
+    val receiverWALConf = new SparkConf().set("spark.streaming.receiver.writeAheadLog.class",
+      classOf[MockWriteAheadLog0].getName())
+    assertDriverLogClass[FileBasedWriteAheadLog](receiverWALConf)
+    assertReceiverLogClass[MockWriteAheadLog0](receiverWALConf)
+
+    // Verify setting receiver WAL class with 1-arg constructor
+    val receiverWALConf2 = new SparkConf().set("spark.streaming.receiver.writeAheadLog.class",
+      classOf[MockWriteAheadLog1].getName())
+    assertReceiverLogClass[MockWriteAheadLog1](receiverWALConf2)
+
+    // Verify failure setting receiver WAL class with 2-arg constructor
+    intercept[SparkException] {
+      val receiverWALConf3 = new SparkConf().set("spark.streaming.receiver.writeAheadLog.class",
+        classOf[MockWriteAheadLog2].getName())
+      assertReceiverLogClass[MockWriteAheadLog1](receiverWALConf3)
+    }
+  }
+
+  test("wrap WriteAheadLog in BatchedWriteAheadLog when batching is enabled") {
+    def getBatchedSparkConf: SparkConf =
+      new SparkConf().set("spark.streaming.driver.writeAheadLog.allowBatching", "true")
+
+    val justBatchingConf = getBatchedSparkConf
+    assertDriverLogClass[FileBasedWriteAheadLog](justBatchingConf, isBatched = true)
+    assertReceiverLogClass[FileBasedWriteAheadLog](justBatchingConf)
+
+    // Verify setting driver WAL class
+    val driverWALConf = getBatchedSparkConf.set("spark.streaming.driver.writeAheadLog.class",
+      classOf[MockWriteAheadLog0].getName())
+    assertDriverLogClass[MockWriteAheadLog0](driverWALConf, isBatched = true)
+    assertReceiverLogClass[FileBasedWriteAheadLog](driverWALConf)
+
+    // Verify receivers are not wrapped
+    val receiverWALConf = getBatchedSparkConf.set("spark.streaming.receiver.writeAheadLog.class",
+      classOf[MockWriteAheadLog0].getName())
+    assertDriverLogClass[FileBasedWriteAheadLog](receiverWALConf, isBatched = true)
+    assertReceiverLogClass[MockWriteAheadLog0](receiverWALConf)
+  }
+}
+
+object WriteAheadLogUtilsSuite {
+
+  class MockWriteAheadLog0() extends WriteAheadLog {
+    override def write(record: ByteBuffer, time: Long): WriteAheadLogRecordHandle = { null }
+    override def read(handle: WriteAheadLogRecordHandle): ByteBuffer = { null }
+    override def readAll(): util.Iterator[ByteBuffer] = { null }
+    override def clean(threshTime: Long, waitForCompletion: Boolean): Unit = { }
+    override def close(): Unit = { }
+  }
+
+  class MockWriteAheadLog1(val conf: SparkConf) extends MockWriteAheadLog0()
+
+  class MockWriteAheadLog2(val conf: SparkConf, x: Int) extends MockWriteAheadLog0()
+}

From 1f0f14efe35f986e338ee2cbc1ef2a9ce7395c00 Mon Sep 17 00:00:00 2001
From: zsxwing <zsxwing@gmail.com>
Date: Mon, 9 Nov 2015 17:38:19 -0800
Subject: [PATCH 147/518] [SPARK-11462][STREAMING] Add JavaStreamingListener

Currently, StreamingListener is not Java friendly because it exposes some Scala collections to Java users directly, such as Option, Map.

This PR added a Java version of StreamingListener and a bunch of Java friendly classes for Java users.

Author: zsxwing <zsxwing@gmail.com>
Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9420 from zsxwing/java-streaming-listener.
---
 .../api/java/JavaStreamingListener.scala      | 168 ++++++++++
 .../java/JavaStreamingListenerWrapper.scala   | 122 ++++++++
 .../JavaStreamingListenerAPISuite.java        |  85 +++++
 .../JavaStreamingListenerWrapperSuite.scala   | 290 ++++++++++++++++++
 4 files changed, 665 insertions(+)
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapper.scala
 create mode 100644 streaming/src/test/java/org/apache/spark/streaming/JavaStreamingListenerAPISuite.java
 create mode 100644 streaming/src/test/java/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapperSuite.scala

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
new file mode 100644
index 0000000000000..c86c7101ff6d5
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
@@ -0,0 +1,168 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.api.java
+
+import org.apache.spark.streaming.Time
+
+/**
+ * A listener interface for receiving information about an ongoing streaming  computation.
+ */
+private[streaming] class JavaStreamingListener {
+
+  /** Called when a receiver has been started */
+  def onReceiverStarted(receiverStarted: JavaStreamingListenerReceiverStarted): Unit = { }
+
+  /** Called when a receiver has reported an error */
+  def onReceiverError(receiverError: JavaStreamingListenerReceiverError): Unit = { }
+
+  /** Called when a receiver has been stopped */
+  def onReceiverStopped(receiverStopped: JavaStreamingListenerReceiverStopped): Unit = { }
+
+  /** Called when a batch of jobs has been submitted for processing. */
+  def onBatchSubmitted(batchSubmitted: JavaStreamingListenerBatchSubmitted): Unit = { }
+
+  /** Called when processing of a batch of jobs has started.  */
+  def onBatchStarted(batchStarted: JavaStreamingListenerBatchStarted): Unit = { }
+
+  /** Called when processing of a batch of jobs has completed. */
+  def onBatchCompleted(batchCompleted: JavaStreamingListenerBatchCompleted): Unit = { }
+
+  /** Called when processing of a job of a batch has started. */
+  def onOutputOperationStarted(
+      outputOperationStarted: JavaStreamingListenerOutputOperationStarted): Unit = { }
+
+  /** Called when processing of a job of a batch has completed. */
+  def onOutputOperationCompleted(
+      outputOperationCompleted: JavaStreamingListenerOutputOperationCompleted): Unit = { }
+}
+
+/**
+ * Base trait for events related to JavaStreamingListener
+ */
+private[streaming] sealed trait JavaStreamingListenerEvent
+
+private[streaming] class JavaStreamingListenerBatchSubmitted(val batchInfo: JavaBatchInfo)
+  extends JavaStreamingListenerEvent
+
+private[streaming] class JavaStreamingListenerBatchCompleted(val batchInfo: JavaBatchInfo)
+  extends JavaStreamingListenerEvent
+
+private[streaming] class JavaStreamingListenerBatchStarted(val batchInfo: JavaBatchInfo)
+  extends JavaStreamingListenerEvent
+
+private[streaming] class JavaStreamingListenerOutputOperationStarted(
+    val outputOperationInfo: JavaOutputOperationInfo) extends JavaStreamingListenerEvent
+
+private[streaming] class JavaStreamingListenerOutputOperationCompleted(
+    val outputOperationInfo: JavaOutputOperationInfo) extends JavaStreamingListenerEvent
+
+private[streaming] class JavaStreamingListenerReceiverStarted(val receiverInfo: JavaReceiverInfo)
+  extends JavaStreamingListenerEvent
+
+private[streaming] class JavaStreamingListenerReceiverError(val receiverInfo: JavaReceiverInfo)
+  extends JavaStreamingListenerEvent
+
+private[streaming] class JavaStreamingListenerReceiverStopped(val receiverInfo: JavaReceiverInfo)
+  extends JavaStreamingListenerEvent
+
+/**
+ * Class having information on batches.
+ *
+ * @param batchTime Time of the batch
+ * @param streamIdToInputInfo A map of input stream id to its input info
+ * @param submissionTime Clock time of when jobs of this batch was submitted to the streaming
+ *                       scheduler queue
+ * @param processingStartTime Clock time of when the first job of this batch started processing.
+ *                            `-1` means the batch has not yet started
+ * @param processingEndTime Clock time of when the last job of this batch finished processing. `-1`
+ *                          means the batch has not yet completed.
+ * @param schedulingDelay Time taken for the first job of this batch to start processing from the
+ *                        time this batch was submitted to the streaming scheduler. Essentially, it
+ *                        is `processingStartTime` - `submissionTime`. `-1` means the batch has not
+ *                        yet started
+ * @param processingDelay Time taken for the all jobs of this batch to finish processing from the
+ *                        time they started processing. Essentially, it is
+ *                        `processingEndTime` - `processingStartTime`. `-1` means the batch has not
+ *                        yet completed.
+ * @param totalDelay Time taken for all the jobs of this batch to finish processing from the time
+ *                   they were submitted.  Essentially, it is `processingDelay` + `schedulingDelay`.
+ *                   `-1` means the batch has not yet completed.
+ * @param numRecords The number of recorders received by the receivers in this batch
+ * @param outputOperationInfos The output operations in this batch
+ */
+private[streaming] case class JavaBatchInfo(
+    batchTime: Time,
+    streamIdToInputInfo: java.util.Map[Int, JavaStreamInputInfo],
+    submissionTime: Long,
+    processingStartTime: Long,
+    processingEndTime: Long,
+    schedulingDelay: Long,
+    processingDelay: Long,
+    totalDelay: Long,
+    numRecords: Long,
+    outputOperationInfos: java.util.Map[Int, JavaOutputOperationInfo])
+
+/**
+ * Track the information of input stream at specified batch time.
+ *
+ * @param inputStreamId the input stream id
+ * @param numRecords the number of records in a batch
+ * @param metadata metadata for this batch. It should contain at least one standard field named
+ *                 "Description" which maps to the content that will be shown in the UI.
+ * @param metadataDescription description of this input stream
+ */
+private[streaming] case class JavaStreamInputInfo(
+    inputStreamId: Int,
+    numRecords: Long,
+    metadata: java.util.Map[String, Any],
+    metadataDescription: String)
+
+/**
+ * Class having information about a receiver
+ */
+private[streaming] case class JavaReceiverInfo(
+    streamId: Int,
+    name: String,
+    active: Boolean,
+    location: String,
+    lastErrorMessage: String,
+    lastError: String,
+    lastErrorTime: Long)
+
+/**
+ * Class having information on output operations.
+ *
+ * @param batchTime Time of the batch
+ * @param id Id of this output operation. Different output operations have different ids in a batch.
+ * @param name The name of this output operation.
+ * @param description The description of this output operation.
+ * @param startTime Clock time of when the output operation started processing. `-1` means the
+ *                  output operation has not yet started
+ * @param endTime Clock time of when the output operation started processing. `-1` means the output
+ *                operation has not yet completed
+ * @param failureReason Failure reason if this output operation fails. If the output operation is
+ *                      successful, this field is `null`.
+ */
+private[streaming] case class JavaOutputOperationInfo(
+    batchTime: Time,
+    id: Int,
+    name: String,
+    description: String,
+    startTime: Long,
+    endTime: Long,
+    failureReason: String)
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapper.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapper.scala
new file mode 100644
index 0000000000000..2c60b396a6616
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapper.scala
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.api.java
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.streaming.scheduler._
+
+/**
+ * A wrapper to convert a [[JavaStreamingListener]] to a [[StreamingListener]].
+ */
+private[streaming] class JavaStreamingListenerWrapper(javaStreamingListener: JavaStreamingListener)
+  extends StreamingListener {
+
+  private def toJavaReceiverInfo(receiverInfo: ReceiverInfo): JavaReceiverInfo = {
+    JavaReceiverInfo(
+      receiverInfo.streamId,
+      receiverInfo.name,
+      receiverInfo.active,
+      receiverInfo.location,
+      receiverInfo.lastErrorMessage,
+      receiverInfo.lastError,
+      receiverInfo.lastErrorTime
+    )
+  }
+
+  private def toJavaStreamInputInfo(streamInputInfo: StreamInputInfo): JavaStreamInputInfo = {
+    JavaStreamInputInfo(
+      streamInputInfo.inputStreamId,
+      streamInputInfo.numRecords: Long,
+      streamInputInfo.metadata.asJava,
+      streamInputInfo.metadataDescription.orNull
+    )
+  }
+
+  private def toJavaOutputOperationInfo(
+      outputOperationInfo: OutputOperationInfo): JavaOutputOperationInfo = {
+    JavaOutputOperationInfo(
+      outputOperationInfo.batchTime,
+      outputOperationInfo.id,
+      outputOperationInfo.name,
+      outputOperationInfo.description: String,
+      outputOperationInfo.startTime.getOrElse(-1),
+      outputOperationInfo.endTime.getOrElse(-1),
+      outputOperationInfo.failureReason.orNull
+    )
+  }
+
+  private def toJavaBatchInfo(batchInfo: BatchInfo): JavaBatchInfo = {
+    JavaBatchInfo(
+      batchInfo.batchTime,
+      batchInfo.streamIdToInputInfo.mapValues(toJavaStreamInputInfo(_)).asJava,
+      batchInfo.submissionTime,
+      batchInfo.processingStartTime.getOrElse(-1),
+      batchInfo.processingEndTime.getOrElse(-1),
+      batchInfo.schedulingDelay.getOrElse(-1),
+      batchInfo.processingDelay.getOrElse(-1),
+      batchInfo.totalDelay.getOrElse(-1),
+      batchInfo.numRecords,
+      batchInfo.outputOperationInfos.mapValues(toJavaOutputOperationInfo(_)).asJava
+    )
+  }
+
+  override def onReceiverStarted(receiverStarted: StreamingListenerReceiverStarted): Unit = {
+    javaStreamingListener.onReceiverStarted(
+      new JavaStreamingListenerReceiverStarted(toJavaReceiverInfo(receiverStarted.receiverInfo)))
+  }
+
+  override def onReceiverError(receiverError: StreamingListenerReceiverError): Unit = {
+    javaStreamingListener.onReceiverError(
+      new JavaStreamingListenerReceiverError(toJavaReceiverInfo(receiverError.receiverInfo)))
+  }
+
+  override def onReceiverStopped(receiverStopped: StreamingListenerReceiverStopped): Unit = {
+    javaStreamingListener.onReceiverStopped(
+      new JavaStreamingListenerReceiverStopped(toJavaReceiverInfo(receiverStopped.receiverInfo)))
+  }
+
+  override def onBatchSubmitted(batchSubmitted: StreamingListenerBatchSubmitted): Unit = {
+    javaStreamingListener.onBatchSubmitted(
+      new JavaStreamingListenerBatchSubmitted(toJavaBatchInfo(batchSubmitted.batchInfo)))
+  }
+
+  override def onBatchStarted(batchStarted: StreamingListenerBatchStarted): Unit = {
+    javaStreamingListener.onBatchStarted(
+      new JavaStreamingListenerBatchStarted(toJavaBatchInfo(batchStarted.batchInfo)))
+  }
+
+  override def onBatchCompleted(batchCompleted: StreamingListenerBatchCompleted): Unit = {
+    javaStreamingListener.onBatchCompleted(
+      new JavaStreamingListenerBatchCompleted(toJavaBatchInfo(batchCompleted.batchInfo)))
+  }
+
+  override def onOutputOperationStarted(
+      outputOperationStarted: StreamingListenerOutputOperationStarted): Unit = {
+    javaStreamingListener.onOutputOperationStarted(new JavaStreamingListenerOutputOperationStarted(
+      toJavaOutputOperationInfo(outputOperationStarted.outputOperationInfo)))
+  }
+
+  override def onOutputOperationCompleted(
+      outputOperationCompleted: StreamingListenerOutputOperationCompleted): Unit = {
+    javaStreamingListener.onOutputOperationCompleted(
+      new JavaStreamingListenerOutputOperationCompleted(
+        toJavaOutputOperationInfo(outputOperationCompleted.outputOperationInfo)))
+  }
+
+}
diff --git a/streaming/src/test/java/org/apache/spark/streaming/JavaStreamingListenerAPISuite.java b/streaming/src/test/java/org/apache/spark/streaming/JavaStreamingListenerAPISuite.java
new file mode 100644
index 0000000000000..8cc285aa7fb34
--- /dev/null
+++ b/streaming/src/test/java/org/apache/spark/streaming/JavaStreamingListenerAPISuite.java
@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+package org.apache.spark.streaming;
+
+import org.apache.spark.streaming.api.java.*;
+
+public class JavaStreamingListenerAPISuite extends JavaStreamingListener {
+
+  @Override
+  public void onReceiverStarted(JavaStreamingListenerReceiverStarted receiverStarted) {
+    JavaReceiverInfo receiverInfo = receiverStarted.receiverInfo();
+    receiverInfo.streamId();
+    receiverInfo.name();
+    receiverInfo.active();
+    receiverInfo.location();
+    receiverInfo.lastErrorMessage();
+    receiverInfo.lastError();
+    receiverInfo.lastErrorTime();
+  }
+
+  @Override
+  public void onReceiverError(JavaStreamingListenerReceiverError receiverError) {
+    JavaReceiverInfo receiverInfo = receiverError.receiverInfo();
+    receiverInfo.streamId();
+    receiverInfo.name();
+    receiverInfo.active();
+    receiverInfo.location();
+    receiverInfo.lastErrorMessage();
+    receiverInfo.lastError();
+    receiverInfo.lastErrorTime();
+  }
+
+  @Override
+  public void onReceiverStopped(JavaStreamingListenerReceiverStopped receiverStopped) {
+    JavaReceiverInfo receiverInfo = receiverStopped.receiverInfo();
+    receiverInfo.streamId();
+    receiverInfo.name();
+    receiverInfo.active();
+    receiverInfo.location();
+    receiverInfo.lastErrorMessage();
+    receiverInfo.lastError();
+    receiverInfo.lastErrorTime();
+  }
+
+  @Override
+  public void onBatchSubmitted(JavaStreamingListenerBatchSubmitted batchSubmitted) {
+    super.onBatchSubmitted(batchSubmitted);
+  }
+
+  @Override
+  public void onBatchStarted(JavaStreamingListenerBatchStarted batchStarted) {
+    super.onBatchStarted(batchStarted);
+  }
+
+  @Override
+  public void onBatchCompleted(JavaStreamingListenerBatchCompleted batchCompleted) {
+    super.onBatchCompleted(batchCompleted);
+  }
+
+  @Override
+  public void onOutputOperationStarted(JavaStreamingListenerOutputOperationStarted outputOperationStarted) {
+    super.onOutputOperationStarted(outputOperationStarted);
+  }
+
+  @Override
+  public void onOutputOperationCompleted(JavaStreamingListenerOutputOperationCompleted outputOperationCompleted) {
+    super.onOutputOperationCompleted(outputOperationCompleted);
+  }
+}
diff --git a/streaming/src/test/java/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapperSuite.scala b/streaming/src/test/java/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapperSuite.scala
new file mode 100644
index 0000000000000..6d6d61e70cafc
--- /dev/null
+++ b/streaming/src/test/java/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapperSuite.scala
@@ -0,0 +1,290 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.api.java
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.streaming.Time
+import org.apache.spark.streaming.scheduler._
+
+class JavaStreamingListenerWrapperSuite extends SparkFunSuite {
+
+  test("basic") {
+    val listener = new TestJavaStreamingListener()
+    val listenerWrapper = new JavaStreamingListenerWrapper(listener)
+
+    val receiverStarted = StreamingListenerReceiverStarted(ReceiverInfo(
+      streamId = 2,
+      name = "test",
+      active = true,
+      location = "localhost"
+    ))
+    listenerWrapper.onReceiverStarted(receiverStarted)
+    assertReceiverInfo(listener.receiverStarted.receiverInfo, receiverStarted.receiverInfo)
+
+    val receiverStopped = StreamingListenerReceiverStopped(ReceiverInfo(
+      streamId = 2,
+      name = "test",
+      active = false,
+      location = "localhost"
+    ))
+    listenerWrapper.onReceiverStopped(receiverStopped)
+    assertReceiverInfo(listener.receiverStopped.receiverInfo, receiverStopped.receiverInfo)
+
+    val receiverError = StreamingListenerReceiverError(ReceiverInfo(
+      streamId = 2,
+      name = "test",
+      active = false,
+      location = "localhost",
+      lastErrorMessage = "failed",
+      lastError = "failed",
+      lastErrorTime = System.currentTimeMillis()
+    ))
+    listenerWrapper.onReceiverError(receiverError)
+    assertReceiverInfo(listener.receiverError.receiverInfo, receiverError.receiverInfo)
+
+    val batchSubmitted = StreamingListenerBatchSubmitted(BatchInfo(
+      batchTime = Time(1000L),
+      streamIdToInputInfo = Map(
+        0 -> StreamInputInfo(
+          inputStreamId = 0,
+          numRecords = 1000,
+          metadata = Map(StreamInputInfo.METADATA_KEY_DESCRIPTION -> "receiver1")),
+        1 -> StreamInputInfo(
+          inputStreamId = 1,
+          numRecords = 2000,
+          metadata = Map(StreamInputInfo.METADATA_KEY_DESCRIPTION -> "receiver2"))),
+      submissionTime = 1001L,
+      None,
+      None,
+      outputOperationInfos = Map(
+        0 -> OutputOperationInfo(
+          batchTime = Time(1000L),
+          id = 0,
+          name = "op1",
+          description = "operation1",
+          startTime = None,
+          endTime = None,
+          failureReason = None),
+        1 -> OutputOperationInfo(
+          batchTime = Time(1000L),
+          id = 1,
+          name = "op2",
+          description = "operation2",
+          startTime = None,
+          endTime = None,
+          failureReason = None))
+    ))
+    listenerWrapper.onBatchSubmitted(batchSubmitted)
+    assertBatchInfo(listener.batchSubmitted.batchInfo, batchSubmitted.batchInfo)
+
+    val batchStarted = StreamingListenerBatchStarted(BatchInfo(
+      batchTime = Time(1000L),
+      streamIdToInputInfo = Map(
+        0 -> StreamInputInfo(
+          inputStreamId = 0,
+          numRecords = 1000,
+          metadata = Map(StreamInputInfo.METADATA_KEY_DESCRIPTION -> "receiver1")),
+        1 -> StreamInputInfo(
+          inputStreamId = 1,
+          numRecords = 2000,
+          metadata = Map(StreamInputInfo.METADATA_KEY_DESCRIPTION -> "receiver2"))),
+      submissionTime = 1001L,
+      Some(1002L),
+      None,
+      outputOperationInfos = Map(
+        0 -> OutputOperationInfo(
+          batchTime = Time(1000L),
+          id = 0,
+          name = "op1",
+          description = "operation1",
+          startTime = Some(1003L),
+          endTime = None,
+          failureReason = None),
+        1 -> OutputOperationInfo(
+          batchTime = Time(1000L),
+          id = 1,
+          name = "op2",
+          description = "operation2",
+          startTime = Some(1005L),
+          endTime = None,
+          failureReason = None))
+    ))
+    listenerWrapper.onBatchStarted(batchStarted)
+    assertBatchInfo(listener.batchStarted.batchInfo, batchStarted.batchInfo)
+
+    val batchCompleted = StreamingListenerBatchCompleted(BatchInfo(
+      batchTime = Time(1000L),
+      streamIdToInputInfo = Map(
+        0 -> StreamInputInfo(
+          inputStreamId = 0,
+          numRecords = 1000,
+          metadata = Map(StreamInputInfo.METADATA_KEY_DESCRIPTION -> "receiver1")),
+        1 -> StreamInputInfo(
+          inputStreamId = 1,
+          numRecords = 2000,
+          metadata = Map(StreamInputInfo.METADATA_KEY_DESCRIPTION -> "receiver2"))),
+      submissionTime = 1001L,
+      Some(1002L),
+      Some(1010L),
+      outputOperationInfos = Map(
+        0 -> OutputOperationInfo(
+          batchTime = Time(1000L),
+          id = 0,
+          name = "op1",
+          description = "operation1",
+          startTime = Some(1003L),
+          endTime = Some(1004L),
+          failureReason = None),
+        1 -> OutputOperationInfo(
+          batchTime = Time(1000L),
+          id = 1,
+          name = "op2",
+          description = "operation2",
+          startTime = Some(1005L),
+          endTime = Some(1010L),
+          failureReason = None))
+    ))
+    listenerWrapper.onBatchCompleted(batchCompleted)
+    assertBatchInfo(listener.batchCompleted.batchInfo, batchCompleted.batchInfo)
+
+    val outputOperationStarted = StreamingListenerOutputOperationStarted(OutputOperationInfo(
+      batchTime = Time(1000L),
+      id = 0,
+      name = "op1",
+      description = "operation1",
+      startTime = Some(1003L),
+      endTime = None,
+      failureReason = None
+    ))
+    listenerWrapper.onOutputOperationStarted(outputOperationStarted)
+    assertOutputOperationInfo(listener.outputOperationStarted.outputOperationInfo,
+      outputOperationStarted.outputOperationInfo)
+
+    val outputOperationCompleted = StreamingListenerOutputOperationCompleted(OutputOperationInfo(
+      batchTime = Time(1000L),
+      id = 0,
+      name = "op1",
+      description = "operation1",
+      startTime = Some(1003L),
+      endTime = Some(1004L),
+      failureReason = None
+    ))
+    listenerWrapper.onOutputOperationCompleted(outputOperationCompleted)
+    assertOutputOperationInfo(listener.outputOperationCompleted.outputOperationInfo,
+      outputOperationCompleted.outputOperationInfo)
+  }
+
+  private def assertReceiverInfo(
+      javaReceiverInfo: JavaReceiverInfo, receiverInfo: ReceiverInfo): Unit = {
+    assert(javaReceiverInfo.streamId === receiverInfo.streamId)
+    assert(javaReceiverInfo.name === receiverInfo.name)
+    assert(javaReceiverInfo.active === receiverInfo.active)
+    assert(javaReceiverInfo.location === receiverInfo.location)
+    assert(javaReceiverInfo.lastErrorMessage === receiverInfo.lastErrorMessage)
+    assert(javaReceiverInfo.lastError === receiverInfo.lastError)
+    assert(javaReceiverInfo.lastErrorTime === receiverInfo.lastErrorTime)
+  }
+
+  private def assertBatchInfo(javaBatchInfo: JavaBatchInfo, batchInfo: BatchInfo): Unit = {
+    assert(javaBatchInfo.batchTime === batchInfo.batchTime)
+    assert(javaBatchInfo.streamIdToInputInfo.size === batchInfo.streamIdToInputInfo.size)
+    batchInfo.streamIdToInputInfo.foreach { case (streamId, streamInputInfo) =>
+      assertStreamingInfo(javaBatchInfo.streamIdToInputInfo.get(streamId), streamInputInfo)
+    }
+    assert(javaBatchInfo.submissionTime === batchInfo.submissionTime)
+    assert(javaBatchInfo.processingStartTime === batchInfo.processingStartTime.getOrElse(-1))
+    assert(javaBatchInfo.processingEndTime === batchInfo.processingEndTime.getOrElse(-1))
+    assert(javaBatchInfo.schedulingDelay === batchInfo.schedulingDelay.getOrElse(-1))
+    assert(javaBatchInfo.processingDelay === batchInfo.processingDelay.getOrElse(-1))
+    assert(javaBatchInfo.totalDelay === batchInfo.totalDelay.getOrElse(-1))
+    assert(javaBatchInfo.numRecords === batchInfo.numRecords)
+    assert(javaBatchInfo.outputOperationInfos.size === batchInfo.outputOperationInfos.size)
+    batchInfo.outputOperationInfos.foreach { case (outputOperationId, outputOperationInfo) =>
+      assertOutputOperationInfo(
+        javaBatchInfo.outputOperationInfos.get(outputOperationId), outputOperationInfo)
+    }
+  }
+
+  private def assertStreamingInfo(
+      javaStreamInputInfo: JavaStreamInputInfo, streamInputInfo: StreamInputInfo): Unit = {
+    assert(javaStreamInputInfo.inputStreamId === streamInputInfo.inputStreamId)
+    assert(javaStreamInputInfo.numRecords === streamInputInfo.numRecords)
+    assert(javaStreamInputInfo.metadata === streamInputInfo.metadata.asJava)
+    assert(javaStreamInputInfo.metadataDescription === streamInputInfo.metadataDescription.orNull)
+  }
+
+  private def assertOutputOperationInfo(
+      javaOutputOperationInfo: JavaOutputOperationInfo,
+      outputOperationInfo: OutputOperationInfo): Unit = {
+    assert(javaOutputOperationInfo.batchTime === outputOperationInfo.batchTime)
+    assert(javaOutputOperationInfo.id === outputOperationInfo.id)
+    assert(javaOutputOperationInfo.name === outputOperationInfo.name)
+    assert(javaOutputOperationInfo.description === outputOperationInfo.description)
+    assert(javaOutputOperationInfo.startTime === outputOperationInfo.startTime.getOrElse(-1))
+    assert(javaOutputOperationInfo.endTime === outputOperationInfo.endTime.getOrElse(-1))
+    assert(javaOutputOperationInfo.failureReason === outputOperationInfo.failureReason.orNull)
+  }
+}
+
+class TestJavaStreamingListener extends JavaStreamingListener {
+
+  var receiverStarted: JavaStreamingListenerReceiverStarted = null
+  var receiverError: JavaStreamingListenerReceiverError = null
+  var receiverStopped: JavaStreamingListenerReceiverStopped = null
+  var batchSubmitted: JavaStreamingListenerBatchSubmitted = null
+  var batchStarted: JavaStreamingListenerBatchStarted = null
+  var batchCompleted: JavaStreamingListenerBatchCompleted = null
+  var outputOperationStarted: JavaStreamingListenerOutputOperationStarted = null
+  var outputOperationCompleted: JavaStreamingListenerOutputOperationCompleted = null
+
+  override def onReceiverStarted(receiverStarted: JavaStreamingListenerReceiverStarted): Unit = {
+    this.receiverStarted = receiverStarted
+  }
+
+  override def onReceiverError(receiverError: JavaStreamingListenerReceiverError): Unit = {
+    this.receiverError = receiverError
+  }
+
+  override def onReceiverStopped(receiverStopped: JavaStreamingListenerReceiverStopped): Unit = {
+    this.receiverStopped = receiverStopped
+  }
+
+  override def onBatchSubmitted(batchSubmitted: JavaStreamingListenerBatchSubmitted): Unit = {
+    this.batchSubmitted = batchSubmitted
+  }
+
+  override def onBatchStarted(batchStarted: JavaStreamingListenerBatchStarted): Unit = {
+    this.batchStarted = batchStarted
+  }
+
+  override def onBatchCompleted(batchCompleted: JavaStreamingListenerBatchCompleted): Unit = {
+    this.batchCompleted = batchCompleted
+  }
+
+  override def onOutputOperationStarted(
+      outputOperationStarted: JavaStreamingListenerOutputOperationStarted): Unit = {
+    this.outputOperationStarted = outputOperationStarted
+  }
+
+  override def onOutputOperationCompleted(
+      outputOperationCompleted: JavaStreamingListenerOutputOperationCompleted): Unit = {
+    this.outputOperationCompleted = outputOperationCompleted
+  }
+}

From 6502944f39893b9dfb472f8406d5f3a02a316eff Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Mon, 9 Nov 2015 18:13:37 -0800
Subject: [PATCH 148/518] [SPARK-11333][STREAMING] Add executorId to
 ReceiverInfo and display it in UI

Expose executorId to `ReceiverInfo` and UI since it's helpful when there are multiple executors running in the same host. Screenshot:

<img width="1058" alt="screen shot 2015-11-02 at 10 52 19 am" src="https://cloud.githubusercontent.com/assets/1000778/10890968/2e2f5512-8150-11e5-8d9d-746e826b69e8.png">

Author: Shixiong Zhu <shixiong@databricks.com>
Author: zsxwing <zsxwing@gmail.com>

Closes #9418 from zsxwing/SPARK-11333.
---
 .../spark/streaming/api/java/JavaStreamingListener.scala  | 1 +
 .../streaming/api/java/JavaStreamingListenerWrapper.scala | 1 +
 .../apache/spark/streaming/scheduler/ReceiverInfo.scala   | 1 +
 .../spark/streaming/scheduler/ReceiverTrackingInfo.scala  | 1 +
 .../org/apache/spark/streaming/ui/StreamingPage.scala     | 8 ++++++--
 .../spark/streaming/JavaStreamingListenerAPISuite.java    | 3 +++
 .../api/java/JavaStreamingListenerWrapperSuite.scala      | 8 ++++++--
 .../streaming/ui/StreamingJobProgressListenerSuite.scala  | 6 +++---
 8 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
index c86c7101ff6d5..34429074fe804 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
@@ -140,6 +140,7 @@ private[streaming] case class JavaReceiverInfo(
     name: String,
     active: Boolean,
     location: String,
+    executorId: String,
     lastErrorMessage: String,
     lastError: String,
     lastErrorTime: Long)
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapper.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapper.scala
index 2c60b396a6616..b109b9f1cbeae 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapper.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapper.scala
@@ -33,6 +33,7 @@ private[streaming] class JavaStreamingListenerWrapper(javaStreamingListener: Jav
       receiverInfo.name,
       receiverInfo.active,
       receiverInfo.location,
+      receiverInfo.executorId,
       receiverInfo.lastErrorMessage,
       receiverInfo.lastError,
       receiverInfo.lastErrorTime
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverInfo.scala b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverInfo.scala
index 59df892397fe0..3b35964114c02 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverInfo.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverInfo.scala
@@ -30,6 +30,7 @@ case class ReceiverInfo(
     name: String,
     active: Boolean,
     location: String,
+    executorId: String,
     lastErrorMessage: String = "",
     lastError: String = "",
     lastErrorTime: Long = -1L
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTrackingInfo.scala b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTrackingInfo.scala
index ab0a84f05214d..4dc5bb9c3bfbe 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTrackingInfo.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/ReceiverTrackingInfo.scala
@@ -49,6 +49,7 @@ private[streaming] case class ReceiverTrackingInfo(
     name.getOrElse(""),
     state == ReceiverState.ACTIVE,
     location = runningExecutor.map(_.host).getOrElse(""),
+    executorId = runningExecutor.map(_.executorId).getOrElse(""),
     lastErrorMessage = errorInfo.map(_.lastErrorMessage).getOrElse(""),
     lastError = errorInfo.map(_.lastError).getOrElse(""),
     lastErrorTime = errorInfo.map(_.lastErrorTime).getOrElse(-1L)
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/ui/StreamingPage.scala b/streaming/src/main/scala/org/apache/spark/streaming/ui/StreamingPage.scala
index 96d943e75d272..4588b2163cd44 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/ui/StreamingPage.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/ui/StreamingPage.scala
@@ -402,7 +402,7 @@ private[ui] class StreamingPage(parent: StreamingTab)
         <tr>
           <th style="width: 151px;"></th>
           <th style="width: 167px; padding: 8px 0 8px 0"><div style="margin: 0 8px 0 8px">Status</div></th>
-          <th style="width: 167px; padding: 8px 0 8px 0"><div style="margin: 0 8px 0 8px">Location</div></th>
+          <th style="width: 167px; padding: 8px 0 8px 0"><div style="margin: 0 8px 0 8px">Executor ID / Host</div></th>
           <th style="width: 166px; padding: 8px 0 8px 0"><div style="margin: 0 8px 0 8px">Last Error Time</div></th>
           <th>Last Error Message</th>
         </tr>
@@ -430,7 +430,11 @@ private[ui] class StreamingPage(parent: StreamingTab)
     val receiverActive = receiverInfo.map { info =>
       if (info.active) "ACTIVE" else "INACTIVE"
     }.getOrElse(emptyCell)
-    val receiverLocation = receiverInfo.map(_.location).getOrElse(emptyCell)
+    val receiverLocation = receiverInfo.map { info =>
+      val executorId = if (info.executorId.isEmpty) emptyCell else info.executorId
+      val location = if (info.location.isEmpty) emptyCell else info.location
+      s"$executorId / $location"
+    }.getOrElse(emptyCell)
     val receiverLastError = receiverInfo.map { info =>
       val msg = s"${info.lastErrorMessage} - ${info.lastError}"
       if (msg.size > 100) msg.take(97) + "..." else msg
diff --git a/streaming/src/test/java/org/apache/spark/streaming/JavaStreamingListenerAPISuite.java b/streaming/src/test/java/org/apache/spark/streaming/JavaStreamingListenerAPISuite.java
index 8cc285aa7fb34..67b2a0703e02b 100644
--- a/streaming/src/test/java/org/apache/spark/streaming/JavaStreamingListenerAPISuite.java
+++ b/streaming/src/test/java/org/apache/spark/streaming/JavaStreamingListenerAPISuite.java
@@ -29,6 +29,7 @@ public void onReceiverStarted(JavaStreamingListenerReceiverStarted receiverStart
     receiverInfo.name();
     receiverInfo.active();
     receiverInfo.location();
+    receiverInfo.executorId();
     receiverInfo.lastErrorMessage();
     receiverInfo.lastError();
     receiverInfo.lastErrorTime();
@@ -41,6 +42,7 @@ public void onReceiverError(JavaStreamingListenerReceiverError receiverError) {
     receiverInfo.name();
     receiverInfo.active();
     receiverInfo.location();
+    receiverInfo.executorId();
     receiverInfo.lastErrorMessage();
     receiverInfo.lastError();
     receiverInfo.lastErrorTime();
@@ -53,6 +55,7 @@ public void onReceiverStopped(JavaStreamingListenerReceiverStopped receiverStopp
     receiverInfo.name();
     receiverInfo.active();
     receiverInfo.location();
+    receiverInfo.executorId();
     receiverInfo.lastErrorMessage();
     receiverInfo.lastError();
     receiverInfo.lastErrorTime();
diff --git a/streaming/src/test/java/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapperSuite.scala b/streaming/src/test/java/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapperSuite.scala
index 6d6d61e70cafc..0295e059f7bc2 100644
--- a/streaming/src/test/java/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapperSuite.scala
+++ b/streaming/src/test/java/org/apache/spark/streaming/api/java/JavaStreamingListenerWrapperSuite.scala
@@ -33,7 +33,8 @@ class JavaStreamingListenerWrapperSuite extends SparkFunSuite {
       streamId = 2,
       name = "test",
       active = true,
-      location = "localhost"
+      location = "localhost",
+      executorId = "1"
     ))
     listenerWrapper.onReceiverStarted(receiverStarted)
     assertReceiverInfo(listener.receiverStarted.receiverInfo, receiverStarted.receiverInfo)
@@ -42,7 +43,8 @@ class JavaStreamingListenerWrapperSuite extends SparkFunSuite {
       streamId = 2,
       name = "test",
       active = false,
-      location = "localhost"
+      location = "localhost",
+      executorId = "1"
     ))
     listenerWrapper.onReceiverStopped(receiverStopped)
     assertReceiverInfo(listener.receiverStopped.receiverInfo, receiverStopped.receiverInfo)
@@ -52,6 +54,7 @@ class JavaStreamingListenerWrapperSuite extends SparkFunSuite {
       name = "test",
       active = false,
       location = "localhost",
+      executorId = "1",
       lastErrorMessage = "failed",
       lastError = "failed",
       lastErrorTime = System.currentTimeMillis()
@@ -197,6 +200,7 @@ class JavaStreamingListenerWrapperSuite extends SparkFunSuite {
     assert(javaReceiverInfo.name === receiverInfo.name)
     assert(javaReceiverInfo.active === receiverInfo.active)
     assert(javaReceiverInfo.location === receiverInfo.location)
+    assert(javaReceiverInfo.executorId === receiverInfo.executorId)
     assert(javaReceiverInfo.lastErrorMessage === receiverInfo.lastErrorMessage)
     assert(javaReceiverInfo.lastError === receiverInfo.lastError)
     assert(javaReceiverInfo.lastErrorTime === receiverInfo.lastErrorTime)
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/ui/StreamingJobProgressListenerSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/ui/StreamingJobProgressListenerSuite.scala
index af4718b4eb705..34cd7435569e1 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/ui/StreamingJobProgressListenerSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/ui/StreamingJobProgressListenerSuite.scala
@@ -130,20 +130,20 @@ class StreamingJobProgressListenerSuite extends TestSuiteBase with Matchers {
     listener.numTotalReceivedRecords should be (600)
 
     // onReceiverStarted
-    val receiverInfoStarted = ReceiverInfo(0, "test", true, "localhost")
+    val receiverInfoStarted = ReceiverInfo(0, "test", true, "localhost", "0")
     listener.onReceiverStarted(StreamingListenerReceiverStarted(receiverInfoStarted))
     listener.receiverInfo(0) should be (Some(receiverInfoStarted))
     listener.receiverInfo(1) should be (None)
 
     // onReceiverError
-    val receiverInfoError = ReceiverInfo(1, "test", true, "localhost")
+    val receiverInfoError = ReceiverInfo(1, "test", true, "localhost", "1")
     listener.onReceiverError(StreamingListenerReceiverError(receiverInfoError))
     listener.receiverInfo(0) should be (Some(receiverInfoStarted))
     listener.receiverInfo(1) should be (Some(receiverInfoError))
     listener.receiverInfo(2) should be (None)
 
     // onReceiverStopped
-    val receiverInfoStopped = ReceiverInfo(2, "test", true, "localhost")
+    val receiverInfoStopped = ReceiverInfo(2, "test", true, "localhost", "2")
     listener.onReceiverStopped(StreamingListenerReceiverStopped(receiverInfoStopped))
     listener.receiverInfo(0) should be (Some(receiverInfoStarted))
     listener.receiverInfo(1) should be (Some(receiverInfoError))

From 1431319e5bc46c7225a8edeeec482816d14a83b8 Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Mon, 9 Nov 2015 18:53:57 -0800
Subject: [PATCH 149/518] Add mockito as an explicit test dependency to
 spark-streaming

While sbt successfully compiles as it properly pulls the mockito dependency, maven builds have broken. We need this in ASAP.
tdas

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9584 from brkyvz/fix-master.
---
 streaming/pom.xml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/streaming/pom.xml b/streaming/pom.xml
index 145c8a7321c05..435e16db13ab4 100644
--- a/streaming/pom.xml
+++ b/streaming/pom.xml
@@ -93,6 +93,11 @@
       <artifactId>selenium-java</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <build>
     <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>

From c4e19b3819df4cd7a1c495a00bd2844cf55f4dbd Mon Sep 17 00:00:00 2001
From: Shivaram Venkataraman <shivaram@cs.berkeley.edu>
Date: Mon, 9 Nov 2015 21:06:01 -0800
Subject: [PATCH 150/518] [SPARK-11587][SPARKR] Fix the summary generic to
 match base R

The signature is summary(object, ...) as defined in
https://stat.ethz.ch/R-manual/R-devel/library/base/html/summary.html

Author: Shivaram Venkataraman <shivaram@cs.berkeley.edu>

Closes #9582 from shivaram/summary-fix.
---
 R/pkg/R/DataFrame.R           |  6 +++---
 R/pkg/R/generics.R            |  2 +-
 R/pkg/R/mllib.R               | 12 ++++++------
 R/pkg/inst/tests/test_mllib.R |  6 ++++++
 4 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index 44ce9414da5cf..e9013aa34a84f 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -1944,9 +1944,9 @@ setMethod("describe",
 #' @rdname summary
 #' @name summary
 setMethod("summary",
-          signature(x = "DataFrame"),
-          function(x) {
-            describe(x)
+          signature(object = "DataFrame"),
+          function(object, ...) {
+            describe(object)
           })
 
 
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 083d37fee28a4..efef7d66b522c 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -561,7 +561,7 @@ setGeneric("summarize", function(x,...) { standardGeneric("summarize") })
 
 #' @rdname summary
 #' @export
-setGeneric("summary", function(x, ...) { standardGeneric("summary") })
+setGeneric("summary", function(object, ...) { standardGeneric("summary") })
 
 # @rdname tojson
 # @export
diff --git a/R/pkg/R/mllib.R b/R/pkg/R/mllib.R
index 7ff859741b4a0..7126b7cde4bd7 100644
--- a/R/pkg/R/mllib.R
+++ b/R/pkg/R/mllib.R
@@ -89,17 +89,17 @@ setMethod("predict", signature(object = "PipelineModel"),
 #' model <- glm(y ~ x, trainingData)
 #' summary(model)
 #'}
-setMethod("summary", signature(x = "PipelineModel"),
-          function(x, ...) {
+setMethod("summary", signature(object = "PipelineModel"),
+          function(object, ...) {
             modelName <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
-                                   "getModelName", x@model)
+                                   "getModelName", object@model)
             features <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
-                                   "getModelFeatures", x@model)
+                                   "getModelFeatures", object@model)
             coefficients <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
-                                   "getModelCoefficients", x@model)
+                                   "getModelCoefficients", object@model)
             if (modelName == "LinearRegressionModel") {
               devianceResiduals <- callJStatic("org.apache.spark.ml.api.r.SparkRWrappers",
-                                               "getModelDevianceResiduals", x@model)
+                                               "getModelDevianceResiduals", object@model)
               devianceResiduals <- matrix(devianceResiduals, nrow = 1)
               colnames(devianceResiduals) <- c("Min", "Max")
               rownames(devianceResiduals) <- rep("", times = 1)
diff --git a/R/pkg/inst/tests/test_mllib.R b/R/pkg/inst/tests/test_mllib.R
index 2606407bdcb44..42287ea19adc5 100644
--- a/R/pkg/inst/tests/test_mllib.R
+++ b/R/pkg/inst/tests/test_mllib.R
@@ -113,3 +113,9 @@ test_that("summary coefficients match with native glm of family 'binomial'", {
     rownames(stats$Coefficients) ==
     c("(Intercept)", "Sepal_Length", "Sepal_Width")))
 })
+
+test_that("summary works on base GLM models", {
+  baseModel <- stats::glm(Sepal.Width ~ Sepal.Length + Species, data = iris)
+  baseSummary <- summary(baseModel)
+  expect_true(abs(baseSummary$deviance - 12.19313) < 1e-4)
+})

From d6cd3a18e720e8f6f1f307e0dffad3512952d997 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Mon, 9 Nov 2015 23:27:36 -0800
Subject: [PATCH 151/518] [SPARK-11599] [SQL] fix NPE when resolve Hive UDF in
 SQLParser

The DataFrame APIs that takes a SQL expression always use SQLParser, then the HiveFunctionRegistry will called outside of Hive state, cause NPE if there is not a active Session State for current thread (in PySpark).

cc rxin yhuai

Author: Davies Liu <davies@databricks.com>

Closes #9576 from davies/hive_udf.
---
 .../apache/spark/sql/hive/HiveContext.scala   | 10 +++++-
 .../sql/hive/execution/HiveQuerySuite.scala   | 33 ++++++++++++++-----
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index 2d72b959af134..c5f69657f5293 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -454,7 +454,15 @@ class HiveContext private[hive](
   // Note that HiveUDFs will be overridden by functions registered in this context.
   @transient
   override protected[sql] lazy val functionRegistry: FunctionRegistry =
-    new HiveFunctionRegistry(FunctionRegistry.builtin.copy())
+    new HiveFunctionRegistry(FunctionRegistry.builtin.copy()) {
+      override def lookupFunction(name: String, children: Seq[Expression]): Expression = {
+        // Hive Registry need current database to lookup function
+        // TODO: the current database of executionHive should be consistent with metadataHive
+        executionHive.withHiveState {
+          super.lookupFunction(name, children)
+        }
+      }
+    }
 
   // The Hive UDF current_database() is foldable, will be evaluated by optimizer, but the optimizer
   // can't access the SessionState of metadataHive.
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
index 78378c8b69c7a..f0a7a6cc7a1e3 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
@@ -20,22 +20,19 @@ package org.apache.spark.sql.hive.execution
 import java.io.File
 import java.util.{Locale, TimeZone}
 
-import org.apache.spark.sql.execution.joins.BroadcastNestedLoopJoin
-
 import scala.util.Try
 
-import org.scalatest.BeforeAndAfter
-
 import org.apache.hadoop.hive.conf.HiveConf.ConfVars
+import org.scalatest.BeforeAndAfter
 
-import org.apache.spark.{SparkFiles, SparkException}
-import org.apache.spark.sql.{AnalysisException, DataFrame, Row}
 import org.apache.spark.sql.catalyst.expressions.Cast
 import org.apache.spark.sql.catalyst.plans.logical.Project
+import org.apache.spark.sql.execution.joins.BroadcastNestedLoopJoin
 import org.apache.spark.sql.hive._
-import org.apache.spark.sql.hive.test.TestHiveContext
-import org.apache.spark.sql.hive.test.TestHive
 import org.apache.spark.sql.hive.test.TestHive._
+import org.apache.spark.sql.hive.test.{TestHive, TestHiveContext}
+import org.apache.spark.sql.{AnalysisException, DataFrame, Row}
+import org.apache.spark.{SparkException, SparkFiles}
 
 case class TestData(a: Int, b: String)
 
@@ -1237,6 +1234,26 @@ class HiveQuerySuite extends HiveComparisonTest with BeforeAndAfter {
 
   }
 
+  test("lookup hive UDF in another thread") {
+    val e = intercept[AnalysisException] {
+      range(1).selectExpr("not_a_udf()")
+    }
+    assert(e.getMessage.contains("undefined function not_a_udf"))
+    var success = false
+    val t = new Thread("test") {
+      override def run(): Unit = {
+        val e = intercept[AnalysisException] {
+          range(1).selectExpr("not_a_udf()")
+        }
+        assert(e.getMessage.contains("undefined function not_a_udf"))
+        success = true
+      }
+    }
+    t.start()
+    t.join()
+    assert(success)
+  }
+
   createQueryTest("select from thrift based table",
     "SELECT * from src_thrift")
 

From 521b3cae118d1e22c170e2aad43f9baa162db55e Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Mon, 9 Nov 2015 23:28:32 -0800
Subject: [PATCH 152/518] [SPARK-11598] [SQL] enable tests for
 ShuffledHashOuterJoin

Author: Davies Liu <davies@databricks.com>

Closes #9573 from davies/join_condition.
---
 .../org/apache/spark/sql/JoinSuite.scala      | 435 ++++++++++--------
 1 file changed, 231 insertions(+), 204 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
index a9ca46cab067d..3f3b837f7581c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
@@ -237,214 +237,241 @@ class JoinSuite extends QueryTest with SharedSQLContext {
         Row(2, 2, 2, 2) :: Nil)
   }
 
-  test("left outer join") {
-    checkAnswer(
-      upperCaseData.join(lowerCaseData, $"n" === $"N", "left"),
-      Row(1, "A", 1, "a") ::
-        Row(2, "B", 2, "b") ::
-        Row(3, "C", 3, "c") ::
-        Row(4, "D", 4, "d") ::
-        Row(5, "E", null, null) ::
-        Row(6, "F", null, null) :: Nil)
-
-    checkAnswer(
-      upperCaseData.join(lowerCaseData, $"n" === $"N" && $"n" > 1, "left"),
-      Row(1, "A", null, null) ::
-        Row(2, "B", 2, "b") ::
-        Row(3, "C", 3, "c") ::
-        Row(4, "D", 4, "d") ::
-        Row(5, "E", null, null) ::
-        Row(6, "F", null, null) :: Nil)
-
-    checkAnswer(
-      upperCaseData.join(lowerCaseData, $"n" === $"N" && $"N" > 1, "left"),
-      Row(1, "A", null, null) ::
-        Row(2, "B", 2, "b") ::
-        Row(3, "C", 3, "c") ::
-        Row(4, "D", 4, "d") ::
-        Row(5, "E", null, null) ::
-        Row(6, "F", null, null) :: Nil)
-
-    checkAnswer(
-      upperCaseData.join(lowerCaseData, $"n" === $"N" && $"l" > $"L", "left"),
-      Row(1, "A", 1, "a") ::
-        Row(2, "B", 2, "b") ::
-        Row(3, "C", 3, "c") ::
-        Row(4, "D", 4, "d") ::
-        Row(5, "E", null, null) ::
-        Row(6, "F", null, null) :: Nil)
-
-    // Make sure we are choosing left.outputPartitioning as the
-    // outputPartitioning for the outer join operator.
-    checkAnswer(
-      sql(
-        """
-          |SELECT l.N, count(*)
-          |FROM upperCaseData l LEFT OUTER JOIN allNulls r ON (l.N = r.a)
-          |GROUP BY l.N
-        """.stripMargin),
-      Row(1, 1) ::
-        Row(2, 1) ::
-        Row(3, 1) ::
-        Row(4, 1) ::
-        Row(5, 1) ::
-        Row(6, 1) :: Nil)
-
-    checkAnswer(
-      sql(
-        """
-          |SELECT r.a, count(*)
-          |FROM upperCaseData l LEFT OUTER JOIN allNulls r ON (l.N = r.a)
-          |GROUP BY r.a
-        """.stripMargin),
-      Row(null, 6) :: Nil)
-  }
+  def test_outer_join(useSMJ: Boolean): Unit = {
+
+    val algo = if (useSMJ) "SortMergeOuterJoin" else "ShuffledHashOuterJoin"
+
+    test("left outer join: " + algo) {
+      withSQLConf(SQLConf.SORTMERGE_JOIN.key -> useSMJ.toString) {
+
+        checkAnswer(
+          upperCaseData.join(lowerCaseData, $"n" === $"N", "left"),
+          Row(1, "A", 1, "a") ::
+            Row(2, "B", 2, "b") ::
+            Row(3, "C", 3, "c") ::
+            Row(4, "D", 4, "d") ::
+            Row(5, "E", null, null) ::
+            Row(6, "F", null, null) :: Nil)
+
+        checkAnswer(
+          upperCaseData.join(lowerCaseData, $"n" === $"N" && $"n" > 1, "left"),
+          Row(1, "A", null, null) ::
+            Row(2, "B", 2, "b") ::
+            Row(3, "C", 3, "c") ::
+            Row(4, "D", 4, "d") ::
+            Row(5, "E", null, null) ::
+            Row(6, "F", null, null) :: Nil)
+
+        checkAnswer(
+          upperCaseData.join(lowerCaseData, $"n" === $"N" && $"N" > 1, "left"),
+          Row(1, "A", null, null) ::
+            Row(2, "B", 2, "b") ::
+            Row(3, "C", 3, "c") ::
+            Row(4, "D", 4, "d") ::
+            Row(5, "E", null, null) ::
+            Row(6, "F", null, null) :: Nil)
+
+        checkAnswer(
+          upperCaseData.join(lowerCaseData, $"n" === $"N" && $"l" > $"L", "left"),
+          Row(1, "A", 1, "a") ::
+            Row(2, "B", 2, "b") ::
+            Row(3, "C", 3, "c") ::
+            Row(4, "D", 4, "d") ::
+            Row(5, "E", null, null) ::
+            Row(6, "F", null, null) :: Nil)
+
+        // Make sure we are choosing left.outputPartitioning as the
+        // outputPartitioning for the outer join operator.
+        checkAnswer(
+          sql(
+            """
+            |SELECT l.N, count(*)
+            |FROM upperCaseData l LEFT OUTER JOIN allNulls r ON (l.N = r.a)
+            |GROUP BY l.N
+          """.
+              stripMargin),
+        Row(1, 1) ::
+          Row(2, 1) ::
+          Row(3, 1) ::
+          Row(4, 1) ::
+          Row(5, 1) ::
+          Row(6, 1) :: Nil)
+
+        checkAnswer(
+          sql(
+            """
+              |SELECT r.a, count(*)
+              |FROM upperCaseData l LEFT OUTER JOIN allNulls r ON (l.N = r.a)
+              |GROUP BY r.a
+            """.stripMargin),
+          Row(null, 6) :: Nil)
+      }
+    }
 
-  test("right outer join") {
-    checkAnswer(
-      lowerCaseData.join(upperCaseData, $"n" === $"N", "right"),
-      Row(1, "a", 1, "A") ::
-        Row(2, "b", 2, "B") ::
-        Row(3, "c", 3, "C") ::
-        Row(4, "d", 4, "D") ::
-        Row(null, null, 5, "E") ::
-        Row(null, null, 6, "F") :: Nil)
-    checkAnswer(
-      lowerCaseData.join(upperCaseData, $"n" === $"N" && $"n" > 1, "right"),
-      Row(null, null, 1, "A") ::
-        Row(2, "b", 2, "B") ::
-        Row(3, "c", 3, "C") ::
-        Row(4, "d", 4, "D") ::
-        Row(null, null, 5, "E") ::
-        Row(null, null, 6, "F") :: Nil)
-    checkAnswer(
-      lowerCaseData.join(upperCaseData, $"n" === $"N" && $"N" > 1, "right"),
-      Row(null, null, 1, "A") ::
-        Row(2, "b", 2, "B") ::
-        Row(3, "c", 3, "C") ::
-        Row(4, "d", 4, "D") ::
-        Row(null, null, 5, "E") ::
-        Row(null, null, 6, "F") :: Nil)
-    checkAnswer(
-      lowerCaseData.join(upperCaseData, $"n" === $"N" && $"l" > $"L", "right"),
-      Row(1, "a", 1, "A") ::
-        Row(2, "b", 2, "B") ::
-        Row(3, "c", 3, "C") ::
-        Row(4, "d", 4, "D") ::
-        Row(null, null, 5, "E") ::
-        Row(null, null, 6, "F") :: Nil)
-
-    // Make sure we are choosing right.outputPartitioning as the
-    // outputPartitioning for the outer join operator.
-    checkAnswer(
-      sql(
-        """
-          |SELECT l.a, count(*)
-          |FROM allNulls l RIGHT OUTER JOIN upperCaseData r ON (l.a = r.N)
-          |GROUP BY l.a
-        """.stripMargin),
-      Row(null, 6))
+    test("right outer join: " + algo) {
+      withSQLConf(SQLConf.SORTMERGE_JOIN.key -> useSMJ.toString) {
+        checkAnswer(
+          lowerCaseData.join(upperCaseData, $"n" === $"N", "right"),
+          Row(1, "a", 1, "A") ::
+            Row(2, "b", 2, "B") ::
+            Row(3, "c", 3, "C") ::
+            Row(4, "d", 4, "D") ::
+            Row(null, null, 5, "E") ::
+            Row(null, null, 6, "F") :: Nil)
+        checkAnswer(
+          lowerCaseData.join(upperCaseData, $"n" === $"N" && $"n" > 1, "right"),
+          Row(null, null, 1, "A") ::
+            Row(2, "b", 2, "B") ::
+            Row(3, "c", 3, "C") ::
+            Row(4, "d", 4, "D") ::
+            Row(null, null, 5, "E") ::
+            Row(null, null, 6, "F") :: Nil)
+        checkAnswer(
+          lowerCaseData.join(upperCaseData, $"n" === $"N" && $"N" > 1, "right"),
+          Row(null, null, 1, "A") ::
+            Row(2, "b", 2, "B") ::
+            Row(3, "c", 3, "C") ::
+            Row(4, "d", 4, "D") ::
+            Row(null, null, 5, "E") ::
+            Row(null, null, 6, "F") :: Nil)
+        checkAnswer(
+          lowerCaseData.join(upperCaseData, $"n" === $"N" && $"l" > $"L", "right"),
+          Row(1, "a", 1, "A") ::
+            Row(2, "b", 2, "B") ::
+            Row(3, "c", 3, "C") ::
+            Row(4, "d", 4, "D") ::
+            Row(null, null, 5, "E") ::
+            Row(null, null, 6, "F") :: Nil)
+
+        // Make sure we are choosing right.outputPartitioning as the
+        // outputPartitioning for the outer join operator.
+        checkAnswer(
+          sql(
+            """
+              |SELECT l.a, count(*)
+              |FROM allNulls l RIGHT OUTER JOIN upperCaseData r ON (l.a = r.N)
+              |GROUP BY l.a
+            """.stripMargin),
+          Row(null,
+            6))
+
+        checkAnswer(
+          sql(
+            """
+              |SELECT r.N, count(*)
+              |FROM allNulls l RIGHT OUTER JOIN upperCaseData r ON (l.a = r.N)
+              |GROUP BY r.N
+            """.stripMargin),
+          Row(1
+            , 1) ::
+            Row(2, 1) ::
+            Row(3, 1) ::
+            Row(4, 1) ::
+            Row(5, 1) ::
+            Row(6, 1) :: Nil)
+        }
+      }
 
-    checkAnswer(
-      sql(
-        """
-          |SELECT r.N, count(*)
-          |FROM allNulls l RIGHT OUTER JOIN upperCaseData r ON (l.a = r.N)
-          |GROUP BY r.N
-        """.stripMargin),
-      Row(1, 1) ::
-        Row(2, 1) ::
-        Row(3, 1) ::
-        Row(4, 1) ::
-        Row(5, 1) ::
-        Row(6, 1) :: Nil)
+    test("full outer join: " + algo) {
+      withSQLConf(SQLConf.SORTMERGE_JOIN.key -> useSMJ.toString) {
+
+        upperCaseData.where('N <= 4).registerTempTable("left")
+        upperCaseData.where('N >= 3).registerTempTable("right")
+
+        val left = UnresolvedRelation(TableIdentifier("left"), None)
+        val right = UnresolvedRelation(TableIdentifier("right"), None)
+
+        checkAnswer(
+          left.join(right, $"left.N" === $"right.N", "full"),
+          Row(1, "A", null, null) ::
+            Row(2, "B", null, null) ::
+            Row(3, "C", 3, "C") ::
+            Row(4, "D", 4, "D") ::
+            Row(null, null, 5, "E") ::
+            Row(null, null, 6, "F") :: Nil)
+
+        checkAnswer(
+          left.join(right, ($"left.N" === $"right.N") && ($"left.N" !== 3), "full"),
+          Row(1, "A", null, null) ::
+            Row(2, "B", null, null) ::
+            Row(3, "C", null, null) ::
+            Row(null, null, 3, "C") ::
+            Row(4, "D", 4, "D") ::
+            Row(null, null, 5, "E") ::
+            Row(null, null, 6, "F") :: Nil)
+
+        checkAnswer(
+          left.join(right, ($"left.N" === $"right.N") && ($"right.N" !== 3), "full"),
+          Row(1, "A", null, null) ::
+            Row(2, "B", null, null) ::
+            Row(3, "C", null, null) ::
+            Row(null, null, 3, "C") ::
+            Row(4, "D", 4, "D") ::
+            Row(null, null, 5, "E") ::
+            Row(null, null, 6, "F") :: Nil)
+
+        // Make sure we are UnknownPartitioning as the outputPartitioning for the outer join
+        // operator.
+        checkAnswer(
+          sql(
+            """
+            |SELECT l.a, count(*)
+            |FROM allNulls l FULL OUTER JOIN upperCaseData r ON (l.a = r.N)
+            |GROUP BY l.a
+          """.
+              stripMargin),
+        Row(
+          null, 10))
+
+        checkAnswer(
+          sql(
+            """
+              |SELECT r.N, count(*)
+              |FROM allNulls l FULL OUTER JOIN upperCaseData r ON (l.a = r.N)
+              |GROUP BY r.N
+            """.stripMargin),
+          Row
+            (1, 1) ::
+            Row(2, 1) ::
+            Row(3, 1) ::
+            Row(4, 1) ::
+            Row(5, 1) ::
+            Row(6, 1) ::
+            Row(null, 4) :: Nil)
+
+        checkAnswer(
+          sql(
+            """
+              |SELECT l.N, count(*)
+              |FROM upperCaseData l FULL OUTER JOIN allNulls r ON (l.N = r.a)
+              |GROUP BY l.N
+            """.stripMargin),
+          Row(1
+            , 1) ::
+            Row(2, 1) ::
+            Row(3, 1) ::
+            Row(4, 1) ::
+            Row(5, 1) ::
+            Row(6, 1) ::
+            Row(null, 4) :: Nil)
+
+        checkAnswer(
+          sql(
+            """
+            |SELECT r.a, count(*)
+            |FROM upperCaseData l FULL OUTER JOIN allNulls r ON (l.N = r.a)
+            |GROUP BY r.a
+          """.
+              stripMargin),
+          Row(null, 10))
+      }
+    }
   }
 
-  test("full outer join") {
-    upperCaseData.where('N <= 4).registerTempTable("left")
-    upperCaseData.where('N >= 3).registerTempTable("right")
-
-    val left = UnresolvedRelation(TableIdentifier("left"), None)
-    val right = UnresolvedRelation(TableIdentifier("right"), None)
-
-    checkAnswer(
-      left.join(right, $"left.N" === $"right.N", "full"),
-      Row(1, "A", null, null) ::
-        Row(2, "B", null, null) ::
-        Row(3, "C", 3, "C") ::
-        Row(4, "D", 4, "D") ::
-        Row(null, null, 5, "E") ::
-        Row(null, null, 6, "F") :: Nil)
-
-    checkAnswer(
-      left.join(right, ($"left.N" === $"right.N") && ($"left.N" !== 3), "full"),
-      Row(1, "A", null, null) ::
-        Row(2, "B", null, null) ::
-        Row(3, "C", null, null) ::
-        Row(null, null, 3, "C") ::
-        Row(4, "D", 4, "D") ::
-        Row(null, null, 5, "E") ::
-        Row(null, null, 6, "F") :: Nil)
-
-    checkAnswer(
-      left.join(right, ($"left.N" === $"right.N") && ($"right.N" !== 3), "full"),
-      Row(1, "A", null, null) ::
-        Row(2, "B", null, null) ::
-        Row(3, "C", null, null) ::
-        Row(null, null, 3, "C") ::
-        Row(4, "D", 4, "D") ::
-        Row(null, null, 5, "E") ::
-        Row(null, null, 6, "F") :: Nil)
-
-    // Make sure we are UnknownPartitioning as the outputPartitioning for the outer join operator.
-    checkAnswer(
-      sql(
-        """
-          |SELECT l.a, count(*)
-          |FROM allNulls l FULL OUTER JOIN upperCaseData r ON (l.a = r.N)
-          |GROUP BY l.a
-        """.stripMargin),
-      Row(null, 10))
-
-    checkAnswer(
-      sql(
-        """
-          |SELECT r.N, count(*)
-          |FROM allNulls l FULL OUTER JOIN upperCaseData r ON (l.a = r.N)
-          |GROUP BY r.N
-        """.stripMargin),
-      Row(1, 1) ::
-        Row(2, 1) ::
-        Row(3, 1) ::
-        Row(4, 1) ::
-        Row(5, 1) ::
-        Row(6, 1) ::
-        Row(null, 4) :: Nil)
-
-    checkAnswer(
-      sql(
-        """
-          |SELECT l.N, count(*)
-          |FROM upperCaseData l FULL OUTER JOIN allNulls r ON (l.N = r.a)
-          |GROUP BY l.N
-        """.stripMargin),
-      Row(1, 1) ::
-        Row(2, 1) ::
-        Row(3, 1) ::
-        Row(4, 1) ::
-        Row(5, 1) ::
-        Row(6, 1) ::
-        Row(null, 4) :: Nil)
-
-    checkAnswer(
-      sql(
-        """
-          |SELECT r.a, count(*)
-          |FROM upperCaseData l FULL OUTER JOIN allNulls r ON (l.N = r.a)
-          |GROUP BY r.a
-        """.stripMargin),
-      Row(null, 10))
-  }
+  // test SortMergeOuterJoin
+  test_outer_join(true)
+  // test ShuffledHashOuterJoin
+  test_outer_join(false)
 
   test("broadcasted left semi join operator selection") {
     sqlContext.cacheManager.clearCache()

From 5507a9d0935aa42d65c3a4fa65da680b5af14faf Mon Sep 17 00:00:00 2001
From: Paul Chandler <pestilence669@users.noreply.github.com>
Date: Tue, 10 Nov 2015 12:59:53 +0100
Subject: [PATCH 153/518] Fix typo in driver page

"Comamnd property" => "Command property"

Author: Paul Chandler <pestilence669@users.noreply.github.com>

Closes #9578 from pestilence669/fix_spelling.
---
 .../scala/org/apache/spark/deploy/mesos/ui/DriverPage.scala     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/mesos/ui/DriverPage.scala b/core/src/main/scala/org/apache/spark/deploy/mesos/ui/DriverPage.scala
index e8ef60bd5428a..bc67fd460d9a9 100644
--- a/core/src/main/scala/org/apache/spark/deploy/mesos/ui/DriverPage.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/mesos/ui/DriverPage.scala
@@ -46,7 +46,7 @@ private[ui] class DriverPage(parent: MesosClusterUI) extends WebUIPage("driver")
     val schedulerHeaders = Seq("Scheduler property", "Value")
     val commandEnvHeaders = Seq("Command environment variable", "Value")
     val launchedHeaders = Seq("Launched property", "Value")
-    val commandHeaders = Seq("Comamnd property", "Value")
+    val commandHeaders = Seq("Command property", "Value")
     val retryHeaders = Seq("Last failed status", "Next retry time", "Retry count")
     val driverDescription = Iterable.apply(driverState.description)
     val submissionState = Iterable.apply(driverState.submissionState)

From a81f47ff7498e7063c855ccf75bba81ab101b43e Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Tue, 10 Nov 2015 10:05:53 -0800
Subject: [PATCH 154/518] [SPARK-11382] Replace example code in
 mllib-decision-tree.md using include_example

https://issues.apache.org/jira/browse/SPARK-11382

B.T.W. I fix an error in naive_bayes_example.py.

Author: Xusen Yin <yinxusen@gmail.com>

Closes #9596 from yinxusen/SPARK-11382.
---
 docs/mllib-decision-tree.md                   | 253 +-----------------
 ...JavaDecisionTreeClassificationExample.java |  91 +++++++
 .../JavaDecisionTreeRegressionExample.java    |  96 +++++++
 .../decision_tree_classification_example.py   |  55 ++++
 .../mllib/decision_tree_regression_example.py |  56 ++++
 .../main/python/mllib/naive_bayes_example.py  |   1 +
 .../DecisionTreeClassificationExample.scala   |  67 +++++
 .../mllib/DecisionTreeRegressionExample.scala |  66 +++++
 8 files changed, 438 insertions(+), 247 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeClassificationExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeRegressionExample.java
 create mode 100644 examples/src/main/python/mllib/decision_tree_classification_example.py
 create mode 100644 examples/src/main/python/mllib/decision_tree_regression_example.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/DecisionTreeClassificationExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/DecisionTreeRegressionExample.scala

diff --git a/docs/mllib-decision-tree.md b/docs/mllib-decision-tree.md
index b5b454bc69245..77ce34e91af3c 100644
--- a/docs/mllib-decision-tree.md
+++ b/docs/mllib-decision-tree.md
@@ -194,137 +194,19 @@ maximum tree depth of 5. The test error is calculated to measure the algorithm a
 <div data-lang="scala" markdown="1">
 Refer to the [`DecisionTree` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.DecisionTree) and [`DecisionTreeModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.model.DecisionTreeModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.tree.DecisionTree
-import org.apache.spark.mllib.tree.model.DecisionTreeModel
-import org.apache.spark.mllib.util.MLUtils
-
-// Load and parse the data file.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-// Split the data into training and test sets (30% held out for testing)
-val splits = data.randomSplit(Array(0.7, 0.3))
-val (trainingData, testData) = (splits(0), splits(1))
-
-// Train a DecisionTree model.
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-val numClasses = 2
-val categoricalFeaturesInfo = Map[Int, Int]()
-val impurity = "gini"
-val maxDepth = 5
-val maxBins = 32
-
-val model = DecisionTree.trainClassifier(trainingData, numClasses, categoricalFeaturesInfo,
-  impurity, maxDepth, maxBins)
-
-// Evaluate model on test instances and compute test error
-val labelAndPreds = testData.map { point =>
-  val prediction = model.predict(point.features)
-  (point.label, prediction)
-}
-val testErr = labelAndPreds.filter(r => r._1 != r._2).count.toDouble / testData.count()
-println("Test Error = " + testErr)
-println("Learned classification tree model:\n" + model.toDebugString)
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = DecisionTreeModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/DecisionTreeClassificationExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`DecisionTree` Java docs](api/java/org/apache/spark/mllib/tree/DecisionTree.html) and [`DecisionTreeModel` Java docs](api/java/org/apache/spark/mllib/tree/model/DecisionTreeModel.html) for details on the API.
 
-{% highlight java %}
-import java.util.HashMap;
-import scala.Tuple2;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.DecisionTree;
-import org.apache.spark.mllib.tree.model.DecisionTreeModel;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.SparkConf;
-
-SparkConf sparkConf = new SparkConf().setAppName("JavaDecisionTree");
-JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-// Load and parse the data file.
-String datapath = "data/mllib/sample_libsvm_data.txt";
-JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD();
-// Split the data into training and test sets (30% held out for testing)
-JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
-JavaRDD<LabeledPoint> trainingData = splits[0];
-JavaRDD<LabeledPoint> testData = splits[1];
-
-// Set parameters.
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-Integer numClasses = 2;
-Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-String impurity = "gini";
-Integer maxDepth = 5;
-Integer maxBins = 32;
-
-// Train a DecisionTree model for classification.
-final DecisionTreeModel model = DecisionTree.trainClassifier(trainingData, numClasses,
-  categoricalFeaturesInfo, impurity, maxDepth, maxBins);
-
-// Evaluate model on test instances and compute test error
-JavaPairRDD<Double, Double> predictionAndLabel =
-  testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-    @Override
-    public Tuple2<Double, Double> call(LabeledPoint p) {
-      return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-    }
-  });
-Double testErr =
-  1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
-    @Override
-    public Boolean call(Tuple2<Double, Double> pl) {
-      return !pl._1().equals(pl._2());
-    }
-  }).count() / testData.count();
-System.out.println("Test Error: " + testErr);
-System.out.println("Learned classification tree model:\n" + model.toDebugString());
-
-// Save and load model
-model.save(sc.sc(), "myModelPath");
-DecisionTreeModel sameModel = DecisionTreeModel.load(sc.sc(), "myModelPath");
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaDecisionTreeClassificationExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`DecisionTree` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.DecisionTree) and [`DecisionTreeModel` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.DecisionTreeModel) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.regression import LabeledPoint
-from pyspark.mllib.tree import DecisionTree, DecisionTreeModel
-from pyspark.mllib.util import MLUtils
-
-# Load and parse the data file into an RDD of LabeledPoint.
-data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a DecisionTree model.
-#  Empty categoricalFeaturesInfo indicates all features are continuous.
-model = DecisionTree.trainClassifier(trainingData, numClasses=2, categoricalFeaturesInfo={},
-                                     impurity='gini', maxDepth=5, maxBins=32)
-
-# Evaluate model on test instances and compute test error
-predictions = model.predict(testData.map(lambda x: x.features))
-labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-testErr = labelsAndPredictions.filter(lambda (v, p): v != p).count() / float(testData.count())
-print('Test Error = ' + str(testErr))
-print('Learned classification tree model:')
-print(model.toDebugString())
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = DecisionTreeModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example python/mllib/decision_tree_classification_example.py %}
 </div>
 
 </div>
@@ -343,142 +225,19 @@ depth of 5. The Mean Squared Error (MSE) is computed at the end to evaluate
 <div data-lang="scala" markdown="1">
 Refer to the [`DecisionTree` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.DecisionTree) and [`DecisionTreeModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.model.DecisionTreeModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.tree.DecisionTree
-import org.apache.spark.mllib.tree.model.DecisionTreeModel
-import org.apache.spark.mllib.util.MLUtils
-
-// Load and parse the data file.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-// Split the data into training and test sets (30% held out for testing)
-val splits = data.randomSplit(Array(0.7, 0.3))
-val (trainingData, testData) = (splits(0), splits(1))
-
-// Train a DecisionTree model.
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-val categoricalFeaturesInfo = Map[Int, Int]()
-val impurity = "variance"
-val maxDepth = 5
-val maxBins = 32
-
-val model = DecisionTree.trainRegressor(trainingData, categoricalFeaturesInfo, impurity,
-  maxDepth, maxBins)
-
-// Evaluate model on test instances and compute test error
-val labelsAndPredictions = testData.map { point =>
-  val prediction = model.predict(point.features)
-  (point.label, prediction)
-}
-val testMSE = labelsAndPredictions.map{ case(v, p) => math.pow((v - p), 2)}.mean()
-println("Test Mean Squared Error = " + testMSE)
-println("Learned regression tree model:\n" + model.toDebugString)
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = DecisionTreeModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/DecisionTreeRegressionExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`DecisionTree` Java docs](api/java/org/apache/spark/mllib/tree/DecisionTree.html) and [`DecisionTreeModel` Java docs](api/java/org/apache/spark/mllib/tree/model/DecisionTreeModel.html) for details on the API.
 
-{% highlight java %}
-import java.util.HashMap;
-import scala.Tuple2;
-import org.apache.spark.api.java.function.Function2;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.DecisionTree;
-import org.apache.spark.mllib.tree.model.DecisionTreeModel;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.SparkConf;
-
-SparkConf sparkConf = new SparkConf().setAppName("JavaDecisionTree");
-JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-// Load and parse the data file.
-String datapath = "data/mllib/sample_libsvm_data.txt";
-JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD();
-// Split the data into training and test sets (30% held out for testing)
-JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
-JavaRDD<LabeledPoint> trainingData = splits[0];
-JavaRDD<LabeledPoint> testData = splits[1];
-
-// Set parameters.
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-String impurity = "variance";
-Integer maxDepth = 5;
-Integer maxBins = 32;
-
-// Train a DecisionTree model.
-final DecisionTreeModel model = DecisionTree.trainRegressor(trainingData,
-  categoricalFeaturesInfo, impurity, maxDepth, maxBins);
-
-// Evaluate model on test instances and compute test error
-JavaPairRDD<Double, Double> predictionAndLabel =
-  testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-    @Override
-    public Tuple2<Double, Double> call(LabeledPoint p) {
-      return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-    }
-  });
-Double testMSE =
-  predictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
-    @Override
-    public Double call(Tuple2<Double, Double> pl) {
-      Double diff = pl._1() - pl._2();
-      return diff * diff;
-    }
-  }).reduce(new Function2<Double, Double, Double>() {
-    @Override
-    public Double call(Double a, Double b) {
-      return a + b;
-    }
-  }) / testData.count();
-System.out.println("Test Mean Squared Error: " + testMSE);
-System.out.println("Learned regression tree model:\n" + model.toDebugString());
-
-// Save and load model
-model.save(sc.sc(), "myModelPath");
-DecisionTreeModel sameModel = DecisionTreeModel.load(sc.sc(), "myModelPath");
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaDecisionTreeRegressionExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`DecisionTree` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.DecisionTree) and [`DecisionTreeModel` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.DecisionTreeModel) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.regression import LabeledPoint
-from pyspark.mllib.tree import DecisionTree, DecisionTreeModel
-from pyspark.mllib.util import MLUtils
-
-# Load and parse the data file into an RDD of LabeledPoint.
-data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a DecisionTree model.
-#  Empty categoricalFeaturesInfo indicates all features are continuous.
-model = DecisionTree.trainRegressor(trainingData, categoricalFeaturesInfo={},
-                                    impurity='variance', maxDepth=5, maxBins=32)
-
-# Evaluate model on test instances and compute test error
-predictions = model.predict(testData.map(lambda x: x.features))
-labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-testMSE = labelsAndPredictions.map(lambda (v, p): (v - p) * (v - p)).sum() / float(testData.count())
-print('Test Mean Squared Error = ' + str(testMSE))
-print('Learned regression tree model:')
-print(model.toDebugString())
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = DecisionTreeModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example python/mllib/decision_tree_regression_example.py %}
 </div>
 
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeClassificationExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeClassificationExample.java
new file mode 100644
index 0000000000000..5839b0cf8a8f8
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeClassificationExample.java
@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.HashMap;
+import java.util.Map;
+
+import scala.Tuple2;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.tree.DecisionTree;
+import org.apache.spark.mllib.tree.model.DecisionTreeModel;
+import org.apache.spark.mllib.util.MLUtils;
+// $example off$
+
+class JavaDecisionTreeClassificationExample {
+
+  public static void main(String[] args) {
+
+    // $example on$
+    SparkConf sparkConf = new SparkConf().setAppName("JavaDecisionTreeClassificationExample");
+    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
+
+    // Load and parse the data file.
+    String datapath = "data/mllib/sample_libsvm_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD();
+    // Split the data into training and test sets (30% held out for testing)
+    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
+    JavaRDD<LabeledPoint> trainingData = splits[0];
+    JavaRDD<LabeledPoint> testData = splits[1];
+
+    // Set parameters.
+    //  Empty categoricalFeaturesInfo indicates all features are continuous.
+    Integer numClasses = 2;
+    Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
+    String impurity = "gini";
+    Integer maxDepth = 5;
+    Integer maxBins = 32;
+
+    // Train a DecisionTree model for classification.
+    final DecisionTreeModel model = DecisionTree.trainClassifier(trainingData, numClasses,
+      categoricalFeaturesInfo, impurity, maxDepth, maxBins);
+
+    // Evaluate model on test instances and compute test error
+    JavaPairRDD<Double, Double> predictionAndLabel =
+      testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
+        @Override
+        public Tuple2<Double, Double> call(LabeledPoint p) {
+          return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
+        }
+      });
+    Double testErr =
+      1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
+        @Override
+        public Boolean call(Tuple2<Double, Double> pl) {
+          return !pl._1().equals(pl._2());
+        }
+      }).count() / testData.count();
+
+    System.out.println("Test Error: " + testErr);
+    System.out.println("Learned classification tree model:\n" + model.toDebugString());
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myDecisionTreeClassificationModel");
+    DecisionTreeModel sameModel = DecisionTreeModel
+      .load(jsc.sc(), "target/tmp/myDecisionTreeClassificationModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeRegressionExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeRegressionExample.java
new file mode 100644
index 0000000000000..ccde578249f7c
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTreeRegressionExample.java
@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.HashMap;
+import java.util.Map;
+
+import scala.Tuple2;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.Function2;
+import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.tree.DecisionTree;
+import org.apache.spark.mllib.tree.model.DecisionTreeModel;
+import org.apache.spark.mllib.util.MLUtils;
+// $example off$
+
+class JavaDecisionTreeRegressionExample {
+
+  public static void main(String[] args) {
+
+    // $example on$
+    SparkConf sparkConf = new SparkConf().setAppName("JavaDecisionTreeRegressionExample");
+    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
+
+    // Load and parse the data file.
+    String datapath = "data/mllib/sample_libsvm_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD();
+    // Split the data into training and test sets (30% held out for testing)
+    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
+    JavaRDD<LabeledPoint> trainingData = splits[0];
+    JavaRDD<LabeledPoint> testData = splits[1];
+
+    // Set parameters.
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
+    String impurity = "variance";
+    Integer maxDepth = 5;
+    Integer maxBins = 32;
+
+    // Train a DecisionTree model.
+    final DecisionTreeModel model = DecisionTree.trainRegressor(trainingData,
+      categoricalFeaturesInfo, impurity, maxDepth, maxBins);
+
+    // Evaluate model on test instances and compute test error
+    JavaPairRDD<Double, Double> predictionAndLabel =
+      testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
+      @Override
+      public Tuple2<Double, Double> call(LabeledPoint p) {
+        return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
+      }
+    });
+    Double testMSE =
+      predictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
+        @Override
+        public Double call(Tuple2<Double, Double> pl) {
+          Double diff = pl._1() - pl._2();
+          return diff * diff;
+        }
+      }).reduce(new Function2<Double, Double, Double>() {
+        @Override
+        public Double call(Double a, Double b) {
+          return a + b;
+        }
+      }) / data.count();
+    System.out.println("Test Mean Squared Error: " + testMSE);
+    System.out.println("Learned regression tree model:\n" + model.toDebugString());
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myDecisionTreeRegressionModel");
+    DecisionTreeModel sameModel = DecisionTreeModel
+      .load(jsc.sc(), "target/tmp/myDecisionTreeRegressionModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/python/mllib/decision_tree_classification_example.py b/examples/src/main/python/mllib/decision_tree_classification_example.py
new file mode 100644
index 0000000000000..1b529768b6c62
--- /dev/null
+++ b/examples/src/main/python/mllib/decision_tree_classification_example.py
@@ -0,0 +1,55 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Decision Tree Classification Example.
+"""
+from __future__ import print_function
+
+from pyspark import SparkContext
+# $example on$
+from pyspark.mllib.tree import DecisionTree, DecisionTreeModel
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+
+    sc = SparkContext(appName="PythonDecisionTreeClassificationExample")
+
+    # $example on$
+    # Load and parse the data file into an RDD of LabeledPoint.
+    data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a DecisionTree model.
+    #  Empty categoricalFeaturesInfo indicates all features are continuous.
+    model = DecisionTree.trainClassifier(trainingData, numClasses=2, categoricalFeaturesInfo={},
+                                         impurity='gini', maxDepth=5, maxBins=32)
+
+    # Evaluate model on test instances and compute test error
+    predictions = model.predict(testData.map(lambda x: x.features))
+    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
+    testErr = labelsAndPredictions.filter(lambda (v, p): v != p).count() / float(testData.count())
+    print('Test Error = ' + str(testErr))
+    print('Learned classification tree model:')
+    print(model.toDebugString())
+
+    # Save and load model
+    model.save(sc, "target/tmp/myDecisionTreeClassificationModel")
+    sameModel = DecisionTreeModel.load(sc, "target/tmp/myDecisionTreeClassificationModel")
+    # $example off$
diff --git a/examples/src/main/python/mllib/decision_tree_regression_example.py b/examples/src/main/python/mllib/decision_tree_regression_example.py
new file mode 100644
index 0000000000000..cf518eac67e81
--- /dev/null
+++ b/examples/src/main/python/mllib/decision_tree_regression_example.py
@@ -0,0 +1,56 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Decision Tree Regression Example.
+"""
+from __future__ import print_function
+
+from pyspark import SparkContext
+# $example on$
+from pyspark.mllib.tree import DecisionTree, DecisionTreeModel
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+
+    sc = SparkContext(appName="PythonDecisionTreeRegressionExample")
+
+    # $example on$
+    # Load and parse the data file into an RDD of LabeledPoint.
+    data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a DecisionTree model.
+    #  Empty categoricalFeaturesInfo indicates all features are continuous.
+    model = DecisionTree.trainRegressor(trainingData, categoricalFeaturesInfo={},
+                                        impurity='variance', maxDepth=5, maxBins=32)
+
+    # Evaluate model on test instances and compute test error
+    predictions = model.predict(testData.map(lambda x: x.features))
+    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
+    testMSE = labelsAndPredictions.map(lambda (v, p): (v - p) * (v - p)).sum() /\
+        float(testData.count())
+    print('Test Mean Squared Error = ' + str(testMSE))
+    print('Learned regression tree model:')
+    print(model.toDebugString())
+
+    # Save and load model
+    model.save(sc, "target/tmp/myDecisionTreeRegressionModel")
+    sameModel = DecisionTreeModel.load(sc, "target/tmp/myDecisionTreeRegressionModel")
+    # $example off$
diff --git a/examples/src/main/python/mllib/naive_bayes_example.py b/examples/src/main/python/mllib/naive_bayes_example.py
index a2e7dacf25491..f5e120c678fcf 100644
--- a/examples/src/main/python/mllib/naive_bayes_example.py
+++ b/examples/src/main/python/mllib/naive_bayes_example.py
@@ -20,6 +20,7 @@
 """
 from __future__ import print_function
 
+from pyspark import SparkContext
 # $example on$
 from pyspark.mllib.classification import NaiveBayes, NaiveBayesModel
 from pyspark.mllib.linalg import Vectors
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/DecisionTreeClassificationExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/DecisionTreeClassificationExample.scala
new file mode 100644
index 0000000000000..d427bbadaa0c1
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/DecisionTreeClassificationExample.scala
@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.tree.DecisionTree
+import org.apache.spark.mllib.tree.model.DecisionTreeModel
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+import org.apache.spark.{SparkConf, SparkContext}
+
+object DecisionTreeClassificationExample {
+
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("DecisionTreeClassificationExample")
+    val sc = new SparkContext(conf)
+
+    // $example on$
+    // Load and parse the data file.
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    // Split the data into training and test sets (30% held out for testing)
+    val splits = data.randomSplit(Array(0.7, 0.3))
+    val (trainingData, testData) = (splits(0), splits(1))
+
+    // Train a DecisionTree model.
+    //  Empty categoricalFeaturesInfo indicates all features are continuous.
+    val numClasses = 2
+    val categoricalFeaturesInfo = Map[Int, Int]()
+    val impurity = "gini"
+    val maxDepth = 5
+    val maxBins = 32
+
+    val model = DecisionTree.trainClassifier(trainingData, numClasses, categoricalFeaturesInfo,
+      impurity, maxDepth, maxBins)
+
+    // Evaluate model on test instances and compute test error
+    val labelAndPreds = testData.map { point =>
+      val prediction = model.predict(point.features)
+      (point.label, prediction)
+    }
+    val testErr = labelAndPreds.filter(r => r._1 != r._2).count().toDouble / testData.count()
+    println("Test Error = " + testErr)
+    println("Learned classification tree model:\n" + model.toDebugString)
+
+    // Save and load model
+    model.save(sc, "target/tmp/myDecisionTreeClassificationModel")
+    val sameModel = DecisionTreeModel.load(sc, "target/tmp/myDecisionTreeClassificationModel")
+    // $example off$
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/DecisionTreeRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/DecisionTreeRegressionExample.scala
new file mode 100644
index 0000000000000..fb05e7d9c5065
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/DecisionTreeRegressionExample.scala
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.tree.DecisionTree
+import org.apache.spark.mllib.tree.model.DecisionTreeModel
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+import org.apache.spark.{SparkConf, SparkContext}
+
+object DecisionTreeRegressionExample {
+
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("DecisionTreeRegressionExample")
+    val sc = new SparkContext(conf)
+
+    // $example on$
+    // Load and parse the data file.
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    // Split the data into training and test sets (30% held out for testing)
+    val splits = data.randomSplit(Array(0.7, 0.3))
+    val (trainingData, testData) = (splits(0), splits(1))
+
+    // Train a DecisionTree model.
+    //  Empty categoricalFeaturesInfo indicates all features are continuous.
+    val categoricalFeaturesInfo = Map[Int, Int]()
+    val impurity = "variance"
+    val maxDepth = 5
+    val maxBins = 32
+
+    val model = DecisionTree.trainRegressor(trainingData, categoricalFeaturesInfo, impurity,
+      maxDepth, maxBins)
+
+    // Evaluate model on test instances and compute test error
+    val labelsAndPredictions = testData.map { point =>
+      val prediction = model.predict(point.features)
+      (point.label, prediction)
+    }
+    val testMSE = labelsAndPredictions.map{ case (v, p) => math.pow(v - p, 2) }.mean()
+    println("Test Mean Squared Error = " + testMSE)
+    println("Learned regression tree model:\n" + model.toDebugString)
+
+    // Save and load model
+    model.save(sc, "target/tmp/myDecisionTreeRegressionModel")
+    val sameModel = DecisionTreeModel.load(sc, "target/tmp/myDecisionTreeRegressionModel")
+    // $example off$
+  }
+}
+// scalastyle:on println

From 689386b1c60997e4505749915f7005a52c207de2 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Tue, 10 Nov 2015 10:14:19 -0800
Subject: [PATCH 155/518] [SPARK-7841][BUILD] Stop using retrieveManaged to
 retrieve dependencies in SBT

This patch modifies Spark's SBT build so that it no longer uses `retrieveManaged` / `lib_managed` to store its dependencies. The motivations for this change are nicely described on the JIRA ticket ([SPARK-7841](https://issues.apache.org/jira/browse/SPARK-7841)); my personal interest in doing this stems from the fact that `lib_managed` has caused me some pain while debugging dependency issues in another PR of mine.

Removing our use of `lib_managed` would be trivial except for one snag: the Datanucleus JARs, required by Spark SQL's Hive integration, cannot be included in assembly JARs due to problems with merging OSGI `plugin.xml` files. As a result, several places in the packaging and deployment pipeline assume that these Datanucleus JARs are copied to `lib_managed/jars`. In the interest of maintaining compatibility, I have chosen to retain the `lib_managed/jars` directory _only_ for these Datanucleus JARs and have added custom code to `SparkBuild.scala` to automatically copy those JARs to that folder as part of the `assembly` task.

`dev/mima` also depended on `lib_managed` in a hacky way in order to set classpaths when generating MiMa excludes; I've updated this to obtain the classpaths directly from SBT instead.

/cc dragos marmbrus pwendell srowen

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9575 from JoshRosen/SPARK-7841.
---
 dev/mima                 |  2 +-
 project/SparkBuild.scala | 22 +++++++++++++++++-----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/dev/mima b/dev/mima
index 2952fa65d42ff..d5baffc6ef8a3 100755
--- a/dev/mima
+++ b/dev/mima
@@ -38,7 +38,7 @@ generate_mima_ignore() {
 # it did not process the new classes (which are in assembly jar).
 generate_mima_ignore
 
-export SPARK_CLASSPATH="`find lib_managed \( -name '*spark*jar' -a -type f \) | tr "\\n" ":"`"
+export SPARK_CLASSPATH="$(build/sbt "export oldDeps/fullClasspath" | tail -n1)"
 echo "SPARK_CLASSPATH=$SPARK_CLASSPATH"
 
 generate_mima_ignore
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index b75ed13a78c68..a9fb741d75933 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -16,6 +16,7 @@
  */
 
 import java.io._
+import java.nio.file.Files
 
 import scala.util.Properties
 import scala.collection.JavaConverters._
@@ -135,8 +136,6 @@ object SparkBuild extends PomBuild {
       .orElse(sys.props.get("java.home").map { p => new File(p).getParentFile().getAbsolutePath() })
       .map(file),
     incOptions := incOptions.value.withNameHashing(true),
-    retrieveManaged := true,
-    retrievePattern := "[type]s/[artifact](-[revision])(-[classifier]).[ext]",
     publishMavenStyle := true,
     unidocGenjavadocVersion := "0.9-spark0",
 
@@ -326,8 +325,6 @@ object OldDeps {
   def oldDepsSettings() = Defaults.coreDefaultSettings ++ Seq(
     name := "old-deps",
     scalaVersion := "2.10.5",
-    retrieveManaged := true,
-    retrievePattern := "[type]s/[artifact](-[revision])(-[classifier]).[ext]",
     libraryDependencies := Seq("spark-streaming-mqtt", "spark-streaming-zeromq",
       "spark-streaming-flume", "spark-streaming-kafka", "spark-streaming-twitter",
       "spark-streaming", "spark-mllib", "spark-bagel", "spark-graphx",
@@ -404,6 +401,8 @@ object Assembly {
 
   val hadoopVersion = taskKey[String]("The version of hadoop that spark is compiled against.")
 
+  val deployDatanucleusJars = taskKey[Unit]("Deploy datanucleus jars to the spark/lib_managed/jars directory")
+
   lazy val settings = assemblySettings ++ Seq(
     test in assembly := {},
     hadoopVersion := {
@@ -429,7 +428,20 @@ object Assembly {
       case m if m.toLowerCase.startsWith("meta-inf/services/") => MergeStrategy.filterDistinctLines
       case "reference.conf"                                    => MergeStrategy.concat
       case _                                                   => MergeStrategy.first
-    }
+    },
+    deployDatanucleusJars := {
+      val jars: Seq[File] = (fullClasspath in assembly).value.map(_.data)
+        .filter(_.getPath.contains("org.datanucleus"))
+      var libManagedJars = new File(BuildCommons.sparkHome, "lib_managed/jars")
+      libManagedJars.mkdirs()
+      jars.foreach { jar =>
+        val dest = new File(libManagedJars, jar.getName)
+        if (!dest.exists()) {
+          Files.copy(jar.toPath, dest.toPath)
+        }
+      }
+    },
+    assembly <<= assembly.dependsOn(deployDatanucleusJars)
   )
 }
 

From 6e5fc37883ed81c3ee2338145a48de3036d19399 Mon Sep 17 00:00:00 2001
From: Lianhui Wang <lianhuiwang09@gmail.com>
Date: Tue, 10 Nov 2015 10:40:08 -0800
Subject: [PATCH 156/518] [SPARK-11252][NETWORK] ShuffleClient should release
 connection after fetching blocks had been completed for external shuffle

with yarn's external shuffle, ExternalShuffleClient of executors reserve its connections for yarn's NodeManager until application has been completed. so it will make NodeManager and executors have many socket connections.
in order to reduce network pressure of NodeManager's shuffleService, after registerWithShuffleServer or fetchBlocks have been completed in ExternalShuffleClient, connection for NM's shuffleService needs to be closed.andrewor14 rxin vanzin

Author: Lianhui Wang <lianhuiwang09@gmail.com>

Closes #9227 from lianhuiwang/spark-11252.
---
 .../spark/deploy/ExternalShuffleService.scala |  3 +-
 .../spark/network/TransportContext.java       | 11 +++++-
 .../client/TransportClientFactory.java        | 10 ++++++
 .../server/TransportChannelHandler.java       | 26 +++++++++-----
 .../network/TransportClientFactorySuite.java  | 34 +++++++++++++++++++
 .../shuffle/ExternalShuffleClient.java        | 12 ++++---
 6 files changed, 81 insertions(+), 15 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala b/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala
index 6840a3ae831f0..a039d543c35e7 100644
--- a/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala
@@ -47,7 +47,8 @@ class ExternalShuffleService(sparkConf: SparkConf, securityManager: SecurityMana
 
   private val transportConf = SparkTransportConf.fromSparkConf(sparkConf, numUsableCores = 0)
   private val blockHandler = newShuffleBlockHandler(transportConf)
-  private val transportContext: TransportContext = new TransportContext(transportConf, blockHandler)
+  private val transportContext: TransportContext =
+    new TransportContext(transportConf, blockHandler, true)
 
   private var server: TransportServer = _
 
diff --git a/network/common/src/main/java/org/apache/spark/network/TransportContext.java b/network/common/src/main/java/org/apache/spark/network/TransportContext.java
index 43900e6f2c972..1b64b863a9fe5 100644
--- a/network/common/src/main/java/org/apache/spark/network/TransportContext.java
+++ b/network/common/src/main/java/org/apache/spark/network/TransportContext.java
@@ -59,15 +59,24 @@ public class TransportContext {
 
   private final TransportConf conf;
   private final RpcHandler rpcHandler;
+  private final boolean closeIdleConnections;
 
   private final MessageEncoder encoder;
   private final MessageDecoder decoder;
 
   public TransportContext(TransportConf conf, RpcHandler rpcHandler) {
+    this(conf, rpcHandler, false);
+  }
+
+  public TransportContext(
+      TransportConf conf,
+      RpcHandler rpcHandler,
+      boolean closeIdleConnections) {
     this.conf = conf;
     this.rpcHandler = rpcHandler;
     this.encoder = new MessageEncoder();
     this.decoder = new MessageDecoder();
+    this.closeIdleConnections = closeIdleConnections;
   }
 
   /**
@@ -144,7 +153,7 @@ private TransportChannelHandler createChannelHandler(Channel channel, RpcHandler
     TransportRequestHandler requestHandler = new TransportRequestHandler(channel, client,
       rpcHandler);
     return new TransportChannelHandler(client, responseHandler, requestHandler,
-      conf.connectionTimeoutMs());
+      conf.connectionTimeoutMs(), closeIdleConnections);
   }
 
   public TransportConf getConf() { return conf; }
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java b/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
index 4952ffb44bb8b..42a4f664e697c 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
@@ -158,6 +158,16 @@ public TransportClient createClient(String remoteHost, int remotePort) throws IO
     }
   }
 
+  /**
+   * Create a completely new {@link TransportClient} to the given remote host / port
+   * But this connection is not pooled.
+   */
+  public TransportClient createUnmanagedClient(String remoteHost, int remotePort)
+      throws IOException {
+    final InetSocketAddress address = new InetSocketAddress(remoteHost, remotePort);
+    return createClient(address);
+  }
+
   /** Create a completely new {@link TransportClient} to the remote address. */
   private TransportClient createClient(InetSocketAddress address) throws IOException {
     logger.debug("Creating new connection to " + address);
diff --git a/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java b/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
index 8e0ee709e38e3..f8fcd1c3d7d76 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
@@ -55,16 +55,19 @@ public class TransportChannelHandler extends SimpleChannelInboundHandler<Message
   private final TransportResponseHandler responseHandler;
   private final TransportRequestHandler requestHandler;
   private final long requestTimeoutNs;
+  private final boolean closeIdleConnections;
 
   public TransportChannelHandler(
       TransportClient client,
       TransportResponseHandler responseHandler,
       TransportRequestHandler requestHandler,
-      long requestTimeoutMs) {
+      long requestTimeoutMs,
+      boolean closeIdleConnections) {
     this.client = client;
     this.responseHandler = responseHandler;
     this.requestHandler = requestHandler;
     this.requestTimeoutNs = requestTimeoutMs * 1000L * 1000;
+    this.closeIdleConnections = closeIdleConnections;
   }
 
   public TransportClient getClient() {
@@ -111,16 +114,21 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exc
       IdleStateEvent e = (IdleStateEvent) evt;
       // See class comment for timeout semantics. In addition to ensuring we only timeout while
       // there are outstanding requests, we also do a secondary consistency check to ensure
-      // there's no race between the idle timeout and incrementing the numOutstandingRequests.
-      boolean hasInFlightRequests = responseHandler.numOutstandingRequests() > 0;
+      // there's no race between the idle timeout and incrementing the numOutstandingRequests
+      // (see SPARK-7003).
       boolean isActuallyOverdue =
         System.nanoTime() - responseHandler.getTimeOfLastRequestNs() > requestTimeoutNs;
-      if (e.state() == IdleState.ALL_IDLE && hasInFlightRequests && isActuallyOverdue) {
-        String address = NettyUtils.getRemoteAddress(ctx.channel());
-        logger.error("Connection to {} has been quiet for {} ms while there are outstanding " +
-          "requests. Assuming connection is dead; please adjust spark.network.timeout if this " +
-          "is wrong.", address, requestTimeoutNs / 1000 / 1000);
-        ctx.close();
+      if (e.state() == IdleState.ALL_IDLE && isActuallyOverdue) {
+        if (responseHandler.numOutstandingRequests() > 0) {
+          String address = NettyUtils.getRemoteAddress(ctx.channel());
+          logger.error("Connection to {} has been quiet for {} ms while there are outstanding " +
+            "requests. Assuming connection is dead; please adjust spark.network.timeout if this " +
+            "is wrong.", address, requestTimeoutNs / 1000 / 1000);
+          ctx.close();
+        } else if (closeIdleConnections) {
+          // While CloseIdleConnections is enable, we also close idle connection
+          ctx.close();
+        }
       }
     }
   }
diff --git a/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java b/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java
index 35de5e57ccb98..f447137419306 100644
--- a/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java
@@ -21,6 +21,7 @@
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -37,6 +38,7 @@
 import org.apache.spark.network.server.NoOpRpcHandler;
 import org.apache.spark.network.server.RpcHandler;
 import org.apache.spark.network.server.TransportServer;
+import org.apache.spark.network.util.ConfigProvider;
 import org.apache.spark.network.util.SystemPropertyConfigProvider;
 import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.MapConfigProvider;
@@ -177,4 +179,36 @@ public void closeBlockClientsWithFactory() throws IOException {
     assertFalse(c1.isActive());
     assertFalse(c2.isActive());
   }
+
+  @Test
+  public void closeIdleConnectionForRequestTimeOut() throws IOException, InterruptedException {
+    TransportConf conf = new TransportConf(new ConfigProvider() {
+
+      @Override
+      public String get(String name) {
+        if ("spark.shuffle.io.connectionTimeout".equals(name)) {
+          // We should make sure there is enough time for us to observe the channel is active
+          return "1s";
+        }
+        String value = System.getProperty(name);
+        if (value == null) {
+          throw new NoSuchElementException(name);
+        }
+        return value;
+      }
+    });
+    TransportContext context = new TransportContext(conf, new NoOpRpcHandler(), true);
+    TransportClientFactory factory = context.createClientFactory();
+    try {
+      TransportClient c1 = factory.createClient(TestUtils.getLocalHost(), server1.getPort());
+      assertTrue(c1.isActive());
+      long expiredTime = System.currentTimeMillis() + 10000; // 10 seconds
+      while (c1.isActive() && System.currentTimeMillis() < expiredTime) {
+        Thread.sleep(10);
+      }
+      assertFalse(c1.isActive());
+    } finally {
+      factory.close();
+    }
+  }
 }
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java
index ea6d248d66be3..ef3a9dcc8711f 100644
--- a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java
+++ b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java
@@ -78,7 +78,7 @@ protected void checkInit() {
   @Override
   public void init(String appId) {
     this.appId = appId;
-    TransportContext context = new TransportContext(conf, new NoOpRpcHandler());
+    TransportContext context = new TransportContext(conf, new NoOpRpcHandler(), true);
     List<TransportClientBootstrap> bootstraps = Lists.newArrayList();
     if (saslEnabled) {
       bootstraps.add(new SaslClientBootstrap(conf, appId, secretKeyHolder, saslEncryptionEnabled));
@@ -137,9 +137,13 @@ public void registerWithShuffleServer(
       String execId,
       ExecutorShuffleInfo executorInfo) throws IOException {
     checkInit();
-    TransportClient client = clientFactory.createClient(host, port);
-    byte[] registerMessage = new RegisterExecutor(appId, execId, executorInfo).toByteArray();
-    client.sendRpcSync(registerMessage, 5000 /* timeoutMs */);
+    TransportClient client = clientFactory.createUnmanagedClient(host, port);
+    try {
+      byte[] registerMessage = new RegisterExecutor(appId, execId, executorInfo).toByteArray();
+      client.sendRpcSync(registerMessage, 5000 /* timeoutMs */);
+    } finally {
+      client.close();
+    }
   }
 
   @Override

From e0701c75601c43f69ed27fc7c252321703db51f2 Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Tue, 10 Nov 2015 11:06:29 -0800
Subject: [PATCH 157/518] [SPARK-9830][SQL] Remove AggregateExpression1 and
 Aggregate Operator used to evaluate AggregateExpression1s

https://issues.apache.org/jira/browse/SPARK-9830

This PR contains the following main changes.
* Removing `AggregateExpression1`.
* Removing `Aggregate` operator, which is used to evaluate `AggregateExpression1`.
* Removing planner rule used to plan `Aggregate`.
* Linking `MultipleDistinctRewriter` to analyzer.
* Renaming `AggregateExpression2` to `AggregateExpression` and `AggregateFunction2` to `AggregateFunction`.
* Updating places where we create aggregate expression. The way to create aggregate expressions is `AggregateExpression(aggregateFunction, mode, isDistinct)`.
* Changing `val`s in `DeclarativeAggregate`s that touch children of this function to `lazy val`s (when we create aggregate expression in DataFrame API, children of an aggregate function can be unresolved).

Author: Yin Huai <yhuai@databricks.com>

Closes #9556 from yhuai/removeAgg1.
---
 R/pkg/R/functions.R                           |    2 +-
 python/pyspark/sql/dataframe.py               |    2 +-
 python/pyspark/sql/functions.py               |    2 +-
 python/pyspark/sql/tests.py                   |    2 +-
 .../spark/sql/catalyst/CatalystConf.scala     |   10 +-
 .../apache/spark/sql/catalyst/SqlParser.scala |   14 +-
 .../sql/catalyst/analysis/Analyzer.scala      |   26 +-
 .../sql/catalyst/analysis/CheckAnalysis.scala |   46 +-
 .../DistinctAggregationRewriter.scala}        |  235 +---
 .../catalyst/analysis/FunctionRegistry.scala  |    2 +
 .../catalyst/analysis/HiveTypeCoercion.scala  |   20 +-
 .../sql/catalyst/analysis/unresolved.scala    |    4 +
 .../spark/sql/catalyst/dsl/package.scala      |   22 +-
 .../expressions/aggregate/Average.scala       |   31 +-
 .../aggregate/CentralMomentAgg.scala          |   13 +-
 .../catalyst/expressions/aggregate/Corr.scala |   15 +
 .../expressions/aggregate/Count.scala         |   28 +-
 .../expressions/aggregate/First.scala         |   14 +-
 .../aggregate/HyperLogLogPlusPlus.scala       |   17 +
 .../expressions/aggregate/Kurtosis.scala      |    2 +
 .../catalyst/expressions/aggregate/Last.scala |   12 +-
 .../catalyst/expressions/aggregate/Max.scala  |   17 +-
 .../catalyst/expressions/aggregate/Min.scala  |   17 +-
 .../expressions/aggregate/Skewness.scala      |    2 +
 .../expressions/aggregate/Stddev.scala        |   31 +-
 .../catalyst/expressions/aggregate/Sum.scala  |   29 +-
 .../expressions/aggregate/Variance.scala      |    7 +-
 .../expressions/aggregate/interfaces.scala    |   57 +-
 .../sql/catalyst/expressions/aggregates.scala | 1073 -----------------
 .../sql/catalyst/optimizer/Optimizer.scala    |   23 +-
 .../sql/catalyst/planning/patterns.scala      |   74 --
 .../spark/sql/catalyst/plans/QueryPlan.scala  |   12 +-
 .../plans/logical/basicOperators.scala        |    4 +-
 .../analysis/AnalysisErrorSuite.scala         |   23 +-
 .../sql/catalyst/analysis/AnalysisSuite.scala |    2 +-
 .../analysis/DecimalPrecisionSuite.scala      |    1 +
 .../ExpressionTypeCheckingSuite.scala         |    6 +-
 .../optimizer/ConstantFoldingSuite.scala      |    4 +-
 .../optimizer/FilterPushdownSuite.scala       |   14 +-
 .../org/apache/spark/sql/DataFrame.scala      |   13 +-
 .../org/apache/spark/sql/GroupedData.scala    |   45 +-
 .../scala/org/apache/spark/sql/SQLConf.scala  |   20 +-
 .../spark/sql/execution/Aggregate.scala       |  205 ----
 .../apache/spark/sql/execution/Expand.scala   |    3 +
 .../spark/sql/execution/SparkPlanner.scala    |    1 -
 .../spark/sql/execution/SparkStrategies.scala |  238 ++--
 .../aggregate/AggregationIterator.scala       |   28 +-
 .../aggregate/SortBasedAggregate.scala        |    4 +-
 .../SortBasedAggregationIterator.scala        |    8 +-
 .../aggregate/TungstenAggregate.scala         |    6 +-
 .../TungstenAggregationIterator.scala         |   36 +-
 .../spark/sql/execution/aggregate/udaf.scala  |    2 +-
 .../spark/sql/execution/aggregate/utils.scala |   20 +-
 .../spark/sql/expressions/Aggregator.scala    |    5 +-
 .../spark/sql/expressions/WindowSpec.scala    |   82 +-
 .../apache/spark/sql/expressions/udaf.scala   |    6 +-
 .../org/apache/spark/sql/functions.scala      |   53 +-
 .../org/apache/spark/sql/SQLQuerySuite.scala  |   69 +-
 .../spark/sql/UserDefinedTypeSuite.scala      |   15 +-
 .../spark/sql/execution/PlannerSuite.scala    |    2 +-
 .../execution/metric/SQLMetricsSuite.scala    |   30 -
 .../apache/spark/sql/hive/HiveContext.scala   |    1 -
 .../org/apache/spark/sql/hive/HiveQl.scala    |    8 +-
 .../execution/AggregationQuerySuite.scala     |  188 ++-
 64 files changed, 743 insertions(+), 2260 deletions(-)
 rename sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/{expressions/aggregate/Utils.scala => analysis/DistinctAggregationRewriter.scala} (58%)
 delete mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
 delete mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/Aggregate.scala

diff --git a/R/pkg/R/functions.R b/R/pkg/R/functions.R
index d7fd279279137..0b280870295a2 100644
--- a/R/pkg/R/functions.R
+++ b/R/pkg/R/functions.R
@@ -1339,7 +1339,7 @@ setMethod("pmod", signature(y = "Column"),
 #' @export
 setMethod("approxCountDistinct",
           signature(x = "Column"),
-          function(x, rsd = 0.95) {
+          function(x, rsd = 0.05) {
             jc <- callJStatic("org.apache.spark.sql.functions", "approxCountDistinct", x@jc, rsd)
             column(jc)
           })
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index b97c94dad834a..0dd75ba7ca820 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -866,7 +866,7 @@ def selectExpr(self, *expr):
         This is a variant of :func:`select` that accepts SQL expressions.
 
         >>> df.selectExpr("age * 2", "abs(age)").collect()
-        [Row((age * 2)=4, 'abs(age)=2), Row((age * 2)=10, 'abs(age)=5)]
+        [Row((age * 2)=4, abs(age)=2), Row((age * 2)=10, abs(age)=5)]
         """
         if len(expr) == 1 and isinstance(expr[0], list):
             expr = expr[0]
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 962f676d406d8..6e1cbde4239f3 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -382,7 +382,7 @@ def expr(str):
     """Parses the expression string into the column that it represents
 
     >>> df.select(expr("length(name)")).collect()
-    [Row('length(name)=5), Row('length(name)=3)]
+    [Row(length(name)=5), Row(length(name)=3)]
     """
     sc = SparkContext._active_spark_context
     return Column(sc._jvm.functions.expr(str))
diff --git a/python/pyspark/sql/tests.py b/python/pyspark/sql/tests.py
index e224574bcb301..9f5f7cfdf7a69 100644
--- a/python/pyspark/sql/tests.py
+++ b/python/pyspark/sql/tests.py
@@ -1017,7 +1017,7 @@ def test_expr(self):
         row = Row(a="length string", b=75)
         df = self.sqlCtx.createDataFrame([row])
         result = df.select(functions.expr("length(a)")).collect()[0].asDict()
-        self.assertEqual(13, result["'length(a)"])
+        self.assertEqual(13, result["length(a)"])
 
     def test_replace(self):
         schema = StructType([
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/CatalystConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/CatalystConf.scala
index 3f351b07b37df..7c2b8a9407884 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/CatalystConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/CatalystConf.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.catalyst
 
 private[spark] trait CatalystConf {
   def caseSensitiveAnalysis: Boolean
+
+  protected[spark] def specializeSingleDistinctAggPlanning: Boolean
 }
 
 /**
@@ -29,7 +31,13 @@ object EmptyConf extends CatalystConf {
   override def caseSensitiveAnalysis: Boolean = {
     throw new UnsupportedOperationException
   }
+
+  protected[spark] override def specializeSingleDistinctAggPlanning: Boolean = {
+    throw new UnsupportedOperationException
+  }
 }
 
 /** A CatalystConf that can be used for local testing. */
-case class SimpleCatalystConf(caseSensitiveAnalysis: Boolean) extends CatalystConf
+case class SimpleCatalystConf(caseSensitiveAnalysis: Boolean) extends CatalystConf {
+  protected[spark] override def specializeSingleDistinctAggPlanning: Boolean = true
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
index cd717c09f8e5e..2a132d8b82bef 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala
@@ -22,6 +22,7 @@ import scala.language.implicitConversions
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.util.DataTypeParser
@@ -272,7 +273,7 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
   protected lazy val function: Parser[Expression] =
     ( ident <~ ("(" ~ "*" ~ ")") ^^ { case udfName =>
       if (lexical.normalizeKeyword(udfName) == "count") {
-        Count(Literal(1))
+        AggregateExpression(Count(Literal(1)), mode = Complete, isDistinct = false)
       } else {
         throw new AnalysisException(s"invalid expression $udfName(*)")
       }
@@ -281,14 +282,14 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
       { case udfName ~ exprs => UnresolvedFunction(udfName, exprs, isDistinct = false) }
     | ident ~ ("(" ~ DISTINCT ~> repsep(expression, ",")) <~ ")" ^^ { case udfName ~ exprs =>
       lexical.normalizeKeyword(udfName) match {
-        case "sum" => SumDistinct(exprs.head)
-        case "count" => CountDistinct(exprs)
+        case "count" =>
+          aggregate.Count(exprs).toAggregateExpression(isDistinct = true)
         case _ => UnresolvedFunction(udfName, exprs, isDistinct = true)
       }
     }
     | APPROXIMATE ~> ident ~ ("(" ~ DISTINCT ~> expression <~ ")") ^^ { case udfName ~ exp =>
       if (lexical.normalizeKeyword(udfName) == "count") {
-        ApproxCountDistinct(exp)
+        AggregateExpression(new HyperLogLogPlusPlus(exp), mode = Complete, isDistinct = false)
       } else {
         throw new AnalysisException(s"invalid function approximate $udfName")
       }
@@ -296,7 +297,10 @@ object SqlParser extends AbstractSparkSQLParser with DataTypeParser {
     | APPROXIMATE ~> "(" ~> unsignedFloat ~ ")" ~ ident ~ "(" ~ DISTINCT ~ expression <~ ")" ^^
       { case s ~ _ ~ udfName ~ _ ~ _ ~ exp =>
         if (lexical.normalizeKeyword(udfName) == "count") {
-          ApproxCountDistinct(exp, s.toDouble)
+          AggregateExpression(
+            HyperLogLogPlusPlus(exp, s.toDouble, 0, 0),
+            mode = Complete,
+            isDistinct = false)
         } else {
           throw new AnalysisException(s"invalid function approximate($s) $udfName")
         }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index 899ee67352df4..b1e14390b7dc0 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -20,8 +20,8 @@ package org.apache.spark.sql.catalyst.analysis
 import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, AggregateExpression2, AggregateFunction2}
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
 import org.apache.spark.sql.catalyst.trees.TreeNodeRef
@@ -79,6 +79,7 @@ class Analyzer(
       ExtractWindowExpressions ::
       GlobalAggregates ::
       ResolveAggregateFunctions ::
+      DistinctAggregationRewriter(conf) ::
       HiveTypeCoercion.typeCoercionRules ++
       extendedResolutionRules : _*),
     Batch("Nondeterministic", Once,
@@ -525,21 +526,14 @@ class Analyzer(
           case u @ UnresolvedFunction(name, children, isDistinct) =>
             withPosition(u) {
               registry.lookupFunction(name, children) match {
-                // We get an aggregate function built based on AggregateFunction2 interface.
-                // So, we wrap it in AggregateExpression2.
-                case agg2: AggregateFunction2 => AggregateExpression2(agg2, Complete, isDistinct)
-                // Currently, our old aggregate function interface supports SUM(DISTINCT ...)
-                // and COUTN(DISTINCT ...).
-                case sumDistinct: SumDistinct => sumDistinct
-                case countDistinct: CountDistinct => countDistinct
-                // DISTINCT is not meaningful with Max and Min.
-                case max: Max if isDistinct => max
-                case min: Min if isDistinct => min
-                // For other aggregate functions, DISTINCT keyword is not supported for now.
-                // Once we converted to the new code path, we will allow using DISTINCT keyword.
-                case other: AggregateExpression1 if isDistinct =>
-                  failAnalysis(s"$name does not support DISTINCT keyword.")
-                // If it does not have DISTINCT keyword, we will return it as is.
+                // DISTINCT is not meaningful for a Max or a Min.
+                case max: Max if isDistinct =>
+                  AggregateExpression(max, Complete, isDistinct = false)
+                case min: Min if isDistinct =>
+                  AggregateExpression(min, Complete, isDistinct = false)
+                // We get an aggregate function, we need to wrap it in an AggregateExpression.
+                case agg2: AggregateFunction => AggregateExpression(agg2, Complete, isDistinct)
+                // This function is not an aggregate function, just return the resolved one.
                 case other => other
               }
             }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
index 98d6637c0601b..8322e9930cd5a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.analysis
 
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateFunction, AggregateExpression}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.types._
 
@@ -108,7 +109,19 @@ trait CheckAnalysis {
 
           case Aggregate(groupingExprs, aggregateExprs, child) =>
             def checkValidAggregateExpression(expr: Expression): Unit = expr match {
-              case _: AggregateExpression => // OK
+              case aggExpr: AggregateExpression =>
+                // TODO: Is it possible that the child of a agg function is another
+                // agg function?
+                aggExpr.aggregateFunction.children.foreach {
+                  // This is just a sanity check, our analysis rule PullOutNondeterministic should
+                  // already pull out those nondeterministic expressions and evaluate them in
+                  // a Project node.
+                  case child if !child.deterministic =>
+                    failAnalysis(
+                      s"nondeterministic expression ${expr.prettyString} should not " +
+                        s"appear in the arguments of an aggregate function.")
+                  case child => // OK
+                }
               case e: Attribute if !groupingExprs.exists(_.semanticEquals(e)) =>
                 failAnalysis(
                   s"expression '${e.prettyString}' is neither present in the group by, " +
@@ -120,14 +133,26 @@ trait CheckAnalysis {
               case e => e.children.foreach(checkValidAggregateExpression)
             }
 
-            def checkValidGroupingExprs(expr: Expression): Unit = expr.dataType match {
-              case BinaryType =>
-                failAnalysis(s"binary type expression ${expr.prettyString} cannot be used " +
-                  "in grouping expression")
-              case m: MapType =>
-                failAnalysis(s"map type expression ${expr.prettyString} cannot be used " +
-                  "in grouping expression")
-              case _ => // OK
+            def checkValidGroupingExprs(expr: Expression): Unit = {
+              expr.dataType match {
+                case BinaryType =>
+                  failAnalysis(s"binary type expression ${expr.prettyString} cannot be used " +
+                    "in grouping expression")
+                case a: ArrayType =>
+                  failAnalysis(s"array type expression ${expr.prettyString} cannot be used " +
+                    "in grouping expression")
+                case m: MapType =>
+                  failAnalysis(s"map type expression ${expr.prettyString} cannot be used " +
+                    "in grouping expression")
+                case _ => // OK
+              }
+              if (!expr.deterministic) {
+                // This is just a sanity check, our analysis rule PullOutNondeterministic should
+                // already pull out those nondeterministic expressions and evaluate them in
+                // a Project node.
+                failAnalysis(s"nondeterministic expression ${expr.prettyString} should not " +
+                  s"appear in grouping expression.")
+              }
             }
 
             aggregateExprs.foreach(checkValidAggregateExpression)
@@ -179,7 +204,8 @@ trait CheckAnalysis {
               s"unresolved operator ${operator.simpleString}")
 
           case o if o.expressions.exists(!_.deterministic) &&
-            !o.isInstanceOf[Project] && !o.isInstanceOf[Filter] =>
+            !o.isInstanceOf[Project] && !o.isInstanceOf[Filter] & !o.isInstanceOf[Aggregate] =>
+            // The rule above is used to check Aggregate operator.
             failAnalysis(
               s"""nondeterministic expressions are only allowed in Project or Filter, found:
                  | ${o.expressions.map(_.prettyString).mkString(",")}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
similarity index 58%
rename from sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
rename to sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
index 9b22ce2619731..397eff05686b6 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Utils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
@@ -15,215 +15,17 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.catalyst.expressions.aggregate
+package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst._
+import org.apache.spark.sql.catalyst.CatalystConf
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.logical.{Expand, Aggregate, LogicalPlan}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, Complete}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Expand, LogicalPlan}
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.types._
+import org.apache.spark.sql.types.IntegerType
 
 /**
- * Utility functions used by the query planner to convert our plan to new aggregation code path.
- */
-object Utils {
-
-  // Check if the DataType given cannot be part of a group by clause.
-  private def isUnGroupable(dt: DataType): Boolean = dt match {
-    case _: ArrayType | _: MapType => true
-    case s: StructType => s.fields.exists(f => isUnGroupable(f.dataType))
-    case _ => false
-  }
-
-  // Right now, we do not support complex types in the grouping key schema.
-  private def supportsGroupingKeySchema(aggregate: Aggregate): Boolean =
-    !aggregate.groupingExpressions.exists(e => isUnGroupable(e.dataType))
-
-  private def doConvert(plan: LogicalPlan): Option[Aggregate] = plan match {
-    case p: Aggregate if supportsGroupingKeySchema(p) =>
-
-      val converted = MultipleDistinctRewriter.rewrite(p.transformExpressionsDown {
-        case expressions.Average(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Average(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.Count(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Count(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.CountDistinct(children) =>
-          val child = if (children.size > 1) {
-            DropAnyNull(CreateStruct(children))
-          } else {
-            children.head
-          }
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Count(child),
-            mode = aggregate.Complete,
-            isDistinct = true)
-
-        case expressions.First(child, ignoreNulls) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.First(child, ignoreNulls),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.Kurtosis(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Kurtosis(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.Last(child, ignoreNulls) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Last(child, ignoreNulls),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.Max(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Max(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.Min(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Min(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.Skewness(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Skewness(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.StddevPop(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.StddevPop(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.StddevSamp(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.StddevSamp(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.Sum(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Sum(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.SumDistinct(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Sum(child),
-            mode = aggregate.Complete,
-            isDistinct = true)
-
-        case expressions.Corr(left, right) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.Corr(left, right),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.ApproxCountDistinct(child, rsd) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.HyperLogLogPlusPlus(child, rsd),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.VariancePop(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.VariancePop(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-
-        case expressions.VarianceSamp(child) =>
-          aggregate.AggregateExpression2(
-            aggregateFunction = aggregate.VarianceSamp(child),
-            mode = aggregate.Complete,
-            isDistinct = false)
-      })
-
-      // Check if there is any expressions.AggregateExpression1 left.
-      // If so, we cannot convert this plan.
-      val hasAggregateExpression1 = converted.aggregateExpressions.exists { expr =>
-        // For every expressions, check if it contains AggregateExpression1.
-        expr.find {
-          case agg: expressions.AggregateExpression1 => true
-          case other => false
-        }.isDefined
-      }
-
-      // Check if there are multiple distinct columns.
-      // TODO remove this.
-      val aggregateExpressions = converted.aggregateExpressions.flatMap { expr =>
-        expr.collect {
-          case agg: AggregateExpression2 => agg
-        }
-      }.toSet.toSeq
-      val functionsWithDistinct = aggregateExpressions.filter(_.isDistinct)
-      val hasMultipleDistinctColumnSets =
-        if (functionsWithDistinct.map(_.aggregateFunction.children).distinct.length > 1) {
-          true
-        } else {
-          false
-        }
-
-      if (!hasAggregateExpression1 && !hasMultipleDistinctColumnSets) Some(converted) else None
-
-    case other => None
-  }
-
-  def checkInvalidAggregateFunction2(aggregate: Aggregate): Unit = {
-    // If the plan cannot be converted, we will do a final round check to see if the original
-    // logical.Aggregate contains both AggregateExpression1 and AggregateExpression2. If so,
-    // we need to throw an exception.
-    val aggregateFunction2s = aggregate.aggregateExpressions.flatMap { expr =>
-      expr.collect {
-        case agg: AggregateExpression2 => agg.aggregateFunction
-      }
-    }.distinct
-    if (aggregateFunction2s.nonEmpty) {
-      // For functions implemented based on the new interface, prepare a list of function names.
-      val invalidFunctions = {
-        if (aggregateFunction2s.length > 1) {
-          s"${aggregateFunction2s.tail.map(_.nodeName).mkString(",")} " +
-            s"and ${aggregateFunction2s.head.nodeName} are"
-        } else {
-          s"${aggregateFunction2s.head.nodeName} is"
-        }
-      }
-      val errorMessage =
-        s"${invalidFunctions} implemented based on the new Aggregate Function " +
-          s"interface and it cannot be used with functions implemented based on " +
-          s"the old Aggregate Function interface."
-      throw new AnalysisException(errorMessage)
-    }
-  }
-
-  def tryConvert(plan: LogicalPlan): Option[Aggregate] = plan match {
-    case p: Aggregate =>
-      val converted = doConvert(p)
-      if (converted.isDefined) {
-        converted
-      } else {
-        checkInvalidAggregateFunction2(p)
-        None
-      }
-    case other => None
-  }
-}
-
-/**
- * This rule rewrites an aggregate query with multiple distinct clauses into an expanded double
+ * This rule rewrites an aggregate query with distinct aggregations into an expanded double
  * aggregation in which the regular aggregation expressions and every distinct clause is aggregated
  * in a separate group. The results are then combined in a second aggregate.
  *
@@ -298,9 +100,11 @@ object Utils {
  * we could improve this in the current rule by applying more advanced expression cannocalization
  * techniques.
  */
-object MultipleDistinctRewriter extends Rule[LogicalPlan] {
+case class DistinctAggregationRewriter(conf: CatalystConf) extends Rule[LogicalPlan] {
 
   def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
+    case p if !p.resolved => p
+    // We need to wait until this Aggregate operator is resolved.
     case a: Aggregate => rewrite(a)
     case p => p
   }
@@ -310,7 +114,7 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
     // Collect all aggregate expressions.
     val aggExpressions = a.aggregateExpressions.flatMap { e =>
       e.collect {
-        case ae: AggregateExpression2 => ae
+        case ae: AggregateExpression => ae
       }
     }
 
@@ -319,8 +123,15 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
       .filter(_.isDistinct)
       .groupBy(_.aggregateFunction.children.toSet)
 
-    // Only continue to rewrite if there is more than one distinct group.
-    if (distinctAggGroups.size > 1) {
+    val shouldRewrite = if (conf.specializeSingleDistinctAggPlanning) {
+      // When the flag is set to specialize single distinct agg planning,
+      // we will rely on our Aggregation strategy to handle queries with a single
+      // distinct column and this aggregate operator does have grouping expressions.
+      distinctAggGroups.size > 1 || (distinctAggGroups.size == 1 && a.groupingExpressions.isEmpty)
+    } else {
+      distinctAggGroups.size >= 1
+    }
+    if (shouldRewrite) {
       // Create the attributes for the grouping id and the group by clause.
       val gid = new AttributeReference("gid", IntegerType, false)()
       val groupByMap = a.groupingExpressions.collect {
@@ -332,11 +143,11 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
       // Functions used to modify aggregate functions and their inputs.
       def evalWithinGroup(id: Literal, e: Expression) = If(EqualTo(gid, id), e, nullify(e))
       def patchAggregateFunctionChildren(
-          af: AggregateFunction2)(
-          attrs: Expression => Expression): AggregateFunction2 = {
+          af: AggregateFunction)(
+          attrs: Expression => Expression): AggregateFunction = {
         af.withNewChildren(af.children.map {
           case afc => attrs(afc)
-        }).asInstanceOf[AggregateFunction2]
+        }).asInstanceOf[AggregateFunction]
       }
 
       // Setup unique distinct aggregate children.
@@ -381,7 +192,7 @@ object MultipleDistinctRewriter extends Rule[LogicalPlan] {
         val operator = Alias(e.copy(aggregateFunction = af), e.prettyString)()
 
         // Select the result of the first aggregate in the last aggregate.
-        val result = AggregateExpression2(
+        val result = AggregateExpression(
           aggregate.First(evalWithinGroup(regularGroupId, operator.toAttribute), Literal(true)),
           mode = Complete,
           isDistinct = false)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index d4334d16289a5..dfa749d1afa5b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -24,6 +24,7 @@ import scala.util.{Failure, Success, Try}
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FunctionBuilder
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.util.StringKeyHashMap
 
 
@@ -177,6 +178,7 @@ object FunctionRegistry {
     expression[ToRadians]("radians"),
 
     // aggregate functions
+    expression[HyperLogLogPlusPlus]("approx_count_distinct"),
     expression[Average]("avg"),
     expression[Corr]("corr"),
     expression[Count]("count"),
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
index 84e2b1366f626..bf2bff0243fa3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.analysis
 import javax.annotation.Nullable
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.types._
@@ -295,14 +296,17 @@ object HiveTypeCoercion {
         i.makeCopy(Array(Cast(a, StringType), b.map(Cast(_, StringType))))
 
       case Sum(e @ StringType()) => Sum(Cast(e, DoubleType))
-      case SumDistinct(e @ StringType()) => Sum(Cast(e, DoubleType))
       case Average(e @ StringType()) => Average(Cast(e, DoubleType))
       case StddevPop(e @ StringType()) => StddevPop(Cast(e, DoubleType))
       case StddevSamp(e @ StringType()) => StddevSamp(Cast(e, DoubleType))
-      case VariancePop(e @ StringType()) => VariancePop(Cast(e, DoubleType))
-      case VarianceSamp(e @ StringType()) => VarianceSamp(Cast(e, DoubleType))
-      case Skewness(e @ StringType()) => Skewness(Cast(e, DoubleType))
-      case Kurtosis(e @ StringType()) => Kurtosis(Cast(e, DoubleType))
+      case VariancePop(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
+        VariancePop(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
+      case VarianceSamp(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
+        VarianceSamp(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
+      case Skewness(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
+        Skewness(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
+      case Kurtosis(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
+        Kurtosis(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
     }
   }
 
@@ -562,12 +566,6 @@ object HiveTypeCoercion {
       case Sum(e @ IntegralType()) if e.dataType != LongType => Sum(Cast(e, LongType))
       case Sum(e @ FractionalType()) if e.dataType != DoubleType => Sum(Cast(e, DoubleType))
 
-      case s @ SumDistinct(e @ DecimalType()) => s // Decimal is already the biggest.
-      case SumDistinct(e @ IntegralType()) if e.dataType != LongType =>
-        SumDistinct(Cast(e, LongType))
-      case SumDistinct(e @ FractionalType()) if e.dataType != DoubleType =>
-        SumDistinct(Cast(e, DoubleType))
-
       case s @ Average(e @ DecimalType()) => s // Decimal is already the biggest.
       case Average(e @ IntegralType()) if e.dataType != LongType =>
         Average(Cast(e, LongType))
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
index eae17c86ddc7a..6485bdfb30234 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
@@ -141,6 +141,10 @@ case class UnresolvedFunction(
   override def nullable: Boolean = throw new UnresolvedException(this, "nullable")
   override lazy val resolved = false
 
+  override def prettyString: String = {
+    s"${name}(${children.map(_.prettyString).mkString(",")})"
+  }
+
   override def toString: String = s"'$name(${children.mkString(",")})"
 }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
index d8df66430a695..af594c25c54cb 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
@@ -23,6 +23,7 @@ import scala.language.implicitConversions
 
 import org.apache.spark.sql.catalyst.analysis.{EliminateSubQueries, UnresolvedExtractValue, UnresolvedAttribute}
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.{Inner, JoinType}
 import org.apache.spark.sql.types._
@@ -144,17 +145,18 @@ package object dsl {
       }
     }
 
-    def sum(e: Expression): Expression = Sum(e)
-    def sumDistinct(e: Expression): Expression = SumDistinct(e)
-    def count(e: Expression): Expression = Count(e)
-    def countDistinct(e: Expression*): Expression = CountDistinct(e)
+    def sum(e: Expression): Expression = Sum(e).toAggregateExpression()
+    def sumDistinct(e: Expression): Expression = Sum(e).toAggregateExpression(isDistinct = true)
+    def count(e: Expression): Expression = Count(e).toAggregateExpression()
+    def countDistinct(e: Expression*): Expression =
+      Count(e).toAggregateExpression(isDistinct = true)
     def approxCountDistinct(e: Expression, rsd: Double = 0.05): Expression =
-      ApproxCountDistinct(e, rsd)
-    def avg(e: Expression): Expression = Average(e)
-    def first(e: Expression): Expression = First(e)
-    def last(e: Expression): Expression = Last(e)
-    def min(e: Expression): Expression = Min(e)
-    def max(e: Expression): Expression = Max(e)
+      HyperLogLogPlusPlus(e, rsd).toAggregateExpression()
+    def avg(e: Expression): Expression = Average(e).toAggregateExpression()
+    def first(e: Expression): Expression = new First(e).toAggregateExpression()
+    def last(e: Expression): Expression = new Last(e).toAggregateExpression()
+    def min(e: Expression): Expression = Min(e).toAggregateExpression()
+    def max(e: Expression): Expression = Max(e).toAggregateExpression()
     def upper(e: Expression): Expression = Upper(e)
     def lower(e: Expression): Expression = Lower(e)
     def sqrt(e: Expression): Expression = Sqrt(e)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
index c8c20ada5fbc7..7f9e5034702e9 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
@@ -17,8 +17,10 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
 
 case class Average(child: Expression) extends DeclarativeAggregate {
@@ -32,36 +34,33 @@ case class Average(child: Expression) extends DeclarativeAggregate {
   // Return data type.
   override def dataType: DataType = resultType
 
-  // Expected input data type.
-  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
-  // new version at planning time (after analysis phase). For now, NullType is added at here
-  // to make it resolved when we have cases like `select avg(null)`.
-  // We can use our analyzer to cast NullType to the default data type of the NumericType once
-  // we remove the old aggregate functions. Then, we will not need NullType at here.
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
+  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType))
 
-  private val resultType = child.dataType match {
+  override def checkInputDataTypes(): TypeCheckResult =
+    TypeUtils.checkForNumericExpr(child.dataType, "function average")
+
+  private lazy val resultType = child.dataType match {
     case DecimalType.Fixed(p, s) =>
       DecimalType.bounded(p + 4, s + 4)
     case _ => DoubleType
   }
 
-  private val sumDataType = child.dataType match {
+  private lazy val sumDataType = child.dataType match {
     case _ @ DecimalType.Fixed(p, s) => DecimalType.bounded(p + 10, s)
     case _ => DoubleType
   }
 
-  private val sum = AttributeReference("sum", sumDataType)()
-  private val count = AttributeReference("count", LongType)()
+  private lazy val sum = AttributeReference("sum", sumDataType)()
+  private lazy val count = AttributeReference("count", LongType)()
 
-  override val aggBufferAttributes = sum :: count :: Nil
+  override lazy val aggBufferAttributes = sum :: count :: Nil
 
-  override val initialValues = Seq(
+  override lazy val initialValues = Seq(
     /* sum = */ Cast(Literal(0), sumDataType),
     /* count = */ Literal(0L)
   )
 
-  override val updateExpressions = Seq(
+  override lazy val updateExpressions = Seq(
     /* sum = */
     Add(
       sum,
@@ -69,13 +68,13 @@ case class Average(child: Expression) extends DeclarativeAggregate {
     /* count = */ If(IsNull(child), count, count + 1L)
   )
 
-  override val mergeExpressions = Seq(
+  override lazy val mergeExpressions = Seq(
     /* sum = */ sum.left + sum.right,
     /* count = */ count.left + count.right
   )
 
   // If all input are nulls, count will be 0 and we will get null after the division.
-  override val evaluateExpression = child.dataType match {
+  override lazy val evaluateExpression = child.dataType match {
     case DecimalType.Fixed(p, s) =>
       // increase the precision and scale to prevent precision loss
       val dt = DecimalType.bounded(p + 14, s + 4)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
index ef08b025ff556..984ce7f24dacc 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
@@ -18,7 +18,9 @@
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
 import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
 
 /**
@@ -55,13 +57,10 @@ abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate w
 
   override def dataType: DataType = DoubleType
 
-  // Expected input data type.
-  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
-  // new version at planning time (after analysis phase). For now, NullType is added at here
-  // to make it resolved when we have cases like `select avg(null)`.
-  // We can use our analyzer to cast NullType to the default data type of the NumericType once
-  // we remove the old aggregate functions. Then, we will not need NullType at here.
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
+  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType))
+
+  override def checkInputDataTypes(): TypeCheckResult =
+    TypeUtils.checkForNumericExpr(child.dataType, s"function $prettyName")
 
   override def aggBufferSchema: StructType = StructType.fromAttributes(aggBufferAttributes)
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
index 832338378fb38..00d7436b710d2 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
@@ -18,7 +18,9 @@
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
 import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
 
 /**
@@ -35,6 +37,9 @@ case class Corr(
     inputAggBufferOffset: Int = 0)
   extends ImperativeAggregate {
 
+  def this(left: Expression, right: Expression) =
+    this(left, right, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+
   override def children: Seq[Expression] = Seq(left, right)
 
   override def nullable: Boolean = false
@@ -43,6 +48,16 @@ case class Corr(
 
   override def inputTypes: Seq[AbstractDataType] = Seq(DoubleType, DoubleType)
 
+  override def checkInputDataTypes(): TypeCheckResult = {
+    if (left.dataType.isInstanceOf[DoubleType] && right.dataType.isInstanceOf[DoubleType]) {
+      TypeCheckResult.TypeCheckSuccess
+    } else {
+      TypeCheckResult.TypeCheckFailure(
+        s"corr requires that both arguments are double type, " +
+          s"not (${left.dataType}, ${right.dataType}).")
+    }
+  }
+
   override def aggBufferSchema: StructType = StructType.fromAttributes(aggBufferAttributes)
 
   override def inputAggBufferAttributes: Seq[AttributeReference] = {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
index ec0c8b483a909..09a1da9200df0 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
@@ -32,23 +32,39 @@ case class Count(child: Expression) extends DeclarativeAggregate {
   // Expected input data type.
   override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
 
-  private val count = AttributeReference("count", LongType)()
+  private lazy val count = AttributeReference("count", LongType)()
 
-  override val aggBufferAttributes = count :: Nil
+  override lazy val aggBufferAttributes = count :: Nil
 
-  override val initialValues = Seq(
+  override lazy val initialValues = Seq(
     /* count = */ Literal(0L)
   )
 
-  override val updateExpressions = Seq(
+  override lazy val updateExpressions = Seq(
     /* count = */ If(IsNull(child), count, count + 1L)
   )
 
-  override val mergeExpressions = Seq(
+  override lazy val mergeExpressions = Seq(
     /* count = */ count.left + count.right
   )
 
-  override val evaluateExpression = Cast(count, LongType)
+  override lazy val evaluateExpression = Cast(count, LongType)
 
   override def defaultResult: Option[Literal] = Option(Literal(0L))
 }
+
+object Count {
+  def apply(children: Seq[Expression]): Count = {
+    // This is used to deal with COUNT DISTINCT. When we have multiple
+    // children (COUNT(DISTINCT col1, col2, ...)), we wrap them in a STRUCT (i.e. a Row).
+    // Also, the semantic of COUNT(DISTINCT col1, col2, ...) is that if there is any
+    // null in the arguments, we will not count that row. So, we use DropAnyNull at here
+    // to return a null when any field of the created STRUCT is null.
+    val child = if (children.size > 1) {
+      DropAnyNull(CreateStruct(children))
+    } else {
+      children.head
+    }
+    Count(child)
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala
index 9028143015853..35f57426feaf2 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/First.scala
@@ -51,18 +51,18 @@ case class First(child: Expression, ignoreNullsExpr: Expression) extends Declara
   // Expected input data type.
   override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
 
-  private val first = AttributeReference("first", child.dataType)()
+  private lazy val first = AttributeReference("first", child.dataType)()
 
-  private val valueSet = AttributeReference("valueSet", BooleanType)()
+  private lazy val valueSet = AttributeReference("valueSet", BooleanType)()
 
-  override val aggBufferAttributes: Seq[AttributeReference] = first :: valueSet :: Nil
+  override lazy val aggBufferAttributes: Seq[AttributeReference] = first :: valueSet :: Nil
 
-  override val initialValues: Seq[Literal] = Seq(
+  override lazy val initialValues: Seq[Literal] = Seq(
     /* first = */ Literal.create(null, child.dataType),
     /* valueSet = */ Literal.create(false, BooleanType)
   )
 
-  override val updateExpressions: Seq[Expression] = {
+  override lazy val updateExpressions: Seq[Expression] = {
     if (ignoreNulls) {
       Seq(
         /* first = */ If(Or(valueSet, IsNull(child)), first, child),
@@ -76,7 +76,7 @@ case class First(child: Expression, ignoreNullsExpr: Expression) extends Declara
     }
   }
 
-  override val mergeExpressions: Seq[Expression] = {
+  override lazy val mergeExpressions: Seq[Expression] = {
     // For first, we can just check if valueSet.left is set to true. If it is set
     // to true, we use first.right. If not, we use first.right (even if valueSet.right is
     // false, we are safe to do so because first.right will be null in this case).
@@ -86,7 +86,7 @@ case class First(child: Expression, ignoreNullsExpr: Expression) extends Declara
     )
   }
 
-  override val evaluateExpression: AttributeReference = first
+  override lazy val evaluateExpression: AttributeReference = first
 
   override def toString: String = s"first($child)${if (ignoreNulls) " ignore nulls"}"
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
index 8d341ee630bdb..8a95c541f1e86 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
@@ -22,6 +22,7 @@ import java.util
 
 import com.clearspring.analytics.hash.MurmurHash
 
+import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.types._
@@ -55,6 +56,22 @@ case class HyperLogLogPlusPlus(
   extends ImperativeAggregate {
   import HyperLogLogPlusPlus._
 
+  def this(child: Expression) = {
+    this(child = child, relativeSD = 0.05, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+  }
+
+  def this(child: Expression, relativeSD: Expression) = {
+    this(
+      child = child,
+      relativeSD = relativeSD match {
+        case Literal(d: Double, DoubleType) => d
+        case _ =>
+          throw new AnalysisException("The second argument should be a double literal.")
+      },
+      mutableAggBufferOffset = 0,
+      inputAggBufferOffset = 0)
+  }
+
   override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
     copy(mutableAggBufferOffset = newMutableAggBufferOffset)
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
index 6da39e7143447..bae78d98493b8 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
@@ -24,6 +24,8 @@ case class Kurtosis(child: Expression,
     inputAggBufferOffset: Int = 0)
   extends CentralMomentAgg(child) {
 
+  def this(child: Expression) = this(child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+
   override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
     copy(mutableAggBufferOffset = newMutableAggBufferOffset)
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Last.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Last.scala
index 8636bfe8d07aa..be7e12d7a2336 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Last.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Last.scala
@@ -51,15 +51,15 @@ case class Last(child: Expression, ignoreNullsExpr: Expression) extends Declarat
   // Expected input data type.
   override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
 
-  private val last = AttributeReference("last", child.dataType)()
+  private lazy val last = AttributeReference("last", child.dataType)()
 
-  override val aggBufferAttributes: Seq[AttributeReference] = last :: Nil
+  override lazy val aggBufferAttributes: Seq[AttributeReference] = last :: Nil
 
-  override val initialValues: Seq[Literal] = Seq(
+  override lazy val initialValues: Seq[Literal] = Seq(
     /* last = */ Literal.create(null, child.dataType)
   )
 
-  override val updateExpressions: Seq[Expression] = {
+  override lazy val updateExpressions: Seq[Expression] = {
     if (ignoreNulls) {
       Seq(
         /* last = */ If(IsNull(child), last, child)
@@ -71,7 +71,7 @@ case class Last(child: Expression, ignoreNullsExpr: Expression) extends Declarat
     }
   }
 
-  override val mergeExpressions: Seq[Expression] = {
+  override lazy val mergeExpressions: Seq[Expression] = {
     if (ignoreNulls) {
       Seq(
         /* last = */ If(IsNull(last.right), last.left, last.right)
@@ -83,7 +83,7 @@ case class Last(child: Expression, ignoreNullsExpr: Expression) extends Declarat
     }
   }
 
-  override val evaluateExpression: AttributeReference = last
+  override lazy val evaluateExpression: AttributeReference = last
 
   override def toString: String = s"last($child)${if (ignoreNulls) " ignore nulls"}"
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
index b9d75ad452838..61cae44cd0f5b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
@@ -17,7 +17,9 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
 
 case class Max(child: Expression) extends DeclarativeAggregate {
@@ -32,24 +34,27 @@ case class Max(child: Expression) extends DeclarativeAggregate {
   // Expected input data type.
   override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
 
-  private val max = AttributeReference("max", child.dataType)()
+  override def checkInputDataTypes(): TypeCheckResult =
+    TypeUtils.checkForOrderingExpr(child.dataType, "function max")
 
-  override val aggBufferAttributes: Seq[AttributeReference] = max :: Nil
+  private lazy val max = AttributeReference("max", child.dataType)()
 
-  override val initialValues: Seq[Literal] = Seq(
+  override lazy val aggBufferAttributes: Seq[AttributeReference] = max :: Nil
+
+  override lazy val initialValues: Seq[Literal] = Seq(
     /* max = */ Literal.create(null, child.dataType)
   )
 
-  override val updateExpressions: Seq[Expression] = Seq(
+  override lazy val updateExpressions: Seq[Expression] = Seq(
     /* max = */ If(IsNull(child), max, If(IsNull(max), child, Greatest(Seq(max, child))))
   )
 
-  override val mergeExpressions: Seq[Expression] = {
+  override lazy val mergeExpressions: Seq[Expression] = {
     val greatest = Greatest(Seq(max.left, max.right))
     Seq(
       /* max = */ If(IsNull(max.right), max.left, If(IsNull(max.left), max.right, greatest))
     )
   }
 
-  override val evaluateExpression: AttributeReference = max
+  override lazy val evaluateExpression: AttributeReference = max
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
index 5ed9cd348daba..242456d9e2e18 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
@@ -17,7 +17,9 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
 
 
@@ -33,24 +35,27 @@ case class Min(child: Expression) extends DeclarativeAggregate {
   // Expected input data type.
   override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
 
-  private val min = AttributeReference("min", child.dataType)()
+  override def checkInputDataTypes(): TypeCheckResult =
+    TypeUtils.checkForOrderingExpr(child.dataType, "function min")
 
-  override val aggBufferAttributes: Seq[AttributeReference] = min :: Nil
+  private lazy val min = AttributeReference("min", child.dataType)()
 
-  override val initialValues: Seq[Expression] = Seq(
+  override lazy val aggBufferAttributes: Seq[AttributeReference] = min :: Nil
+
+  override lazy val initialValues: Seq[Expression] = Seq(
     /* min = */ Literal.create(null, child.dataType)
   )
 
-  override val updateExpressions: Seq[Expression] = Seq(
+  override lazy val updateExpressions: Seq[Expression] = Seq(
     /* min = */ If(IsNull(child), min, If(IsNull(min), child, Least(Seq(min, child))))
   )
 
-  override val mergeExpressions: Seq[Expression] = {
+  override lazy val mergeExpressions: Seq[Expression] = {
     val least = Least(Seq(min.left, min.right))
     Seq(
       /* min = */ If(IsNull(min.right), min.left, If(IsNull(min.left), min.right, least))
     )
   }
 
-  override val evaluateExpression: AttributeReference = min
+  override lazy val evaluateExpression: AttributeReference = min
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
index 0def7ddfd9d3d..c593074fa2479 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
@@ -24,6 +24,8 @@ case class Skewness(child: Expression,
     inputAggBufferOffset: Int = 0)
   extends CentralMomentAgg(child) {
 
+  def this(child: Expression) = this(child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+
   override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
     copy(mutableAggBufferOffset = newMutableAggBufferOffset)
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
index 3f47ffe13cbc8..5b9eb7ae02f25 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
@@ -17,8 +17,10 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
 
 
@@ -48,29 +50,26 @@ abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
 
   override def dataType: DataType = resultType
 
-  // Expected input data type.
-  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
-  // new version at planning time (after analysis phase). For now, NullType is added at here
-  // to make it resolved when we have cases like `select stddev(null)`.
-  // We can use our analyzer to cast NullType to the default data type of the NumericType once
-  // we remove the old aggregate functions. Then, we will not need NullType at here.
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType, NullType))
+  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType))
 
-  private val resultType = DoubleType
+  override def checkInputDataTypes(): TypeCheckResult =
+    TypeUtils.checkForNumericExpr(child.dataType, "function stddev")
 
-  private val count = AttributeReference("count", resultType)()
-  private val avg = AttributeReference("avg", resultType)()
-  private val mk = AttributeReference("mk", resultType)()
+  private lazy val resultType = DoubleType
 
-  override val aggBufferAttributes = count :: avg :: mk :: Nil
+  private lazy val count = AttributeReference("count", resultType)()
+  private lazy val avg = AttributeReference("avg", resultType)()
+  private lazy val mk = AttributeReference("mk", resultType)()
 
-  override val initialValues: Seq[Expression] = Seq(
+  override lazy val aggBufferAttributes = count :: avg :: mk :: Nil
+
+  override lazy val initialValues: Seq[Expression] = Seq(
     /* count = */ Cast(Literal(0), resultType),
     /* avg = */ Cast(Literal(0), resultType),
     /* mk = */ Cast(Literal(0), resultType)
   )
 
-  override val updateExpressions: Seq[Expression] = {
+  override lazy val updateExpressions: Seq[Expression] = {
     val value = Cast(child, resultType)
     val newCount = count + Cast(Literal(1), resultType)
 
@@ -89,7 +88,7 @@ abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
     )
   }
 
-  override val mergeExpressions: Seq[Expression] = {
+  override lazy val mergeExpressions: Seq[Expression] = {
 
     // count merge
     val newCount = count.left + count.right
@@ -114,7 +113,7 @@ abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
     )
   }
 
-  override val evaluateExpression: Expression = {
+  override lazy val evaluateExpression: Expression = {
     // when count == 0, return null
     // when count == 1, return 0
     // when count >1
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
index 7f8adbc56ad1d..c005ec9657211 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
@@ -17,7 +17,9 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
 
 case class Sum(child: Expression) extends DeclarativeAggregate {
@@ -29,16 +31,13 @@ case class Sum(child: Expression) extends DeclarativeAggregate {
   // Return data type.
   override def dataType: DataType = resultType
 
-  // Expected input data type.
-  // TODO: Right now, we replace old aggregate functions (based on AggregateExpression1) to the
-  // new version at planning time (after analysis phase). For now, NullType is added at here
-  // to make it resolved when we have cases like `select sum(null)`.
-  // We can use our analyzer to cast NullType to the default data type of the NumericType once
-  // we remove the old aggregate functions. Then, we will not need NullType at here.
   override def inputTypes: Seq[AbstractDataType] =
     Seq(TypeCollection(LongType, DoubleType, DecimalType, NullType))
 
-  private val resultType = child.dataType match {
+  override def checkInputDataTypes(): TypeCheckResult =
+    TypeUtils.checkForNumericExpr(child.dataType, "function sum")
+
+  private lazy val resultType = child.dataType match {
     case DecimalType.Fixed(precision, scale) =>
       DecimalType.bounded(precision + 10, scale)
     // TODO: Remove this line once we remove the NullType from inputTypes.
@@ -46,24 +45,24 @@ case class Sum(child: Expression) extends DeclarativeAggregate {
     case _ => child.dataType
   }
 
-  private val sumDataType = resultType
+  private lazy val sumDataType = resultType
 
-  private val sum = AttributeReference("sum", sumDataType)()
+  private lazy val sum = AttributeReference("sum", sumDataType)()
 
-  private val zero = Cast(Literal(0), sumDataType)
+  private lazy val zero = Cast(Literal(0), sumDataType)
 
-  override val aggBufferAttributes = sum :: Nil
+  override lazy val aggBufferAttributes = sum :: Nil
 
-  override val initialValues: Seq[Expression] = Seq(
+  override lazy val initialValues: Seq[Expression] = Seq(
     /* sum = */ Literal.create(null, sumDataType)
   )
 
-  override val updateExpressions: Seq[Expression] = Seq(
+  override lazy val updateExpressions: Seq[Expression] = Seq(
     /* sum = */
     Coalesce(Seq(Add(Coalesce(Seq(sum, zero)), Cast(child, sumDataType)), sum))
   )
 
-  override val mergeExpressions: Seq[Expression] = {
+  override lazy val mergeExpressions: Seq[Expression] = {
     val add = Add(Coalesce(Seq(sum.left, zero)), Cast(sum.right, sumDataType))
     Seq(
       /* sum = */
@@ -71,5 +70,5 @@ case class Sum(child: Expression) extends DeclarativeAggregate {
     )
   }
 
-  override val evaluateExpression: Expression = Cast(sum, resultType)
+  override lazy val evaluateExpression: Expression = Cast(sum, resultType)
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala
index ec63534e5290a..ede2da2805966 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala
@@ -24,6 +24,8 @@ case class VarianceSamp(child: Expression,
     inputAggBufferOffset: Int = 0)
   extends CentralMomentAgg(child) {
 
+  def this(child: Expression) = this(child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+
   override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
     copy(mutableAggBufferOffset = newMutableAggBufferOffset)
 
@@ -42,11 +44,14 @@ case class VarianceSamp(child: Expression,
   }
 }
 
-case class VariancePop(child: Expression,
+case class VariancePop(
+    child: Expression,
     mutableAggBufferOffset: Int = 0,
     inputAggBufferOffset: Int = 0)
   extends CentralMomentAgg(child) {
 
+  def this(child: Expression) = this(child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
+
   override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
     copy(mutableAggBufferOffset = newMutableAggBufferOffset)
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala
index 5c5b3d1ccd3cd..3b441de34a49f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala
@@ -17,23 +17,24 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.{GeneratedExpressionCode, CodeGenContext}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.types._
 
-/** The mode of an [[AggregateFunction2]]. */
+/** The mode of an [[AggregateFunction]]. */
 private[sql] sealed trait AggregateMode
 
 /**
- * An [[AggregateFunction2]] with [[Partial]] mode is used for partial aggregation.
+ * An [[AggregateFunction]] with [[Partial]] mode is used for partial aggregation.
  * This function updates the given aggregation buffer with the original input of this
  * function. When it has processed all input rows, the aggregation buffer is returned.
  */
 private[sql] case object Partial extends AggregateMode
 
 /**
- * An [[AggregateFunction2]] with [[PartialMerge]] mode is used to merge aggregation buffers
+ * An [[AggregateFunction]] with [[PartialMerge]] mode is used to merge aggregation buffers
  * containing intermediate results for this function.
  * This function updates the given aggregation buffer by merging multiple aggregation buffers.
  * When it has processed all input rows, the aggregation buffer is returned.
@@ -41,7 +42,7 @@ private[sql] case object Partial extends AggregateMode
 private[sql] case object PartialMerge extends AggregateMode
 
 /**
- * An [[AggregateFunction2]] with [[Final]] mode is used to merge aggregation buffers
+ * An [[AggregateFunction]] with [[Final]] mode is used to merge aggregation buffers
  * containing intermediate results for this function and then generate final result.
  * This function updates the given aggregation buffer by merging multiple aggregation buffers.
  * When it has processed all input rows, the final result of this function is returned.
@@ -49,7 +50,7 @@ private[sql] case object PartialMerge extends AggregateMode
 private[sql] case object Final extends AggregateMode
 
 /**
- * An [[AggregateFunction2]] with [[Complete]] mode is used to evaluate this function directly
+ * An [[AggregateFunction]] with [[Complete]] mode is used to evaluate this function directly
  * from original input rows without any partial aggregation.
  * This function updates the given aggregation buffer with the original input of this
  * function. When it has processed all input rows, the final result of this function is returned.
@@ -67,13 +68,15 @@ private[sql] case object NoOp extends Expression with Unevaluable {
 }
 
 /**
- * A container for an [[AggregateFunction2]] with its [[AggregateMode]] and a field
+ * A container for an [[AggregateFunction]] with its [[AggregateMode]] and a field
  * (`isDistinct`) indicating if DISTINCT keyword is specified for this function.
  */
-private[sql] case class AggregateExpression2(
-    aggregateFunction: AggregateFunction2,
+private[sql] case class AggregateExpression(
+    aggregateFunction: AggregateFunction,
     mode: AggregateMode,
-    isDistinct: Boolean) extends AggregateExpression {
+    isDistinct: Boolean)
+  extends Expression
+  with Unevaluable {
 
   override def children: Seq[Expression] = aggregateFunction :: Nil
   override def dataType: DataType = aggregateFunction.dataType
@@ -89,6 +92,8 @@ private[sql] case class AggregateExpression2(
     AttributeSet(childReferences)
   }
 
+  override def prettyString: String = aggregateFunction.prettyString
+
   override def toString: String = s"(${aggregateFunction},mode=$mode,isDistinct=$isDistinct)"
 }
 
@@ -106,10 +111,10 @@ private[sql] case class AggregateExpression2(
  * combined aggregation buffer which concatenates the aggregation buffers of the individual
  * aggregate functions.
  *
- * Code which accepts [[AggregateFunction2]] instances should be prepared to handle both types of
+ * Code which accepts [[AggregateFunction]] instances should be prepared to handle both types of
  * aggregate functions.
  */
-sealed abstract class AggregateFunction2 extends Expression with ImplicitCastInputTypes {
+sealed abstract class AggregateFunction extends Expression with ImplicitCastInputTypes {
 
   /** An aggregate function is not foldable. */
   final override def foldable: Boolean = false
@@ -141,6 +146,27 @@ sealed abstract class AggregateFunction2 extends Expression with ImplicitCastInp
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String =
     throw new UnsupportedOperationException(s"Cannot evaluate expression: $this")
+
+  /**
+   * Wraps this [[AggregateFunction]] in an [[AggregateExpression]] because
+   * [[AggregateExpression]] is the container of an [[AggregateFunction]], aggregation mode,
+   * and the flag indicating if this aggregation is distinct aggregation or not.
+   * An [[AggregateFunction]] should not be used without being wrapped in
+   * an [[AggregateExpression]].
+   */
+  def toAggregateExpression(): AggregateExpression = toAggregateExpression(isDistinct = false)
+
+  /**
+   * Wraps this [[AggregateFunction]] in an [[AggregateExpression]] and set isDistinct
+   * field of the [[AggregateExpression]] to the given value because
+   * [[AggregateExpression]] is the container of an [[AggregateFunction]], aggregation mode,
+   * and the flag indicating if this aggregation is distinct aggregation or not.
+   * An [[AggregateFunction]] should not be used without being wrapped in
+   * an [[AggregateExpression]].
+   */
+  def toAggregateExpression(isDistinct: Boolean): AggregateExpression = {
+    AggregateExpression(aggregateFunction = this, mode = Complete, isDistinct = isDistinct)
+  }
 }
 
 /**
@@ -161,7 +187,7 @@ sealed abstract class AggregateFunction2 extends Expression with ImplicitCastInp
  * `inputAggBufferOffset`, but not on the correctness of the attribute ids in `aggBufferAttributes`
  * and `inputAggBufferAttributes`.
  */
-abstract class ImperativeAggregate extends AggregateFunction2 {
+abstract class ImperativeAggregate extends AggregateFunction {
 
   /**
    * The offset of this function's first buffer value in the underlying shared mutable aggregation
@@ -258,9 +284,14 @@ abstract class ImperativeAggregate extends AggregateFunction2 {
  * `bufferAttributes`, defining attributes for the fields of the mutable aggregation buffer. You
  * can then use these attributes when defining `updateExpressions`, `mergeExpressions`, and
  * `evaluateExpressions`.
+ *
+ * Please note that children of an aggregate function can be unresolved (it will happen when
+ * we create this function in DataFrame API). So, if there is any fields in
+ * the implemented class that need to access fields of its children, please make
+ * those fields `lazy val`s.
  */
 abstract class DeclarativeAggregate
-  extends AggregateFunction2
+  extends AggregateFunction
   with Serializable
   with Unevaluable {
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
deleted file mode 100644
index 3dcf7915d77b3..0000000000000
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregates.scala
+++ /dev/null
@@ -1,1073 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.catalyst.expressions
-
-import com.clearspring.analytics.stream.cardinality.HyperLogLog
-
-import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
-import org.apache.spark.sql.catalyst.expressions.codegen.{CodeGenContext, GeneratedExpressionCode}
-import org.apache.spark.sql.catalyst.util.{GenericArrayData, ArrayData, TypeUtils}
-import org.apache.spark.sql.types._
-import org.apache.spark.util.collection.OpenHashSet
-
-
-trait AggregateExpression extends Expression with Unevaluable
-
-trait AggregateExpression1 extends AggregateExpression {
-
-  /**
-   * Aggregate expressions should not be foldable.
-   */
-  override def foldable: Boolean = false
-
-  /**
-   * Creates a new instance that can be used to compute this aggregate expression for a group
-   * of input rows/
-   */
-  def newInstance(): AggregateFunction1
-}
-
-/**
- * Represents an aggregation that has been rewritten to be performed in two steps.
- *
- * @param finalEvaluation an aggregate expression that evaluates to same final result as the
- *                        original aggregation.
- * @param partialEvaluations A sequence of [[NamedExpression]]s that can be computed on partial
- *                           data sets and are required to compute the `finalEvaluation`.
- */
-case class SplitEvaluation(
-    finalEvaluation: Expression,
-    partialEvaluations: Seq[NamedExpression])
-
-/**
- * An [[AggregateExpression1]] that can be partially computed without seeing all relevant tuples.
- * These partial evaluations can then be combined to compute the actual answer.
- */
-trait PartialAggregate1 extends AggregateExpression1 {
-
-  /**
-   * Returns a [[SplitEvaluation]] that computes this aggregation using partial aggregation.
-   */
-  def asPartial: SplitEvaluation
-}
-
-/**
- * A specific implementation of an aggregate function. Used to wrap a generic
- * [[AggregateExpression1]] with an algorithm that will be used to compute one specific result.
- */
-abstract class AggregateFunction1 extends LeafExpression with Serializable {
-
-  /** Base should return the generic aggregate expression that this function is computing */
-  val base: AggregateExpression1
-
-  override def nullable: Boolean = base.nullable
-  override def dataType: DataType = base.dataType
-
-  def update(input: InternalRow): Unit
-
-  override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    throw new UnsupportedOperationException(
-      "AggregateFunction1 should not be used for generated aggregates")
-  }
-}
-
-case class Min(child: Expression) extends UnaryExpression with PartialAggregate1 {
-
-  override def nullable: Boolean = true
-  override def dataType: DataType = child.dataType
-
-  override def asPartial: SplitEvaluation = {
-    val partialMin = Alias(Min(child), "PartialMin")()
-    SplitEvaluation(Min(partialMin.toAttribute), partialMin :: Nil)
-  }
-
-  override def newInstance(): MinFunction = new MinFunction(child, this)
-
-  override def checkInputDataTypes(): TypeCheckResult =
-    TypeUtils.checkForOrderingExpr(child.dataType, "function min")
-}
-
-case class MinFunction(expr: Expression, base: AggregateExpression1) extends AggregateFunction1 {
-  def this() = this(null, null) // Required for serialization.
-
-  val currentMin: MutableLiteral = MutableLiteral(null, expr.dataType)
-  val cmp = GreaterThan(currentMin, expr)
-
-  override def update(input: InternalRow): Unit = {
-    if (currentMin.value == null) {
-      currentMin.value = expr.eval(input)
-    } else if (cmp.eval(input) == true) {
-      currentMin.value = expr.eval(input)
-    }
-  }
-
-  override def eval(input: InternalRow): Any = currentMin.value
-}
-
-case class Max(child: Expression) extends UnaryExpression with PartialAggregate1 {
-
-  override def nullable: Boolean = true
-  override def dataType: DataType = child.dataType
-
-  override def asPartial: SplitEvaluation = {
-    val partialMax = Alias(Max(child), "PartialMax")()
-    SplitEvaluation(Max(partialMax.toAttribute), partialMax :: Nil)
-  }
-
-  override def newInstance(): MaxFunction = new MaxFunction(child, this)
-
-  override def checkInputDataTypes(): TypeCheckResult =
-    TypeUtils.checkForOrderingExpr(child.dataType, "function max")
-}
-
-case class MaxFunction(expr: Expression, base: AggregateExpression1) extends AggregateFunction1 {
-  def this() = this(null, null) // Required for serialization.
-
-  val currentMax: MutableLiteral = MutableLiteral(null, expr.dataType)
-  val cmp = LessThan(currentMax, expr)
-
-  override def update(input: InternalRow): Unit = {
-    if (currentMax.value == null) {
-      currentMax.value = expr.eval(input)
-    } else if (cmp.eval(input) == true) {
-      currentMax.value = expr.eval(input)
-    }
-  }
-
-  override def eval(input: InternalRow): Any = currentMax.value
-}
-
-case class Count(child: Expression) extends UnaryExpression with PartialAggregate1 {
-
-  override def nullable: Boolean = false
-  override def dataType: LongType.type = LongType
-
-  override def asPartial: SplitEvaluation = {
-    val partialCount = Alias(Count(child), "PartialCount")()
-    SplitEvaluation(Coalesce(Seq(Sum(partialCount.toAttribute), Literal(0L))), partialCount :: Nil)
-  }
-
-  override def newInstance(): CountFunction = new CountFunction(child, this)
-}
-
-case class CountFunction(expr: Expression, base: AggregateExpression1) extends AggregateFunction1 {
-  def this() = this(null, null) // Required for serialization.
-
-  var count: Long = _
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = expr.eval(input)
-    if (evaluatedExpr != null) {
-      count += 1L
-    }
-  }
-
-  override def eval(input: InternalRow): Any = count
-}
-
-case class CountDistinct(expressions: Seq[Expression]) extends PartialAggregate1 {
-  def this() = this(null)
-
-  override def children: Seq[Expression] = expressions
-
-  override def nullable: Boolean = false
-  override def dataType: DataType = LongType
-  override def toString: String = s"COUNT(DISTINCT ${expressions.mkString(",")})"
-  override def newInstance(): CountDistinctFunction = new CountDistinctFunction(expressions, this)
-
-  override def asPartial: SplitEvaluation = {
-    val partialSet = Alias(CollectHashSet(expressions), "partialSets")()
-    SplitEvaluation(
-      CombineSetsAndCount(partialSet.toAttribute),
-      partialSet :: Nil)
-  }
-}
-
-case class CountDistinctFunction(
-    @transient expr: Seq[Expression],
-    @transient base: AggregateExpression1)
-  extends AggregateFunction1 {
-
-  def this() = this(null, null) // Required for serialization.
-
-  val seen = new OpenHashSet[Any]()
-
-  @transient
-  val distinctValue = new InterpretedProjection(expr)
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = distinctValue(input)
-    if (!evaluatedExpr.anyNull) {
-      seen.add(evaluatedExpr)
-    }
-  }
-
-  override def eval(input: InternalRow): Any = seen.size.toLong
-}
-
-case class CollectHashSet(expressions: Seq[Expression]) extends AggregateExpression1 {
-  def this() = this(null)
-
-  override def children: Seq[Expression] = expressions
-  override def nullable: Boolean = false
-  override def dataType: OpenHashSetUDT = new OpenHashSetUDT(expressions.head.dataType)
-  override def toString: String = s"AddToHashSet(${expressions.mkString(",")})"
-  override def newInstance(): CollectHashSetFunction =
-    new CollectHashSetFunction(expressions, this)
-}
-
-case class CollectHashSetFunction(
-    @transient expr: Seq[Expression],
-    @transient base: AggregateExpression1)
-  extends AggregateFunction1 {
-
-  def this() = this(null, null) // Required for serialization.
-
-  val seen = new OpenHashSet[Any]()
-
-  @transient
-  val distinctValue = new InterpretedProjection(expr)
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = distinctValue(input)
-    if (!evaluatedExpr.anyNull) {
-      seen.add(evaluatedExpr)
-    }
-  }
-
-  override def eval(input: InternalRow): Any = {
-    seen
-  }
-}
-
-case class CombineSetsAndCount(inputSet: Expression) extends AggregateExpression1 {
-  def this() = this(null)
-
-  override def children: Seq[Expression] = inputSet :: Nil
-  override def nullable: Boolean = false
-  override def dataType: DataType = LongType
-  override def toString: String = s"CombineAndCount($inputSet)"
-  override def newInstance(): CombineSetsAndCountFunction = {
-    new CombineSetsAndCountFunction(inputSet, this)
-  }
-}
-
-case class CombineSetsAndCountFunction(
-    @transient inputSet: Expression,
-    @transient base: AggregateExpression1)
-  extends AggregateFunction1 {
-
-  def this() = this(null, null) // Required for serialization.
-
-  val seen = new OpenHashSet[Any]()
-
-  override def update(input: InternalRow): Unit = {
-    val inputSetEval = inputSet.eval(input).asInstanceOf[OpenHashSet[Any]]
-    val inputIterator = inputSetEval.iterator
-    while (inputIterator.hasNext) {
-      seen.add(inputIterator.next)
-    }
-  }
-
-  override def eval(input: InternalRow): Any = seen.size.toLong
-}
-
-/** The data type of ApproxCountDistinctPartition since its output is a HyperLogLog object. */
-private[sql] case object HyperLogLogUDT extends UserDefinedType[HyperLogLog] {
-
-  override def sqlType: DataType = BinaryType
-
-  /** Since we are using HyperLogLog internally, usually it will not be called. */
-  override def serialize(obj: Any): Array[Byte] =
-    obj.asInstanceOf[HyperLogLog].getBytes
-
-
-  /** Since we are using HyperLogLog internally, usually it will not be called. */
-  override def deserialize(datum: Any): HyperLogLog =
-    HyperLogLog.Builder.build(datum.asInstanceOf[Array[Byte]])
-
-  override def userClass: Class[HyperLogLog] = classOf[HyperLogLog]
-}
-
-case class ApproxCountDistinctPartition(child: Expression, relativeSD: Double)
-  extends UnaryExpression with AggregateExpression1 {
-
-  override def nullable: Boolean = false
-  override def dataType: DataType = HyperLogLogUDT
-  override def toString: String = s"APPROXIMATE COUNT(DISTINCT $child)"
-  override def newInstance(): ApproxCountDistinctPartitionFunction = {
-    new ApproxCountDistinctPartitionFunction(child, this, relativeSD)
-  }
-}
-
-case class ApproxCountDistinctPartitionFunction(
-    expr: Expression,
-    base: AggregateExpression1,
-    relativeSD: Double)
-  extends AggregateFunction1 {
-  def this() = this(null, null, 0) // Required for serialization.
-
-  private val hyperLogLog = new HyperLogLog(relativeSD)
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = expr.eval(input)
-    if (evaluatedExpr != null) {
-      hyperLogLog.offer(evaluatedExpr)
-    }
-  }
-
-  override def eval(input: InternalRow): Any = hyperLogLog
-}
-
-case class ApproxCountDistinctMerge(child: Expression, relativeSD: Double)
-  extends UnaryExpression with AggregateExpression1 {
-
-  override def nullable: Boolean = false
-  override def dataType: LongType.type = LongType
-  override def toString: String = s"APPROXIMATE COUNT(DISTINCT $child)"
-  override def newInstance(): ApproxCountDistinctMergeFunction = {
-    new ApproxCountDistinctMergeFunction(child, this, relativeSD)
-  }
-}
-
-case class ApproxCountDistinctMergeFunction(
-    expr: Expression,
-    base: AggregateExpression1,
-    relativeSD: Double)
-  extends AggregateFunction1 {
-  def this() = this(null, null, 0) // Required for serialization.
-
-  private val hyperLogLog = new HyperLogLog(relativeSD)
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = expr.eval(input)
-    hyperLogLog.addAll(evaluatedExpr.asInstanceOf[HyperLogLog])
-  }
-
-  override def eval(input: InternalRow): Any = hyperLogLog.cardinality()
-}
-
-case class ApproxCountDistinct(child: Expression, relativeSD: Double = 0.05)
-  extends UnaryExpression with PartialAggregate1 {
-
-  override def nullable: Boolean = false
-  override def dataType: LongType.type = LongType
-  override def toString: String = s"APPROXIMATE COUNT(DISTINCT $child)"
-
-  override def asPartial: SplitEvaluation = {
-    val partialCount =
-      Alias(ApproxCountDistinctPartition(child, relativeSD), "PartialApproxCountDistinct")()
-
-    SplitEvaluation(
-      ApproxCountDistinctMerge(partialCount.toAttribute, relativeSD),
-      partialCount :: Nil)
-  }
-
-  override def newInstance(): CountDistinctFunction = new CountDistinctFunction(child :: Nil, this)
-}
-
-case class Average(child: Expression) extends UnaryExpression with PartialAggregate1 {
-
-  override def prettyName: String = "avg"
-
-  override def nullable: Boolean = true
-
-  override def dataType: DataType = child.dataType match {
-    case DecimalType.Fixed(precision, scale) =>
-      // Add 4 digits after decimal point, like Hive
-      DecimalType.bounded(precision + 4, scale + 4)
-    case _ =>
-      DoubleType
-  }
-
-  override def asPartial: SplitEvaluation = {
-    child.dataType match {
-      case DecimalType.Fixed(precision, scale) =>
-        val partialSum = Alias(Sum(child), "PartialSum")()
-        val partialCount = Alias(Count(child), "PartialCount")()
-
-        // partialSum already increase the precision by 10
-        val castedSum = Cast(Sum(partialSum.toAttribute), partialSum.dataType)
-        val castedCount = Cast(Sum(partialCount.toAttribute), partialSum.dataType)
-        SplitEvaluation(
-          Cast(Divide(castedSum, castedCount), dataType),
-          partialCount :: partialSum :: Nil)
-
-      case _ =>
-        val partialSum = Alias(Sum(child), "PartialSum")()
-        val partialCount = Alias(Count(child), "PartialCount")()
-
-        val castedSum = Cast(Sum(partialSum.toAttribute), dataType)
-        val castedCount = Cast(Sum(partialCount.toAttribute), dataType)
-        SplitEvaluation(
-          Divide(castedSum, castedCount),
-          partialCount :: partialSum :: Nil)
-    }
-  }
-
-  override def newInstance(): AverageFunction = new AverageFunction(child, this)
-
-  override def checkInputDataTypes(): TypeCheckResult =
-    TypeUtils.checkForNumericExpr(child.dataType, "function average")
-}
-
-case class AverageFunction(expr: Expression, base: AggregateExpression1)
-  extends AggregateFunction1 {
-
-  def this() = this(null, null) // Required for serialization.
-
-  private val calcType =
-    expr.dataType match {
-      case DecimalType.Fixed(precision, scale) =>
-        DecimalType.bounded(precision + 10, scale)
-      case _ =>
-        expr.dataType
-    }
-
-  private val zero = Cast(Literal(0), calcType)
-
-  private var count: Long = _
-  private val sum = MutableLiteral(zero.eval(null), calcType)
-
-  private def addFunction(value: Any) = Add(sum,
-    Cast(Literal.create(value, expr.dataType), calcType))
-
-  override def eval(input: InternalRow): Any = {
-    if (count == 0L) {
-      null
-    } else {
-      expr.dataType match {
-        case DecimalType.Fixed(precision, scale) =>
-          val dt = DecimalType.bounded(precision + 14, scale + 4)
-          Cast(Divide(Cast(sum, dt), Cast(Literal(count), dt)), dataType).eval(null)
-        case _ =>
-          Divide(
-            Cast(sum, dataType),
-            Cast(Literal(count), dataType)).eval(null)
-      }
-    }
-  }
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = expr.eval(input)
-    if (evaluatedExpr != null) {
-      count += 1
-      sum.update(addFunction(evaluatedExpr), input)
-    }
-  }
-}
-
-case class Sum(child: Expression) extends UnaryExpression with PartialAggregate1 {
-
-  override def nullable: Boolean = true
-
-  override def dataType: DataType = child.dataType match {
-    case DecimalType.Fixed(precision, scale) =>
-      // Add 10 digits left of decimal point, like Hive
-      DecimalType.bounded(precision + 10, scale)
-    case _ =>
-      child.dataType
-  }
-
-  override def asPartial: SplitEvaluation = {
-    child.dataType match {
-      case DecimalType.Fixed(_, _) =>
-        val partialSum = Alias(Sum(child), "PartialSum")()
-        SplitEvaluation(
-          Cast(Sum(partialSum.toAttribute), dataType),
-          partialSum :: Nil)
-
-      case _ =>
-        val partialSum = Alias(Sum(child), "PartialSum")()
-        SplitEvaluation(
-          Sum(partialSum.toAttribute),
-          partialSum :: Nil)
-    }
-  }
-
-  override def newInstance(): SumFunction = new SumFunction(child, this)
-
-  override def checkInputDataTypes(): TypeCheckResult =
-    TypeUtils.checkForNumericExpr(child.dataType, "function sum")
-}
-
-case class SumFunction(expr: Expression, base: AggregateExpression1) extends AggregateFunction1 {
-  def this() = this(null, null) // Required for serialization.
-
-  private val calcType =
-    expr.dataType match {
-      case DecimalType.Fixed(precision, scale) =>
-        DecimalType.bounded(precision + 10, scale)
-      case _ =>
-        expr.dataType
-    }
-
-  private val zero = Cast(Literal(0), calcType)
-
-  private val sum = MutableLiteral(null, calcType)
-
-  private val addFunction = Coalesce(Seq(Add(Coalesce(Seq(sum, zero)), Cast(expr, calcType)), sum))
-
-  override def update(input: InternalRow): Unit = {
-    sum.update(addFunction, input)
-  }
-
-  override def eval(input: InternalRow): Any = {
-    expr.dataType match {
-      case DecimalType.Fixed(_, _) =>
-        Cast(sum, dataType).eval(null)
-      case _ => sum.eval(null)
-    }
-  }
-}
-
-case class SumDistinct(child: Expression) extends UnaryExpression with PartialAggregate1 {
-
-  def this() = this(null)
-  override def nullable: Boolean = true
-  override def dataType: DataType = child.dataType match {
-    case DecimalType.Fixed(precision, scale) =>
-      // Add 10 digits left of decimal point, like Hive
-      DecimalType.bounded(precision + 10, scale)
-    case _ =>
-      child.dataType
-  }
-  override def toString: String = s"sum(distinct $child)"
-  override def newInstance(): SumDistinctFunction = new SumDistinctFunction(child, this)
-
-  override def asPartial: SplitEvaluation = {
-    val partialSet = Alias(CollectHashSet(child :: Nil), "partialSets")()
-    SplitEvaluation(
-      CombineSetsAndSum(partialSet.toAttribute, this),
-      partialSet :: Nil)
-  }
-
-  override def checkInputDataTypes(): TypeCheckResult =
-    TypeUtils.checkForNumericExpr(child.dataType, "function sumDistinct")
-}
-
-case class SumDistinctFunction(expr: Expression, base: AggregateExpression1)
-  extends AggregateFunction1 {
-
-  def this() = this(null, null) // Required for serialization.
-
-  private val seen = new scala.collection.mutable.HashSet[Any]()
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = expr.eval(input)
-    if (evaluatedExpr != null) {
-      seen += evaluatedExpr
-    }
-  }
-
-  override def eval(input: InternalRow): Any = {
-    if (seen.size == 0) {
-      null
-    } else {
-      Cast(Literal(
-        seen.reduceLeft(
-          dataType.asInstanceOf[NumericType].numeric.asInstanceOf[Numeric[Any]].plus)),
-        dataType).eval(null)
-    }
-  }
-}
-
-case class CombineSetsAndSum(inputSet: Expression, base: Expression) extends AggregateExpression1 {
-  def this() = this(null, null)
-
-  override def children: Seq[Expression] = inputSet :: Nil
-  override def nullable: Boolean = true
-  override def dataType: DataType = base.dataType
-  override def toString: String = s"CombineAndSum($inputSet)"
-  override def newInstance(): CombineSetsAndSumFunction = {
-    new CombineSetsAndSumFunction(inputSet, this)
-  }
-}
-
-case class CombineSetsAndSumFunction(
-    @transient inputSet: Expression,
-    @transient base: AggregateExpression1)
-  extends AggregateFunction1 {
-
-  def this() = this(null, null) // Required for serialization.
-
-  val seen = new OpenHashSet[Any]()
-
-  override def update(input: InternalRow): Unit = {
-    val inputSetEval = inputSet.eval(input).asInstanceOf[OpenHashSet[Any]]
-    val inputIterator = inputSetEval.iterator
-    while (inputIterator.hasNext) {
-      seen.add(inputIterator.next())
-    }
-  }
-
-  override def eval(input: InternalRow): Any = {
-    val casted = seen.asInstanceOf[OpenHashSet[InternalRow]]
-    if (casted.size == 0) {
-      null
-    } else {
-      Cast(Literal(
-        casted.iterator.map(f => f.get(0, null)).reduceLeft(
-          base.dataType.asInstanceOf[NumericType].numeric.asInstanceOf[Numeric[Any]].plus)),
-        base.dataType).eval(null)
-    }
-  }
-}
-
-case class First(
-    child: Expression,
-    ignoreNullsExpr: Expression)
-  extends UnaryExpression with PartialAggregate1 {
-
-  def this(child: Expression) = this(child, Literal.create(false, BooleanType))
-
-  private val ignoreNulls: Boolean = ignoreNullsExpr match {
-    case Literal(b: Boolean, BooleanType) => b
-    case _ =>
-      throw new AnalysisException("The second argument of First should be a boolean literal.")
-  }
-
-  override def nullable: Boolean = true
-  override def dataType: DataType = child.dataType
-  override def toString: String = s"first(${child}${if (ignoreNulls) " ignore nulls"})"
-
-  override def asPartial: SplitEvaluation = {
-    val partialFirst = Alias(First(child, ignoreNulls), "PartialFirst")()
-    SplitEvaluation(
-      First(partialFirst.toAttribute, ignoreNulls),
-      partialFirst :: Nil)
-  }
-  override def newInstance(): FirstFunction = new FirstFunction(child, ignoreNulls, this)
-}
-
-object First {
-  def apply(child: Expression): First = First(child, ignoreNulls = false)
-
-  def apply(child: Expression, ignoreNulls: Boolean): First =
-    First(child, Literal.create(ignoreNulls, BooleanType))
-}
-
-case class FirstFunction(
-    expr: Expression,
-    ignoreNulls: Boolean,
-    base: AggregateExpression1)
-  extends AggregateFunction1 {
-
-  def this() = this(null, null.asInstanceOf[Boolean], null) // Required for serialization.
-
-  private[this] var result: Any = null
-
-  private[this] var valueSet: Boolean = false
-
-  override def update(input: InternalRow): Unit = {
-    if (!valueSet) {
-      val value = expr.eval(input)
-      // When we have not set the result, we will set the result if we respect nulls
-      // (i.e. ignoreNulls is false), or we ignore nulls and the evaluated value is not null.
-      if (!ignoreNulls || (ignoreNulls && value != null)) {
-        result = value
-        valueSet = true
-      }
-    }
-  }
-
-  override def eval(input: InternalRow): Any = result
-}
-
-case class Last(
-    child: Expression,
-    ignoreNullsExpr: Expression)
-  extends UnaryExpression with PartialAggregate1 {
-
-  def this(child: Expression) = this(child, Literal.create(false, BooleanType))
-
-  private val ignoreNulls: Boolean = ignoreNullsExpr match {
-    case Literal(b: Boolean, BooleanType) => b
-    case _ =>
-      throw new AnalysisException("The second argument of First should be a boolean literal.")
-  }
-
-  override def references: AttributeSet = child.references
-  override def nullable: Boolean = true
-  override def dataType: DataType = child.dataType
-  override def toString: String = s"last($child)${if (ignoreNulls) " ignore nulls"}"
-
-  override def asPartial: SplitEvaluation = {
-    val partialLast = Alias(Last(child, ignoreNulls), "PartialLast")()
-    SplitEvaluation(
-      Last(partialLast.toAttribute, ignoreNulls),
-      partialLast :: Nil)
-  }
-  override def newInstance(): LastFunction = new LastFunction(child, ignoreNulls, this)
-}
-
-object Last {
-  def apply(child: Expression): Last = Last(child, ignoreNulls = false)
-
-  def apply(child: Expression, ignoreNulls: Boolean): Last =
-    Last(child, Literal.create(ignoreNulls, BooleanType))
-}
-
-case class LastFunction(
-    expr: Expression,
-    ignoreNulls: Boolean,
-    base: AggregateExpression1)
-  extends AggregateFunction1 {
-
-  def this() = this(null, null.asInstanceOf[Boolean], null) // Required for serialization.
-
-  var result: Any = null
-
-  override def update(input: InternalRow): Unit = {
-    val value = expr.eval(input)
-    if (!ignoreNulls || (ignoreNulls && value != null)) {
-      result = value
-    }
-  }
-
-  override def eval(input: InternalRow): Any = {
-    result
-  }
-}
-
-/**
- * Calculate Pearson Correlation Coefficient for the given columns.
- * Only support AggregateExpression2.
- *
- */
-case class Corr(left: Expression, right: Expression)
-    extends BinaryExpression with AggregateExpression1 with ImplicitCastInputTypes {
-  override def nullable: Boolean = false
-  override def dataType: DoubleType.type = DoubleType
-  override def toString: String = s"corr($left, $right)"
-  override def inputTypes: Seq[AbstractDataType] = Seq(DoubleType, DoubleType)
-  override def newInstance(): AggregateFunction1 = {
-    throw new UnsupportedOperationException(
-      "Corr only supports the new AggregateExpression2 and can only be used " +
-        "when spark.sql.useAggregate2 = true")
-  }
-}
-
-// Compute standard deviation based on online algorithm specified here:
-// http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
-abstract class StddevAgg1(child: Expression) extends UnaryExpression with PartialAggregate1 {
-  override def nullable: Boolean = true
-  override def dataType: DataType = DoubleType
-
-  def isSample: Boolean
-
-  override def asPartial: SplitEvaluation = {
-    val partialStd = Alias(ComputePartialStd(child), "PartialStddev")()
-    SplitEvaluation(MergePartialStd(partialStd.toAttribute, isSample), partialStd :: Nil)
-  }
-
-  override def newInstance(): StddevFunction = new StddevFunction(child, this, isSample)
-
-  override def checkInputDataTypes(): TypeCheckResult =
-    TypeUtils.checkForNumericExpr(child.dataType, "function stddev")
-
-}
-
-// Compute the population standard deviation of a column
-case class StddevPop(child: Expression) extends StddevAgg1(child) {
-
-  override def toString: String = s"stddev_pop($child)"
-  override def isSample: Boolean = false
-}
-
-// Compute the sample standard deviation of a column
-case class StddevSamp(child: Expression) extends StddevAgg1(child) {
-
-  override def toString: String = s"stddev_samp($child)"
-  override def isSample: Boolean = true
-}
-
-case class ComputePartialStd(child: Expression) extends UnaryExpression with AggregateExpression1 {
-  def this() = this(null)
-
-  override def children: Seq[Expression] = child :: Nil
-  override def nullable: Boolean = false
-  override def dataType: DataType = ArrayType(DoubleType)
-  override def toString: String = s"computePartialStddev($child)"
-  override def newInstance(): ComputePartialStdFunction =
-    new ComputePartialStdFunction(child, this)
-}
-
-case class ComputePartialStdFunction (
-    expr: Expression,
-    base: AggregateExpression1
-  ) extends AggregateFunction1 {
-
-  def this() = this(null, null)  // Required for serialization
-
-  private val computeType = DoubleType
-  private val zero = Cast(Literal(0), computeType)
-  private var partialCount: Long = 0L
-
-  // the mean of data processed so far
-  private val partialAvg: MutableLiteral = MutableLiteral(zero.eval(null), computeType)
-
-  // update average based on this formula:
-  // avg = avg + (value - avg)/count
-  private def avgAddFunction (value: Literal): Expression = {
-    val delta = Subtract(Cast(value, computeType), partialAvg)
-    Add(partialAvg, Divide(delta, Cast(Literal(partialCount), computeType)))
-  }
-
-  // the sum of squares of difference from mean
-  private val partialMk: MutableLiteral = MutableLiteral(zero.eval(null), computeType)
-
-  // update sum of square of difference from mean based on following formula:
-  // Mk = Mk + (value - preAvg) * (value - updatedAvg)
-  private def mkAddFunction(value: Literal, prePartialAvg: MutableLiteral): Expression = {
-    val delta1 = Subtract(Cast(value, computeType), prePartialAvg)
-    val delta2 = Subtract(Cast(value, computeType), partialAvg)
-    Add(partialMk, Multiply(delta1, delta2))
-  }
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = expr.eval(input)
-    if (evaluatedExpr != null) {
-      val exprValue = Literal.create(evaluatedExpr, expr.dataType)
-      val prePartialAvg = partialAvg.copy()
-      partialCount += 1
-      partialAvg.update(avgAddFunction(exprValue), input)
-      partialMk.update(mkAddFunction(exprValue, prePartialAvg), input)
-    }
-  }
-
-  override def eval(input: InternalRow): Any = {
-    new GenericArrayData(Array(Cast(Literal(partialCount), computeType).eval(null),
-        partialAvg.eval(null),
-        partialMk.eval(null)))
-  }
-}
-
-case class MergePartialStd(
-    child: Expression,
-    isSample: Boolean
-) extends UnaryExpression with AggregateExpression1 {
-  def this() = this(null, false) // required for serialization
-
-  override def children: Seq[Expression] = child:: Nil
-  override def nullable: Boolean = false
-  override def dataType: DataType = DoubleType
-  override def toString: String = s"MergePartialStd($child)"
-  override def newInstance(): MergePartialStdFunction = {
-    new MergePartialStdFunction(child, this, isSample)
-  }
-}
-
-case class MergePartialStdFunction(
-    expr: Expression,
-    base: AggregateExpression1,
-    isSample: Boolean
-) extends AggregateFunction1 {
-  def this() = this (null, null, false) // Required for serialization
-
-  private val computeType = DoubleType
-  private val zero = Cast(Literal(0), computeType)
-  private val combineCount = MutableLiteral(zero.eval(null), computeType)
-  private val combineAvg = MutableLiteral(zero.eval(null), computeType)
-  private val combineMk = MutableLiteral(zero.eval(null), computeType)
-
-  private def avgUpdateFunction(preCount: Expression,
-                                partialCount: Expression,
-                                partialAvg: Expression): Expression = {
-    Divide(Add(Multiply(combineAvg, preCount),
-               Multiply(partialAvg, partialCount)),
-           Add(preCount, partialCount))
-  }
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = expr.eval(input).asInstanceOf[ArrayData]
-
-    if (evaluatedExpr != null) {
-      val exprValue = evaluatedExpr.toArray(computeType)
-      val (partialCount, partialAvg, partialMk) =
-        (Literal.create(exprValue(0), computeType),
-         Literal.create(exprValue(1), computeType),
-         Literal.create(exprValue(2), computeType))
-
-      if (Cast(partialCount, LongType).eval(null).asInstanceOf[Long] > 0) {
-        val preCount = combineCount.copy()
-        combineCount.update(Add(combineCount, partialCount), input)
-
-        val preAvg = combineAvg.copy()
-        val avgDelta = Subtract(partialAvg, preAvg)
-        val mkDelta = Multiply(Multiply(avgDelta, avgDelta),
-                               Divide(Multiply(preCount, partialCount),
-                                      combineCount))
-
-        // update average based on following formula
-        // (combineAvg * preCount + partialAvg * partialCount) / (preCount + partialCount)
-        combineAvg.update(avgUpdateFunction(preCount, partialCount, partialAvg), input)
-
-        // update sum of square differences from mean based on following formula
-        // (combineMk + partialMk + (avgDelta * avgDelta) * (preCount * partialCount/combineCount)
-        combineMk.update(Add(combineMk, Add(partialMk, mkDelta)), input)
-      }
-    }
-  }
-
-  override def eval(input: InternalRow): Any = {
-    val count: Long = Cast(combineCount, LongType).eval(null).asInstanceOf[Long]
-
-    if (count == 0) null
-    else if (count < 2) zero.eval(null)
-    else {
-      // when total count > 2
-      // stddev_samp = sqrt (combineMk/(combineCount -1))
-      // stddev_pop = sqrt (combineMk/combineCount)
-      val varCol = {
-        if (isSample) {
-          Divide(combineMk, Cast(Literal(count - 1), computeType))
-        }
-        else {
-          Divide(combineMk, Cast(Literal(count), computeType))
-        }
-      }
-      Sqrt(varCol).eval(null)
-    }
-  }
-}
-
-case class StddevFunction(
-    expr: Expression,
-    base: AggregateExpression1,
-    isSample: Boolean
-) extends AggregateFunction1 {
-
-  def this() = this(null, null, false) // Required for serialization
-
-  private val computeType = DoubleType
-  private var curCount: Long = 0L
-  private val zero = Cast(Literal(0), computeType)
-  private val curAvg = MutableLiteral(zero.eval(null), computeType)
-  private val curMk = MutableLiteral(zero.eval(null), computeType)
-
-  private def curAvgAddFunction(value: Literal): Expression = {
-    val delta = Subtract(Cast(value, computeType), curAvg)
-    Add(curAvg, Divide(delta, Cast(Literal(curCount), computeType)))
-  }
-  private def curMkAddFunction(value: Literal, preAvg: MutableLiteral): Expression = {
-    val delta1 = Subtract(Cast(value, computeType), preAvg)
-    val delta2 = Subtract(Cast(value, computeType), curAvg)
-    Add(curMk, Multiply(delta1, delta2))
-  }
-
-  override def update(input: InternalRow): Unit = {
-    val evaluatedExpr = expr.eval(input)
-    if (evaluatedExpr != null) {
-      val preAvg: MutableLiteral = curAvg.copy()
-      val exprValue = Literal.create(evaluatedExpr, expr.dataType)
-      curCount += 1L
-      curAvg.update(curAvgAddFunction(exprValue), input)
-      curMk.update(curMkAddFunction(exprValue, preAvg), input)
-    }
-  }
-
-  override def eval(input: InternalRow): Any = {
-    if (curCount == 0) null
-    else if (curCount < 2) zero.eval(null)
-    else {
-      // when total count > 2,
-      // stddev_samp = sqrt(curMk/(curCount - 1))
-      // stddev_pop = sqrt(curMk/curCount)
-      val varCol = {
-        if (isSample) {
-          Divide(curMk, Cast(Literal(curCount - 1), computeType))
-        }
-        else {
-          Divide(curMk, Cast(Literal(curCount), computeType))
-        }
-      }
-      Sqrt(varCol).eval(null)
-    }
-  }
-}
-
-// placeholder
-case class Kurtosis(child: Expression) extends UnaryExpression with AggregateExpression1 {
-
-  override def newInstance(): AggregateFunction1 = {
-    throw new UnsupportedOperationException("AggregateExpression1 is no longer supported, " +
-      "please set spark.sql.useAggregate2 = true")
-  }
-
-  override def nullable: Boolean = false
-
-  override def dataType: DoubleType.type = DoubleType
-
-  override def foldable: Boolean = false
-
-  override def prettyName: String = "kurtosis"
-}
-
-// placeholder
-case class Skewness(child: Expression) extends UnaryExpression with AggregateExpression1 {
-
-  override def newInstance(): AggregateFunction1 = {
-    throw new UnsupportedOperationException("AggregateExpression1 is no longer supported, " +
-      "please set spark.sql.useAggregate2 = true")
-  }
-
-  override def nullable: Boolean = false
-
-  override def dataType: DoubleType.type = DoubleType
-
-  override def foldable: Boolean = false
-
-  override def prettyName: String = "skewness"
-}
-
-// placeholder
-case class VariancePop(child: Expression) extends UnaryExpression with AggregateExpression1 {
-
-  override def newInstance(): AggregateFunction1 = {
-    throw new UnsupportedOperationException("AggregateExpression1 is no longer supported, " +
-      "please set spark.sql.useAggregate2 = true")
-  }
-
-  override def nullable: Boolean = false
-
-  override def dataType: DoubleType.type = DoubleType
-
-  override def foldable: Boolean = false
-
-  override def prettyName: String = "var_pop"
-}
-
-// placeholder
-case class VarianceSamp(child: Expression) extends UnaryExpression with AggregateExpression1 {
-
-  override def newInstance(): AggregateFunction1 = {
-    throw new UnsupportedOperationException("AggregateExpression1 is no longer supported, " +
-      "please set spark.sql.useAggregate2 = true")
-  }
-
-  override def nullable: Boolean = false
-
-  override def dataType: DoubleType.type = DoubleType
-
-  override def foldable: Boolean = false
-
-  override def prettyName: String = "var_samp"
-}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index d222dfa33ad8a..f4dba67f13b54 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.optimizer
 import scala.collection.immutable.HashSet
 import org.apache.spark.sql.catalyst.analysis.{CleanupAliases, EliminateSubQueries}
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.Inner
 import org.apache.spark.sql.catalyst.plans.FullOuter
 import org.apache.spark.sql.catalyst.plans.LeftOuter
@@ -201,8 +202,8 @@ object SetOperationPushDown extends Rule[LogicalPlan] with PredicateHelper {
 object ColumnPruning extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case a @ Aggregate(_, _, e @ Expand(_, _, child))
-      if (child.outputSet -- AttributeSet(e.output) -- a.references).nonEmpty =>
-      a.copy(child = e.copy(child = prunedChild(child, AttributeSet(e.output) ++ a.references)))
+      if (child.outputSet -- e.references -- a.references).nonEmpty =>
+      a.copy(child = e.copy(child = prunedChild(child, e.references ++ a.references)))
 
     // Eliminate attributes that are not needed to calculate the specified aggregates.
     case a @ Aggregate(_, _, child) if (child.outputSet -- a.references).nonEmpty =>
@@ -363,7 +364,8 @@ object LikeSimplification extends Rule[LogicalPlan] {
 object NullPropagation extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case q: LogicalPlan => q transformExpressionsUp {
-      case e @ Count(Literal(null, _)) => Cast(Literal(0L), e.dataType)
+      case e @ AggregateExpression(Count(Literal(null, _)), _, _) =>
+        Cast(Literal(0L), e.dataType)
       case e @ IsNull(c) if !c.nullable => Literal.create(false, BooleanType)
       case e @ IsNotNull(c) if !c.nullable => Literal.create(true, BooleanType)
       case e @ GetArrayItem(Literal(null, _), _) => Literal.create(null, e.dataType)
@@ -375,7 +377,9 @@ object NullPropagation extends Rule[LogicalPlan] {
         Literal.create(null, e.dataType)
       case e @ EqualNullSafe(Literal(null, _), r) => IsNull(r)
       case e @ EqualNullSafe(l, Literal(null, _)) => IsNull(l)
-      case e @ Count(expr) if !expr.nullable => Count(Literal(1))
+      case e @ AggregateExpression(Count(expr), mode, false) if !expr.nullable =>
+        // This rule should be only triggered when isDistinct field is false.
+        AggregateExpression(Count(Literal(1)), mode, isDistinct = false)
 
       // For Coalesce, remove null literals.
       case e @ Coalesce(children) =>
@@ -857,12 +861,15 @@ object DecimalAggregates extends Rule[LogicalPlan] {
   private val MAX_DOUBLE_DIGITS = 15
 
   def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
-    case Sum(e @ DecimalType.Expression(prec, scale)) if prec + 10 <= MAX_LONG_DIGITS =>
-      MakeDecimal(Sum(UnscaledValue(e)), prec + 10, scale)
+    case AggregateExpression(Sum(e @ DecimalType.Expression(prec, scale)), mode, isDistinct)
+      if prec + 10 <= MAX_LONG_DIGITS =>
+      MakeDecimal(AggregateExpression(Sum(UnscaledValue(e)), mode, isDistinct), prec + 10, scale)
 
-    case Average(e @ DecimalType.Expression(prec, scale)) if prec + 4 <= MAX_DOUBLE_DIGITS =>
+    case AggregateExpression(Average(e @ DecimalType.Expression(prec, scale)), mode, isDistinct)
+      if prec + 4 <= MAX_DOUBLE_DIGITS =>
+      val newAggExpr = AggregateExpression(Average(UnscaledValue(e)), mode, isDistinct)
       Cast(
-        Divide(Average(UnscaledValue(e)), Literal.create(math.pow(10.0, scale), DoubleType)),
+        Divide(newAggExpr, Literal.create(math.pow(10.0, scale), DoubleType)),
         DecimalType(prec + 4, scale + 4))
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala
index 3b975b904a332..6f4f11406d7c4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala
@@ -84,80 +84,6 @@ object PhysicalOperation extends PredicateHelper {
   }
 }
 
-/**
- * Matches a logical aggregation that can be performed on distributed data in two steps.  The first
- * operates on the data in each partition performing partial aggregation for each group.  The second
- * occurs after the shuffle and completes the aggregation.
- *
- * This pattern will only match if all aggregate expressions can be computed partially and will
- * return the rewritten aggregation expressions for both phases.
- *
- * The returned values for this match are as follows:
- *  - Grouping attributes for the final aggregation.
- *  - Aggregates for the final aggregation.
- *  - Grouping expressions for the partial aggregation.
- *  - Partial aggregate expressions.
- *  - Input to the aggregation.
- */
-object PartialAggregation {
-  type ReturnType =
-    (Seq[Attribute], Seq[NamedExpression], Seq[Expression], Seq[NamedExpression], LogicalPlan)
-
-  def unapply(plan: LogicalPlan): Option[ReturnType] = plan match {
-    case logical.Aggregate(groupingExpressions, aggregateExpressions, child) =>
-      // Collect all aggregate expressions.
-      val allAggregates =
-        aggregateExpressions.flatMap(_ collect { case a: AggregateExpression1 => a})
-      // Collect all aggregate expressions that can be computed partially.
-      val partialAggregates =
-        aggregateExpressions.flatMap(_ collect { case p: PartialAggregate1 => p})
-
-      // Only do partial aggregation if supported by all aggregate expressions.
-      if (allAggregates.size == partialAggregates.size) {
-        // Create a map of expressions to their partial evaluations for all aggregate expressions.
-        val partialEvaluations: Map[TreeNodeRef, SplitEvaluation] =
-          partialAggregates.map(a => (new TreeNodeRef(a), a.asPartial)).toMap
-
-        // We need to pass all grouping expressions though so the grouping can happen a second
-        // time. However some of them might be unnamed so we alias them allowing them to be
-        // referenced in the second aggregation.
-        val namedGroupingExpressions: Seq[(Expression, NamedExpression)] =
-          groupingExpressions.map {
-            case n: NamedExpression => (n, n)
-            case other => (other, Alias(other, "PartialGroup")())
-          }
-
-        // Replace aggregations with a new expression that computes the result from the already
-        // computed partial evaluations and grouping values.
-        val rewrittenAggregateExpressions = aggregateExpressions.map(_.transformDown {
-          case e: Expression if partialEvaluations.contains(new TreeNodeRef(e)) =>
-            partialEvaluations(new TreeNodeRef(e)).finalEvaluation
-
-          case e: Expression =>
-            namedGroupingExpressions.collectFirst {
-              case (expr, ne) if expr semanticEquals e => ne.toAttribute
-            }.getOrElse(e)
-        }).asInstanceOf[Seq[NamedExpression]]
-
-        val partialComputation = namedGroupingExpressions.map(_._2) ++
-          partialEvaluations.values.flatMap(_.partialEvaluations)
-
-        val namedGroupingAttributes = namedGroupingExpressions.map(_._2.toAttribute)
-
-        Some(
-          (namedGroupingAttributes,
-           rewrittenAggregateExpressions,
-           groupingExpressions,
-           partialComputation,
-           child))
-      } else {
-        None
-      }
-    case _ => None
-  }
-}
-
-
 /**
  * A pattern that finds joins with equality conditions that can be evaluated using equi-join.
  *
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
index 0ec9f08571082..b9db7838db08a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
@@ -137,13 +137,17 @@ abstract class QueryPlan[PlanType <: TreeNode[PlanType]] extends TreeNode[PlanTy
 
   /** Returns all of the expressions present in this query plan operator. */
   def expressions: Seq[Expression] = {
+    // Recursively find all expressions from a traversable.
+    def seqToExpressions(seq: Traversable[Any]): Traversable[Expression] = seq.flatMap {
+      case e: Expression => e :: Nil
+      case s: Traversable[_] => seqToExpressions(s)
+      case other => Nil
+    }
+
     productIterator.flatMap {
       case e: Expression => e :: Nil
       case Some(e: Expression) => e :: Nil
-      case seq: Traversable[_] => seq.flatMap {
-        case e: Expression => e :: Nil
-        case other => Nil
-      }
+      case seq: Traversable[_] => seqToExpressions(seq)
       case other => Nil
     }.toSeq
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index d771088d69dea..764f8aaebddf1 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.plans.logical
 
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.Utils
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.types._
 import org.apache.spark.util.collection.OpenHashSet
@@ -219,8 +219,6 @@ case class Aggregate(
     !expressions.exists(!_.resolved) && childrenResolved && !hasWindowExpressions
   }
 
-  lazy val newAggregation: Option[Aggregate] = Utils.tryConvert(this)
-
   override def output: Seq[Attribute] = aggregateExpressions.map(_.toAttribute)
 }
 
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
index fbdd3a7776f50..5a2368e329976 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
@@ -171,16 +171,18 @@ class AnalysisErrorSuite extends AnalysisTest {
 
   test("SPARK-6452 regression test") {
     // CheckAnalysis should throw AnalysisException when Aggregate contains missing attribute(s)
+    // Since we manually construct the logical plan at here and Sum only accetp
+    // LongType, DoubleType, and DecimalType. We use LongType as the type of a.
     val plan =
       Aggregate(
         Nil,
-        Alias(Sum(AttributeReference("a", IntegerType)(exprId = ExprId(1))), "b")() :: Nil,
+        Alias(sum(AttributeReference("a", LongType)(exprId = ExprId(1))), "b")() :: Nil,
         LocalRelation(
-          AttributeReference("a", IntegerType)(exprId = ExprId(2))))
+          AttributeReference("a", LongType)(exprId = ExprId(2))))
 
     assert(plan.resolved)
 
-    assertAnalysisError(plan, "resolved attribute(s) a#1 missing from a#2" :: Nil)
+    assertAnalysisError(plan, "resolved attribute(s) a#1L missing from a#2L" :: Nil)
   }
 
   test("error test for self-join") {
@@ -196,7 +198,7 @@ class AnalysisErrorSuite extends AnalysisTest {
     val plan =
       Aggregate(
         AttributeReference("a", BinaryType)(exprId = ExprId(2)) :: Nil,
-        Alias(Sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
+        Alias(sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
         LocalRelation(
           AttributeReference("a", BinaryType)(exprId = ExprId(2)),
           AttributeReference("b", IntegerType)(exprId = ExprId(1))))
@@ -207,13 +209,24 @@ class AnalysisErrorSuite extends AnalysisTest {
     val plan2 =
       Aggregate(
         AttributeReference("a", MapType(IntegerType, StringType))(exprId = ExprId(2)) :: Nil,
-        Alias(Sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
+        Alias(sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
         LocalRelation(
           AttributeReference("a", MapType(IntegerType, StringType))(exprId = ExprId(2)),
           AttributeReference("b", IntegerType)(exprId = ExprId(1))))
 
     assertAnalysisError(plan2,
       "map type expression a cannot be used in grouping expression" :: Nil)
+
+    val plan3 =
+      Aggregate(
+        AttributeReference("a", ArrayType(IntegerType))(exprId = ExprId(2)) :: Nil,
+        Alias(sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
+        LocalRelation(
+          AttributeReference("a", ArrayType(IntegerType))(exprId = ExprId(2)),
+          AttributeReference("b", IntegerType)(exprId = ExprId(1))))
+
+    assertAnalysisError(plan3,
+      "array type expression a cannot be used in grouping expression" :: Nil)
   }
 
   test("Join can't work on binary and map types") {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
index 71d2939ecffe6..65f09b46afae1 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
@@ -45,7 +45,7 @@ class AnalysisSuite extends AnalysisTest {
     val explode = Explode(AttributeReference("a", IntegerType, nullable = true)())
     assert(!Project(Seq(Alias(explode, "explode")()), testRelation).resolved)
 
-    assert(!Project(Seq(Alias(Count(Literal(1)), "count")()), testRelation).resolved)
+    assert(!Project(Seq(Alias(count(Literal(1)), "count")()), testRelation).resolved)
   }
 
   test("analyze project") {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/DecimalPrecisionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/DecimalPrecisionSuite.scala
index 40c4ae7920918..fed591fd90a9a 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/DecimalPrecisionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/DecimalPrecisionSuite.scala
@@ -21,6 +21,7 @@ import org.scalatest.BeforeAndAfter
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.{Union, Project, LocalRelation}
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.catalyst.{TableIdentifier, SimpleCatalystConf}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala
index c9bcc68f02030..b902982add8ff 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala
@@ -22,6 +22,7 @@ import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
 import org.apache.spark.sql.types.{TypeCollection, StringType}
 
@@ -140,15 +141,16 @@ class ExpressionTypeCheckingSuite extends SparkFunSuite {
   }
 
   test("check types for aggregates") {
+    // We use AggregateFunction directly at here because the error will be thrown from it
+    // instead of from AggregateExpression, which is the wrapper of an AggregateFunction.
+
     // We will cast String to Double for sum and average
     assertSuccess(Sum('stringField))
-    assertSuccess(SumDistinct('stringField))
     assertSuccess(Average('stringField))
 
     assertError(Min('complexField), "min does not support ordering on type")
     assertError(Max('complexField), "max does not support ordering on type")
     assertError(Sum('booleanField), "function sum requires numeric type")
-    assertError(SumDistinct('booleanField), "function sumDistinct requires numeric type")
     assertError(Average('booleanField), "function average requires numeric type")
   }
 
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConstantFoldingSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConstantFoldingSuite.scala
index e67606288f514..8aaefa84937c2 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConstantFoldingSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConstantFoldingSuite.scala
@@ -162,7 +162,7 @@ class ConstantFoldingSuite extends PlanTest {
       testRelation
         .select(
           Rand(5L) + Literal(1) as Symbol("c1"),
-          Sum('a) as Symbol("c2"))
+          sum('a) as Symbol("c2"))
 
     val optimized = Optimize.execute(originalQuery.analyze)
 
@@ -170,7 +170,7 @@ class ConstantFoldingSuite extends PlanTest {
       testRelation
         .select(
           Rand(5L) + Literal(1.0) as Symbol("c1"),
-          Sum('a) as Symbol("c2"))
+          sum('a) as Symbol("c2"))
         .analyze
 
     comparePlans(optimized, correctAnswer)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
index ed810a12808f0..0290fafe879f6 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
@@ -68,7 +68,7 @@ class FilterPushdownSuite extends PlanTest {
   test("column pruning for group") {
     val originalQuery =
       testRelation
-        .groupBy('a)('a, Count('b))
+        .groupBy('a)('a, count('b))
         .select('a)
 
     val optimized = Optimize.execute(originalQuery.analyze)
@@ -84,7 +84,7 @@ class FilterPushdownSuite extends PlanTest {
   test("column pruning for group with alias") {
     val originalQuery =
       testRelation
-        .groupBy('a)('a as 'c, Count('b))
+        .groupBy('a)('a as 'c, count('b))
         .select('c)
 
     val optimized = Optimize.execute(originalQuery.analyze)
@@ -656,7 +656,7 @@ class FilterPushdownSuite extends PlanTest {
 
   test("aggregate: push down filter when filter on group by expression") {
     val originalQuery = testRelation
-                        .groupBy('a)('a, Count('b) as 'c)
+                        .groupBy('a)('a, count('b) as 'c)
                         .select('a, 'c)
                         .where('a === 2)
 
@@ -664,7 +664,7 @@ class FilterPushdownSuite extends PlanTest {
 
     val correctAnswer = testRelation
                         .where('a === 2)
-                        .groupBy('a)('a, Count('b) as 'c)
+                        .groupBy('a)('a, count('b) as 'c)
                         .analyze
     comparePlans(optimized, correctAnswer)
   }
@@ -672,7 +672,7 @@ class FilterPushdownSuite extends PlanTest {
   test("aggregate: don't push down filter when filter not on group by expression") {
     val originalQuery = testRelation
                         .select('a, 'b)
-                        .groupBy('a)('a, Count('b) as 'c)
+                        .groupBy('a)('a, count('b) as 'c)
                         .where('c === 2L)
 
     val optimized = Optimize.execute(originalQuery.analyze)
@@ -683,7 +683,7 @@ class FilterPushdownSuite extends PlanTest {
   test("aggregate: push down filters partially which are subset of group by expressions") {
     val originalQuery = testRelation
                         .select('a, 'b)
-                        .groupBy('a)('a, Count('b) as 'c)
+                        .groupBy('a)('a, count('b) as 'c)
                         .where('c === 2L && 'a === 3)
 
     val optimized = Optimize.execute(originalQuery.analyze)
@@ -691,7 +691,7 @@ class FilterPushdownSuite extends PlanTest {
     val correctAnswer = testRelation
                         .select('a, 'b)
                         .where('a === 3)
-                        .groupBy('a)('a, Count('b) as 'c)
+                        .groupBy('a)('a, count('b) as 'c)
                         .where('c === 2L)
                         .analyze
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index d25807cf8d09c..3b69247dc54ef 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -34,6 +34,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.encoders.Encoder
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.{Inner, JoinType}
@@ -1338,7 +1339,7 @@ class DataFrame private[sql](
       if (groupColExprIds.contains(attr.exprId)) {
         attr
       } else {
-        Alias(First(attr), attr.name)()
+        Alias(new First(attr).toAggregateExpression(), attr.name)()
       }
     }
     Aggregate(groupCols, aggCols, logicalPlan)
@@ -1381,11 +1382,11 @@ class DataFrame private[sql](
 
     // The list of summary statistics to compute, in the form of expressions.
     val statistics = List[(String, Expression => Expression)](
-      "count" -> Count,
-      "mean" -> Average,
-      "stddev" -> StddevSamp,
-      "min" -> Min,
-      "max" -> Max)
+      "count" -> ((child: Expression) => Count(child).toAggregateExpression()),
+      "mean" -> ((child: Expression) => Average(child).toAggregateExpression()),
+      "stddev" -> ((child: Expression) => StddevSamp(child).toAggregateExpression()),
+      "min" -> ((child: Expression) => Min(child).toAggregateExpression()),
+      "max" -> ((child: Expression) => Max(child).toAggregateExpression()))
 
     val outputCols = (if (cols.isEmpty) numericColumns.map(_.prettyString) else cols).toList
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
index f9eab5c2e965b..5babf2cc0ca25 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
@@ -21,8 +21,9 @@ import scala.collection.JavaConverters._
 import scala.language.implicitConversions
 
 import org.apache.spark.annotation.Experimental
-import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttribute, Star}
+import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunction, UnresolvedAlias, UnresolvedAttribute, Star}
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.{Rollup, Cube, Aggregate}
 import org.apache.spark.sql.types.NumericType
 
@@ -70,7 +71,7 @@ class GroupedData protected[sql](
     }
   }
 
-  private[this] def aggregateNumericColumns(colNames: String*)(f: Expression => Expression)
+  private[this] def aggregateNumericColumns(colNames: String*)(f: Expression => AggregateFunction)
     : DataFrame = {
 
     val columnExprs = if (colNames.isEmpty) {
@@ -88,30 +89,28 @@ class GroupedData protected[sql](
         namedExpr
       }
     }
-    toDF(columnExprs.map(f))
+    toDF(columnExprs.map(expr => f(expr).toAggregateExpression()))
   }
 
   private[this] def strToExpr(expr: String): (Expression => Expression) = {
-    expr.toLowerCase match {
-      case "avg" | "average" | "mean" => Average
-      case "max" => Max
-      case "min" => Min
-      case "stddev" | "std" => StddevSamp
-      case "stddev_pop" => StddevPop
-      case "stddev_samp" => StddevSamp
-      case "variance" => VarianceSamp
-      case "var_pop" => VariancePop
-      case "var_samp" => VarianceSamp
-      case "sum" => Sum
-      case "skewness" => Skewness
-      case "kurtosis" => Kurtosis
-      case "count" | "size" =>
-        // Turn count(*) into count(1)
-        (inputExpr: Expression) => inputExpr match {
-          case s: Star => Count(Literal(1))
-          case _ => Count(inputExpr)
-        }
+    val exprToFunc: (Expression => Expression) = {
+      (inputExpr: Expression) => expr.toLowerCase match {
+        // We special handle a few cases that have alias that are not in function registry.
+        case "avg" | "average" | "mean" =>
+          UnresolvedFunction("avg", inputExpr :: Nil, isDistinct = false)
+        case "stddev" | "std" =>
+          UnresolvedFunction("stddev", inputExpr :: Nil, isDistinct = false)
+        // Also special handle count because we need to take care count(*).
+        case "count" | "size" =>
+          // Turn count(*) into count(1)
+          inputExpr match {
+            case s: Star => Count(Literal(1)).toAggregateExpression()
+            case _ => Count(inputExpr).toAggregateExpression()
+          }
+        case name => UnresolvedFunction(name, inputExpr :: Nil, isDistinct = false)
+      }
     }
+    (inputExpr: Expression) => exprToFunc(inputExpr)
   }
 
   /**
@@ -213,7 +212,7 @@ class GroupedData protected[sql](
    *
    * @since 1.3.0
    */
-  def count(): DataFrame = toDF(Seq(Alias(Count(Literal(1)), "count")()))
+  def count(): DataFrame = toDF(Seq(Alias(Count(Literal(1)).toAggregateExpression(), "count")()))
 
   /**
    * Compute the average value for each numeric columns for each group. This is an alias for `avg`.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
index ed8b634ad5630..b7314189b5403 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
@@ -448,15 +448,24 @@ private[spark] object SQLConf {
     defaultValue = Some(true),
     isPublic = false)
 
-  val USE_SQL_AGGREGATE2 = booleanConf("spark.sql.useAggregate2",
-    defaultValue = Some(true), doc = "<TODO>")
-
   val RUN_SQL_ON_FILES = booleanConf("spark.sql.runSQLOnFiles",
     defaultValue = Some(true),
     isPublic = false,
     doc = "When true, we could use `datasource`.`path` as table in SQL query"
   )
 
+  val SPECIALIZE_SINGLE_DISTINCT_AGG_PLANNING =
+    booleanConf("spark.sql.specializeSingleDistinctAggPlanning",
+      defaultValue = Some(true),
+      isPublic = false,
+      doc = "When true, if a query only has a single distinct column and it has " +
+        "grouping expressions, we will use our planner rule to handle this distinct " +
+        "column (other cases are handled by DistinctAggregationRewriter). " +
+        "When false, we will always use DistinctAggregationRewriter to plan " +
+        "aggregation queries with DISTINCT keyword. This is an internal flag that is " +
+        "used to benchmark the performance impact of using DistinctAggregationRewriter to " +
+        "plan aggregation queries with a single distinct column.")
+
   object Deprecated {
     val MAPRED_REDUCE_TASKS = "mapred.reduce.tasks"
     val EXTERNAL_SORT = "spark.sql.planner.externalSort"
@@ -532,8 +541,6 @@ private[sql] class SQLConf extends Serializable with CatalystConf {
 
   private[spark] def unsafeEnabled: Boolean = getConf(UNSAFE_ENABLED, getConf(TUNGSTEN_ENABLED))
 
-  private[spark] def useSqlAggregate2: Boolean = getConf(USE_SQL_AGGREGATE2)
-
   private[spark] def autoBroadcastJoinThreshold: Int = getConf(AUTO_BROADCASTJOIN_THRESHOLD)
 
   private[spark] def defaultSizeInBytes: Long =
@@ -575,6 +582,9 @@ private[sql] class SQLConf extends Serializable with CatalystConf {
 
   private[spark] def runSQLOnFile: Boolean = getConf(RUN_SQL_ON_FILES)
 
+  protected[spark] override def specializeSingleDistinctAggPlanning: Boolean =
+    getConf(SPECIALIZE_SINGLE_DISTINCT_AGG_PLANNING)
+
   /** ********************** SQLConf functionality methods ************ */
 
   /** Set Spark SQL configuration properties. */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Aggregate.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Aggregate.scala
deleted file mode 100644
index 6f3f1bd97ad52..0000000000000
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Aggregate.scala
+++ /dev/null
@@ -1,205 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.execution
-
-import java.util.HashMap
-
-import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.errors._
-import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.physical._
-import org.apache.spark.sql.execution.metric.SQLMetrics
-
-/**
- * Groups input data by `groupingExpressions` and computes the `aggregateExpressions` for each
- * group.
- *
- * @param partial if true then aggregation is done partially on local data without shuffling to
- *                ensure all values where `groupingExpressions` are equal are present.
- * @param groupingExpressions expressions that are evaluated to determine grouping.
- * @param aggregateExpressions expressions that are computed for each group.
- * @param child the input data source.
- */
-case class Aggregate(
-    partial: Boolean,
-    groupingExpressions: Seq[Expression],
-    aggregateExpressions: Seq[NamedExpression],
-    child: SparkPlan)
-  extends UnaryNode {
-
-  override private[sql] lazy val metrics = Map(
-    "numInputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of input rows"),
-    "numOutputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of output rows"))
-
-  override def requiredChildDistribution: List[Distribution] = {
-    if (partial) {
-      UnspecifiedDistribution :: Nil
-    } else {
-      if (groupingExpressions == Nil) {
-        AllTuples :: Nil
-      } else {
-        ClusteredDistribution(groupingExpressions) :: Nil
-      }
-    }
-  }
-
-  override def output: Seq[Attribute] = aggregateExpressions.map(_.toAttribute)
-
-  /**
-   * An aggregate that needs to be computed for each row in a group.
-   *
-   * @param unbound Unbound version of this aggregate, used for result substitution.
-   * @param aggregate A bound copy of this aggregate used to create a new aggregation buffer.
-   * @param resultAttribute An attribute used to refer to the result of this aggregate in the final
-   *                        output.
-   */
-  case class ComputedAggregate(
-      unbound: AggregateExpression1,
-      aggregate: AggregateExpression1,
-      resultAttribute: AttributeReference)
-
-  /** A list of aggregates that need to be computed for each group. */
-  private[this] val computedAggregates = aggregateExpressions.flatMap { agg =>
-    agg.collect {
-      case a: AggregateExpression1 =>
-        ComputedAggregate(
-          a,
-          BindReferences.bindReference(a, child.output),
-          AttributeReference(s"aggResult:$a", a.dataType, a.nullable)())
-    }
-  }.toArray
-
-  /** The schema of the result of all aggregate evaluations */
-  private[this] val computedSchema = computedAggregates.map(_.resultAttribute)
-
-  /** Creates a new aggregate buffer for a group. */
-  private[this] def newAggregateBuffer(): Array[AggregateFunction1] = {
-    val buffer = new Array[AggregateFunction1](computedAggregates.length)
-    var i = 0
-    while (i < computedAggregates.length) {
-      buffer(i) = computedAggregates(i).aggregate.newInstance()
-      i += 1
-    }
-    buffer
-  }
-
-  /** Named attributes used to substitute grouping attributes into the final result. */
-  private[this] val namedGroups = groupingExpressions.map {
-    case ne: NamedExpression => ne -> ne.toAttribute
-    case e => e -> Alias(e, s"groupingExpr:$e")().toAttribute
-  }
-
-  /**
-   * A map of substitutions that are used to insert the aggregate expressions and grouping
-   * expression into the final result expression.
-   */
-  private[this] val resultMap =
-    (computedAggregates.map { agg => agg.unbound -> agg.resultAttribute } ++ namedGroups).toMap
-
-  /**
-   * Substituted version of aggregateExpressions expressions which are used to compute final
-   * output rows given a group and the result of all aggregate computations.
-   */
-  private[this] val resultExpressions = aggregateExpressions.map { agg =>
-    agg.transform {
-      case e: Expression if resultMap.contains(e) => resultMap(e)
-    }
-  }
-
-  protected override def doExecute(): RDD[InternalRow] = attachTree(this, "execute") {
-    val numInputRows = longMetric("numInputRows")
-    val numOutputRows = longMetric("numOutputRows")
-    if (groupingExpressions.isEmpty) {
-      child.execute().mapPartitions { iter =>
-        val buffer = newAggregateBuffer()
-        var currentRow: InternalRow = null
-        while (iter.hasNext) {
-          currentRow = iter.next()
-          numInputRows += 1
-          var i = 0
-          while (i < buffer.length) {
-            buffer(i).update(currentRow)
-            i += 1
-          }
-        }
-        val resultProjection = new InterpretedProjection(resultExpressions, computedSchema)
-        val aggregateResults = new GenericMutableRow(computedAggregates.length)
-
-        var i = 0
-        while (i < buffer.length) {
-          aggregateResults(i) = buffer(i).eval(EmptyRow)
-          i += 1
-        }
-
-        numOutputRows += 1
-        Iterator(resultProjection(aggregateResults))
-      }
-    } else {
-      child.execute().mapPartitions { iter =>
-        val hashTable = new HashMap[InternalRow, Array[AggregateFunction1]]
-        val groupingProjection = new InterpretedMutableProjection(groupingExpressions, child.output)
-
-        var currentRow: InternalRow = null
-        while (iter.hasNext) {
-          currentRow = iter.next()
-          numInputRows += 1
-          val currentGroup = groupingProjection(currentRow)
-          var currentBuffer = hashTable.get(currentGroup)
-          if (currentBuffer == null) {
-            currentBuffer = newAggregateBuffer()
-            hashTable.put(currentGroup.copy(), currentBuffer)
-          }
-
-          var i = 0
-          while (i < currentBuffer.length) {
-            currentBuffer(i).update(currentRow)
-            i += 1
-          }
-        }
-
-        new Iterator[InternalRow] {
-          private[this] val hashTableIter = hashTable.entrySet().iterator()
-          private[this] val aggregateResults = new GenericMutableRow(computedAggregates.length)
-          private[this] val resultProjection =
-            new InterpretedMutableProjection(
-              resultExpressions, computedSchema ++ namedGroups.map(_._2))
-          private[this] val joinedRow = new JoinedRow
-
-          override final def hasNext: Boolean = hashTableIter.hasNext
-
-          override final def next(): InternalRow = {
-            val currentEntry = hashTableIter.next()
-            val currentGroup = currentEntry.getKey
-            val currentBuffer = currentEntry.getValue
-            numOutputRows += 1
-
-            var i = 0
-            while (i < currentBuffer.length) {
-              // Evaluating an aggregate buffer returns the result.  No row is required since we
-              // already added all rows in the group using update.
-              aggregateResults(i) = currentBuffer(i).eval(EmptyRow)
-              i += 1
-            }
-            resultProjection(joinedRow(aggregateResults, currentGroup))
-          }
-        }
-      }
-    }
-  }
-}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Expand.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Expand.scala
index 55e95769d3faa..91530bd63798a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Expand.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Expand.scala
@@ -45,6 +45,9 @@ case class Expand(
   override def canProcessUnsafeRows: Boolean = true
   override def canProcessSafeRows: Boolean = true
 
+  override def references: AttributeSet =
+    AttributeSet(projections.flatten.flatMap(_.references))
+
   private[this] val projection = {
     if (outputsUnsafeRows) {
       (exprs: Seq[Expression]) => UnsafeProjection.create(exprs, child.output)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
index 0f98fe88b2101..a10d1edcc91aa 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
@@ -38,7 +38,6 @@ class SparkPlanner(val sqlContext: SQLContext) extends SparkStrategies {
       DataSourceStrategy ::
       DDLStrategy ::
       TakeOrderedAndProject ::
-      HashAggregation ::
       Aggregation ::
       LeftSemiJoin ::
       EquiJoinSelection ::
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index dd3bb33c57287..d65cb1bae7fb5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.execution
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression2, Utils}
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.planning._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical.{BroadcastHint, LogicalPlan}
@@ -146,148 +146,104 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
     }
   }
 
-  object HashAggregation extends Strategy {
-    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
-      // Aggregations that can be performed in two phases, before and after the shuffle.
-      case PartialAggregation(
-          namedGroupingAttributes,
-          rewrittenAggregateExpressions,
-          groupingExpressions,
-          partialComputation,
-          child) if !canBeConvertedToNewAggregation(plan) =>
-        execution.Aggregate(
-          partial = false,
-          namedGroupingAttributes,
-          rewrittenAggregateExpressions,
-          execution.Aggregate(
-            partial = true,
-            groupingExpressions,
-            partialComputation,
-            planLater(child))) :: Nil
-
-      case _ => Nil
-    }
-
-    def canBeConvertedToNewAggregation(plan: LogicalPlan): Boolean = plan match {
-      case a: logical.Aggregate =>
-        if (sqlContext.conf.useSqlAggregate2 && sqlContext.conf.codegenEnabled) {
-          a.newAggregation.isDefined
-        } else {
-          Utils.checkInvalidAggregateFunction2(a)
-          false
-        }
-      case _ => false
-    }
-
-    def allAggregates(exprs: Seq[Expression]): Seq[AggregateExpression1] =
-      exprs.flatMap(_.collect { case a: AggregateExpression1 => a })
-  }
-
   /**
    * Used to plan the aggregate operator for expressions based on the AggregateFunction2 interface.
    */
   object Aggregation extends Strategy {
     def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
-      case p: logical.Aggregate if sqlContext.conf.useSqlAggregate2 &&
-          sqlContext.conf.codegenEnabled =>
-        val converted = p.newAggregation
-        converted match {
-          case None => Nil // Cannot convert to new aggregation code path.
-          case Some(logical.Aggregate(groupingExpressions, resultExpressions, child)) =>
-            // A single aggregate expression might appear multiple times in resultExpressions.
-            // In order to avoid evaluating an individual aggregate function multiple times, we'll
-            // build a set of the distinct aggregate expressions and build a function which can
-            // be used to re-write expressions so that they reference the single copy of the
-            // aggregate function which actually gets computed.
-            val aggregateExpressions = resultExpressions.flatMap { expr =>
-              expr.collect {
-                case agg: AggregateExpression2 => agg
-              }
-            }.distinct
-            // For those distinct aggregate expressions, we create a map from the
-            // aggregate function to the corresponding attribute of the function.
-            val aggregateFunctionToAttribute = aggregateExpressions.map { agg =>
-              val aggregateFunction = agg.aggregateFunction
-              val attribute = Alias(aggregateFunction, aggregateFunction.toString)().toAttribute
-              (aggregateFunction, agg.isDistinct) -> attribute
-            }.toMap
-
-            val (functionsWithDistinct, functionsWithoutDistinct) =
-              aggregateExpressions.partition(_.isDistinct)
-            if (functionsWithDistinct.map(_.aggregateFunction.children).distinct.length > 1) {
-              // This is a sanity check. We should not reach here when we have multiple distinct
-              // column sets (aggregate.NewAggregation will not match).
-              sys.error(
-                "Multiple distinct column sets are not supported by the new aggregation" +
-                  "code path.")
-            }
+      case logical.Aggregate(groupingExpressions, resultExpressions, child) =>
+        // A single aggregate expression might appear multiple times in resultExpressions.
+        // In order to avoid evaluating an individual aggregate function multiple times, we'll
+        // build a set of the distinct aggregate expressions and build a function which can
+        // be used to re-write expressions so that they reference the single copy of the
+        // aggregate function which actually gets computed.
+        val aggregateExpressions = resultExpressions.flatMap { expr =>
+          expr.collect {
+            case agg: AggregateExpression => agg
+          }
+        }.distinct
+        // For those distinct aggregate expressions, we create a map from the
+        // aggregate function to the corresponding attribute of the function.
+        val aggregateFunctionToAttribute = aggregateExpressions.map { agg =>
+          val aggregateFunction = agg.aggregateFunction
+          val attribute = Alias(aggregateFunction, aggregateFunction.toString)().toAttribute
+          (aggregateFunction, agg.isDistinct) -> attribute
+        }.toMap
+
+        val (functionsWithDistinct, functionsWithoutDistinct) =
+          aggregateExpressions.partition(_.isDistinct)
+        if (functionsWithDistinct.map(_.aggregateFunction.children).distinct.length > 1) {
+          // This is a sanity check. We should not reach here when we have multiple distinct
+          // column sets. Our MultipleDistinctRewriter should take care this case.
+          sys.error("You hit a query analyzer bug. Please report your query to " +
+            "Spark user mailing list.")
+        }
 
-            val namedGroupingExpressions = groupingExpressions.map {
-              case ne: NamedExpression => ne -> ne
-              // If the expression is not a NamedExpressions, we add an alias.
-              // So, when we generate the result of the operator, the Aggregate Operator
-              // can directly get the Seq of attributes representing the grouping expressions.
-              case other =>
-                val withAlias = Alias(other, other.toString)()
-                other -> withAlias
-            }
-            val groupExpressionMap = namedGroupingExpressions.toMap
-
-            // The original `resultExpressions` are a set of expressions which may reference
-            // aggregate expressions, grouping column values, and constants. When aggregate operator
-            // emits output rows, we will use `resultExpressions` to generate an output projection
-            // which takes the grouping columns and final aggregate result buffer as input.
-            // Thus, we must re-write the result expressions so that their attributes match up with
-            // the attributes of the final result projection's input row:
-            val rewrittenResultExpressions = resultExpressions.map { expr =>
-              expr.transformDown {
-                case AggregateExpression2(aggregateFunction, _, isDistinct) =>
-                  // The final aggregation buffer's attributes will be `finalAggregationAttributes`,
-                  // so replace each aggregate expression by its corresponding attribute in the set:
-                  aggregateFunctionToAttribute(aggregateFunction, isDistinct)
-                case expression =>
-                  // Since we're using `namedGroupingAttributes` to extract the grouping key
-                  // columns, we need to replace grouping key expressions with their corresponding
-                  // attributes. We do not rely on the equality check at here since attributes may
-                  // differ cosmetically. Instead, we use semanticEquals.
-                  groupExpressionMap.collectFirst {
-                    case (expr, ne) if expr semanticEquals expression => ne.toAttribute
-                  }.getOrElse(expression)
-              }.asInstanceOf[NamedExpression]
+        val namedGroupingExpressions = groupingExpressions.map {
+          case ne: NamedExpression => ne -> ne
+          // If the expression is not a NamedExpressions, we add an alias.
+          // So, when we generate the result of the operator, the Aggregate Operator
+          // can directly get the Seq of attributes representing the grouping expressions.
+          case other =>
+            val withAlias = Alias(other, other.toString)()
+            other -> withAlias
+        }
+        val groupExpressionMap = namedGroupingExpressions.toMap
+
+        // The original `resultExpressions` are a set of expressions which may reference
+        // aggregate expressions, grouping column values, and constants. When aggregate operator
+        // emits output rows, we will use `resultExpressions` to generate an output projection
+        // which takes the grouping columns and final aggregate result buffer as input.
+        // Thus, we must re-write the result expressions so that their attributes match up with
+        // the attributes of the final result projection's input row:
+        val rewrittenResultExpressions = resultExpressions.map { expr =>
+          expr.transformDown {
+            case AggregateExpression(aggregateFunction, _, isDistinct) =>
+              // The final aggregation buffer's attributes will be `finalAggregationAttributes`,
+              // so replace each aggregate expression by its corresponding attribute in the set:
+              aggregateFunctionToAttribute(aggregateFunction, isDistinct)
+            case expression =>
+              // Since we're using `namedGroupingAttributes` to extract the grouping key
+              // columns, we need to replace grouping key expressions with their corresponding
+              // attributes. We do not rely on the equality check at here since attributes may
+              // differ cosmetically. Instead, we use semanticEquals.
+              groupExpressionMap.collectFirst {
+                case (expr, ne) if expr semanticEquals expression => ne.toAttribute
+              }.getOrElse(expression)
+          }.asInstanceOf[NamedExpression]
+        }
+
+        val aggregateOperator =
+          if (aggregateExpressions.map(_.aggregateFunction).exists(!_.supportsPartial)) {
+            if (functionsWithDistinct.nonEmpty) {
+              sys.error("Distinct columns cannot exist in Aggregate operator containing " +
+                "aggregate functions which don't support partial aggregation.")
+            } else {
+              aggregate.Utils.planAggregateWithoutPartial(
+                namedGroupingExpressions.map(_._2),
+                aggregateExpressions,
+                aggregateFunctionToAttribute,
+                rewrittenResultExpressions,
+                planLater(child))
             }
+          } else if (functionsWithDistinct.isEmpty) {
+            aggregate.Utils.planAggregateWithoutDistinct(
+              namedGroupingExpressions.map(_._2),
+              aggregateExpressions,
+              aggregateFunctionToAttribute,
+              rewrittenResultExpressions,
+              planLater(child))
+          } else {
+            aggregate.Utils.planAggregateWithOneDistinct(
+              namedGroupingExpressions.map(_._2),
+              functionsWithDistinct,
+              functionsWithoutDistinct,
+              aggregateFunctionToAttribute,
+              rewrittenResultExpressions,
+              planLater(child))
+          }
 
-            val aggregateOperator =
-              if (aggregateExpressions.map(_.aggregateFunction).exists(!_.supportsPartial)) {
-                if (functionsWithDistinct.nonEmpty) {
-                  sys.error("Distinct columns cannot exist in Aggregate operator containing " +
-                    "aggregate functions which don't support partial aggregation.")
-                } else {
-                  aggregate.Utils.planAggregateWithoutPartial(
-                    namedGroupingExpressions.map(_._2),
-                    aggregateExpressions,
-                    aggregateFunctionToAttribute,
-                    rewrittenResultExpressions,
-                    planLater(child))
-                }
-              } else if (functionsWithDistinct.isEmpty) {
-                aggregate.Utils.planAggregateWithoutDistinct(
-                  namedGroupingExpressions.map(_._2),
-                  aggregateExpressions,
-                  aggregateFunctionToAttribute,
-                  rewrittenResultExpressions,
-                  planLater(child))
-              } else {
-                aggregate.Utils.planAggregateWithOneDistinct(
-                  namedGroupingExpressions.map(_._2),
-                  functionsWithDistinct,
-                  functionsWithoutDistinct,
-                  aggregateFunctionToAttribute,
-                  rewrittenResultExpressions,
-                  planLater(child))
-              }
-
-            aggregateOperator
-        }
+        aggregateOperator
 
       case _ => Nil
     }
@@ -422,18 +378,6 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         execution.Filter(condition, planLater(child)) :: Nil
       case e @ logical.Expand(_, _, child) =>
         execution.Expand(e.projections, e.output, planLater(child)) :: Nil
-      case a @ logical.Aggregate(group, agg, child) => {
-        val useNewAggregation = sqlContext.conf.useSqlAggregate2 && sqlContext.conf.codegenEnabled
-        if (useNewAggregation && a.newAggregation.isDefined) {
-          // If this logical.Aggregate can be planned to use new aggregation code path
-          // (i.e. it can be planned by the Strategy Aggregation), we will not use the old
-          // aggregation code path.
-          Nil
-        } else {
-          Utils.checkInvalidAggregateFunction2(a)
-          execution.Aggregate(partial = false, group, agg, planLater(child)) :: Nil
-        }
-      }
       case logical.Window(projectList, windowExprs, partitionSpec, orderSpec, child) =>
         execution.Window(
           projectList, windowExprs, partitionSpec, orderSpec, planLater(child)) :: Nil
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregationIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregationIterator.scala
index 99fb7a40b72e1..008478a6a0e17 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregationIterator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggregationIterator.scala
@@ -35,9 +35,9 @@ import scala.collection.mutable.ArrayBuffer
 abstract class AggregationIterator(
     groupingKeyAttributes: Seq[Attribute],
     valueAttributes: Seq[Attribute],
-    nonCompleteAggregateExpressions: Seq[AggregateExpression2],
+    nonCompleteAggregateExpressions: Seq[AggregateExpression],
     nonCompleteAggregateAttributes: Seq[Attribute],
-    completeAggregateExpressions: Seq[AggregateExpression2],
+    completeAggregateExpressions: Seq[AggregateExpression],
     completeAggregateAttributes: Seq[Attribute],
     initialInputBufferOffset: Int,
     resultExpressions: Seq[NamedExpression],
@@ -76,14 +76,14 @@ abstract class AggregationIterator(
 
   // Initialize all AggregateFunctions by binding references if necessary,
   // and set inputBufferOffset and mutableBufferOffset.
-  protected val allAggregateFunctions: Array[AggregateFunction2] = {
+  protected val allAggregateFunctions: Array[AggregateFunction] = {
     var mutableBufferOffset = 0
     var inputBufferOffset: Int = initialInputBufferOffset
-    val functions = new Array[AggregateFunction2](allAggregateExpressions.length)
+    val functions = new Array[AggregateFunction](allAggregateExpressions.length)
     var i = 0
     while (i < allAggregateExpressions.length) {
       val func = allAggregateExpressions(i).aggregateFunction
-      val funcWithBoundReferences: AggregateFunction2 = allAggregateExpressions(i).mode match {
+      val funcWithBoundReferences: AggregateFunction = allAggregateExpressions(i).mode match {
         case Partial | Complete if func.isInstanceOf[ImperativeAggregate] =>
           // We need to create BoundReferences if the function is not an
           // expression-based aggregate function (it does not support code-gen) and the mode of
@@ -135,7 +135,7 @@ abstract class AggregationIterator(
   }
 
   // All AggregateFunctions functions with mode Partial, PartialMerge, or Final.
-  private[this] val nonCompleteAggregateFunctions: Array[AggregateFunction2] =
+  private[this] val nonCompleteAggregateFunctions: Array[AggregateFunction] =
     allAggregateFunctions.take(nonCompleteAggregateExpressions.length)
 
   // All imperative aggregate functions with mode Partial, PartialMerge, or Final.
@@ -172,7 +172,7 @@ abstract class AggregationIterator(
       case (Some(Partial), None) =>
         val updateExpressions = nonCompleteAggregateFunctions.flatMap {
           case ae: DeclarativeAggregate => ae.updateExpressions
-          case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+          case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
         }
         val expressionAggUpdateProjection =
           newMutableProjection(updateExpressions, aggregationBufferSchema ++ valueAttributes)()
@@ -204,7 +204,7 @@ abstract class AggregationIterator(
         //    allAggregateFunctions.flatMap(_.cloneBufferAttributes)
         val mergeExpressions = nonCompleteAggregateFunctions.flatMap {
           case ae: DeclarativeAggregate => ae.mergeExpressions
-          case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+          case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
         }
         // This projection is used to merge buffer values for all expression-based aggregates.
         val expressionAggMergeProjection =
@@ -225,7 +225,7 @@ abstract class AggregationIterator(
 
       // Final-Complete
       case (Some(Final), Some(Complete)) =>
-        val completeAggregateFunctions: Array[AggregateFunction2] =
+        val completeAggregateFunctions: Array[AggregateFunction] =
           allAggregateFunctions.takeRight(completeAggregateExpressions.length)
         // All imperative aggregate functions with mode Complete.
         val completeImperativeAggregateFunctions: Array[ImperativeAggregate] =
@@ -248,7 +248,7 @@ abstract class AggregationIterator(
         val mergeExpressions =
           nonCompleteAggregateFunctions.flatMap {
             case ae: DeclarativeAggregate => ae.mergeExpressions
-            case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+            case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
           } ++ completeOffsetExpressions
         val finalExpressionAggMergeProjection =
           newMutableProjection(mergeExpressions, mergeInputSchema)()
@@ -256,7 +256,7 @@ abstract class AggregationIterator(
         val updateExpressions =
           finalOffsetExpressions ++ completeAggregateFunctions.flatMap {
             case ae: DeclarativeAggregate => ae.updateExpressions
-            case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+            case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
           }
         val completeExpressionAggUpdateProjection =
           newMutableProjection(updateExpressions, aggregationBufferSchema ++ valueAttributes)()
@@ -282,7 +282,7 @@ abstract class AggregationIterator(
 
       // Complete-only
       case (None, Some(Complete)) =>
-        val completeAggregateFunctions: Array[AggregateFunction2] =
+        val completeAggregateFunctions: Array[AggregateFunction] =
           allAggregateFunctions.takeRight(completeAggregateExpressions.length)
         // All imperative aggregate functions with mode Complete.
         val completeImperativeAggregateFunctions: Array[ImperativeAggregate] =
@@ -291,7 +291,7 @@ abstract class AggregationIterator(
         val updateExpressions =
           completeAggregateFunctions.flatMap {
             case ae: DeclarativeAggregate => ae.updateExpressions
-            case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+            case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
           }
         val completeExpressionAggUpdateProjection =
           newMutableProjection(updateExpressions, aggregationBufferSchema ++ valueAttributes)()
@@ -353,7 +353,7 @@ abstract class AggregationIterator(
           allAggregateFunctions.flatMap(_.aggBufferAttributes)
         val evalExpressions = allAggregateFunctions.map {
           case ae: DeclarativeAggregate => ae.evaluateExpression
-          case agg: AggregateFunction2 => NoOp
+          case agg: AggregateFunction => NoOp
         }
         val expressionAggEvalProjection = newMutableProjection(evalExpressions, bufferSchemata)()
         val aggregateResultSchema = nonCompleteAggregateAttributes ++ completeAggregateAttributes
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
index 4d37106e007f5..fb7f30c2aec99 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
@@ -29,9 +29,9 @@ import org.apache.spark.sql.execution.metric.SQLMetrics
 case class SortBasedAggregate(
     requiredChildDistributionExpressions: Option[Seq[Expression]],
     groupingExpressions: Seq[NamedExpression],
-    nonCompleteAggregateExpressions: Seq[AggregateExpression2],
+    nonCompleteAggregateExpressions: Seq[AggregateExpression],
     nonCompleteAggregateAttributes: Seq[Attribute],
-    completeAggregateExpressions: Seq[AggregateExpression2],
+    completeAggregateExpressions: Seq[AggregateExpression],
     completeAggregateAttributes: Seq[Attribute],
     initialInputBufferOffset: Int,
     resultExpressions: Seq[NamedExpression],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregationIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregationIterator.scala
index 64c673064f576..fe5c3195f867b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregationIterator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregationIterator.scala
@@ -19,11 +19,11 @@ package org.apache.spark.sql.execution.aggregate
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression2, AggregateFunction2}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction}
 import org.apache.spark.sql.execution.metric.LongSQLMetric
 
 /**
- * An iterator used to evaluate [[AggregateFunction2]]. It assumes the input rows have been
+ * An iterator used to evaluate [[AggregateFunction]]. It assumes the input rows have been
  * sorted by values of [[groupingKeyAttributes]].
  */
 class SortBasedAggregationIterator(
@@ -31,9 +31,9 @@ class SortBasedAggregationIterator(
     groupingKeyAttributes: Seq[Attribute],
     valueAttributes: Seq[Attribute],
     inputIterator: Iterator[InternalRow],
-    nonCompleteAggregateExpressions: Seq[AggregateExpression2],
+    nonCompleteAggregateExpressions: Seq[AggregateExpression],
     nonCompleteAggregateAttributes: Seq[Attribute],
-    completeAggregateExpressions: Seq[AggregateExpression2],
+    completeAggregateExpressions: Seq[AggregateExpression],
     completeAggregateAttributes: Seq[Attribute],
     initialInputBufferOffset: Int,
     resultExpressions: Seq[NamedExpression],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregate.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregate.scala
index 15616915f7364..1edde1e5a16d9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregate.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregate.scala
@@ -21,7 +21,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.errors._
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression2
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans.physical._
 import org.apache.spark.sql.execution.metric.SQLMetrics
 import org.apache.spark.sql.execution.{SparkPlan, UnaryNode, UnsafeFixedWidthAggregationMap}
@@ -30,9 +30,9 @@ import org.apache.spark.sql.types.StructType
 case class TungstenAggregate(
     requiredChildDistributionExpressions: Option[Seq[Expression]],
     groupingExpressions: Seq[NamedExpression],
-    nonCompleteAggregateExpressions: Seq[AggregateExpression2],
+    nonCompleteAggregateExpressions: Seq[AggregateExpression],
     nonCompleteAggregateAttributes: Seq[Attribute],
-    completeAggregateExpressions: Seq[AggregateExpression2],
+    completeAggregateExpressions: Seq[AggregateExpression],
     completeAggregateAttributes: Seq[Attribute],
     initialInputBufferOffset: Int,
     resultExpressions: Seq[NamedExpression],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala
index ce8d592c368ee..04391443920ac 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala
@@ -64,12 +64,12 @@ import org.apache.spark.sql.types.StructType
  * @param groupingExpressions
  *   expressions for grouping keys
  * @param nonCompleteAggregateExpressions
- *   [[AggregateExpression2]] containing [[AggregateFunction2]]s with mode [[Partial]],
- *   [[PartialMerge]], or [[Final]].
+ * [[AggregateExpression]] containing [[AggregateFunction]]s with mode [[Partial]],
+ * [[PartialMerge]], or [[Final]].
  * @param nonCompleteAggregateAttributes the attributes of the nonCompleteAggregateExpressions'
  *   outputs when they are stored in the final aggregation buffer.
  * @param completeAggregateExpressions
- *   [[AggregateExpression2]] containing [[AggregateFunction2]]s with mode [[Complete]].
+ * [[AggregateExpression]] containing [[AggregateFunction]]s with mode [[Complete]].
  * @param completeAggregateAttributes the attributes of completeAggregateExpressions' outputs
  *   when they are stored in the final aggregation buffer.
  * @param resultExpressions
@@ -83,9 +83,9 @@ import org.apache.spark.sql.types.StructType
  */
 class TungstenAggregationIterator(
     groupingExpressions: Seq[NamedExpression],
-    nonCompleteAggregateExpressions: Seq[AggregateExpression2],
+    nonCompleteAggregateExpressions: Seq[AggregateExpression],
     nonCompleteAggregateAttributes: Seq[Attribute],
-    completeAggregateExpressions: Seq[AggregateExpression2],
+    completeAggregateExpressions: Seq[AggregateExpression],
     completeAggregateAttributes: Seq[Attribute],
     initialInputBufferOffset: Int,
     resultExpressions: Seq[NamedExpression],
@@ -106,7 +106,7 @@ class TungstenAggregationIterator(
   // A Seq containing all AggregateExpressions.
   // It is important that all AggregateExpressions with the mode Partial, PartialMerge or Final
   // are at the beginning of the allAggregateExpressions.
-  private[this] val allAggregateExpressions: Seq[AggregateExpression2] =
+  private[this] val allAggregateExpressions: Seq[AggregateExpression] =
     nonCompleteAggregateExpressions ++ completeAggregateExpressions
 
   // Check to make sure we do not have more than three modes in our AggregateExpressions.
@@ -150,10 +150,10 @@ class TungstenAggregationIterator(
   // Initialize all AggregateFunctions by binding references, if necessary,
   // and setting inputBufferOffset and mutableBufferOffset.
   private def initializeAllAggregateFunctions(
-      startingInputBufferOffset: Int): Array[AggregateFunction2] = {
+      startingInputBufferOffset: Int): Array[AggregateFunction] = {
     var mutableBufferOffset = 0
     var inputBufferOffset: Int = startingInputBufferOffset
-    val functions = new Array[AggregateFunction2](allAggregateExpressions.length)
+    val functions = new Array[AggregateFunction](allAggregateExpressions.length)
     var i = 0
     while (i < allAggregateExpressions.length) {
       val func = allAggregateExpressions(i).aggregateFunction
@@ -195,7 +195,7 @@ class TungstenAggregationIterator(
     functions
   }
 
-  private[this] var allAggregateFunctions: Array[AggregateFunction2] =
+  private[this] var allAggregateFunctions: Array[AggregateFunction] =
     initializeAllAggregateFunctions(initialInputBufferOffset)
 
   // Positions of those imperative aggregate functions in allAggregateFunctions.
@@ -263,7 +263,7 @@ class TungstenAggregationIterator(
       case (Some(Partial), None) =>
         val updateExpressions = allAggregateFunctions.flatMap {
           case ae: DeclarativeAggregate => ae.updateExpressions
-          case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+          case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
         }
         val imperativeAggregateFunctions: Array[ImperativeAggregate] =
           allAggregateFunctions.collect { case func: ImperativeAggregate => func}
@@ -286,7 +286,7 @@ class TungstenAggregationIterator(
       case (Some(PartialMerge), None) | (Some(Final), None) =>
         val mergeExpressions = allAggregateFunctions.flatMap {
           case ae: DeclarativeAggregate => ae.mergeExpressions
-          case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+          case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
         }
         val imperativeAggregateFunctions: Array[ImperativeAggregate] =
           allAggregateFunctions.collect { case func: ImperativeAggregate => func}
@@ -307,11 +307,11 @@ class TungstenAggregationIterator(
 
       // Final-Complete
       case (Some(Final), Some(Complete)) =>
-        val completeAggregateFunctions: Array[AggregateFunction2] =
+        val completeAggregateFunctions: Array[AggregateFunction] =
           allAggregateFunctions.takeRight(completeAggregateExpressions.length)
         val completeImperativeAggregateFunctions: Array[ImperativeAggregate] =
           completeAggregateFunctions.collect { case func: ImperativeAggregate => func }
-        val nonCompleteAggregateFunctions: Array[AggregateFunction2] =
+        val nonCompleteAggregateFunctions: Array[AggregateFunction] =
           allAggregateFunctions.take(nonCompleteAggregateExpressions.length)
         val nonCompleteImperativeAggregateFunctions: Array[ImperativeAggregate] =
           nonCompleteAggregateFunctions.collect { case func: ImperativeAggregate => func }
@@ -321,7 +321,7 @@ class TungstenAggregationIterator(
         val mergeExpressions =
           nonCompleteAggregateFunctions.flatMap {
             case ae: DeclarativeAggregate => ae.mergeExpressions
-            case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+            case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
           } ++ completeOffsetExpressions
         val finalMergeProjection =
           newMutableProjection(mergeExpressions, aggregationBufferAttributes ++ inputAttributes)()
@@ -331,7 +331,7 @@ class TungstenAggregationIterator(
           Seq.fill(nonCompleteAggregateFunctions.map(_.aggBufferAttributes.length).sum)(NoOp)
         val updateExpressions = finalOffsetExpressions ++ completeAggregateFunctions.flatMap {
           case ae: DeclarativeAggregate => ae.updateExpressions
-          case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+          case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
         }
         val completeUpdateProjection =
           newMutableProjection(updateExpressions, aggregationBufferAttributes ++ inputAttributes)()
@@ -358,7 +358,7 @@ class TungstenAggregationIterator(
 
       // Complete-only
       case (None, Some(Complete)) =>
-        val completeAggregateFunctions: Array[AggregateFunction2] =
+        val completeAggregateFunctions: Array[AggregateFunction] =
           allAggregateFunctions.takeRight(completeAggregateExpressions.length)
         // All imperative aggregate functions with mode Complete.
         val completeImperativeAggregateFunctions: Array[ImperativeAggregate] =
@@ -366,7 +366,7 @@ class TungstenAggregationIterator(
 
         val updateExpressions = completeAggregateFunctions.flatMap {
           case ae: DeclarativeAggregate => ae.updateExpressions
-          case agg: AggregateFunction2 => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
+          case agg: AggregateFunction => Seq.fill(agg.aggBufferAttributes.length)(NoOp)
         }
         val completeExpressionAggUpdateProjection =
           newMutableProjection(updateExpressions, aggregationBufferAttributes ++ inputAttributes)()
@@ -414,7 +414,7 @@ class TungstenAggregationIterator(
         val joinedRow = new JoinedRow()
         val evalExpressions = allAggregateFunctions.map {
           case ae: DeclarativeAggregate => ae.evaluateExpression
-          case agg: AggregateFunction2 => NoOp
+          case agg: AggregateFunction => NoOp
         }
         val expressionAggEvalProjection = newMutableProjection(evalExpressions, bufferAttributes)()
         // These are the attributes of the row produced by `expressionAggEvalProjection`
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
index d2f56e0fc14a4..20359c1e540e2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/udaf.scala
@@ -22,7 +22,7 @@ import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.{InternalRow, CatalystTypeConverters}
 import org.apache.spark.sql.catalyst.expressions.codegen.GenerateMutableProjection
 import org.apache.spark.sql.catalyst.expressions.{MutableRow, InterpretedMutableProjection, AttributeReference, Expression}
-import org.apache.spark.sql.catalyst.expressions.aggregate.{ImperativeAggregate, AggregateFunction2}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{ImperativeAggregate, AggregateFunction}
 import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
 import org.apache.spark.sql.types._
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
index eaafd83158a15..79abf2d5929be 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
@@ -28,8 +28,8 @@ object Utils {
 
   def planAggregateWithoutPartial(
       groupingExpressions: Seq[NamedExpression],
-      aggregateExpressions: Seq[AggregateExpression2],
-      aggregateFunctionToAttribute: Map[(AggregateFunction2, Boolean), Attribute],
+      aggregateExpressions: Seq[AggregateExpression],
+      aggregateFunctionToAttribute: Map[(AggregateFunction, Boolean), Attribute],
       resultExpressions: Seq[NamedExpression],
       child: SparkPlan): Seq[SparkPlan] = {
 
@@ -54,8 +54,8 @@ object Utils {
 
   def planAggregateWithoutDistinct(
       groupingExpressions: Seq[NamedExpression],
-      aggregateExpressions: Seq[AggregateExpression2],
-      aggregateFunctionToAttribute: Map[(AggregateFunction2, Boolean), Attribute],
+      aggregateExpressions: Seq[AggregateExpression],
+      aggregateFunctionToAttribute: Map[(AggregateFunction, Boolean), Attribute],
       resultExpressions: Seq[NamedExpression],
       child: SparkPlan): Seq[SparkPlan] = {
     // Check if we can use TungstenAggregate.
@@ -137,9 +137,9 @@ object Utils {
 
   def planAggregateWithOneDistinct(
       groupingExpressions: Seq[NamedExpression],
-      functionsWithDistinct: Seq[AggregateExpression2],
-      functionsWithoutDistinct: Seq[AggregateExpression2],
-      aggregateFunctionToAttribute: Map[(AggregateFunction2, Boolean), Attribute],
+      functionsWithDistinct: Seq[AggregateExpression],
+      functionsWithoutDistinct: Seq[AggregateExpression],
+      aggregateFunctionToAttribute: Map[(AggregateFunction, Boolean), Attribute],
       resultExpressions: Seq[NamedExpression],
       child: SparkPlan): Seq[SparkPlan] = {
 
@@ -253,16 +253,16 @@ object Utils {
         // Children of an AggregateFunction with DISTINCT keyword has already
         // been evaluated. At here, we need to replace original children
         // to AttributeReferences.
-        case agg @ AggregateExpression2(aggregateFunction, mode, true) =>
+        case agg @ AggregateExpression(aggregateFunction, mode, true) =>
           val rewrittenAggregateFunction = aggregateFunction.transformDown {
             case expr if expr == distinctColumnExpression => distinctColumnAttribute
-          }.asInstanceOf[AggregateFunction2]
+          }.asInstanceOf[AggregateFunction]
           // We rewrite the aggregate function to a non-distinct aggregation because
           // its input will have distinct arguments.
           // We just keep the isDistinct setting to true, so when users look at the query plan,
           // they still can see distinct aggregations.
           val rewrittenAggregateExpression =
-            AggregateExpression2(rewrittenAggregateFunction, Complete, isDistinct = true)
+            AggregateExpression(rewrittenAggregateFunction, Complete, isDistinct = true)
 
           val aggregateFunctionAttribute = aggregateFunctionToAttribute(agg.aggregateFunction, true)
           (rewrittenAggregateExpression, aggregateFunctionAttribute)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
index 0b3192a6da9d8..8cc25c2440633 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
@@ -18,7 +18,7 @@
 package org.apache.spark.sql.expressions
 
 import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, AggregateExpression2}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete}
 import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
 import org.apache.spark.sql.{Dataset, DataFrame, TypedColumn}
 
@@ -70,7 +70,7 @@ abstract class Aggregator[-A, B, C] {
       implicit bEncoder: Encoder[B],
       cEncoder: Encoder[C]): TypedColumn[A, C] = {
     val expr =
-      new AggregateExpression2(
+      new AggregateExpression(
         TypedAggregateExpression(this),
         Complete,
         false)
@@ -78,4 +78,3 @@ abstract class Aggregator[-A, B, C] {
     new TypedColumn[A, C](expr, encoderFor[C])
   }
 }
-
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/WindowSpec.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/WindowSpec.scala
index 8b9247adea200..fc873c04f88f0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/WindowSpec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/WindowSpec.scala
@@ -18,9 +18,9 @@
 package org.apache.spark.sql.expressions
 
 import org.apache.spark.annotation.Experimental
-import org.apache.spark.sql.types.BooleanType
 import org.apache.spark.sql.{Column, catalyst}
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 
 
 /**
@@ -141,40 +141,56 @@ class WindowSpec private[sql](
    */
   private[sql] def withAggregate(aggregate: Column): Column = {
     val windowExpr = aggregate.expr match {
-      case Average(child) => WindowExpression(
-        UnresolvedWindowFunction("avg", child :: Nil),
-        WindowSpecDefinition(partitionSpec, orderSpec, frame))
-      case Sum(child) => WindowExpression(
-        UnresolvedWindowFunction("sum", child :: Nil),
-        WindowSpecDefinition(partitionSpec, orderSpec, frame))
-      case Count(child) => WindowExpression(
-        UnresolvedWindowFunction("count", child :: Nil),
-        WindowSpecDefinition(partitionSpec, orderSpec, frame))
-      case First(child, ignoreNulls) => WindowExpression(
-        // TODO this is a hack for Hive UDAF first_value
-        UnresolvedWindowFunction(
-          "first_value",
-          child :: ignoreNulls :: Nil),
-        WindowSpecDefinition(partitionSpec, orderSpec, frame))
-      case Last(child, ignoreNulls) => WindowExpression(
-        // TODO this is a hack for Hive UDAF last_value
-        UnresolvedWindowFunction(
-          "last_value",
-          child :: ignoreNulls :: Nil),
-        WindowSpecDefinition(partitionSpec, orderSpec, frame))
-      case Min(child) => WindowExpression(
-        UnresolvedWindowFunction("min", child :: Nil),
-        WindowSpecDefinition(partitionSpec, orderSpec, frame))
-      case Max(child) => WindowExpression(
-        UnresolvedWindowFunction("max", child :: Nil),
-        WindowSpecDefinition(partitionSpec, orderSpec, frame))
-      case wf: WindowFunction => WindowExpression(
-        wf,
-        WindowSpecDefinition(partitionSpec, orderSpec, frame))
+      // First, we check if we get an aggregate function without the DISTINCT keyword.
+      // Right now, we do not support using a DISTINCT aggregate function as a
+      // window function.
+      case AggregateExpression(aggregateFunction, _, isDistinct) if !isDistinct =>
+        aggregateFunction match {
+          case Average(child) => WindowExpression(
+            UnresolvedWindowFunction("avg", child :: Nil),
+            WindowSpecDefinition(partitionSpec, orderSpec, frame))
+          case Sum(child) => WindowExpression(
+            UnresolvedWindowFunction("sum", child :: Nil),
+            WindowSpecDefinition(partitionSpec, orderSpec, frame))
+          case Count(child) => WindowExpression(
+            UnresolvedWindowFunction("count", child :: Nil),
+            WindowSpecDefinition(partitionSpec, orderSpec, frame))
+          case First(child, ignoreNulls) => WindowExpression(
+            // TODO this is a hack for Hive UDAF first_value
+            UnresolvedWindowFunction(
+              "first_value",
+              child :: ignoreNulls :: Nil),
+            WindowSpecDefinition(partitionSpec, orderSpec, frame))
+          case Last(child, ignoreNulls) => WindowExpression(
+            // TODO this is a hack for Hive UDAF last_value
+            UnresolvedWindowFunction(
+              "last_value",
+              child :: ignoreNulls :: Nil),
+            WindowSpecDefinition(partitionSpec, orderSpec, frame))
+          case Min(child) => WindowExpression(
+            UnresolvedWindowFunction("min", child :: Nil),
+            WindowSpecDefinition(partitionSpec, orderSpec, frame))
+          case Max(child) => WindowExpression(
+            UnresolvedWindowFunction("max", child :: Nil),
+            WindowSpecDefinition(partitionSpec, orderSpec, frame))
+          case x =>
+            throw new UnsupportedOperationException(s"$x is not supported in a window operation.")
+        }
+
+      case AggregateExpression(aggregateFunction, _, isDistinct) if isDistinct =>
+        throw new UnsupportedOperationException(
+          s"Distinct aggregate function ${aggregateFunction} is not supported " +
+            s"in window operation.")
+
+      case wf: WindowFunction =>
+        WindowExpression(
+          wf,
+          WindowSpecDefinition(partitionSpec, orderSpec, frame))
+
       case x =>
-        throw new UnsupportedOperationException(s"$x is not supported in window operation.")
+        throw new UnsupportedOperationException(s"$x is not supported in a window operation.")
     }
+
     new Column(windowExpr)
   }
-
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/udaf.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/udaf.scala
index 258afadc76951..11dbf391cff98 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/udaf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/udaf.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.expressions
 
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, AggregateExpression2}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, AggregateExpression}
 import org.apache.spark.sql.execution.aggregate.ScalaUDAF
 import org.apache.spark.sql.{Column, Row}
 import org.apache.spark.sql.types._
@@ -109,7 +109,7 @@ abstract class UserDefinedAggregateFunction extends Serializable {
   @scala.annotation.varargs
   def apply(exprs: Column*): Column = {
     val aggregateExpression =
-      AggregateExpression2(
+      AggregateExpression(
         ScalaUDAF(exprs.map(_.expr), this),
         Complete,
         isDistinct = false)
@@ -123,7 +123,7 @@ abstract class UserDefinedAggregateFunction extends Serializable {
   @scala.annotation.varargs
   def distinct(exprs: Column*): Column = {
     val aggregateExpression =
-      AggregateExpression2(
+      AggregateExpression(
         ScalaUDAF(exprs.map(_.expr), this),
         Complete,
         isDistinct = true)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 6d56542ee0875..22104e4d48617 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.{SqlParser, ScalaReflection}
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunction, Star}
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, Encoder}
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.BroadcastHint
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
@@ -76,6 +77,12 @@ object functions extends LegacyFunctions {
 
   private def withExpr(expr: Expression): Column = Column(expr)
 
+  private def withAggregateFunction(
+    func: AggregateFunction,
+    isDistinct: Boolean = false): Column = {
+    Column(func.toAggregateExpression(isDistinct))
+  }
+
   private implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder[Long](flat = true)
 
 
@@ -154,7 +161,9 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def approxCountDistinct(e: Column): Column = withExpr { ApproxCountDistinct(e.expr) }
+  def approxCountDistinct(e: Column): Column = withAggregateFunction {
+    HyperLogLogPlusPlus(e.expr)
+  }
 
   /**
    * Aggregate function: returns the approximate number of distinct items in a group.
@@ -170,8 +179,8 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def approxCountDistinct(e: Column, rsd: Double): Column = withExpr {
-    ApproxCountDistinct(e.expr, rsd)
+  def approxCountDistinct(e: Column, rsd: Double): Column = withAggregateFunction {
+    HyperLogLogPlusPlus(e.expr, rsd, 0, 0)
   }
 
   /**
@@ -190,7 +199,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def avg(e: Column): Column = withExpr { Average(e.expr) }
+  def avg(e: Column): Column = withAggregateFunction { Average(e.expr) }
 
   /**
    * Aggregate function: returns the average of the values in a group.
@@ -226,7 +235,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def corr(column1: Column, column2: Column): Column = withExpr {
+  def corr(column1: Column, column2: Column): Column = withAggregateFunction {
     Corr(column1.expr, column2.expr)
   }
 
@@ -246,7 +255,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def count(e: Column): Column = withExpr {
+  def count(e: Column): Column = withAggregateFunction {
     e.expr match {
       // Turn count(*) into count(1)
       case s: Star => Count(Literal(1))
@@ -269,8 +278,8 @@ object functions extends LegacyFunctions {
    * @since 1.3.0
    */
   @scala.annotation.varargs
-  def countDistinct(expr: Column, exprs: Column*): Column = withExpr {
-    CountDistinct((expr +: exprs).map(_.expr))
+  def countDistinct(expr: Column, exprs: Column*): Column = {
+    withAggregateFunction(Count.apply((expr +: exprs).map(_.expr)), isDistinct = true)
   }
 
   /**
@@ -289,7 +298,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def first(e: Column): Column = withExpr { First(e.expr) }
+  def first(e: Column): Column = withAggregateFunction { new First(e.expr) }
 
   /**
    * Aggregate function: returns the first value of a column in a group.
@@ -305,7 +314,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def kurtosis(e: Column): Column = withExpr { Kurtosis(e.expr) }
+  def kurtosis(e: Column): Column = withAggregateFunction { Kurtosis(e.expr) }
 
   /**
    * Aggregate function: returns the last value in a group.
@@ -313,7 +322,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def last(e: Column): Column = withExpr { Last(e.expr) }
+  def last(e: Column): Column = withAggregateFunction { new Last(e.expr) }
 
   /**
    * Aggregate function: returns the last value of the column in a group.
@@ -329,7 +338,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def max(e: Column): Column = withExpr { Max(e.expr) }
+  def max(e: Column): Column = withAggregateFunction { Max(e.expr) }
 
   /**
    * Aggregate function: returns the maximum value of the column in a group.
@@ -363,7 +372,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def min(e: Column): Column = withExpr { Min(e.expr) }
+  def min(e: Column): Column = withAggregateFunction { Min(e.expr) }
 
   /**
    * Aggregate function: returns the minimum value of the column in a group.
@@ -379,7 +388,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def skewness(e: Column): Column = withExpr { Skewness(e.expr) }
+  def skewness(e: Column): Column = withAggregateFunction { Skewness(e.expr) }
 
   /**
    * Aggregate function: alias for [[stddev_samp]].
@@ -387,7 +396,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def stddev(e: Column): Column = withExpr { StddevSamp(e.expr) }
+  def stddev(e: Column): Column = withAggregateFunction { StddevSamp(e.expr) }
 
   /**
    * Aggregate function: returns the unbiased sample standard deviation of
@@ -396,7 +405,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def stddev_samp(e: Column): Column = withExpr { StddevSamp(e.expr) }
+  def stddev_samp(e: Column): Column = withAggregateFunction { StddevSamp(e.expr) }
 
   /**
    * Aggregate function: returns the population standard deviation of
@@ -405,7 +414,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def stddev_pop(e: Column): Column = withExpr { StddevPop(e.expr) }
+  def stddev_pop(e: Column): Column = withAggregateFunction { StddevPop(e.expr) }
 
   /**
    * Aggregate function: returns the sum of all values in the expression.
@@ -413,7 +422,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def sum(e: Column): Column = withExpr { Sum(e.expr) }
+  def sum(e: Column): Column = withAggregateFunction { Sum(e.expr) }
 
   /**
    * Aggregate function: returns the sum of all values in the given column.
@@ -429,7 +438,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def sumDistinct(e: Column): Column = withExpr { SumDistinct(e.expr) }
+  def sumDistinct(e: Column): Column = withAggregateFunction(Sum(e.expr), isDistinct = true)
 
   /**
    * Aggregate function: returns the sum of distinct values in the expression.
@@ -445,7 +454,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def variance(e: Column): Column = withExpr { VarianceSamp(e.expr) }
+  def variance(e: Column): Column = withAggregateFunction { VarianceSamp(e.expr) }
 
   /**
    * Aggregate function: returns the unbiased variance of the values in a group.
@@ -453,7 +462,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def var_samp(e: Column): Column = withExpr { VarianceSamp(e.expr) }
+  def var_samp(e: Column): Column = withAggregateFunction { VarianceSamp(e.expr) }
 
   /**
    * Aggregate function: returns the population variance of the values in a group.
@@ -461,7 +470,7 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.6.0
    */
-  def var_pop(e: Column): Column = withExpr { VariancePop(e.expr) }
+  def var_pop(e: Column): Column = withAggregateFunction { VariancePop(e.expr) }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
   // Window functions
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index 3de277a79a52c..441a0c6d0e36e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -237,34 +237,10 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
   }
 
   test("SPARK-8828 sum should return null if all input values are null") {
-    withSQLConf(SQLConf.USE_SQL_AGGREGATE2.key -> "true") {
-      withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "true") {
-        checkAnswer(
-          sql("select sum(a), avg(a) from allNulls"),
-          Seq(Row(null, null))
-        )
-      }
-      withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "false") {
-        checkAnswer(
-          sql("select sum(a), avg(a) from allNulls"),
-          Seq(Row(null, null))
-        )
-      }
-    }
-    withSQLConf(SQLConf.USE_SQL_AGGREGATE2.key -> "false") {
-      withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "true") {
-        checkAnswer(
-          sql("select sum(a), avg(a) from allNulls"),
-          Seq(Row(null, null))
-        )
-      }
-      withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "false") {
-        checkAnswer(
-          sql("select sum(a), avg(a) from allNulls"),
-          Seq(Row(null, null))
-        )
-      }
-    }
+    checkAnswer(
+      sql("select sum(a), avg(a) from allNulls"),
+      Seq(Row(null, null))
+    )
   }
 
   private def testCodeGen(sqlText: String, expectedResults: Seq[Row]): Unit = {
@@ -507,29 +483,22 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
   }
 
   test("literal in agg grouping expressions") {
-    def literalInAggTest(): Unit = {
-      checkAnswer(
-        sql("SELECT a, count(1) FROM testData2 GROUP BY a, 1"),
-        Seq(Row(1, 2), Row(2, 2), Row(3, 2)))
-      checkAnswer(
-        sql("SELECT a, count(2) FROM testData2 GROUP BY a, 2"),
-        Seq(Row(1, 2), Row(2, 2), Row(3, 2)))
-
-      checkAnswer(
-        sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a, 1"),
-        sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a"))
-      checkAnswer(
-        sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a, 1 + 2"),
-        sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a"))
-      checkAnswer(
-        sql("SELECT 1, 2, sum(b) FROM testData2 GROUP BY 1, 2"),
-        sql("SELECT 1, 2, sum(b) FROM testData2"))
-    }
+    checkAnswer(
+      sql("SELECT a, count(1) FROM testData2 GROUP BY a, 1"),
+      Seq(Row(1, 2), Row(2, 2), Row(3, 2)))
+    checkAnswer(
+      sql("SELECT a, count(2) FROM testData2 GROUP BY a, 2"),
+      Seq(Row(1, 2), Row(2, 2), Row(3, 2)))
 
-    literalInAggTest()
-    withSQLConf(SQLConf.USE_SQL_AGGREGATE2.key -> "false") {
-      literalInAggTest()
-    }
+    checkAnswer(
+      sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a, 1"),
+      sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a"))
+    checkAnswer(
+      sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a, 1 + 2"),
+      sql("SELECT a, 1, sum(b) FROM testData2 GROUP BY a"))
+    checkAnswer(
+      sql("SELECT 1, 2, sum(b) FROM testData2 GROUP BY 1, 2"),
+      sql("SELECT 1, 2, sum(b) FROM testData2"))
   }
 
   test("aggregates with nulls") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
index a229e5814df89..e31c528f3a633 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
@@ -21,16 +21,13 @@ import org.apache.spark.sql.catalyst.util.{GenericArrayData, ArrayData}
 
 import scala.beans.{BeanInfo, BeanProperty}
 
-import com.clearspring.analytics.stream.cardinality.HyperLogLog
-
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.CatalystTypeConverters
-import org.apache.spark.sql.catalyst.expressions.{OpenHashSetUDT, HyperLogLogUDT}
+import org.apache.spark.sql.catalyst.expressions.OpenHashSetUDT
 import org.apache.spark.sql.execution.datasources.parquet.ParquetTest
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
-import org.apache.spark.util.Utils
 import org.apache.spark.util.collection.OpenHashSet
 
 
@@ -134,16 +131,6 @@ class UserDefinedTypeSuite extends QueryTest with SharedSQLContext with ParquetT
     df.orderBy('int).limit(1).groupBy('int).agg(first('vec)).collect()(0).getAs[MyDenseVector](0)
   }
 
-  test("HyperLogLogUDT") {
-    val hyperLogLogUDT = HyperLogLogUDT
-    val hyperLogLog = new HyperLogLog(0.4)
-    (1 to 10).foreach(i => hyperLogLog.offer(Row(i)))
-
-    val actual = hyperLogLogUDT.deserialize(hyperLogLogUDT.serialize(hyperLogLog))
-    assert(actual.cardinality() === hyperLogLog.cardinality())
-    assert(java.util.Arrays.equals(actual.getBytes, hyperLogLog.getBytes))
-  }
-
   test("OpenHashSetUDT") {
     val openHashSetUDT = new OpenHashSetUDT(IntegerType)
     val set = new OpenHashSet[Int]
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
index 2076c573b56c1..44634dacbde68 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -38,7 +38,7 @@ class PlannerSuite extends SharedSQLContext {
   private def testPartialAggregationPlan(query: LogicalPlan): Unit = {
     val planner = sqlContext.planner
     import planner._
-    val plannedOption = HashAggregation(query).headOption.orElse(Aggregation(query).headOption)
+    val plannedOption = Aggregation(query).headOption
     val planned =
       plannedOption.getOrElse(
         fail(s"Could query play aggregation query $query. Is it an aggregation query?"))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index cdd885ba14203..4b4f5c6c45c7a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -152,36 +152,6 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
     )
   }
 
-  test("Aggregate metrics") {
-    withSQLConf(
-      SQLConf.UNSAFE_ENABLED.key -> "false",
-      SQLConf.CODEGEN_ENABLED.key -> "false",
-      SQLConf.TUNGSTEN_ENABLED.key -> "false") {
-      // Assume the execution plan is
-      // ... -> Aggregate(nodeId = 2) -> TungstenExchange(nodeId = 1) -> Aggregate(nodeId = 0)
-      val df = testData2.groupBy().count() // 2 partitions
-      testSparkPlanMetrics(df, 1, Map(
-        2L -> ("Aggregate", Map(
-          "number of input rows" -> 6L,
-          "number of output rows" -> 2L)),
-        0L -> ("Aggregate", Map(
-          "number of input rows" -> 2L,
-          "number of output rows" -> 1L)))
-      )
-
-      // 2 partitions and each partition contains 2 keys
-      val df2 = testData2.groupBy('a).count()
-      testSparkPlanMetrics(df2, 1, Map(
-        2L -> ("Aggregate", Map(
-          "number of input rows" -> 6L,
-          "number of output rows" -> 4L)),
-        0L -> ("Aggregate", Map(
-          "number of input rows" -> 4L,
-          "number of output rows" -> 3L)))
-      )
-    }
-  }
-
   test("SortBasedAggregate metrics") {
     // Because SortBasedAggregate may skip different rows if the number of partitions is different,
     // this test should use the deterministic number of partitions.
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index c5f69657f5293..ba6204633b9ca 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -584,7 +584,6 @@ class HiveContext private[hive](
       HiveTableScans,
       DataSinks,
       Scripts,
-      HashAggregation,
       Aggregation,
       LeftSemiJoin,
       EquiJoinSelection,
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
index ab88c1e68fd72..6f8ed413a06cd 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
@@ -38,6 +38,7 @@ import org.apache.spark.Logging
 import org.apache.spark.sql.{AnalysisException, catalyst}
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.{logical, _}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.trees.CurrentOrigin
@@ -1508,9 +1509,10 @@ https://cwiki.apache.org/confluence/display/Hive/Enhanced+Aggregation%2C+Cube%2C
       UnresolvedStar(Some(UnresolvedAttribute.parseAttributeName(name)))
 
     /* Aggregate Functions */
-    case Token("TOK_FUNCTIONSTAR", Token(COUNT(), Nil) :: Nil) => Count(Literal(1))
-    case Token("TOK_FUNCTIONDI", Token(COUNT(), Nil) :: args) => CountDistinct(args.map(nodeToExpr))
-    case Token("TOK_FUNCTIONDI", Token(SUM(), Nil) :: arg :: Nil) => SumDistinct(nodeToExpr(arg))
+    case Token("TOK_FUNCTIONDI", Token(COUNT(), Nil) :: args) =>
+      Count(args.map(nodeToExpr)).toAggregateExpression(isDistinct = true)
+    case Token("TOK_FUNCTIONSTAR", Token(COUNT(), Nil) :: Nil) =>
+      Count(Literal(1)).toAggregateExpression()
 
     /* Casts */
     case Token("TOK_FUNCTION", Token("TOK_STRING", Nil) :: arg :: Nil) =>
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index ea36c132bb190..6bf2c53440baf 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -69,11 +69,7 @@ class ScalaAggregateFunction(schema: StructType) extends UserDefinedAggregateFun
 abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
   import testImplicits._
 
-  var originalUseAggregate2: Boolean = _
-
   override def beforeAll(): Unit = {
-    originalUseAggregate2 = sqlContext.conf.useSqlAggregate2
-    sqlContext.setConf(SQLConf.USE_SQL_AGGREGATE2.key, "true")
     val data1 = Seq[(Integer, Integer)](
       (1, 10),
       (null, -60),
@@ -120,7 +116,6 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
     sqlContext.sql("DROP TABLE IF EXISTS agg1")
     sqlContext.sql("DROP TABLE IF EXISTS agg2")
     sqlContext.dropTempTable("emptyTable")
-    sqlContext.setConf(SQLConf.USE_SQL_AGGREGATE2.key, originalUseAggregate2.toString)
   }
 
   test("empty table") {
@@ -447,73 +442,80 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
   }
 
   test("single distinct column set") {
-    // DISTINCT is not meaningful with Max and Min, so we just ignore the DISTINCT keyword.
-    checkAnswer(
-      sqlContext.sql(
-        """
-          |SELECT
-          |  min(distinct value1),
-          |  sum(distinct value1),
-          |  avg(value1),
-          |  avg(value2),
-          |  max(distinct value1)
-          |FROM agg2
-        """.stripMargin),
-      Row(-60, 70.0, 101.0/9.0, 5.6, 100))
-
-    checkAnswer(
-      sqlContext.sql(
-        """
-          |SELECT
-          |  mydoubleavg(distinct value1),
-          |  avg(value1),
-          |  avg(value2),
-          |  key,
-          |  mydoubleavg(value1 - 1),
-          |  mydoubleavg(distinct value1) * 0.1,
-          |  avg(value1 + value2)
-          |FROM agg2
-          |GROUP BY key
-        """.stripMargin),
-      Row(120.0, 70.0/3.0, -10.0/3.0, 1, 67.0/3.0 + 100.0, 12.0, 20.0) ::
-        Row(100.0, 1.0/3.0, 1.0, 2, -2.0/3.0 + 100.0, 10.0, 2.0) ::
-        Row(null, null, 3.0, 3, null, null, null) ::
-        Row(110.0, 10.0, 20.0, null, 109.0, 11.0, 30.0) :: Nil)
-
-    checkAnswer(
-      sqlContext.sql(
-        """
-          |SELECT
-          |  key,
-          |  mydoubleavg(distinct value1),
-          |  mydoublesum(value2),
-          |  mydoublesum(distinct value1),
-          |  mydoubleavg(distinct value1),
-          |  mydoubleavg(value1)
-          |FROM agg2
-          |GROUP BY key
-        """.stripMargin),
-      Row(1, 120.0, -10.0, 40.0, 120.0, 70.0/3.0 + 100.0) ::
-        Row(2, 100.0, 3.0, 0.0, 100.0, 1.0/3.0 + 100.0) ::
-        Row(3, null, 3.0, null, null, null) ::
-        Row(null, 110.0, 60.0, 30.0, 110.0, 110.0) :: Nil)
-
-    checkAnswer(
-      sqlContext.sql(
-        """
-          |SELECT
-          |  count(value1),
-          |  count(*),
-          |  count(1),
-          |  count(DISTINCT value1),
-          |  key
-          |FROM agg2
-          |GROUP BY key
-        """.stripMargin),
-      Row(3, 3, 3, 2, 1) ::
-        Row(3, 4, 4, 2, 2) ::
-        Row(0, 2, 2, 0, 3) ::
-        Row(3, 4, 4, 3, null) :: Nil)
+    Seq(true, false).foreach { specializeSingleDistinctAgg =>
+      val conf =
+        (SQLConf.SPECIALIZE_SINGLE_DISTINCT_AGG_PLANNING.key,
+          specializeSingleDistinctAgg.toString)
+      withSQLConf(conf) {
+        // DISTINCT is not meaningful with Max and Min, so we just ignore the DISTINCT keyword.
+        checkAnswer(
+          sqlContext.sql(
+            """
+              |SELECT
+              |  min(distinct value1),
+              |  sum(distinct value1),
+              |  avg(value1),
+              |  avg(value2),
+              |  max(distinct value1)
+              |FROM agg2
+            """.stripMargin),
+          Row(-60, 70.0, 101.0/9.0, 5.6, 100))
+
+        checkAnswer(
+          sqlContext.sql(
+            """
+              |SELECT
+              |  mydoubleavg(distinct value1),
+              |  avg(value1),
+              |  avg(value2),
+              |  key,
+              |  mydoubleavg(value1 - 1),
+              |  mydoubleavg(distinct value1) * 0.1,
+              |  avg(value1 + value2)
+              |FROM agg2
+              |GROUP BY key
+            """.stripMargin),
+          Row(120.0, 70.0/3.0, -10.0/3.0, 1, 67.0/3.0 + 100.0, 12.0, 20.0) ::
+            Row(100.0, 1.0/3.0, 1.0, 2, -2.0/3.0 + 100.0, 10.0, 2.0) ::
+            Row(null, null, 3.0, 3, null, null, null) ::
+            Row(110.0, 10.0, 20.0, null, 109.0, 11.0, 30.0) :: Nil)
+
+        checkAnswer(
+          sqlContext.sql(
+            """
+              |SELECT
+              |  key,
+              |  mydoubleavg(distinct value1),
+              |  mydoublesum(value2),
+              |  mydoublesum(distinct value1),
+              |  mydoubleavg(distinct value1),
+              |  mydoubleavg(value1)
+              |FROM agg2
+              |GROUP BY key
+            """.stripMargin),
+          Row(1, 120.0, -10.0, 40.0, 120.0, 70.0/3.0 + 100.0) ::
+            Row(2, 100.0, 3.0, 0.0, 100.0, 1.0/3.0 + 100.0) ::
+            Row(3, null, 3.0, null, null, null) ::
+            Row(null, 110.0, 60.0, 30.0, 110.0, 110.0) :: Nil)
+
+        checkAnswer(
+          sqlContext.sql(
+            """
+              |SELECT
+              |  count(value1),
+              |  count(*),
+              |  count(1),
+              |  count(DISTINCT value1),
+              |  key
+              |FROM agg2
+              |GROUP BY key
+            """.stripMargin),
+          Row(3, 3, 3, 2, 1) ::
+            Row(3, 4, 4, 2, 2) ::
+            Row(0, 2, 2, 0, 3) ::
+            Row(3, 4, 4, 3, null) :: Nil)
+      }
+    }
   }
 
   test("single distinct multiple columns set") {
@@ -699,48 +701,6 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
 
     val corr7 = sqlContext.sql("SELECT corr(b, c) FROM covar_tab").collect()(0).getDouble(0)
     assert(math.abs(corr7 - 0.6633880657639323) < 1e-12)
-
-    withSQLConf(SQLConf.USE_SQL_AGGREGATE2.key -> "false") {
-      val errorMessage = intercept[SparkException] {
-        val df = Seq.tabulate(10)(i => (1.0 * i, 2.0 * i, i * -1.0)).toDF("a", "b", "c")
-        val corr1 = df.repartition(2).groupBy().agg(corr("a", "b")).collect()(0).getDouble(0)
-      }.getMessage
-      assert(errorMessage.contains("java.lang.UnsupportedOperationException: " +
-        "Corr only supports the new AggregateExpression2"))
-    }
-  }
-
-  test("test Last implemented based on AggregateExpression1") {
-    // TODO: Remove this test once we remove AggregateExpression1.
-    import org.apache.spark.sql.functions._
-    val df = Seq((1, 1), (2, 2), (3, 3)).toDF("i", "j").repartition(1)
-    withSQLConf(
-      SQLConf.SHUFFLE_PARTITIONS.key -> "1",
-      SQLConf.USE_SQL_AGGREGATE2.key -> "false") {
-
-      checkAnswer(
-        df.groupBy("i").agg(last("j")),
-        df
-      )
-    }
-  }
-
-  test("error handling") {
-    withSQLConf("spark.sql.useAggregate2" -> "false") {
-      val errorMessage = intercept[AnalysisException] {
-        sqlContext.sql(
-          """
-            |SELECT
-            |  key,
-            |  sum(value + 1.5 * key),
-            |  mydoublesum(value),
-            |  mydoubleavg(value)
-            |FROM agg1
-            |GROUP BY key
-          """.stripMargin).collect()
-      }.getMessage
-      assert(errorMessage.contains("implemented based on the new Aggregate Function interface"))
-    }
   }
 
   test("no aggregation function (SPARK-11486)") {

From 47735cdc2a878cfdbe76316d3ff8314a45dabf54 Mon Sep 17 00:00:00 2001
From: "Oscar D. Lara Yejas" <olarayej@mail.usf.edu>
Date: Tue, 10 Nov 2015 11:07:57 -0800
Subject: [PATCH 158/518] [SPARK-10863][SPARKR] Method coltypes() (New version)

This is a follow up on PR #8984, as the corresponding branch for such PR was damaged.

Author: Oscar D. Lara Yejas <olarayej@mail.usf.edu>

Closes #9579 from olarayej/SPARK-10863_NEW14.
---
 R/pkg/DESCRIPTION                |  1 +
 R/pkg/NAMESPACE                  |  6 ++--
 R/pkg/R/DataFrame.R              | 49 ++++++++++++++++++++++++++++++++
 R/pkg/R/generics.R               |  4 +++
 R/pkg/R/schema.R                 | 15 +---------
 R/pkg/R/types.R                  | 43 ++++++++++++++++++++++++++++
 R/pkg/inst/tests/test_sparkSQL.R | 24 +++++++++++++++-
 7 files changed, 124 insertions(+), 18 deletions(-)
 create mode 100644 R/pkg/R/types.R

diff --git a/R/pkg/DESCRIPTION b/R/pkg/DESCRIPTION
index 3d6edb70ec98e..369714f7b99c2 100644
--- a/R/pkg/DESCRIPTION
+++ b/R/pkg/DESCRIPTION
@@ -34,4 +34,5 @@ Collate:
     'serialize.R'
     'sparkR.R'
     'stats.R'
+    'types.R'
     'utils.R'
diff --git a/R/pkg/NAMESPACE b/R/pkg/NAMESPACE
index 56b8ed0bf271b..52fd6c9f76c54 100644
--- a/R/pkg/NAMESPACE
+++ b/R/pkg/NAMESPACE
@@ -23,9 +23,11 @@ export("setJobGroup",
 exportClasses("DataFrame")
 
 exportMethods("arrange",
+              "as.data.frame",
               "attach",
               "cache",
               "collect",
+              "coltypes",
               "columns",
               "count",
               "cov",
@@ -262,6 +264,4 @@ export("structField",
        "structType",
        "structType.jobj",
        "structType.structField",
-       "print.structType")
-
-export("as.data.frame")
+       "print.structType")
\ No newline at end of file
diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index e9013aa34a84f..cc868069d1e5a 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -2152,3 +2152,52 @@ setMethod("with",
             newEnv <- assignNewEnv(data)
             eval(substitute(expr), envir = newEnv, enclos = newEnv)
           })
+
+#' Returns the column types of a DataFrame.
+#' 
+#' @name coltypes
+#' @title Get column types of a DataFrame
+#' @family dataframe_funcs
+#' @param x (DataFrame)
+#' @return value (character) A character vector with the column types of the given DataFrame
+#' @rdname coltypes
+#' @examples \dontrun{
+#' irisDF <- createDataFrame(sqlContext, iris)
+#' coltypes(irisDF)
+#' }
+setMethod("coltypes",
+          signature(x = "DataFrame"),
+          function(x) {
+            # Get the data types of the DataFrame by invoking dtypes() function
+            types <- sapply(dtypes(x), function(x) {x[[2]]})
+
+            # Map Spark data types into R's data types using DATA_TYPES environment
+            rTypes <- sapply(types, USE.NAMES=F, FUN=function(x) {
+
+              # Check for primitive types
+              type <- PRIMITIVE_TYPES[[x]]
+
+              if (is.null(type)) {
+                # Check for complex types
+                for (t in names(COMPLEX_TYPES)) {
+                  if (substring(x, 1, nchar(t)) == t) {
+                    type <- COMPLEX_TYPES[[t]]
+                    break
+                  }
+                }
+
+                if (is.null(type)) {
+                  stop(paste("Unsupported data type: ", x))
+                }
+              }
+              type
+            })
+
+            # Find which types don't have mapping to R
+            naIndices <- which(is.na(rTypes))
+
+            # Assign the original scala data types to the unmatched ones
+            rTypes[naIndices] <- types[naIndices]
+
+            rTypes
+          })
\ No newline at end of file
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index efef7d66b522c..89731affeb898 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -1047,3 +1047,7 @@ setGeneric("attach")
 #' @rdname with
 #' @export
 setGeneric("with")
+
+#' @rdname coltypes
+#' @export
+setGeneric("coltypes", function(x) { standardGeneric("coltypes") })
\ No newline at end of file
diff --git a/R/pkg/R/schema.R b/R/pkg/R/schema.R
index 6f0e9a94e9bfa..c6ddb562270b7 100644
--- a/R/pkg/R/schema.R
+++ b/R/pkg/R/schema.R
@@ -115,20 +115,7 @@ structField.jobj <- function(x) {
 }
 
 checkType <- function(type) {
-  primtiveTypes <- c("byte",
-                     "integer",
-                     "float",
-                     "double",
-                     "numeric",
-                     "character",
-                     "string",
-                     "binary",
-                     "raw",
-                     "logical",
-                     "boolean",
-                     "timestamp",
-                     "date")
-  if (type %in% primtiveTypes) {
+  if (!is.null(PRIMITIVE_TYPES[[type]])) {
     return()
   } else {
     # Check complex types
diff --git a/R/pkg/R/types.R b/R/pkg/R/types.R
new file mode 100644
index 0000000000000..1828c23ab0f6d
--- /dev/null
+++ b/R/pkg/R/types.R
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# types.R. This file handles the data type mapping between Spark and R
+
+# The primitive data types, where names(PRIMITIVE_TYPES) are Scala types whereas
+# values are equivalent R types. This is stored in an environment to allow for
+# more efficient look up (environments use hashmaps).
+PRIMITIVE_TYPES <- as.environment(list(
+  "byte"="integer",
+  "tinyint"="integer",
+  "smallint"="integer",
+  "integer"="integer",
+  "bigint"="numeric",
+  "float"="numeric",
+  "double"="numeric",
+  "decimal"="numeric",
+  "string"="character",
+  "binary"="raw",
+  "boolean"="logical",
+  "timestamp"="POSIXct",
+  "date"="Date"))
+
+# The complex data types. These do not have any direct mapping to R's types.
+COMPLEX_TYPES <- list(
+  "map"=NA,
+  "array"=NA,
+  "struct"=NA)
+
+# The full list of data types.
+DATA_TYPES <- as.environment(c(as.list(PRIMITIVE_TYPES), COMPLEX_TYPES))
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index fbdb9a8f1ef6b..06f52d021cff8 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -1467,8 +1467,9 @@ test_that("SQL error message is returned from JVM", {
   expect_equal(grepl("Table not found: blah", retError), TRUE)
 })
 
+irisDF <- createDataFrame(sqlContext, iris)
+
 test_that("Method as.data.frame as a synonym for collect()", {
-  irisDF <- createDataFrame(sqlContext, iris)
   expect_equal(as.data.frame(irisDF), collect(irisDF))
   irisDF2 <- irisDF[irisDF$Species == "setosa", ]
   expect_equal(as.data.frame(irisDF2), collect(irisDF2))
@@ -1503,6 +1504,27 @@ test_that("with() on a DataFrame", {
   expect_equal(nrow(sum2), 35)
 })
 
+test_that("Method coltypes() to get R's data types of a DataFrame", {
+  expect_equal(coltypes(irisDF), c(rep("numeric", 4), "character"))
+
+  data <- data.frame(c1=c(1,2,3),
+                     c2=c(T,F,T),
+                     c3=c("2015/01/01 10:00:00", "2015/01/02 10:00:00", "2015/01/03 10:00:00"))
+
+  schema <- structType(structField("c1", "byte"),
+                       structField("c3", "boolean"),
+                       structField("c4", "timestamp"))
+
+  # Test primitive types
+  DF <- createDataFrame(sqlContext, data, schema)
+  expect_equal(coltypes(DF), c("integer", "logical", "POSIXct"))
+
+  # Test complex types
+  x <- createDataFrame(sqlContext, list(list(as.environment(
+    list("a"="b", "c"="d", "e"="f")))))
+  expect_equal(coltypes(x), "map<string,string>")
+})
+
 unlink(parquetPath)
 unlink(jsonPath)
 unlink(jsonPathNa)

From dfcfcbcc0448ebc6f02eba6bf0495832a321c87e Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Tue, 10 Nov 2015 11:14:25 -0800
Subject: [PATCH 159/518] [SPARK-11578][SQL][FOLLOW-UP] complete the user
 facing api for typed aggregation

Currently the user facing api for typed aggregation has some limitations:

* the customized typed aggregation must be the first of aggregation list
* the customized typed aggregation can only use long as buffer type
* the customized typed aggregation can only use flat type as result type

This PR tries to remove these limitations.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9599 from cloud-fan/agg.
---
 .../catalyst/encoders/ExpressionEncoder.scala |  6 +++
 .../aggregate/TypedAggregateExpression.scala  | 50 +++++++++++++-----
 .../spark/sql/expressions/Aggregator.scala    |  5 ++
 .../spark/sql/DatasetAggregatorSuite.scala    | 52 +++++++++++++++++++
 4 files changed, 99 insertions(+), 14 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index c287aebeeee05..005c0627f56b8 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -185,6 +185,12 @@ case class ExpressionEncoder[T](
     })
   }
 
+  def shift(delta: Int): ExpressionEncoder[T] = {
+    copy(constructExpression = constructExpression transform {
+      case r: BoundReference => r.copy(ordinal = r.ordinal + delta)
+    })
+  }
+
   /**
    * Returns a copy of this encoder where the expressions used to create an object given an
    * input row have been modified to pull the object out from a nested struct, instead of the
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
index 24d8122b6222b..0e5bc1f9abf28 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
@@ -20,13 +20,13 @@ package org.apache.spark.sql.execution.aggregate
 import scala.language.existentials
 
 import org.apache.spark.Logging
+import org.apache.spark.sql.expressions.Aggregator
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
 import org.apache.spark.sql.catalyst.expressions.aggregate.ImperativeAggregate
-import org.apache.spark.sql.expressions.Aggregator
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.types.{StructType, DataType}
+import org.apache.spark.sql.types._
 
 object TypedAggregateExpression {
   def apply[A, B : Encoder, C : Encoder](
@@ -67,8 +67,11 @@ case class TypedAggregateExpression(
 
   override def nullable: Boolean = true
 
-  // TODO: this assumes flat results...
-  override def dataType: DataType = cEncoder.schema.head.dataType
+  override def dataType: DataType = if (cEncoder.flat) {
+    cEncoder.schema.head.dataType
+  } else {
+    cEncoder.schema
+  }
 
   override def deterministic: Boolean = true
 
@@ -93,32 +96,51 @@ case class TypedAggregateExpression(
       case a: AttributeReference => inputMapping(a)
     })
 
-  // TODO: this probably only works when we are in the first column.
   val bAttributes = bEncoder.schema.toAttributes
   lazy val boundB = bEncoder.resolve(bAttributes).bind(bAttributes)
 
+  private def updateBuffer(buffer: MutableRow, value: InternalRow): Unit = {
+    // todo: need a more neat way to assign the value.
+    var i = 0
+    while (i < aggBufferAttributes.length) {
+      aggBufferSchema(i).dataType match {
+        case IntegerType => buffer.setInt(mutableAggBufferOffset + i, value.getInt(i))
+        case LongType => buffer.setLong(mutableAggBufferOffset + i, value.getLong(i))
+      }
+      i += 1
+    }
+  }
+
   override def initialize(buffer: MutableRow): Unit = {
-    // TODO: We need to either force Aggregator to have a zero or we need to eliminate the need for
-    // this in execution.
-    buffer.setInt(mutableAggBufferOffset, aggregator.zero.asInstanceOf[Int])
+    val zero = bEncoder.toRow(aggregator.zero)
+    updateBuffer(buffer, zero)
   }
 
   override def update(buffer: MutableRow, input: InternalRow): Unit = {
     val inputA = boundA.fromRow(input)
-    val currentB = boundB.fromRow(buffer)
+    val currentB = boundB.shift(mutableAggBufferOffset).fromRow(buffer)
     val merged = aggregator.reduce(currentB, inputA)
     val returned = boundB.toRow(merged)
-    buffer.setInt(mutableAggBufferOffset, returned.getInt(0))
+
+    updateBuffer(buffer, returned)
   }
 
   override def merge(buffer1: MutableRow, buffer2: InternalRow): Unit = {
-    buffer1.setLong(
-      mutableAggBufferOffset,
-      buffer1.getLong(mutableAggBufferOffset) + buffer2.getLong(inputAggBufferOffset))
+    val b1 = boundB.shift(mutableAggBufferOffset).fromRow(buffer1)
+    val b2 = boundB.shift(inputAggBufferOffset).fromRow(buffer2)
+    val merged = aggregator.merge(b1, b2)
+    val returned = boundB.toRow(merged)
+
+    updateBuffer(buffer1, returned)
   }
 
   override def eval(buffer: InternalRow): Any = {
-    buffer.getInt(mutableAggBufferOffset)
+    val b = boundB.shift(mutableAggBufferOffset).fromRow(buffer)
+    val result = cEncoder.toRow(aggregator.present(b))
+    dataType match {
+      case _: StructType => result
+      case _ => result.get(0, dataType)
+    }
   }
 
   override def toString: String = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
index 8cc25c2440633..3c1c457e06d5f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
@@ -57,6 +57,11 @@ abstract class Aggregator[-A, B, C] {
    */
   def reduce(b: B, a: A): B
 
+  /**
+   * Merge two intermediate values
+   */
+  def merge(b1: B, b2: B): B
+
   /**
    * Transform the output of the reduction.
    */
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
index 340470c096b87..206095a519762 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -34,9 +34,41 @@ class SumOf[I, N : Numeric](f: I => N) extends Aggregator[I, N, N] with Serializ
 
   override def reduce(b: N, a: I): N = numeric.plus(b, f(a))
 
+  override def merge(b1: N, b2: N): N = numeric.plus(b1, b2)
+
   override def present(reduction: N): N = reduction
 }
 
+object TypedAverage extends Aggregator[(String, Int), (Long, Long), Double] with Serializable {
+  override def zero: (Long, Long) = (0, 0)
+
+  override def reduce(countAndSum: (Long, Long), input: (String, Int)): (Long, Long) = {
+    (countAndSum._1 + 1, countAndSum._2 + input._2)
+  }
+
+  override def merge(b1: (Long, Long), b2: (Long, Long)): (Long, Long) = {
+    (b1._1 + b2._1, b1._2 + b2._2)
+  }
+
+  override def present(countAndSum: (Long, Long)): Double = countAndSum._2 / countAndSum._1
+}
+
+object ComplexResultAgg extends Aggregator[(String, Int), (Long, Long), (Long, Long)]
+  with Serializable {
+
+  override def zero: (Long, Long) = (0, 0)
+
+  override def reduce(countAndSum: (Long, Long), input: (String, Int)): (Long, Long) = {
+    (countAndSum._1 + 1, countAndSum._2 + input._2)
+  }
+
+  override def merge(b1: (Long, Long), b2: (Long, Long)): (Long, Long) = {
+    (b1._1 + b2._1, b1._2 + b2._2)
+  }
+
+  override def present(reduction: (Long, Long)): (Long, Long) = reduction
+}
+
 class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
 
   import testImplicits._
@@ -62,4 +94,24 @@ class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
         count("*")),
       ("a", 30, 30, 2L), ("b", 3, 3, 2L), ("c", 1, 1, 1L))
   }
+
+  test("typed aggregation: complex case") {
+    val ds = Seq("a" -> 1, "a" -> 3, "b" -> 3).toDS()
+
+    checkAnswer(
+      ds.groupBy(_._1).agg(
+        expr("avg(_2)").as[Double],
+        TypedAverage.toColumn),
+      ("a", 2.0, 2.0), ("b", 3.0, 3.0))
+  }
+
+  test("typed aggregation: complex result type") {
+    val ds = Seq("a" -> 1, "a" -> 3, "b" -> 3).toDS()
+
+    checkAnswer(
+      ds.groupBy(_._1).agg(
+        expr("avg(_2)").as[Double],
+        ComplexResultAgg.toColumn),
+      ("a", 2.0, (2L, 4L)), ("b", 3.0, (1L, 3L)))
+  }
 }

From 53600854c270d4c953fe95fbae528740b5cf6603 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Tue, 10 Nov 2015 11:21:31 -0800
Subject: [PATCH 160/518] [SPARK-11590][SQL] use native json_tuple in lateral
 view

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9562 from cloud-fan/json-tuple.
---
 .../expressions/jsonExpressions.scala         | 23 +++++---------
 .../expressions/JsonExpressionsSuite.scala    | 30 ++++++++++--------
 .../org/apache/spark/sql/DataFrame.scala      |  8 +++--
 .../org/apache/spark/sql/functions.scala      | 12 +++++++
 .../apache/spark/sql/JsonFunctionsSuite.scala | 23 ++++++++------
 .../org/apache/spark/sql/hive/HiveQl.scala    |  4 +++
 .../apache/spark/sql/hive/HiveQlSuite.scala   | 13 ++++++++
 .../sql/hive/execution/SQLQuerySuite.scala    | 31 +++++++++++++++++++
 8 files changed, 104 insertions(+), 40 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
index 8c9853e628d2c..8cd73236a7876 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
@@ -314,7 +314,7 @@ case class GetJsonObject(json: Expression, path: Expression)
 }
 
 case class JsonTuple(children: Seq[Expression])
-  extends Expression with CodegenFallback {
+  extends Generator with CodegenFallback {
 
   import SharedFactory._
 
@@ -324,8 +324,8 @@ case class JsonTuple(children: Seq[Expression])
   }
 
   // if processing fails this shared value will be returned
-  @transient private lazy val nullRow: InternalRow =
-    new GenericInternalRow(Array.ofDim[Any](fieldExpressions.length))
+  @transient private lazy val nullRow: Seq[InternalRow] =
+    new GenericInternalRow(Array.ofDim[Any](fieldExpressions.length)) :: Nil
 
   // the json body is the first child
   @transient private lazy val jsonExpr: Expression = children.head
@@ -344,15 +344,8 @@ case class JsonTuple(children: Seq[Expression])
   // and count the number of foldable fields, we'll use this later to optimize evaluation
   @transient private lazy val constantFields: Int = foldableFieldNames.count(_ != null)
 
-  override lazy val dataType: StructType = {
-    val fields = fieldExpressions.zipWithIndex.map {
-      case (_, idx) => StructField(
-        name = s"c$idx", // mirroring GenericUDTFJSONTuple.initialize
-        dataType = StringType,
-        nullable = true)
-    }
-
-    StructType(fields)
+  override def elementTypes: Seq[(DataType, Boolean, String)] = fieldExpressions.zipWithIndex.map {
+    case (_, idx) => (StringType, true, s"c$idx")
   }
 
   override def prettyName: String = "json_tuple"
@@ -367,7 +360,7 @@ case class JsonTuple(children: Seq[Expression])
     }
   }
 
-  override def eval(input: InternalRow): InternalRow = {
+  override def eval(input: InternalRow): TraversableOnce[InternalRow] = {
     val json = jsonExpr.eval(input).asInstanceOf[UTF8String]
     if (json == null) {
       return nullRow
@@ -383,7 +376,7 @@ case class JsonTuple(children: Seq[Expression])
     }
   }
 
-  private def parseRow(parser: JsonParser, input: InternalRow): InternalRow = {
+  private def parseRow(parser: JsonParser, input: InternalRow): Seq[InternalRow] = {
     // only objects are supported
     if (parser.nextToken() != JsonToken.START_OBJECT) {
       return nullRow
@@ -433,7 +426,7 @@ case class JsonTuple(children: Seq[Expression])
       parser.skipChildren()
     }
 
-    new GenericInternalRow(row)
+    new GenericInternalRow(row) :: Nil
   }
 
   private def copyCurrentStructure(generator: JsonGenerator, parser: JsonParser): Unit = {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala
index f33125f463e14..7b754091f4714 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala
@@ -209,8 +209,12 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     Literal("f5") ::
     Nil
 
+  private def checkJsonTuple(jt: JsonTuple, expected: InternalRow): Unit = {
+    assert(jt.eval(null).toSeq.head === expected)
+  }
+
   test("json_tuple - hive key 1") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(
         Literal("""{"f1": "value1", "f2": "value2", "f3": 3, "f5": 5.23}""") ::
           jsonTupleQuery),
@@ -218,7 +222,7 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("json_tuple - hive key 2") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(
         Literal("""{"f1": "value12", "f3": "value3", "f2": 2, "f4": 4.01}""") ::
           jsonTupleQuery),
@@ -226,7 +230,7 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("json_tuple - hive key 2 (mix of foldable fields)") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal("""{"f1": "value12", "f3": "value3", "f2": 2, "f4": 4.01}""") ::
         Literal("f1") ::
         NonFoldableLiteral("f2") ::
@@ -238,7 +242,7 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("json_tuple - hive key 3") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(
         Literal("""{"f1": "value13", "f4": "value44", "f3": "value33", "f2": 2, "f5": 5.01}""") ::
           jsonTupleQuery),
@@ -247,7 +251,7 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("json_tuple - hive key 3 (nonfoldable json)") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(
         NonFoldableLiteral(
           """{"f1": "value13", "f4": "value44",
@@ -258,7 +262,7 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("json_tuple - hive key 3 (nonfoldable fields)") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal(
         """{"f1": "value13", "f4": "value44",
           | "f3": "value33", "f2": 2, "f5": 5.01}""".stripMargin) ::
@@ -273,43 +277,43 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("json_tuple - hive key 4 - null json") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal(null) :: jsonTupleQuery),
       InternalRow.fromSeq(Seq(null, null, null, null, null)))
   }
 
   test("json_tuple - hive key 5 - null and empty fields") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal("""{"f1": "", "f5": null}""") :: jsonTupleQuery),
       InternalRow.fromSeq(Seq(UTF8String.fromString(""), null, null, null, null)))
   }
 
   test("json_tuple - hive key 6 - invalid json (array)") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal("[invalid JSON string]") :: jsonTupleQuery),
       InternalRow.fromSeq(Seq(null, null, null, null, null)))
   }
 
   test("json_tuple - invalid json (object start only)") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal("{") :: jsonTupleQuery),
       InternalRow.fromSeq(Seq(null, null, null, null, null)))
   }
 
   test("json_tuple - invalid json (no object end)") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal("""{"foo": "bar"""") :: jsonTupleQuery),
       InternalRow.fromSeq(Seq(null, null, null, null, null)))
   }
 
   test("json_tuple - invalid json (invalid json)") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal("\\") :: jsonTupleQuery),
       InternalRow.fromSeq(Seq(null, null, null, null, null)))
   }
 
   test("json_tuple - preserve newlines") {
-    checkEvaluation(
+    checkJsonTuple(
       JsonTuple(Literal("{\"a\":\"b\nc\"}") :: Literal("a") :: Nil),
       InternalRow.fromSeq(Seq(UTF8String.fromString("b\nc"))))
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 3b69247dc54ef..9368435a63c35 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -750,10 +750,14 @@ class DataFrame private[sql](
       // will remove intermediate Alias for ExtractValue chain, and we need to alias it again to
       // make it a NamedExpression.
       case Column(u: UnresolvedAttribute) => UnresolvedAlias(u)
+
       case Column(expr: NamedExpression) => expr
-      // Leave an unaliased explode with an empty list of names since the analyzer will generate the
-      // correct defaults after the nested expression's type has been resolved.
+
+      // Leave an unaliased generator with an empty list of names since the analyzer will generate
+      // the correct defaults after the nested expression's type has been resolved.
       case Column(explode: Explode) => MultiAlias(explode, Nil)
+      case Column(jt: JsonTuple) => MultiAlias(jt, Nil)
+
       case Column(expr: Expression) => Alias(expr, expr.prettyString)()
     }
     Project(namedExpressions.toSeq, logicalPlan)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 22104e4d48617..a59d738010f7c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -2307,6 +2307,18 @@ object functions extends LegacyFunctions {
    */
   def explode(e: Column): Column = withExpr { Explode(e.expr) }
 
+  /**
+   * Creates a new row for a json column according to the given field names.
+   *
+   * @group collection_funcs
+   * @since 1.6.0
+   */
+  @scala.annotation.varargs
+  def json_tuple(json: Column, fields: String*): Column = withExpr {
+    require(fields.length > 0, "at least 1 field name should be given.")
+    JsonTuple(json.expr +: fields.map(Literal.apply))
+  }
+
   /**
    * Returns length of array or map.
    *
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala
index e3531d0d6d799..14fd56fc8c222 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala
@@ -41,23 +41,26 @@ class JsonFunctionsSuite extends QueryTest with SharedSQLContext {
 
   test("json_tuple select") {
     val df: DataFrame = tuples.toDF("key", "jstring")
-    val expected = Row("1", Row("value1", "value2", "3", null, "5.23")) ::
-      Row("2", Row("value12", "2", "value3", "4.01", null)) ::
-      Row("3", Row("value13", "2", "value33", "value44", "5.01")) ::
-      Row("4", Row(null, null, null, null, null)) ::
-      Row("5", Row("", null, null, null, null)) ::
-      Row("6", Row(null, null, null, null, null)) ::
+    val expected =
+      Row("1", "value1", "value2", "3", null, "5.23") ::
+      Row("2", "value12", "2", "value3", "4.01", null) ::
+      Row("3", "value13", "2", "value33", "value44", "5.01") ::
+      Row("4", null, null, null, null, null) ::
+      Row("5", "", null, null, null, null) ::
+      Row("6", null, null, null, null, null) ::
       Nil
 
-    checkAnswer(df.selectExpr("key", "json_tuple(jstring, 'f1', 'f2', 'f3', 'f4', 'f5')"), expected)
+    checkAnswer(
+      df.select($"key", functions.json_tuple($"jstring", "f1", "f2", "f3", "f4", "f5")),
+      expected)
   }
 
   test("json_tuple filter and group") {
     val df: DataFrame = tuples.toDF("key", "jstring")
     val expr = df
-      .selectExpr("json_tuple(jstring, 'f1', 'f2') as jt")
-      .where($"jt.c0".isNotNull)
-      .groupBy($"jt.c1")
+      .select(functions.json_tuple($"jstring", "f1", "f2"))
+      .where($"c0".isNotNull)
+      .groupBy($"c1")
       .count()
 
     val expected = Row(null, 1) ::
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
index 6f8ed413a06cd..091caab921fe9 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveQl.scala
@@ -1821,6 +1821,7 @@ https://cwiki.apache.org/confluence/display/Hive/Enhanced+Aggregation%2C+Cube%2C
   }
 
   val explode = "(?i)explode".r
+  val jsonTuple = "(?i)json_tuple".r
   def nodesToGenerator(nodes: Seq[Node]): (Generator, Seq[String]) = {
     val function = nodes.head
 
@@ -1833,6 +1834,9 @@ https://cwiki.apache.org/confluence/display/Hive/Enhanced+Aggregation%2C+Cube%2C
       case Token("TOK_FUNCTION", Token(explode(), Nil) :: child :: Nil) =>
         (Explode(nodeToExpr(child)), attributes)
 
+      case Token("TOK_FUNCTION", Token(jsonTuple(), Nil) :: children) =>
+        (JsonTuple(children.map(nodeToExpr)), attributes)
+
       case Token("TOK_FUNCTION", Token(functionName, Nil) :: children) =>
         val functionInfo: FunctionInfo =
           Option(FunctionRegistry.getFunctionInfo(functionName.toLowerCase)).getOrElse(
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveQlSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveQlSuite.scala
index 528a7398b10df..a330362b4e1d1 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveQlSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveQlSuite.scala
@@ -18,6 +18,8 @@
 package org.apache.spark.sql.hive
 
 import org.apache.hadoop.hive.serde.serdeConstants
+import org.apache.spark.sql.catalyst.expressions.JsonTuple
+import org.apache.spark.sql.catalyst.plans.logical.Generate
 import org.scalatest.BeforeAndAfterAll
 
 import org.apache.spark.SparkFunSuite
@@ -183,4 +185,15 @@ class HiveQlSuite extends SparkFunSuite with BeforeAndAfterAll {
     assertError("select interval '.1111111111' second",
       "nanosecond 1111111111 outside range")
   }
+
+  test("use native json_tuple instead of hive's UDTF in LATERAL VIEW") {
+    val plan = HiveQl.parseSql(
+      """
+        |SELECT *
+        |FROM (SELECT '{"f1": "value1", "f2": 12}' json) test
+        |LATERAL VIEW json_tuple(json, 'f1', 'f2') jt AS a, b
+      """.stripMargin)
+
+    assert(plan.children.head.asInstanceOf[Generate].generator.isInstanceOf[JsonTuple])
+  }
 }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
index 9a425d7f6b265..3427152b2da02 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
@@ -1448,4 +1448,35 @@ class SQLQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
         Row("1", "10") :: Row("2", "20") :: Row("3", "30") :: Row("4", "40") :: Nil)
     }
   }
+
+  test("SPARK-11590: use native json_tuple in lateral view") {
+    checkAnswer(sql(
+      """
+        |SELECT a, b
+        |FROM (SELECT '{"f1": "value1", "f2": 12}' json) test
+        |LATERAL VIEW json_tuple(json, 'f1', 'f2') jt AS a, b
+      """.stripMargin), Row("value1", "12"))
+
+    // we should use `c0`, `c1`... as the name of fields if no alias is provided, to follow hive.
+    checkAnswer(sql(
+      """
+        |SELECT c0, c1
+        |FROM (SELECT '{"f1": "value1", "f2": 12}' json) test
+        |LATERAL VIEW json_tuple(json, 'f1', 'f2') jt
+      """.stripMargin), Row("value1", "12"))
+
+    // we can also use `json_tuple` in project list.
+    checkAnswer(sql(
+      """
+        |SELECT json_tuple(json, 'f1', 'f2')
+        |FROM (SELECT '{"f1": "value1", "f2": 12}' json) test
+      """.stripMargin), Row("value1", "12"))
+
+    // we can also mix `json_tuple` with other project expressions.
+    checkAnswer(sql(
+      """
+        |SELECT json_tuple(json, 'f1', 'f2'), 3.14, str
+        |FROM (SELECT '{"f1": "value1", "f2": 12}' json, 'hello' as str) test
+      """.stripMargin), Row("value1", "12", 3.14, "hello"))
+  }
 }

From 87aedc48c01dffbd880e6ca84076ed47c68f88d0 Mon Sep 17 00:00:00 2001
From: Nong Li <nong@databricks.com>
Date: Tue, 10 Nov 2015 11:28:53 -0800
Subject: [PATCH 161/518] [SPARK-10371][SQL] Implement subexpr elimination for
 UnsafeProjections

This patch adds the building blocks for codegening subexpr elimination and implements
it end to end for UnsafeProjection. The building blocks can be used to do the same thing
for other operators.

It introduces some utilities to compute common sub expressions. Expressions can be added to
this data structure. The expr and its children will be recursively matched against existing
expressions (ones previously added) and grouped into common groups. This is built using
the existing `semanticEquals`. It does not understand things like commutative or associative
expressions. This can be done as future work.

After building this data structure, the codegen process takes advantage of it by:
  1. Generating a helper function in the generated class that computes the common
     subexpression. This is done for all common subexpressions that have at least
     two occurrences and the expression tree is sufficiently complex.
  2. When generating the apply() function, if the helper function exists, call that
     instead of regenerating the expression tree. Repeated calls to the helper function
     shortcircuit the evaluation logic.

Author: Nong Li <nong@databricks.com>
Author: Nong Li <nongli@gmail.com>

This patch had conflicts when merged, resolved by
Committer: Michael Armbrust <michael@databricks.com>

Closes #9480 from nongli/spark-10371.
---
 .../expressions/EquivalentExpressions.scala   | 106 ++++++++++++
 .../sql/catalyst/expressions/Expression.scala |  50 +++++-
 .../sql/catalyst/expressions/Projection.scala |  16 ++
 .../expressions/codegen/CodeGenerator.scala   | 110 ++++++++++++-
 .../codegen/GenerateUnsafeProjection.scala    |  36 ++++-
 .../expressions/namedExpressions.scala        |   4 +
 .../SubexpressionEliminationSuite.scala       | 153 ++++++++++++++++++
 .../scala/org/apache/spark/sql/SQLConf.scala  |   8 +
 .../spark/sql/execution/SparkPlan.scala       |   5 +
 .../spark/sql/execution/basicOperators.scala  |   3 +-
 .../org/apache/spark/sql/SQLQuerySuite.scala  |  48 ++++++
 11 files changed, 523 insertions(+), 16 deletions(-)
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
 create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
new file mode 100644
index 0000000000000..e7380d21f98af
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions
+
+import scala.collection.mutable
+
+/**
+ * This class is used to compute equality of (sub)expression trees. Expressions can be added
+ * to this class and they subsequently query for expression equality. Expression trees are
+ * considered equal if for the same input(s), the same result is produced.
+ */
+class EquivalentExpressions {
+  /**
+   * Wrapper around an Expression that provides semantic equality.
+   */
+  case class Expr(e: Expression) {
+    val hash = e.semanticHash()
+    override def equals(o: Any): Boolean = o match {
+      case other: Expr => e.semanticEquals(other.e)
+      case _ => false
+    }
+    override def hashCode: Int = hash
+  }
+
+  // For each expression, the set of equivalent expressions.
+  private val equivalenceMap: mutable.HashMap[Expr, mutable.MutableList[Expression]] =
+      new mutable.HashMap[Expr, mutable.MutableList[Expression]]
+
+  /**
+   * Adds each expression to this data structure, grouping them with existing equivalent
+   * expressions. Non-recursive.
+   * Returns if there was already a matching expression.
+   */
+  def addExpr(expr: Expression): Boolean = {
+    if (expr.deterministic) {
+      val e: Expr = Expr(expr)
+      val f = equivalenceMap.get(e)
+      if (f.isDefined) {
+        f.get.+= (expr)
+        true
+      } else {
+        equivalenceMap.put(e, mutable.MutableList(expr))
+        false
+      }
+    } else {
+      false
+    }
+  }
+
+  /**
+   * Adds the expression to this datastructure recursively. Stops if a matching expression
+   * is found. That is, if `expr` has already been added, its children are not added.
+   * If ignoreLeaf is true, leaf nodes are ignored.
+   */
+  def addExprTree(root: Expression, ignoreLeaf: Boolean = true): Unit = {
+    val skip = root.isInstanceOf[LeafExpression] && ignoreLeaf
+    if (!skip && root.deterministic && !addExpr(root)) {
+     root.children.foreach(addExprTree(_, ignoreLeaf))
+    }
+  }
+
+  /**
+   * Returns all fo the expression trees that are equivalent to `e`. Returns
+   * an empty collection if there are none.
+   */
+  def getEquivalentExprs(e: Expression): Seq[Expression] = {
+    equivalenceMap.get(Expr(e)).getOrElse(mutable.MutableList())
+  }
+
+  /**
+   * Returns all the equivalent sets of expressions.
+   */
+  def getAllEquivalentExprs: Seq[Seq[Expression]] = {
+    equivalenceMap.values.map(_.toSeq).toSeq
+  }
+
+  /**
+   * Returns the state of the datastructure as a string. If all is false, skips sets of equivalent
+   * expressions with cardinality 1.
+   */
+  def debugString(all: Boolean = false): String = {
+    val sb: mutable.StringBuilder = new StringBuilder()
+    sb.append("Equivalent expressions:\n")
+    equivalenceMap.foreach { case (k, v) => {
+      if (all || v.length > 1) {
+        sb.append("  " + v.mkString(", ")).append("\n")
+      }
+    }}
+    sb.toString()
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
index 96fcc799e537a..7d5741eefcc7d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
@@ -92,12 +92,24 @@ abstract class Expression extends TreeNode[Expression] {
    * @return [[GeneratedExpressionCode]]
    */
   def gen(ctx: CodeGenContext): GeneratedExpressionCode = {
-    val isNull = ctx.freshName("isNull")
-    val primitive = ctx.freshName("primitive")
-    val ve = GeneratedExpressionCode("", isNull, primitive)
-    ve.code = genCode(ctx, ve)
-    // Add `this` in the comment.
-    ve.copy(s"/* $this */\n" + ve.code)
+    val subExprState = ctx.subExprEliminationExprs.get(this)
+    if (subExprState.isDefined) {
+      // This expression is repeated meaning the code to evaluated has already been added
+      // as a function, `subExprState.fnName`. Just call that.
+      val code =
+        s"""
+           |/* $this */
+           |${subExprState.get.fnName}(${ctx.INPUT_ROW});
+           |""".stripMargin.trim
+      GeneratedExpressionCode(code, subExprState.get.code.isNull, subExprState.get.code.value)
+    } else {
+      val isNull = ctx.freshName("isNull")
+      val primitive = ctx.freshName("primitive")
+      val ve = GeneratedExpressionCode("", isNull, primitive)
+      ve.code = genCode(ctx, ve)
+      // Add `this` in the comment.
+      ve.copy(s"/* $this */\n" + ve.code.trim)
+    }
   }
 
   /**
@@ -145,11 +157,37 @@ abstract class Expression extends TreeNode[Expression] {
         case (i1, i2) => i1 == i2
       }
     }
+    // Non-determinstic expressions cannot be equal
+    if (!deterministic || !other.deterministic) return false
     val elements1 = this.productIterator.toSeq
     val elements2 = other.asInstanceOf[Product].productIterator.toSeq
     checkSemantic(elements1, elements2)
   }
 
+  /**
+   * Returns the hash for this expression. Expressions that compute the same result, even if
+   * they differ cosmetically should return the same hash.
+   */
+  def semanticHash() : Int = {
+    def computeHash(e: Seq[Any]): Int = {
+      // See http://stackoverflow.com/questions/113511/hash-code-implementation
+      var hash: Int = 17
+      e.foreach(i => {
+        val h: Int = i match {
+          case (e: Expression) => e.semanticHash()
+          case (Some(e: Expression)) => e.semanticHash()
+          case (t: Traversable[_]) => computeHash(t.toSeq)
+          case null => 0
+          case (o) => o.hashCode()
+        }
+        hash = hash * 37 + h
+      })
+      hash
+    }
+
+    computeHash(this.productIterator.toSeq)
+  }
+
   /**
    * Checks the input data types, returns `TypeCheckResult.success` if it's valid,
    * or returns a `TypeCheckResult` with an error message if invalid.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
index 79dabe8e925ad..9f0b7821ae74a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
@@ -144,6 +144,22 @@ object UnsafeProjection {
   def create(exprs: Seq[Expression], inputSchema: Seq[Attribute]): UnsafeProjection = {
     create(exprs.map(BindReferences.bindReference(_, inputSchema)))
   }
+
+  /**
+    * Same as other create()'s but allowing enabling/disabling subexpression elimination.
+    * TODO: refactor the plumbing and clean this up.
+    */
+  def create(
+      exprs: Seq[Expression],
+      inputSchema: Seq[Attribute],
+      subexpressionEliminationEnabled: Boolean): UnsafeProjection = {
+    val e = exprs.map(BindReferences.bindReference(_, inputSchema))
+      .map(_ transform {
+        case CreateStruct(children) => CreateStructUnsafe(children)
+        case CreateNamedStruct(children) => CreateNamedStructUnsafe(children)
+    })
+    GenerateUnsafeProjection.generate(e, subexpressionEliminationEnabled)
+  }
 }
 
 /**
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
index f0f7a6cf0cc4d..60a3d6018496c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
@@ -92,6 +92,33 @@ class CodeGenContext {
     addedFunctions += ((funcName, funcCode))
   }
 
+  /**
+   * Holds expressions that are equivalent. Used to perform subexpression elimination
+   * during codegen.
+   *
+   * For expressions that appear more than once, generate additional code to prevent
+   * recomputing the value.
+   *
+   * For example, consider two exprsesion generated from this SQL statement:
+   *  SELECT (col1 + col2), (col1 + col2) / col3.
+   *
+   *  equivalentExpressions will match the tree containing `col1 + col2` and it will only
+   *  be evaluated once.
+   */
+  val equivalentExpressions: EquivalentExpressions = new EquivalentExpressions
+
+  // State used for subexpression elimination.
+  case class SubExprEliminationState(
+    val isLoaded: String, code: GeneratedExpressionCode, val fnName: String)
+
+  // Foreach expression that is participating in subexpression elimination, the state to use.
+  val subExprEliminationExprs: mutable.HashMap[Expression, SubExprEliminationState] =
+    mutable.HashMap[Expression, SubExprEliminationState]()
+
+  // The collection of isLoaded variables that need to be reset on each row.
+  val subExprIsLoadedVariables: mutable.ArrayBuffer[String] =
+    mutable.ArrayBuffer.empty[String]
+
   final val JAVA_BOOLEAN = "boolean"
   final val JAVA_BYTE = "byte"
   final val JAVA_SHORT = "short"
@@ -317,6 +344,87 @@ class CodeGenContext {
       functions.map(name => s"$name($row);").mkString("\n")
     }
   }
+
+  /**
+   * Checks and sets up the state and codegen for subexpression elimination. This finds the
+   * common subexpresses, generates the functions that evaluate those expressions and populates
+   * the mapping of common subexpressions to the generated functions.
+   */
+  private def subexpressionElimination(expressions: Seq[Expression]) = {
+    // Add each expression tree and compute the common subexpressions.
+    expressions.foreach(equivalentExpressions.addExprTree(_))
+
+    // Get all the exprs that appear at least twice and set up the state for subexpression
+    // elimination.
+    val commonExprs = equivalentExpressions.getAllEquivalentExprs.filter(_.size > 1)
+    commonExprs.foreach(e => {
+      val expr = e.head
+      val isLoaded = freshName("isLoaded")
+      val isNull = freshName("isNull")
+      val primitive = freshName("primitive")
+      val fnName = freshName("evalExpr")
+
+      // Generate the code for this expression tree and wrap it in a function.
+      val code = expr.gen(this)
+      val fn =
+        s"""
+           |private void $fnName(InternalRow ${INPUT_ROW}) {
+           |  if (!$isLoaded) {
+           |    ${code.code.trim}
+           |    $isLoaded = true;
+           |    $isNull = ${code.isNull};
+           |    $primitive = ${code.value};
+           |  }
+           |}
+           """.stripMargin
+      code.code = fn
+      code.isNull = isNull
+      code.value = primitive
+
+      addNewFunction(fnName, fn)
+
+      // Add a state and a mapping of the common subexpressions that are associate with this
+      // state. Adding this expression to subExprEliminationExprMap means it will call `fn`
+      // when it is code generated. This decision should be a cost based one.
+      //
+      // The cost of doing subexpression elimination is:
+      //   1. Extra function call, although this is probably *good* as the JIT can decide to
+      //      inline or not.
+      //   2. Extra branch to check isLoaded. This branch is likely to be predicted correctly
+      //      very often. The reason it is not loaded is because of a prior branch.
+      //   3. Extra store into isLoaded.
+      // The benefit doing subexpression elimination is:
+      //   1. Running the expression logic. Even for a simple expression, it is likely more than 3
+      //      above.
+      //   2. Less code.
+      // Currently, we will do this for all non-leaf only expression trees (i.e. expr trees with
+      // at least two nodes) as the cost of doing it is expected to be low.
+
+      // Maintain the loaded value and isNull as member variables. This is necessary if the codegen
+      // function is split across multiple functions.
+      // TODO: maintaining this as a local variable probably allows the compiler to do better
+      // optimizations.
+      addMutableState("boolean", isLoaded, s"$isLoaded = false;")
+      addMutableState("boolean", isNull, s"$isNull = false;")
+      addMutableState(javaType(expr.dataType), primitive,
+        s"$primitive = ${defaultValue(expr.dataType)};")
+      subExprIsLoadedVariables += isLoaded
+
+      val state = SubExprEliminationState(isLoaded, code, fnName)
+      e.foreach(subExprEliminationExprs.put(_, state))
+    })
+  }
+
+  /**
+   * Generates code for expressions. If doSubexpressionElimination is true, subexpression
+   * elimination will be performed. Subexpression elimination assumes that the code will for each
+   * expression will be combined in the `expressions` order.
+   */
+  def generateExpressions(expressions: Seq[Expression],
+      doSubexpressionElimination: Boolean = false): Seq[GeneratedExpressionCode] = {
+    if (doSubexpressionElimination) subexpressionElimination(expressions)
+    expressions.map(e => e.gen(this))
+  }
 }
 
 /**
@@ -349,7 +457,7 @@ abstract class CodeGenerator[InType <: AnyRef, OutType <: AnyRef] extends Loggin
   }
 
   protected def declareAddedFunctions(ctx: CodeGenContext): String = {
-    ctx.addedFunctions.map { case (funcName, funcCode) => funcCode }.mkString("\n")
+    ctx.addedFunctions.map { case (funcName, funcCode) => funcCode }.mkString("\n").trim
   }
 
   /**
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
index 2136f82ba4752..9ef226141421b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
@@ -139,9 +139,9 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
         s"""
           ${input.code}
           if (${input.isNull}) {
-            $setNull
+            ${setNull.trim}
           } else {
-            $writeField
+            ${writeField.trim}
           }
         """
     }
@@ -149,7 +149,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     s"""
       $rowWriter.initialize($bufferHolder, ${inputs.length});
       ${ctx.splitExpressions(row, writeFields)}
-    """
+    """.trim
   }
 
   // TODO: if the nullability of array element is correct, we can use it to save null check.
@@ -275,8 +275,11 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     """
   }
 
-  def createCode(ctx: CodeGenContext, expressions: Seq[Expression]): GeneratedExpressionCode = {
-    val exprEvals = expressions.map(e => e.gen(ctx))
+  def createCode(
+      ctx: CodeGenContext,
+      expressions: Seq[Expression],
+      useSubexprElimination: Boolean = false): GeneratedExpressionCode = {
+    val exprEvals = ctx.generateExpressions(expressions, useSubexprElimination)
     val exprTypes = expressions.map(_.dataType)
 
     val result = ctx.freshName("result")
@@ -285,10 +288,15 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     val holderClass = classOf[BufferHolder].getName
     ctx.addMutableState(holderClass, bufferHolder, s"this.$bufferHolder = new $holderClass();")
 
+    // Reset the isLoaded flag for each row.
+    val subexprReset = ctx.subExprIsLoadedVariables.map { v => s"${v} = false;" }.mkString("\n")
+
     val code =
       s"""
         $bufferHolder.reset();
+        $subexprReset
         ${writeExpressionsToBuffer(ctx, ctx.INPUT_ROW, exprEvals, exprTypes, bufferHolder)}
+
         $result.pointTo($bufferHolder.buffer, ${expressions.length}, $bufferHolder.totalSize());
       """
     GeneratedExpressionCode(code, "false", result)
@@ -300,10 +308,21 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
   protected def bind(in: Seq[Expression], inputSchema: Seq[Attribute]): Seq[Expression] =
     in.map(BindReferences.bindReference(_, inputSchema))
 
+  def generate(
+    expressions: Seq[Expression],
+    subexpressionEliminationEnabled: Boolean): UnsafeProjection = {
+    create(canonicalize(expressions), subexpressionEliminationEnabled)
+  }
+
   protected def create(expressions: Seq[Expression]): UnsafeProjection = {
-    val ctx = newCodeGenContext()
+    create(expressions, false)
+  }
 
-    val eval = createCode(ctx, expressions)
+  private def create(
+      expressions: Seq[Expression],
+      subexpressionEliminationEnabled: Boolean): UnsafeProjection = {
+    val ctx = newCodeGenContext()
+    val eval = createCode(ctx, expressions, subexpressionEliminationEnabled)
 
     val code = s"""
       public Object generate($exprType[] exprs) {
@@ -315,6 +334,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
         private $exprType[] expressions;
 
         ${declareMutableStates(ctx)}
+
         ${declareAddedFunctions(ctx)}
 
         public SpecificUnsafeProjection($exprType[] expressions) {
@@ -328,7 +348,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
         }
 
         public UnsafeRow apply(InternalRow ${ctx.INPUT_ROW}) {
-          ${eval.code}
+          ${eval.code.trim}
           return ${eval.value};
         }
       }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
index 9ab5c299d0f55..f80bcfcb0b0bf 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
@@ -203,6 +203,10 @@ case class AttributeReference(
     case _ => false
   }
 
+  override def semanticHash(): Int = {
+    this.exprId.hashCode()
+  }
+
   override def hashCode: Int = {
     // See http://stackoverflow.com/questions/113511/hash-code-implementation
     var h = 17
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
new file mode 100644
index 0000000000000..9de066e99d637
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.catalyst.expressions
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.types.IntegerType
+
+class SubexpressionEliminationSuite extends SparkFunSuite {
+  test("Semantic equals and hash") {
+    val id = ExprId(1)
+    val a: AttributeReference = AttributeReference("name", IntegerType)()
+    val b1 = a.withName("name2").withExprId(id)
+    val b2 = a.withExprId(id)
+
+    assert(b1 != b2)
+    assert(a != b1)
+    assert(b1.semanticEquals(b2))
+    assert(!b1.semanticEquals(a))
+    assert(a.hashCode != b1.hashCode)
+    assert(b1.hashCode == b2.hashCode)
+    assert(b1.semanticHash() == b2.semanticHash())
+  }
+
+  test("Expression Equivalence - basic") {
+    val equivalence = new EquivalentExpressions
+    assert(equivalence.getAllEquivalentExprs.isEmpty)
+
+    val oneA = Literal(1)
+    val oneB = Literal(1)
+    val twoA = Literal(2)
+    var twoB = Literal(2)
+
+    assert(equivalence.getEquivalentExprs(oneA).isEmpty)
+    assert(equivalence.getEquivalentExprs(twoA).isEmpty)
+
+    // Add oneA and test if it is returned. Since it is a group of one, it does not.
+    assert(!equivalence.addExpr(oneA))
+    assert(equivalence.getEquivalentExprs(oneA).size == 1)
+    assert(equivalence.getEquivalentExprs(twoA).isEmpty)
+    assert(equivalence.addExpr((oneA)))
+    assert(equivalence.getEquivalentExprs(oneA).size == 2)
+
+    // Add B and make sure they can see each other.
+    assert(equivalence.addExpr(oneB))
+    // Use exists and reference equality because of how equals is defined.
+    assert(equivalence.getEquivalentExprs(oneA).exists(_ eq oneB))
+    assert(equivalence.getEquivalentExprs(oneA).exists(_ eq oneA))
+    assert(equivalence.getEquivalentExprs(oneB).exists(_ eq oneA))
+    assert(equivalence.getEquivalentExprs(oneB).exists(_ eq oneB))
+    assert(equivalence.getEquivalentExprs(twoA).isEmpty)
+    assert(equivalence.getAllEquivalentExprs.size == 1)
+    assert(equivalence.getAllEquivalentExprs.head.size == 3)
+    assert(equivalence.getAllEquivalentExprs.head.contains(oneA))
+    assert(equivalence.getAllEquivalentExprs.head.contains(oneB))
+
+    val add1 = Add(oneA, oneB)
+    val add2 = Add(oneA, oneB)
+
+    equivalence.addExpr(add1)
+    equivalence.addExpr(add2)
+
+    assert(equivalence.getAllEquivalentExprs.size == 2)
+    assert(equivalence.getEquivalentExprs(add2).exists(_ eq add1))
+    assert(equivalence.getEquivalentExprs(add2).size == 2)
+    assert(equivalence.getEquivalentExprs(add1).exists(_ eq add2))
+  }
+
+  test("Expression Equivalence - Trees") {
+    val one = Literal(1)
+    val two = Literal(2)
+
+    val add = Add(one, two)
+    val abs = Abs(add)
+    val add2 = Add(add, add)
+
+    var equivalence = new EquivalentExpressions
+    equivalence.addExprTree(add, true)
+    equivalence.addExprTree(abs, true)
+    equivalence.addExprTree(add2, true)
+
+    // Should only have one equivalence for `one + two`
+    assert(equivalence.getAllEquivalentExprs.filter(_.size > 1).size == 1)
+    assert(equivalence.getAllEquivalentExprs.filter(_.size > 1).head.size == 4)
+
+    // Set up the expressions
+    //   one * two,
+    //   (one * two) * (one * two)
+    //   sqrt( (one * two) * (one * two) )
+    //   (one * two) + sqrt( (one * two) * (one * two) )
+    equivalence = new EquivalentExpressions
+    val mul = Multiply(one, two)
+    val mul2 = Multiply(mul, mul)
+    val sqrt = Sqrt(mul2)
+    val sum = Add(mul2, sqrt)
+    equivalence.addExprTree(mul, true)
+    equivalence.addExprTree(mul2, true)
+    equivalence.addExprTree(sqrt, true)
+    equivalence.addExprTree(sum, true)
+
+    // (one * two), (one * two) * (one * two) and sqrt( (one * two) * (one * two) ) should be found
+    assert(equivalence.getAllEquivalentExprs.filter(_.size > 1).size == 3)
+    assert(equivalence.getEquivalentExprs(mul).size == 3)
+    assert(equivalence.getEquivalentExprs(mul2).size == 3)
+    assert(equivalence.getEquivalentExprs(sqrt).size == 2)
+    assert(equivalence.getEquivalentExprs(sum).size == 1)
+
+    // Some expressions inspired by TPCH-Q1
+    // sum(l_quantity) as sum_qty,
+    // sum(l_extendedprice) as sum_base_price,
+    // sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
+    // sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge,
+    // avg(l_extendedprice) as avg_price,
+    // avg(l_discount) as avg_disc
+    equivalence = new EquivalentExpressions
+    val quantity = Literal(1)
+    val price = Literal(1.1)
+    val discount = Literal(.24)
+    val tax = Literal(0.1)
+    equivalence.addExprTree(quantity, false)
+    equivalence.addExprTree(price, false)
+    equivalence.addExprTree(Multiply(price, Subtract(Literal(1), discount)), false)
+    equivalence.addExprTree(
+      Multiply(
+        Multiply(price, Subtract(Literal(1), discount)),
+        Add(Literal(1), tax)), false)
+    equivalence.addExprTree(price, false)
+    equivalence.addExprTree(discount, false)
+    // quantity, price, discount and (price * (1 - discount))
+    assert(equivalence.getAllEquivalentExprs.filter(_.size > 1).size == 4)
+  }
+
+  test("Expression equivalence - non deterministic") {
+    val sum = Add(Rand(0), Rand(0))
+    val equivalence = new EquivalentExpressions
+    equivalence.addExpr(sum)
+    equivalence.addExpr(sum)
+    assert(equivalence.getAllEquivalentExprs.isEmpty)
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
index b7314189b5403..89e196c066007 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
@@ -268,6 +268,11 @@ private[spark] object SQLConf {
     doc = "When true, use the new optimized Tungsten physical execution backend.",
     isPublic = false)
 
+  val SUBEXPRESSION_ELIMINATION_ENABLED = booleanConf("spark.sql.subexpressionElimination.enabled",
+    defaultValue = Some(true),  // use CODEGEN_ENABLED as default
+    doc = "When true, common subexpressions will be eliminated.",
+    isPublic = false)
+
   val DIALECT = stringConf(
     "spark.sql.dialect",
     defaultValue = Some("sql"),
@@ -541,6 +546,9 @@ private[sql] class SQLConf extends Serializable with CatalystConf {
 
   private[spark] def unsafeEnabled: Boolean = getConf(UNSAFE_ENABLED, getConf(TUNGSTEN_ENABLED))
 
+  private[spark] def subexpressionEliminationEnabled: Boolean =
+    getConf(SUBEXPRESSION_ELIMINATION_ENABLED, codegenEnabled)
+
   private[spark] def autoBroadcastJoinThreshold: Int = getConf(AUTO_BROADCASTJOIN_THRESHOLD)
 
   private[spark] def defaultSizeInBytes: Long =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
index 8bb293ae87e64..8650ac500b652 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
@@ -66,6 +66,11 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
   } else {
     false
   }
+  val subexpressionEliminationEnabled: Boolean = if (sqlContext != null) {
+    sqlContext.conf.subexpressionEliminationEnabled
+  } else {
+    false
+  }
 
   /**
    * Whether the "prepare" method is called.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index 145de0db9edaa..303d636164adb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -70,7 +70,8 @@ case class TungstenProject(projectList: Seq[NamedExpression], child: SparkPlan)
   protected override def doExecute(): RDD[InternalRow] = {
     val numRows = longMetric("numRows")
     child.execute().mapPartitions { iter =>
-      val project = UnsafeProjection.create(projectList, child.output)
+      val project = UnsafeProjection.create(projectList, child.output,
+        subexpressionEliminationEnabled)
       iter.map { row =>
         numRows += 1
         project(row)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index 441a0c6d0e36e..19e850a46fdfc 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -1970,4 +1970,52 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
         Row(1, 1) :: Row(1, 2) :: Row(2, 1) :: Row(2, 2) :: Row(3, 1) :: Row(3, 2) :: Nil)
     }
   }
+
+  test("Common subexpression elimination") {
+    // select from a table to prevent constant folding.
+    val df = sql("SELECT a, b from testData2 limit 1")
+    checkAnswer(df, Row(1, 1))
+
+    checkAnswer(df.selectExpr("a + 1", "a + 1"), Row(2, 2))
+    checkAnswer(df.selectExpr("a + 1", "a + 1 + 1"), Row(2, 3))
+
+    // This does not work because the expressions get grouped like (a + a) + 1
+    checkAnswer(df.selectExpr("a + 1", "a + a + 1"), Row(2, 3))
+    checkAnswer(df.selectExpr("a + 1", "a + (a + 1)"), Row(2, 3))
+
+    // Identity udf that tracks the number of times it is called.
+    val countAcc = sparkContext.accumulator(0, "CallCount")
+    sqlContext.udf.register("testUdf", (x: Int) => {
+      countAcc.++=(1)
+      x
+    })
+
+    // Evaluates df, verifying it is equal to the expectedResult and the accumulator's value
+    // is correct.
+    def verifyCallCount(df: DataFrame, expectedResult: Row, expectedCount: Int): Unit = {
+      countAcc.setValue(0)
+      checkAnswer(df, expectedResult)
+      assert(countAcc.value == expectedCount)
+    }
+
+    verifyCallCount(df.selectExpr("testUdf(a)"), Row(1), 1)
+    verifyCallCount(df.selectExpr("testUdf(a)", "testUdf(a)"), Row(1, 1), 1)
+    verifyCallCount(df.selectExpr("testUdf(a + 1)", "testUdf(a + 1)"), Row(2, 2), 1)
+    verifyCallCount(df.selectExpr("testUdf(a + 1)", "testUdf(a)"), Row(2, 1), 2)
+    verifyCallCount(
+      df.selectExpr("testUdf(a + 1) + testUdf(a + 1)", "testUdf(a + 1)"), Row(4, 2), 1)
+
+    verifyCallCount(
+      df.selectExpr("testUdf(a + 1) + testUdf(1 + b)", "testUdf(a + 1)"), Row(4, 2), 2)
+
+    // Would be nice if semantic equals for `+` understood commutative
+    verifyCallCount(
+      df.selectExpr("testUdf(a + 1) + testUdf(1 + a)", "testUdf(a + 1)"), Row(4, 2), 2)
+
+    // Try disabling it via configuration.
+    sqlContext.setConf("spark.sql.subexpressionElimination.enabled", "false")
+    verifyCallCount(df.selectExpr("testUdf(a)", "testUdf(a)"), Row(1, 1), 2)
+    sqlContext.setConf("spark.sql.subexpressionElimination.enabled", "true")
+    verifyCallCount(df.selectExpr("testUdf(a)", "testUdf(a)"), Row(1, 1), 1)
+  }
 }

From f14e95115c0939a77ebcb00209696a87fd651ff9 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Tue, 10 Nov 2015 11:34:36 -0800
Subject: [PATCH 162/518] [ML][R] SparkR::glm summary result to compare with
 native R

Follow up #9561. Due to [SPARK-11587](https://issues.apache.org/jira/browse/SPARK-11587) has been fixed, we should compare SparkR::glm summary result with native R output rather than hard-code one. mengxr

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9590 from yanboliang/glm-r-test.
---
 R/pkg/R/mllib.R               |  2 +-
 R/pkg/inst/tests/test_mllib.R | 31 ++++++++++---------------------
 2 files changed, 11 insertions(+), 22 deletions(-)

diff --git a/R/pkg/R/mllib.R b/R/pkg/R/mllib.R
index 7126b7cde4bd7..f23e1c7f1fce4 100644
--- a/R/pkg/R/mllib.R
+++ b/R/pkg/R/mllib.R
@@ -106,7 +106,7 @@ setMethod("summary", signature(object = "PipelineModel"),
               coefficients <- matrix(coefficients, ncol = 4)
               colnames(coefficients) <- c("Estimate", "Std. Error", "t value", "Pr(>|t|)")
               rownames(coefficients) <- unlist(features)
-              return(list(DevianceResiduals = devianceResiduals, Coefficients = coefficients))
+              return(list(devianceResiduals = devianceResiduals, coefficients = coefficients))
             } else {
               coefficients <- as.matrix(unlist(coefficients))
               colnames(coefficients) <- c("Estimate")
diff --git a/R/pkg/inst/tests/test_mllib.R b/R/pkg/inst/tests/test_mllib.R
index 42287ea19adc5..d497ad8c9daa3 100644
--- a/R/pkg/inst/tests/test_mllib.R
+++ b/R/pkg/inst/tests/test_mllib.R
@@ -72,22 +72,17 @@ test_that("feature interaction vs native glm", {
 test_that("summary coefficients match with native glm", {
   training <- createDataFrame(sqlContext, iris)
   stats <- summary(glm(Sepal_Width ~ Sepal_Length + Species, data = training, solver = "normal"))
-  coefs <- unlist(stats$Coefficients)
-  devianceResiduals <- unlist(stats$DevianceResiduals)
+  coefs <- unlist(stats$coefficients)
+  devianceResiduals <- unlist(stats$devianceResiduals)
 
-  rCoefs <- as.vector(coef(glm(Sepal.Width ~ Sepal.Length + Species, data = iris)))
-  rStdError <- c(0.23536, 0.04630, 0.07207, 0.09331)
-  rTValue <- c(7.123, 7.557, -13.644, -10.798)
-  rPValue <- c(0.0, 0.0, 0.0, 0.0)
+  rStats <- summary(glm(Sepal.Width ~ Sepal.Length + Species, data = iris))
+  rCoefs <- unlist(rStats$coefficients)
   rDevianceResiduals <- c(-0.95096, 0.72918)
 
-  expect_true(all(abs(rCoefs - coefs[1:4]) < 1e-6))
-  expect_true(all(abs(rStdError - coefs[5:8]) < 1e-5))
-  expect_true(all(abs(rTValue - coefs[9:12]) < 1e-3))
-  expect_true(all(abs(rPValue - coefs[13:16]) < 1e-6))
+  expect_true(all(abs(rCoefs - coefs) < 1e-5))
   expect_true(all(abs(rDevianceResiduals - devianceResiduals) < 1e-5))
   expect_true(all(
-    rownames(stats$Coefficients) ==
+    rownames(stats$coefficients) ==
     c("(Intercept)", "Sepal_Length", "Species_versicolor", "Species_virginica")))
 })
 
@@ -96,21 +91,15 @@ test_that("summary coefficients match with native glm of family 'binomial'", {
   training <- filter(df, df$Species != "setosa")
   stats <- summary(glm(Species ~ Sepal_Length + Sepal_Width, data = training,
     family = "binomial"))
-  coefs <- as.vector(stats$Coefficients)
+  coefs <- as.vector(stats$coefficients[,1])
 
   rTraining <- iris[iris$Species %in% c("versicolor","virginica"),]
   rCoefs <- as.vector(coef(glm(Species ~ Sepal.Length + Sepal.Width, data = rTraining,
     family = binomial(link = "logit"))))
-  rStdError <- c(3.0974, 0.5169, 0.8628)
-  rTValue <- c(-4.212, 3.680, 0.469)
-  rPValue <- c(0.000, 0.000, 0.639)
-
-  expect_true(all(abs(rCoefs - coefs[1:3]) < 1e-4))
-  expect_true(all(abs(rStdError - coefs[4:6]) < 1e-4))
-  expect_true(all(abs(rTValue - coefs[7:9]) < 1e-3))
-  expect_true(all(abs(rPValue - coefs[10:12]) < 1e-3))
+
+  expect_true(all(abs(rCoefs - coefs) < 1e-4))
   expect_true(all(
-    rownames(stats$Coefficients) ==
+    rownames(stats$coefficients) ==
     c("(Intercept)", "Sepal_Length", "Sepal_Width")))
 })
 

From 18350a57004eb87cafa9504ff73affab4b818e06 Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Tue, 10 Nov 2015 11:36:43 -0800
Subject: [PATCH 163/518] [SPARK-11618][ML] Minor refactoring of basic ML
 import/export

Refactoring
* separated overwrite and param save logic in DefaultParamsWriter
* added sparkVersion to DefaultParamsWriter

CC: mengxr

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9587 from jkbradley/logreg-io.
---
 .../org/apache/spark/ml/util/ReadWrite.scala  | 57 ++++++++++---------
 1 file changed, 30 insertions(+), 27 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index ea790e0dddc7f..cbdf913ba8dfa 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -51,6 +51,9 @@ private[util] sealed trait BaseReadWrite {
   protected final def sqlContext: SQLContext = optionSQLContext.getOrElse {
     SQLContext.getOrCreate(SparkContext.getOrCreate())
   }
+
+  /** Returns the [[SparkContext]] underlying [[sqlContext]] */
+  protected final def sc: SparkContext = sqlContext.sparkContext
 }
 
 /**
@@ -58,7 +61,7 @@ private[util] sealed trait BaseReadWrite {
  */
 @Experimental
 @Since("1.6.0")
-abstract class Writer extends BaseReadWrite {
+abstract class Writer extends BaseReadWrite with Logging {
 
   protected var shouldOverwrite: Boolean = false
 
@@ -67,7 +70,29 @@ abstract class Writer extends BaseReadWrite {
    */
   @Since("1.6.0")
   @throws[IOException]("If the input path already exists but overwrite is not enabled.")
-  def save(path: String): Unit
+  def save(path: String): Unit = {
+    val hadoopConf = sc.hadoopConfiguration
+    val fs = FileSystem.get(hadoopConf)
+    val p = new Path(path)
+    if (fs.exists(p)) {
+      if (shouldOverwrite) {
+        logInfo(s"Path $path already exists. It will be overwritten.")
+        // TODO: Revert back to the original content if save is not successful.
+        fs.delete(p, true)
+      } else {
+        throw new IOException(
+          s"Path $path already exists. Please use write.overwrite().save(path) to overwrite it.")
+      }
+    }
+    saveImpl(path)
+  }
+
+  /**
+   * [[save()]] handles overwriting and then calls this method.  Subclasses should override this
+   * method to implement the actual saving of the instance.
+   */
+  @Since("1.6.0")
+  protected def saveImpl(path: String): Unit
 
   /**
    * Overwrites if the output path already exists.
@@ -147,28 +172,9 @@ trait Readable[T] {
  * data (e.g., models with coefficients).
  * @param instance object to save
  */
-private[ml] class DefaultParamsWriter(instance: Params) extends Writer with Logging {
-
-  /**
-   * Saves the ML component to the input path.
-   */
-  override def save(path: String): Unit = {
-    val sc = sqlContext.sparkContext
-
-    val hadoopConf = sc.hadoopConfiguration
-    val fs = FileSystem.get(hadoopConf)
-    val p = new Path(path)
-    if (fs.exists(p)) {
-      if (shouldOverwrite) {
-        logInfo(s"Path $path already exists. It will be overwritten.")
-        // TODO: Revert back to the original content if save is not successful.
-        fs.delete(p, true)
-      } else {
-        throw new IOException(
-          s"Path $path already exists. Please use write.overwrite().save(path) to overwrite it.")
-      }
-    }
+private[ml] class DefaultParamsWriter(instance: Params) extends Writer {
 
+  override protected def saveImpl(path: String): Unit = {
     val uid = instance.uid
     val cls = instance.getClass.getName
     val params = instance.extractParamMap().toSeq.asInstanceOf[Seq[ParamPair[Any]]]
@@ -177,6 +183,7 @@ private[ml] class DefaultParamsWriter(instance: Params) extends Writer with Logg
     }.toList
     val metadata = ("class" -> cls) ~
       ("timestamp" -> System.currentTimeMillis()) ~
+      ("sparkVersion" -> sc.version) ~
       ("uid" -> uid) ~
       ("paramMap" -> jsonParams)
     val metadataPath = new Path(path, "metadata").toString
@@ -193,12 +200,8 @@ private[ml] class DefaultParamsWriter(instance: Params) extends Writer with Logg
  */
 private[ml] class DefaultParamsReader[T] extends Reader[T] {
 
-  /**
-   * Loads the ML component from the input path.
-   */
   override def load(path: String): T = {
     implicit val format = DefaultFormats
-    val sc = sqlContext.sparkContext
     val metadataPath = new Path(path, "metadata").toString
     val metadataStr = sc.textFile(metadataPath, 1).first()
     val metadata = parse(metadataStr)

From dba1a62cf1baa9ae1ee665d592e01dfad78331a2 Mon Sep 17 00:00:00 2001
From: unknown <ulanov@ULANOV3.americas.hpqcorp.net>
Date: Tue, 10 Nov 2015 14:25:06 -0800
Subject: [PATCH 164/518] [SPARK-7316][MLLIB] RDD sliding window with step

Implementation of step capability for sliding window function in MLlib's RDD.

Though one can use current sliding window with step 1 and then filter every Nth window, it will take more time and space (N*data.count times more than needed). For example, below are the results for various windows and steps on 10M data points:

Window | Step | Time | Windows produced
------------ | ------------- | ---------- | ----------
128 | 1 |  6.38 | 9999873
128 | 10 | 0.9 | 999988
128 | 100 | 0.41 | 99999
1024 | 1 | 44.67 | 9998977
1024 | 10 | 4.74 | 999898
1024 | 100 | 0.78 | 99990
```
import org.apache.spark.mllib.rdd.RDDFunctions._
val rdd = sc.parallelize(1 to 10000000, 10)
rdd.count
val window = 1024
val step = 1
val t = System.nanoTime(); val windows = rdd.sliding(window, step); println(windows.count); println((System.nanoTime() - t) / 1e9)
```

Author: unknown <ulanov@ULANOV3.americas.hpqcorp.net>
Author: Alexander Ulanov <nashb@yandex.ru>
Author: Xiangrui Meng <meng@databricks.com>

Closes #5855 from avulanov/SPARK-7316-sliding.
---
 .../apache/spark/mllib/rdd/RDDFunctions.scala | 11 ++-
 .../apache/spark/mllib/rdd/SlidingRDD.scala   | 71 ++++++++++---------
 .../spark/mllib/rdd/RDDFunctionsSuite.scala   | 11 +--
 3 files changed, 54 insertions(+), 39 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/rdd/RDDFunctions.scala b/mllib/src/main/scala/org/apache/spark/mllib/rdd/RDDFunctions.scala
index 78172843be56e..19a047ded257c 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/rdd/RDDFunctions.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/rdd/RDDFunctions.scala
@@ -37,15 +37,20 @@ class RDDFunctions[T: ClassTag](self: RDD[T]) extends Serializable {
    * trigger a Spark job if the parent RDD has more than one partitions and the window size is
    * greater than 1.
    */
-  def sliding(windowSize: Int): RDD[Array[T]] = {
+  def sliding(windowSize: Int, step: Int): RDD[Array[T]] = {
     require(windowSize > 0, s"Sliding window size must be positive, but got $windowSize.")
-    if (windowSize == 1) {
+    if (windowSize == 1 && step == 1) {
       self.map(Array(_))
     } else {
-      new SlidingRDD[T](self, windowSize)
+      new SlidingRDD[T](self, windowSize, step)
     }
   }
 
+  /**
+   * [[sliding(Int, Int)*]] with step = 1.
+   */
+  def sliding(windowSize: Int): RDD[Array[T]] = sliding(windowSize, 1)
+
   /**
    * Reduces the elements of this RDD in a multi-level tree pattern.
    *
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/rdd/SlidingRDD.scala b/mllib/src/main/scala/org/apache/spark/mllib/rdd/SlidingRDD.scala
index 1facf83d806d0..ead8db6344998 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/rdd/SlidingRDD.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/rdd/SlidingRDD.scala
@@ -24,13 +24,13 @@ import org.apache.spark.{TaskContext, Partition}
 import org.apache.spark.rdd.RDD
 
 private[mllib]
-class SlidingRDDPartition[T](val idx: Int, val prev: Partition, val tail: Seq[T])
+class SlidingRDDPartition[T](val idx: Int, val prev: Partition, val tail: Seq[T], val offset: Int)
   extends Partition with Serializable {
   override val index: Int = idx
 }
 
 /**
- * Represents a RDD from grouping items of its parent RDD in fixed size blocks by passing a sliding
+ * Represents an RDD from grouping items of its parent RDD in fixed size blocks by passing a sliding
  * window over them. The ordering is first based on the partition index and then the ordering of
  * items within each partition. This is similar to sliding in Scala collections, except that it
  * becomes an empty RDD if the window size is greater than the total number of items. It needs to
@@ -40,19 +40,24 @@ class SlidingRDDPartition[T](val idx: Int, val prev: Partition, val tail: Seq[T]
  *
  * @param parent the parent RDD
  * @param windowSize the window size, must be greater than 1
+ * @param step step size for windows
  *
- * @see [[org.apache.spark.mllib.rdd.RDDFunctions#sliding]]
+ * @see [[org.apache.spark.mllib.rdd.RDDFunctions.sliding(Int, Int)*]]
+ * @see [[scala.collection.IterableLike.sliding(Int, Int)*]]
  */
 private[mllib]
-class SlidingRDD[T: ClassTag](@transient val parent: RDD[T], val windowSize: Int)
+class SlidingRDD[T: ClassTag](@transient val parent: RDD[T], val windowSize: Int, val step: Int)
   extends RDD[Array[T]](parent) {
 
-  require(windowSize > 1, s"Window size must be greater than 1, but got $windowSize.")
+  require(windowSize > 0 && step > 0 && !(windowSize == 1 && step == 1),
+    "Window size and step must be greater than 0, " +
+      s"and they cannot be both 1, but got windowSize = $windowSize and step = $step.")
 
   override def compute(split: Partition, context: TaskContext): Iterator[Array[T]] = {
     val part = split.asInstanceOf[SlidingRDDPartition[T]]
     (firstParent[T].iterator(part.prev, context) ++ part.tail)
-      .sliding(windowSize)
+      .drop(part.offset)
+      .sliding(windowSize, step)
       .withPartial(false)
       .map(_.toArray)
   }
@@ -62,40 +67,42 @@ class SlidingRDD[T: ClassTag](@transient val parent: RDD[T], val windowSize: Int
 
   override def getPartitions: Array[Partition] = {
     val parentPartitions = parent.partitions
-    val n = parentPartitions.size
+    val n = parentPartitions.length
     if (n == 0) {
       Array.empty
     } else if (n == 1) {
-      Array(new SlidingRDDPartition[T](0, parentPartitions(0), Seq.empty))
+      Array(new SlidingRDDPartition[T](0, parentPartitions(0), Seq.empty, 0))
     } else {
-      val n1 = n - 1
       val w1 = windowSize - 1
-      // Get the first w1 items of each partition, starting from the second partition.
-      val nextHeads =
-        parent.context.runJob(parent, (iter: Iterator[T]) => iter.take(w1).toArray, 1 until n)
-      val partitions = mutable.ArrayBuffer[SlidingRDDPartition[T]]()
+      // Get partition sizes and first w1 elements.
+      val (sizes, heads) = parent.mapPartitions { iter =>
+        val w1Array = iter.take(w1).toArray
+        Iterator.single((w1Array.length + iter.length, w1Array))
+      }.collect().unzip
+      val partitions = mutable.ArrayBuffer.empty[SlidingRDDPartition[T]]
       var i = 0
+      var cumSize = 0
       var partitionIndex = 0
-      while (i < n1) {
-        var j = i
-        val tail = mutable.ListBuffer[T]()
-        // Keep appending to the current tail until appended a head of size w1.
-        while (j < n1 && nextHeads(j).size < w1) {
-          tail ++= nextHeads(j)
-          j += 1
+      while (i < n) {
+        val mod = cumSize % step
+        val offset = if (mod == 0) 0 else step - mod
+        val size = sizes(i)
+        if (offset < size) {
+          val tail = mutable.ListBuffer.empty[T]
+          // Keep appending to the current tail until it has w1 elements.
+          var j = i + 1
+          while (j < n && tail.length < w1) {
+            tail ++= heads(j).take(w1 - tail.length)
+            j += 1
+          }
+          if (sizes(i) + tail.length >= offset + windowSize) {
+            partitions +=
+              new SlidingRDDPartition[T](partitionIndex, parentPartitions(i), tail, offset)
+            partitionIndex += 1
+          }
         }
-        if (j < n1) {
-          tail ++= nextHeads(j)
-          j += 1
-        }
-        partitions += new SlidingRDDPartition[T](partitionIndex, parentPartitions(i), tail)
-        partitionIndex += 1
-        // Skip appended heads.
-        i = j
-      }
-      // If the head of last partition has size w1, we also need to add this partition.
-      if (nextHeads.last.size == w1) {
-        partitions += new SlidingRDDPartition[T](partitionIndex, parentPartitions(n1), Seq.empty)
+        cumSize += size
+        i += 1
       }
       partitions.toArray
     }
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/rdd/RDDFunctionsSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/rdd/RDDFunctionsSuite.scala
index bc64172614830..ac93733bab5f5 100644
--- a/mllib/src/test/scala/org/apache/spark/mllib/rdd/RDDFunctionsSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/mllib/rdd/RDDFunctionsSuite.scala
@@ -28,9 +28,12 @@ class RDDFunctionsSuite extends SparkFunSuite with MLlibTestSparkContext {
     for (numPartitions <- 1 to 8) {
       val rdd = sc.parallelize(data, numPartitions)
       for (windowSize <- 1 to 6) {
-        val sliding = rdd.sliding(windowSize).collect().map(_.toList).toList
-        val expected = data.sliding(windowSize).map(_.toList).toList
-        assert(sliding === expected)
+        for (step <- 1 to 3) {
+          val sliding = rdd.sliding(windowSize, step).collect().map(_.toList).toList
+          val expected = data.sliding(windowSize, step)
+            .map(_.toList).toList.filter(l => l.size == windowSize)
+          assert(sliding === expected)
+        }
       }
       assert(rdd.sliding(7).collect().isEmpty,
         "Should return an empty RDD if the window size is greater than the number of items.")
@@ -40,7 +43,7 @@ class RDDFunctionsSuite extends SparkFunSuite with MLlibTestSparkContext {
   test("sliding with empty partitions") {
     val data = Seq(Seq(1, 2, 3), Seq.empty[Int], Seq(4), Seq.empty[Int], Seq(5, 6, 7))
     val rdd = sc.parallelize(data, data.length).flatMap(s => s)
-    assert(rdd.partitions.size === data.length)
+    assert(rdd.partitions.length === data.length)
     val sliding = rdd.sliding(3).collect().toSeq.map(_.toSeq)
     val expected = data.flatMap(x => x).sliding(3).toSeq.map(_.toSeq)
     assert(sliding === expected)

From 724cf7a38c551bf2a79b87a8158bbe1725f9f888 Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Tue, 10 Nov 2015 14:30:19 -0800
Subject: [PATCH 165/518] [SPARK-11616][SQL] Improve toString for Dataset

Author: Michael Armbrust <michael@databricks.com>

Closes #9586 from marmbrus/dataset-toString.
---
 .../org/apache/spark/sql/DataFrame.scala      | 14 ++-----
 .../scala/org/apache/spark/sql/Dataset.scala  |  4 +-
 .../spark/sql/execution/Queryable.scala       | 37 +++++++++++++++++++
 .../org/apache/spark/sql/DatasetSuite.scala   |  5 +++
 4 files changed, 47 insertions(+), 13 deletions(-)
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 9368435a63c35..691b476fff8d6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.encoders.Encoder
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.{Inner, JoinType}
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, ScalaReflection, SqlParser}
-import org.apache.spark.sql.execution.{EvaluatePython, ExplainCommand, FileRelation, LogicalRDD, QueryExecution, SQLExecution}
+import org.apache.spark.sql.execution.{EvaluatePython, ExplainCommand, FileRelation, LogicalRDD, QueryExecution, Queryable, SQLExecution}
 import org.apache.spark.sql.execution.datasources.{CreateTableUsingAsSelect, LogicalRelation}
 import org.apache.spark.sql.execution.datasources.json.JacksonGenerator
 import org.apache.spark.sql.sources.HadoopFsRelation
@@ -116,7 +116,8 @@ private[sql] object DataFrame {
 @Experimental
 class DataFrame private[sql](
     @transient val sqlContext: SQLContext,
-    @DeveloperApi @transient val queryExecution: QueryExecution) extends Serializable {
+    @DeveloperApi @transient val queryExecution: QueryExecution)
+  extends Queryable with Serializable {
 
   // Note for Spark contributors: if adding or updating any action in `DataFrame`, please make sure
   // you wrap it with `withNewExecutionId` if this actions doesn't call other action.
@@ -234,15 +235,6 @@ class DataFrame private[sql](
     sb.toString()
   }
 
-  override def toString: String = {
-    try {
-      schema.map(f => s"${f.name}: ${f.dataType.simpleString}").mkString("[", ", ", "]")
-    } catch {
-      case NonFatal(e) =>
-        s"Invalid tree; ${e.getMessage}:\n$queryExecution"
-    }
-  }
-
   /**
    * Returns the object itself.
    * @group basic
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 6d2968e2881f8..a7e5ab19bf846 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.Inner
 import org.apache.spark.sql.catalyst.plans.logical._
-import org.apache.spark.sql.execution.QueryExecution
+import org.apache.spark.sql.execution.{Queryable, QueryExecution}
 import org.apache.spark.sql.types.StructType
 
 /**
@@ -62,7 +62,7 @@ import org.apache.spark.sql.types.StructType
 class Dataset[T] private[sql](
     @transient val sqlContext: SQLContext,
     @transient val queryExecution: QueryExecution,
-    unresolvedEncoder: Encoder[T]) extends Serializable {
+    unresolvedEncoder: Encoder[T]) extends Queryable with Serializable {
 
   /** The encoder for this [[Dataset]] that has been resolved to its output schema. */
   private[sql] implicit val encoder: ExpressionEncoder[T] = unresolvedEncoder match {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
new file mode 100644
index 0000000000000..9ca383896a09b
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.sql.types.StructType
+
+import scala.util.control.NonFatal
+
+/** A trait that holds shared code between DataFrames and Datasets. */
+private[sql] trait Queryable {
+  def schema: StructType
+  def queryExecution: QueryExecution
+
+  override def toString: String = {
+    try {
+      schema.map(f => s"${f.name}: ${f.dataType.simpleString}").mkString("[", ", ", "]")
+    } catch {
+      case NonFatal(e) =>
+        s"Invalid tree; ${e.getMessage}:\n$queryExecution"
+    }
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index aea5a700d0204..621148528714f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -313,4 +313,9 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     val joined = ds1.joinWith(ds2, $"a.value" === $"b.value")
     checkAnswer(joined, ("2", 2))
   }
+
+  test("toString") {
+    val ds = Seq((1, 2)).toDS()
+    assert(ds.toString == "[_1: int, _2: int]")
+  }
 }

From 638c51d9380081b3b8182be2c2460bd53b8b0a4f Mon Sep 17 00:00:00 2001
From: Pravin Gadakh <pravingadakh177@gmail.com>
Date: Tue, 10 Nov 2015 14:47:04 -0800
Subject: [PATCH 166/518] [SPARK-11550][DOCS] Replace example code in
 mllib-optimization.md using include_example

Author: Pravin Gadakh <pravingadakh177@gmail.com>

Closes #9516 from pravingadakh/SPARK-11550.
---
 docs/mllib-optimization.md                    | 145 +-----------------
 .../examples/mllib/JavaLBFGSExample.java      | 108 +++++++++++++
 .../spark/examples/mllib/LBFGSExample.scala   |  90 +++++++++++
 3 files changed, 200 insertions(+), 143 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaLBFGSExample.java
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/LBFGSExample.scala

diff --git a/docs/mllib-optimization.md b/docs/mllib-optimization.md
index a3bd130ba077c..ad7bcd9bfd407 100644
--- a/docs/mllib-optimization.md
+++ b/docs/mllib-optimization.md
@@ -220,154 +220,13 @@ L-BFGS optimizer.
 <div data-lang="scala" markdown="1">
 Refer to the [`LBFGS` Scala docs](api/scala/index.html#org.apache.spark.mllib.optimization.LBFGS) and [`SquaredL2Updater` Scala docs](api/scala/index.html#org.apache.spark.mllib.optimization.SquaredL2Updater) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.SparkContext
-import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics
-import org.apache.spark.mllib.linalg.Vectors
-import org.apache.spark.mllib.util.MLUtils
-import org.apache.spark.mllib.classification.LogisticRegressionModel
-import org.apache.spark.mllib.optimization.{LBFGS, LogisticGradient, SquaredL2Updater}
-
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-val numFeatures = data.take(1)(0).features.size
-
-// Split data into training (60%) and test (40%).
-val splits = data.randomSplit(Array(0.6, 0.4), seed = 11L)
-
-// Append 1 into the training data as intercept.
-val training = splits(0).map(x => (x.label, MLUtils.appendBias(x.features))).cache()
-
-val test = splits(1)
-
-// Run training algorithm to build the model
-val numCorrections = 10
-val convergenceTol = 1e-4
-val maxNumIterations = 20
-val regParam = 0.1
-val initialWeightsWithIntercept = Vectors.dense(new Array[Double](numFeatures + 1))
-
-val (weightsWithIntercept, loss) = LBFGS.runLBFGS(
-  training,
-  new LogisticGradient(),
-  new SquaredL2Updater(),
-  numCorrections,
-  convergenceTol,
-  maxNumIterations,
-  regParam,
-  initialWeightsWithIntercept)
-
-val model = new LogisticRegressionModel(
-  Vectors.dense(weightsWithIntercept.toArray.slice(0, weightsWithIntercept.size - 1)),
-  weightsWithIntercept(weightsWithIntercept.size - 1))
-
-// Clear the default threshold.
-model.clearThreshold()
-
-// Compute raw scores on the test set.
-val scoreAndLabels = test.map { point =>
-  val score = model.predict(point.features)
-  (score, point.label)
-}
-
-// Get evaluation metrics.
-val metrics = new BinaryClassificationMetrics(scoreAndLabels)
-val auROC = metrics.areaUnderROC()
-
-println("Loss of each step in training process")
-loss.foreach(println)
-println("Area under ROC = " + auROC)
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/LBFGSExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`LBFGS` Java docs](api/java/org/apache/spark/mllib/optimization/LBFGS.html) and [`SquaredL2Updater` Java docs](api/java/org/apache/spark/mllib/optimization/SquaredL2Updater.html) for details on the API.
 
-{% highlight java %}
-import java.util.Arrays;
-import java.util.Random;
-
-import scala.Tuple2;
-
-import org.apache.spark.api.java.*;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.mllib.classification.LogisticRegressionModel;
-import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics;
-import org.apache.spark.mllib.linalg.Vector;
-import org.apache.spark.mllib.linalg.Vectors;
-import org.apache.spark.mllib.optimization.*;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.SparkConf;
-import org.apache.spark.SparkContext;
-
-public class LBFGSExample {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf().setAppName("L-BFGS Example");
-    SparkContext sc = new SparkContext(conf);
-    String path = "data/mllib/sample_libsvm_data.txt";
-    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD();
-    int numFeatures = data.take(1).get(0).features().size();
-    
-    // Split initial RDD into two... [60% training data, 40% testing data].
-    JavaRDD<LabeledPoint> trainingInit = data.sample(false, 0.6, 11L);
-    JavaRDD<LabeledPoint> test = data.subtract(trainingInit);
-    
-    // Append 1 into the training data as intercept.
-    JavaRDD<Tuple2<Object, Vector>> training = data.map(
-      new Function<LabeledPoint, Tuple2<Object, Vector>>() {
-        public Tuple2<Object, Vector> call(LabeledPoint p) {
-          return new Tuple2<Object, Vector>(p.label(), MLUtils.appendBias(p.features()));
-        }
-      });
-    training.cache();
-
-    // Run training algorithm to build the model.
-    int numCorrections = 10;
-    double convergenceTol = 1e-4;
-    int maxNumIterations = 20;
-    double regParam = 0.1;
-    Vector initialWeightsWithIntercept = Vectors.dense(new double[numFeatures + 1]);
-
-    Tuple2<Vector, double[]> result = LBFGS.runLBFGS(
-      training.rdd(),
-      new LogisticGradient(),
-      new SquaredL2Updater(),
-      numCorrections,
-      convergenceTol,
-      maxNumIterations,
-      regParam,
-      initialWeightsWithIntercept);
-    Vector weightsWithIntercept = result._1();
-    double[] loss = result._2();
-
-    final LogisticRegressionModel model = new LogisticRegressionModel(
-      Vectors.dense(Arrays.copyOf(weightsWithIntercept.toArray(), weightsWithIntercept.size() - 1)),
-      (weightsWithIntercept.toArray())[weightsWithIntercept.size() - 1]);
-
-    // Clear the default threshold.
-    model.clearThreshold();
-
-    // Compute raw scores on the test set.
-    JavaRDD<Tuple2<Object, Object>> scoreAndLabels = test.map(
-      new Function<LabeledPoint, Tuple2<Object, Object>>() {
-      public Tuple2<Object, Object> call(LabeledPoint p) {
-        Double score = model.predict(p.features());
-        return new Tuple2<Object, Object>(score, p.label());
-      }
-    });
-
-    // Get evaluation metrics.
-    BinaryClassificationMetrics metrics = 
-      new BinaryClassificationMetrics(scoreAndLabels.rdd());
-    double auROC = metrics.areaUnderROC();
-     
-    System.out.println("Loss of each step in training process");
-    for (double l : loss)
-      System.out.println(l);
-    System.out.println("Area under ROC = " + auROC);
-  }
-}
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaLBFGSExample.java %}
 </div>
 </div>
 
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaLBFGSExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaLBFGSExample.java
new file mode 100644
index 0000000000000..355883f61bd64
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaLBFGSExample.java
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.Arrays;
+
+import scala.Tuple2;
+
+import org.apache.spark.api.java.*;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.mllib.classification.LogisticRegressionModel;
+import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics;
+import org.apache.spark.mllib.linalg.Vector;
+import org.apache.spark.mllib.linalg.Vectors;
+import org.apache.spark.mllib.optimization.*;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.util.MLUtils;
+import org.apache.spark.SparkConf;
+import org.apache.spark.SparkContext;
+// $example off$
+
+public class JavaLBFGSExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("L-BFGS Example");
+    SparkContext sc = new SparkContext(conf);
+
+    // $example on$
+    String path = "data/mllib/sample_libsvm_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD();
+    int numFeatures = data.take(1).get(0).features().size();
+
+    // Split initial RDD into two... [60% training data, 40% testing data].
+    JavaRDD<LabeledPoint> trainingInit = data.sample(false, 0.6, 11L);
+    JavaRDD<LabeledPoint> test = data.subtract(trainingInit);
+
+    // Append 1 into the training data as intercept.
+    JavaRDD<Tuple2<Object, Vector>> training = data.map(
+      new Function<LabeledPoint, Tuple2<Object, Vector>>() {
+        public Tuple2<Object, Vector> call(LabeledPoint p) {
+          return new Tuple2<Object, Vector>(p.label(), MLUtils.appendBias(p.features()));
+        }
+      });
+    training.cache();
+
+    // Run training algorithm to build the model.
+    int numCorrections = 10;
+    double convergenceTol = 1e-4;
+    int maxNumIterations = 20;
+    double regParam = 0.1;
+    Vector initialWeightsWithIntercept = Vectors.dense(new double[numFeatures + 1]);
+
+    Tuple2<Vector, double[]> result = LBFGS.runLBFGS(
+      training.rdd(),
+      new LogisticGradient(),
+      new SquaredL2Updater(),
+      numCorrections,
+      convergenceTol,
+      maxNumIterations,
+      regParam,
+      initialWeightsWithIntercept);
+    Vector weightsWithIntercept = result._1();
+    double[] loss = result._2();
+
+    final LogisticRegressionModel model = new LogisticRegressionModel(
+      Vectors.dense(Arrays.copyOf(weightsWithIntercept.toArray(), weightsWithIntercept.size() - 1)),
+      (weightsWithIntercept.toArray())[weightsWithIntercept.size() - 1]);
+
+    // Clear the default threshold.
+    model.clearThreshold();
+
+    // Compute raw scores on the test set.
+    JavaRDD<Tuple2<Object, Object>> scoreAndLabels = test.map(
+      new Function<LabeledPoint, Tuple2<Object, Object>>() {
+        public Tuple2<Object, Object> call(LabeledPoint p) {
+          Double score = model.predict(p.features());
+          return new Tuple2<Object, Object>(score, p.label());
+        }
+      });
+
+    // Get evaluation metrics.
+    BinaryClassificationMetrics metrics =
+      new BinaryClassificationMetrics(scoreAndLabels.rdd());
+    double auROC = metrics.areaUnderROC();
+
+    System.out.println("Loss of each step in training process");
+    for (double l : loss)
+      System.out.println(l);
+    System.out.println("Area under ROC = " + auROC);
+    // $example off$
+  }
+}
+
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/LBFGSExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/LBFGSExample.scala
new file mode 100644
index 0000000000000..61d2e7715f53d
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/LBFGSExample.scala
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.classification.LogisticRegressionModel
+import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics
+import org.apache.spark.mllib.linalg.Vectors
+import org.apache.spark.mllib.optimization.{LBFGS, LogisticGradient, SquaredL2Updater}
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+
+import org.apache.spark.{SparkConf, SparkContext}
+
+object LBFGSExample {
+
+  def main(args: Array[String]): Unit = {
+
+    val conf = new SparkConf().setAppName("LBFGSExample")
+    val sc = new SparkContext(conf)
+
+    // $example on$
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    val numFeatures = data.take(1)(0).features.size
+
+    // Split data into training (60%) and test (40%).
+    val splits = data.randomSplit(Array(0.6, 0.4), seed = 11L)
+
+    // Append 1 into the training data as intercept.
+    val training = splits(0).map(x => (x.label, MLUtils.appendBias(x.features))).cache()
+
+    val test = splits(1)
+
+    // Run training algorithm to build the model
+    val numCorrections = 10
+    val convergenceTol = 1e-4
+    val maxNumIterations = 20
+    val regParam = 0.1
+    val initialWeightsWithIntercept = Vectors.dense(new Array[Double](numFeatures + 1))
+
+    val (weightsWithIntercept, loss) = LBFGS.runLBFGS(
+      training,
+      new LogisticGradient(),
+      new SquaredL2Updater(),
+      numCorrections,
+      convergenceTol,
+      maxNumIterations,
+      regParam,
+      initialWeightsWithIntercept)
+
+    val model = new LogisticRegressionModel(
+      Vectors.dense(weightsWithIntercept.toArray.slice(0, weightsWithIntercept.size - 1)),
+      weightsWithIntercept(weightsWithIntercept.size - 1))
+
+    // Clear the default threshold.
+    model.clearThreshold()
+
+    // Compute raw scores on the test set.
+    val scoreAndLabels = test.map { point =>
+      val score = model.predict(point.features)
+      (score, point.label)
+    }
+
+    // Get evaluation metrics.
+    val metrics = new BinaryClassificationMetrics(scoreAndLabels)
+    val auROC = metrics.areaUnderROC()
+
+    println("Loss of each step in training process")
+    loss.foreach(println)
+    println("Area under ROC = " + auROC)
+    // $example off$
+  }
+}
+// scalastyle:on println

From 32790fe7249b0efe2cbc5c4ee2df0fb687dcd624 Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Tue, 10 Nov 2015 15:47:10 -0800
Subject: [PATCH 167/518] [SPARK-11567] [PYTHON] Add Python API for corr
 Aggregate function

like `df.agg(corr("col1", "col2")`

davies

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9536 from felixcheung/pyfunc.
---
 python/pyspark/sql/functions.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index 6e1cbde4239f3..c3da513c13897 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -255,6 +255,22 @@ def coalesce(*cols):
     return Column(jc)
 
 
+@since(1.6)
+def corr(col1, col2):
+    """Returns a new :class:`Column` for the Pearson Correlation Coefficient for ``col1``
+    and ``col2``.
+
+    >>> a = [x * x - 2 * x + 3.5 for x in range(20)]
+    >>> b = range(20)
+    >>> corrDf = sqlContext.createDataFrame(zip(a, b))
+    >>> corrDf = corrDf.agg(corr(corrDf._1, corrDf._2).alias('c'))
+    >>> corrDf.selectExpr('abs(c - 0.9572339139475857) < 1e-16 as t').collect()
+    [Row(t=True)]
+    """
+    sc = SparkContext._active_spark_context
+    return Column(sc._jvm.functions.corr(_to_java_column(col1), _to_java_column(col2)))
+
+
 @since(1.3)
 def countDistinct(col, *cols):
     """Returns a new :class:`Column` for distinct count of ``col`` or ``cols``.

From 1dde39d796bbf42336051a86bedf871c7fddd513 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Tue, 10 Nov 2015 15:58:30 -0800
Subject: [PATCH 168/518] [SPARK-9818] Re-enable Docker tests for JDBC data
 source

This patch re-enables tests for the Docker JDBC data source. These tests were reverted in #4872 due to transitive dependency conflicts introduced by the `docker-client` library. This patch should avoid those problems by using a version of `docker-client` which shades its transitive dependencies and by performing some build-magic to work around problems with that shaded JAR.

In addition, I significantly refactored the tests to simplify the setup and teardown code and to fix several Docker networking issues which caused problems when running in `boot2docker`.

Closes #8101.

Author: Josh Rosen <joshrosen@databricks.com>
Author: Yijie Shen <henry.yijieshen@gmail.com>

Closes #9503 from JoshRosen/docker-jdbc-tests.
---
 docker-integration-tests/pom.xml              | 149 ++++++++++++++++
 .../sql/jdbc/DockerJDBCIntegrationSuite.scala | 160 ++++++++++++++++++
 .../sql/jdbc/MySQLIntegrationSuite.scala      | 153 +++++++++++++++++
 .../sql/jdbc/PostgresIntegrationSuite.scala   |  82 +++++++++
 .../org/apache/spark/util/DockerUtils.scala   |  68 ++++++++
 pom.xml                                       |  14 ++
 project/SparkBuild.scala                      |  14 +-
 .../org/apache/spark/tags/DockerTest.java     |  26 +++
 8 files changed, 664 insertions(+), 2 deletions(-)
 create mode 100644 docker-integration-tests/pom.xml
 create mode 100644 docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala
 create mode 100644 docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/MySQLIntegrationSuite.scala
 create mode 100644 docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
 create mode 100644 docker-integration-tests/src/test/scala/org/apache/spark/util/DockerUtils.scala
 create mode 100644 tags/src/main/java/org/apache/spark/tags/DockerTest.java

diff --git a/docker-integration-tests/pom.xml b/docker-integration-tests/pom.xml
new file mode 100644
index 0000000000000..dee0c4aa37ae8
--- /dev/null
+++ b/docker-integration-tests/pom.xml
@@ -0,0 +1,149 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Licensed to the Apache Software Foundation (ASF) under one or more
+  ~ contributor license agreements.  See the NOTICE file distributed with
+  ~ this work for additional information regarding copyright ownership.
+  ~ The ASF licenses this file to You under the Apache License, Version 2.0
+  ~ (the "License"); you may not use this file except in compliance with
+  ~ the License.  You may obtain a copy of the License at
+  ~
+  ~    http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>org.apache.spark</groupId>
+    <artifactId>spark-parent_2.10</artifactId>
+    <version>1.6.0-SNAPSHOT</version>
+    <relativePath>../pom.xml</relativePath>
+  </parent>
+
+  <artifactId>spark-docker-integration-tests_2.10</artifactId>
+  <packaging>jar</packaging>
+  <name>Spark Project Docker Integration Tests</name>
+  <url>http://spark.apache.org/</url>
+  <properties>
+    <sbt.project.name>docker-integration-tests</sbt.project.name>
+  </properties>
+
+  <dependencies>
+    <dependency>
+      <groupId>com.spotify</groupId>
+      <artifactId>docker-client</artifactId>
+      <classifier>shaded</classifier>
+      <scope>test</scope>
+      <!--
+        See https://github.com/spotify/docker-client/pull/272#issuecomment-155249101
+        for an explanation of why these exclusions are (necessarily) a mess.
+      -->
+      <exclusions>
+        <exclusion>
+          <groupId>com.fasterxml.jackson.jaxrs</groupId>
+          <artifactId>jackson-jaxrs-json-provider</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>com.fasterxml.jackson.datatype</groupId>
+          <artifactId>jackson-datatype-guava</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>com.fasterxml.jackson.core</groupId>
+          <artifactId>jackson-databind</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>org.glassfish.jersey.core</groupId>
+          <artifactId>jersey-client</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>org.glassfish.jersey.connectors</groupId>
+          <artifactId>jersey-apache-connector</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>org.glassfish.jersey.media</groupId>
+          <artifactId>jersey-media-json-jackson</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+    <!-- Necessary in order to avoid errors in log messages: -->
+    <dependency>
+      <groupId>com.google.guava</groupId>
+      <artifactId>guava</artifactId>
+      <version>18.0</version>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-core_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-core_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+      <type>test-jar</type>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-sql_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-sql_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+      <type>test-jar</type>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-test-tags_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- Jersey dependencies, used to override version.
+     See https://github.com/apache/spark/pull/9503#issuecomment-154369560 for
+     background on why we need to use a newer Jersey only in this test module;
+     we can remove this once https://github.com/spotify/docker-client/pull/272 is
+     merged and a new docker-client release is published. -->
+    <dependency>
+      <groupId>com.sun.jersey</groupId>
+      <artifactId>jersey-server</artifactId>
+      <version>1.19</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.sun.jersey</groupId>
+      <artifactId>jersey-core</artifactId>
+      <version>1.19</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.sun.jersey</groupId>
+      <artifactId>jersey-servlet</artifactId>
+      <version>1.19</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.sun.jersey</groupId>
+      <artifactId>jersey-json</artifactId>
+      <version>1.19</version>
+      <scope>test</scope>
+      <exclusions>
+        <exclusion>
+          <groupId>stax</groupId>
+          <artifactId>stax-api</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+    <!-- End Jersey dependencies -->
+  </dependencies>
+</project>
diff --git a/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala
new file mode 100644
index 0000000000000..c503c4a13b482
--- /dev/null
+++ b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/DockerJDBCIntegrationSuite.scala
@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import java.net.ServerSocket
+import java.sql.Connection
+
+import scala.collection.JavaConverters._
+import scala.util.control.NonFatal
+
+import com.spotify.docker.client._
+import com.spotify.docker.client.messages.{ContainerConfig, HostConfig, PortBinding}
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.concurrent.Eventually
+import org.scalatest.time.SpanSugar._
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.util.DockerUtils
+import org.apache.spark.sql.test.SharedSQLContext
+
+abstract class DatabaseOnDocker {
+  /**
+   * The docker image to be pulled.
+   */
+  val imageName: String
+
+  /**
+   * Environment variables to set inside of the Docker container while launching it.
+   */
+  val env: Map[String, String]
+
+  /**
+   * The container-internal JDBC port that the database listens on.
+   */
+  val jdbcPort: Int
+
+  /**
+   * Return a JDBC URL that connects to the database running at the given IP address and port.
+   */
+  def getJdbcUrl(ip: String, port: Int): String
+}
+
+abstract class DockerJDBCIntegrationSuite
+  extends SparkFunSuite
+  with BeforeAndAfterAll
+  with Eventually
+  with SharedSQLContext {
+
+  val db: DatabaseOnDocker
+
+  private var docker: DockerClient = _
+  private var containerId: String = _
+  protected var jdbcUrl: String = _
+
+  override def beforeAll() {
+    super.beforeAll()
+    try {
+      docker = DefaultDockerClient.fromEnv.build()
+      // Check that Docker is actually up
+      try {
+        docker.ping()
+      } catch {
+        case NonFatal(e) =>
+          log.error("Exception while connecting to Docker. Check whether Docker is running.")
+          throw e
+      }
+      // Ensure that the Docker image is installed:
+      try {
+        docker.inspectImage(db.imageName)
+      } catch {
+        case e: ImageNotFoundException =>
+          log.warn(s"Docker image ${db.imageName} not found; pulling image from registry")
+          docker.pull(db.imageName)
+      }
+      // Configure networking (necessary for boot2docker / Docker Machine)
+      val externalPort: Int = {
+        val sock = new ServerSocket(0)
+        val port = sock.getLocalPort
+        sock.close()
+        port
+      }
+      val dockerIp = DockerUtils.getDockerIp()
+      val hostConfig: HostConfig = HostConfig.builder()
+        .networkMode("bridge")
+        .portBindings(
+          Map(s"${db.jdbcPort}/tcp" -> List(PortBinding.of(dockerIp, externalPort)).asJava).asJava)
+        .build()
+      // Create the database container:
+      val config = ContainerConfig.builder()
+        .image(db.imageName)
+        .networkDisabled(false)
+        .env(db.env.map { case (k, v) => s"$k=$v" }.toSeq.asJava)
+        .hostConfig(hostConfig)
+        .exposedPorts(s"${db.jdbcPort}/tcp")
+        .build()
+      containerId = docker.createContainer(config).id
+      // Start the container and wait until the database can accept JDBC connections:
+      docker.startContainer(containerId)
+      jdbcUrl = db.getJdbcUrl(dockerIp, externalPort)
+      eventually(timeout(60.seconds), interval(1.seconds)) {
+        val conn = java.sql.DriverManager.getConnection(jdbcUrl)
+        conn.close()
+      }
+      // Run any setup queries:
+      val conn: Connection = java.sql.DriverManager.getConnection(jdbcUrl)
+      try {
+        dataPreparation(conn)
+      } finally {
+        conn.close()
+      }
+    } catch {
+      case NonFatal(e) =>
+        try {
+          afterAll()
+        } finally {
+          throw e
+        }
+    }
+  }
+
+  override def afterAll() {
+    try {
+      if (docker != null) {
+        try {
+          if (containerId != null) {
+            docker.killContainer(containerId)
+            docker.removeContainer(containerId)
+          }
+        } catch {
+          case NonFatal(e) =>
+            logWarning(s"Could not stop container $containerId", e)
+        } finally {
+          docker.close()
+        }
+      }
+    } finally {
+      super.afterAll()
+    }
+  }
+
+  /**
+   * Prepare databases and tables for testing.
+   */
+  def dataPreparation(connection: Connection): Unit
+}
diff --git a/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/MySQLIntegrationSuite.scala b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/MySQLIntegrationSuite.scala
new file mode 100644
index 0000000000000..c68e4dc4933b1
--- /dev/null
+++ b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/MySQLIntegrationSuite.scala
@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import java.math.BigDecimal
+import java.sql.{Connection, Date, Timestamp}
+import java.util.Properties
+
+import org.apache.spark.tags.DockerTest
+
+@DockerTest
+class MySQLIntegrationSuite extends DockerJDBCIntegrationSuite {
+  override val db = new DatabaseOnDocker {
+    override val imageName = "mysql:5.7.9"
+    override val env = Map(
+      "MYSQL_ROOT_PASSWORD" -> "rootpass"
+    )
+    override val jdbcPort: Int = 3306
+    override def getJdbcUrl(ip: String, port: Int): String =
+      s"jdbc:mysql://$ip:$port/mysql?user=root&password=rootpass"
+  }
+
+  override def dataPreparation(conn: Connection): Unit = {
+    conn.prepareStatement("CREATE DATABASE foo").executeUpdate()
+    conn.prepareStatement("CREATE TABLE tbl (x INTEGER, y TEXT(8))").executeUpdate()
+    conn.prepareStatement("INSERT INTO tbl VALUES (42,'fred')").executeUpdate()
+    conn.prepareStatement("INSERT INTO tbl VALUES (17,'dave')").executeUpdate()
+
+    conn.prepareStatement("CREATE TABLE numbers (onebit BIT(1), tenbits BIT(10), "
+      + "small SMALLINT, med MEDIUMINT, nor INT, big BIGINT, deci DECIMAL(40,20), flt FLOAT, "
+      + "dbl DOUBLE)").executeUpdate()
+    conn.prepareStatement("INSERT INTO numbers VALUES (b'0', b'1000100101', "
+      + "17, 77777, 123456789, 123456789012345, 123456789012345.123456789012345, "
+      + "42.75, 1.0000000000000002)").executeUpdate()
+
+    conn.prepareStatement("CREATE TABLE dates (d DATE, t TIME, dt DATETIME, ts TIMESTAMP, "
+      + "yr YEAR)").executeUpdate()
+    conn.prepareStatement("INSERT INTO dates VALUES ('1991-11-09', '13:31:24', "
+      + "'1996-01-01 01:23:45', '2009-02-13 23:31:30', '2001')").executeUpdate()
+
+    // TODO: Test locale conversion for strings.
+    conn.prepareStatement("CREATE TABLE strings (a CHAR(10), b VARCHAR(10), c TINYTEXT, "
+      + "d TEXT, e MEDIUMTEXT, f LONGTEXT, g BINARY(4), h VARBINARY(10), i BLOB)"
+    ).executeUpdate()
+    conn.prepareStatement("INSERT INTO strings VALUES ('the', 'quick', 'brown', 'fox', " +
+      "'jumps', 'over', 'the', 'lazy', 'dog')").executeUpdate()
+  }
+
+  test("Basic test") {
+    val df = sqlContext.read.jdbc(jdbcUrl, "tbl", new Properties)
+    val rows = df.collect()
+    assert(rows.length == 2)
+    val types = rows(0).toSeq.map(x => x.getClass.toString)
+    assert(types.length == 2)
+    assert(types(0).equals("class java.lang.Integer"))
+    assert(types(1).equals("class java.lang.String"))
+  }
+
+  test("Numeric types") {
+    val df = sqlContext.read.jdbc(jdbcUrl, "numbers", new Properties)
+    val rows = df.collect()
+    assert(rows.length == 1)
+    val types = rows(0).toSeq.map(x => x.getClass.toString)
+    assert(types.length == 9)
+    assert(types(0).equals("class java.lang.Boolean"))
+    assert(types(1).equals("class java.lang.Long"))
+    assert(types(2).equals("class java.lang.Integer"))
+    assert(types(3).equals("class java.lang.Integer"))
+    assert(types(4).equals("class java.lang.Integer"))
+    assert(types(5).equals("class java.lang.Long"))
+    assert(types(6).equals("class java.math.BigDecimal"))
+    assert(types(7).equals("class java.lang.Double"))
+    assert(types(8).equals("class java.lang.Double"))
+    assert(rows(0).getBoolean(0) == false)
+    assert(rows(0).getLong(1) == 0x225)
+    assert(rows(0).getInt(2) == 17)
+    assert(rows(0).getInt(3) == 77777)
+    assert(rows(0).getInt(4) == 123456789)
+    assert(rows(0).getLong(5) == 123456789012345L)
+    val bd = new BigDecimal("123456789012345.12345678901234500000")
+    assert(rows(0).getAs[BigDecimal](6).equals(bd))
+    assert(rows(0).getDouble(7) == 42.75)
+    assert(rows(0).getDouble(8) == 1.0000000000000002)
+  }
+
+  test("Date types") {
+    val df = sqlContext.read.jdbc(jdbcUrl, "dates", new Properties)
+    val rows = df.collect()
+    assert(rows.length == 1)
+    val types = rows(0).toSeq.map(x => x.getClass.toString)
+    assert(types.length == 5)
+    assert(types(0).equals("class java.sql.Date"))
+    assert(types(1).equals("class java.sql.Timestamp"))
+    assert(types(2).equals("class java.sql.Timestamp"))
+    assert(types(3).equals("class java.sql.Timestamp"))
+    assert(types(4).equals("class java.sql.Date"))
+    assert(rows(0).getAs[Date](0).equals(Date.valueOf("1991-11-09")))
+    assert(rows(0).getAs[Timestamp](1).equals(Timestamp.valueOf("1970-01-01 13:31:24")))
+    assert(rows(0).getAs[Timestamp](2).equals(Timestamp.valueOf("1996-01-01 01:23:45")))
+    assert(rows(0).getAs[Timestamp](3).equals(Timestamp.valueOf("2009-02-13 23:31:30")))
+    assert(rows(0).getAs[Date](4).equals(Date.valueOf("2001-01-01")))
+  }
+
+  test("String types") {
+    val df = sqlContext.read.jdbc(jdbcUrl, "strings", new Properties)
+    val rows = df.collect()
+    assert(rows.length == 1)
+    val types = rows(0).toSeq.map(x => x.getClass.toString)
+    assert(types.length == 9)
+    assert(types(0).equals("class java.lang.String"))
+    assert(types(1).equals("class java.lang.String"))
+    assert(types(2).equals("class java.lang.String"))
+    assert(types(3).equals("class java.lang.String"))
+    assert(types(4).equals("class java.lang.String"))
+    assert(types(5).equals("class java.lang.String"))
+    assert(types(6).equals("class [B"))
+    assert(types(7).equals("class [B"))
+    assert(types(8).equals("class [B"))
+    assert(rows(0).getString(0).equals("the"))
+    assert(rows(0).getString(1).equals("quick"))
+    assert(rows(0).getString(2).equals("brown"))
+    assert(rows(0).getString(3).equals("fox"))
+    assert(rows(0).getString(4).equals("jumps"))
+    assert(rows(0).getString(5).equals("over"))
+    assert(java.util.Arrays.equals(rows(0).getAs[Array[Byte]](6), Array[Byte](116, 104, 101, 0)))
+    assert(java.util.Arrays.equals(rows(0).getAs[Array[Byte]](7), Array[Byte](108, 97, 122, 121)))
+    assert(java.util.Arrays.equals(rows(0).getAs[Array[Byte]](8), Array[Byte](100, 111, 103)))
+  }
+
+  test("Basic write test") {
+    val df1 = sqlContext.read.jdbc(jdbcUrl, "numbers", new Properties)
+    val df2 = sqlContext.read.jdbc(jdbcUrl, "dates", new Properties)
+    val df3 = sqlContext.read.jdbc(jdbcUrl, "strings", new Properties)
+    df1.write.jdbc(jdbcUrl, "numberscopy", new Properties)
+    df2.write.jdbc(jdbcUrl, "datescopy", new Properties)
+    df3.write.jdbc(jdbcUrl, "stringscopy", new Properties)
+  }
+}
diff --git a/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
new file mode 100644
index 0000000000000..164a7f396280c
--- /dev/null
+++ b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.jdbc
+
+import java.sql.Connection
+import java.util.Properties
+
+import org.apache.spark.tags.DockerTest
+
+@DockerTest
+class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite {
+  override val db = new DatabaseOnDocker {
+    override val imageName = "postgres:9.4.5"
+    override val env = Map(
+      "POSTGRES_PASSWORD" -> "rootpass"
+    )
+    override val jdbcPort = 5432
+    override def getJdbcUrl(ip: String, port: Int): String =
+      s"jdbc:postgresql://$ip:$port/postgres?user=postgres&password=rootpass"
+  }
+
+  override def dataPreparation(conn: Connection): Unit = {
+    conn.prepareStatement("CREATE DATABASE foo").executeUpdate()
+    conn.setCatalog("foo")
+    conn.prepareStatement("CREATE TABLE bar (a text, b integer, c double precision, d bigint, "
+      + "e bit(1), f bit(10), g bytea, h boolean, i inet, j cidr)").executeUpdate()
+    conn.prepareStatement("INSERT INTO bar VALUES ('hello', 42, 1.25, 123456789012345, B'0', "
+      + "B'1000100101', E'\\\\xDEADBEEF', true, '172.16.0.42', '192.168.0.0/16')").executeUpdate()
+  }
+
+  test("Type mapping for various types") {
+    val df = sqlContext.read.jdbc(jdbcUrl, "bar", new Properties)
+    val rows = df.collect()
+    assert(rows.length == 1)
+    val types = rows(0).toSeq.map(x => x.getClass.toString)
+    assert(types.length == 10)
+    assert(types(0).equals("class java.lang.String"))
+    assert(types(1).equals("class java.lang.Integer"))
+    assert(types(2).equals("class java.lang.Double"))
+    assert(types(3).equals("class java.lang.Long"))
+    assert(types(4).equals("class java.lang.Boolean"))
+    assert(types(5).equals("class [B"))
+    assert(types(6).equals("class [B"))
+    assert(types(7).equals("class java.lang.Boolean"))
+    assert(types(8).equals("class java.lang.String"))
+    assert(types(9).equals("class java.lang.String"))
+    assert(rows(0).getString(0).equals("hello"))
+    assert(rows(0).getInt(1) == 42)
+    assert(rows(0).getDouble(2) == 1.25)
+    assert(rows(0).getLong(3) == 123456789012345L)
+    assert(rows(0).getBoolean(4) == false)
+    // BIT(10)'s come back as ASCII strings of ten ASCII 0's and 1's...
+    assert(java.util.Arrays.equals(rows(0).getAs[Array[Byte]](5),
+      Array[Byte](49, 48, 48, 48, 49, 48, 48, 49, 48, 49)))
+    assert(java.util.Arrays.equals(rows(0).getAs[Array[Byte]](6),
+      Array[Byte](0xDE.toByte, 0xAD.toByte, 0xBE.toByte, 0xEF.toByte)))
+    assert(rows(0).getBoolean(7) == true)
+    assert(rows(0).getString(8) == "172.16.0.42")
+    assert(rows(0).getString(9) == "192.168.0.0/16")
+  }
+
+  test("Basic write test") {
+    val df = sqlContext.read.jdbc(jdbcUrl, "bar", new Properties)
+    df.write.jdbc(jdbcUrl, "public.barcopy", new Properties)
+    // Test only that it doesn't crash.
+  }
+}
diff --git a/docker-integration-tests/src/test/scala/org/apache/spark/util/DockerUtils.scala b/docker-integration-tests/src/test/scala/org/apache/spark/util/DockerUtils.scala
new file mode 100644
index 0000000000000..87271776d8564
--- /dev/null
+++ b/docker-integration-tests/src/test/scala/org/apache/spark/util/DockerUtils.scala
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.util
+
+import java.net.{Inet4Address, NetworkInterface, InetAddress}
+
+import scala.collection.JavaConverters._
+import scala.sys.process._
+import scala.util.Try
+
+private[spark] object DockerUtils {
+
+  def getDockerIp(): String = {
+    /** If docker-machine is setup on this box, attempts to find the ip from it. */
+    def findFromDockerMachine(): Option[String] = {
+      sys.env.get("DOCKER_MACHINE_NAME").flatMap { name =>
+        Try(Seq("/bin/bash", "-c", s"docker-machine ip $name 2>/dev/null").!!.trim).toOption
+      }
+    }
+    sys.env.get("DOCKER_IP")
+      .orElse(findFromDockerMachine())
+      .orElse(Try(Seq("/bin/bash", "-c", "boot2docker ip 2>/dev/null").!!.trim).toOption)
+      .getOrElse {
+        // This block of code is based on Utils.findLocalInetAddress(), but is modified to blacklist
+        // certain interfaces.
+        val address = InetAddress.getLocalHost
+        // Address resolves to something like 127.0.1.1, which happens on Debian; try to find
+        // a better address using the local network interfaces
+        // getNetworkInterfaces returns ifs in reverse order compared to ifconfig output order
+        // on unix-like system. On windows, it returns in index order.
+        // It's more proper to pick ip address following system output order.
+        val blackListedIFs = Seq(
+          "vboxnet0",  // Mac
+          "docker0"    // Linux
+        )
+        val activeNetworkIFs = NetworkInterface.getNetworkInterfaces.asScala.toSeq.filter { i =>
+          !blackListedIFs.contains(i.getName)
+        }
+        val reOrderedNetworkIFs = activeNetworkIFs.reverse
+        for (ni <- reOrderedNetworkIFs) {
+          val addresses = ni.getInetAddresses.asScala
+            .filterNot(addr => addr.isLinkLocalAddress || addr.isLoopbackAddress).toSeq
+          if (addresses.nonEmpty) {
+            val addr = addresses.find(_.isInstanceOf[Inet4Address]).getOrElse(addresses.head)
+            // because of Inet6Address.toHostName may add interface at the end if it knows about it
+            val strippedAddress = InetAddress.getByAddress(addr.getAddress)
+            return strippedAddress.getHostAddress
+          }
+        }
+        address.getHostAddress
+      }
+  }
+}
diff --git a/pom.xml b/pom.xml
index fd8c773513881..c499a80aa0f43 100644
--- a/pom.xml
+++ b/pom.xml
@@ -98,6 +98,7 @@
     <module>sql/catalyst</module>
     <module>sql/core</module>
     <module>sql/hive</module>
+    <module>docker-integration-tests</module>
     <module>unsafe</module>
     <module>assembly</module>
     <module>external/twitter</module>
@@ -778,6 +779,19 @@
         <version>0.11</version>
         <scope>test</scope>
       </dependency>
+      <dependency>
+        <groupId>com.spotify</groupId>
+        <artifactId>docker-client</artifactId>
+        <classifier>shaded</classifier>
+        <version>3.2.1</version>
+        <scope>test</scope>
+        <exclusions>
+          <exclusion>
+            <artifactId>guava</artifactId>
+            <groupId>com.google.guava</groupId>
+          </exclusion>
+        </exclusions>
+      </dependency>
       <dependency>
         <groupId>org.apache.curator</groupId>
         <artifactId>curator-recipes</artifactId>
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index a9fb741d75933..b7c619224329f 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -43,8 +43,9 @@ object BuildCommons {
       "streaming-zeromq", "launcher", "unsafe", "test-tags").map(ProjectRef(buildLocation, _))
 
   val optionallyEnabledProjects@Seq(yarn, yarnStable, java8Tests, sparkGangliaLgpl,
-    streamingKinesisAsl) = Seq("yarn", "yarn-stable", "java8-tests", "ganglia-lgpl",
-    "streaming-kinesis-asl").map(ProjectRef(buildLocation, _))
+    streamingKinesisAsl, dockerIntegrationTests) =
+    Seq("yarn", "yarn-stable", "java8-tests", "ganglia-lgpl", "streaming-kinesis-asl",
+      "docker-integration-tests").map(ProjectRef(buildLocation, _))
 
   val assemblyProjects@Seq(assembly, examples, networkYarn, streamingFlumeAssembly, streamingKafkaAssembly, streamingMqttAssembly, streamingKinesisAslAssembly) =
     Seq("assembly", "examples", "network-yarn", "streaming-flume-assembly", "streaming-kafka-assembly", "streaming-mqtt-assembly", "streaming-kinesis-asl-assembly")
@@ -240,6 +241,8 @@ object SparkBuild extends PomBuild {
 
   enable(Flume.settings)(streamingFlumeSink)
 
+  enable(DockerIntegrationTests.settings)(dockerIntegrationTests)
+
 
   /**
    * Adds the ability to run the spark shell directly from SBT without building an assembly
@@ -291,6 +294,13 @@ object Flume {
   lazy val settings = sbtavro.SbtAvro.avroSettings
 }
 
+object DockerIntegrationTests {
+  // This serves to override the override specified in DependencyOverrides:
+  lazy val settings = Seq(
+    dependencyOverrides += "com.google.guava" % "guava" % "18.0"
+  )
+}
+
 /**
  * Overrides to work around sbt's dependency resolution being different from Maven's.
  */
diff --git a/tags/src/main/java/org/apache/spark/tags/DockerTest.java b/tags/src/main/java/org/apache/spark/tags/DockerTest.java
new file mode 100644
index 0000000000000..0fecf3b8f979a
--- /dev/null
+++ b/tags/src/main/java/org/apache/spark/tags/DockerTest.java
@@ -0,0 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.tags;
+
+import java.lang.annotation.*;
+import org.scalatest.TagAnnotation;
+
+@TagAnnotation
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.METHOD, ElementType.TYPE})
+public @interface DockerTest { }

From e281b87398f1298cc3df8e0409c7040acdddce03 Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Tue, 10 Nov 2015 16:20:10 -0800
Subject: [PATCH 169/518] [SPARK-5565][ML] LDA wrapper for Pipelines API

This adds LDA to spark.ml, the Pipelines API.  It follows the design doc in the JIRA: [https://issues.apache.org/jira/browse/SPARK-5565], with one major change:
* I eliminated doc IDs.  These are not necessary with DataFrames since the user can add an ID column as needed.

Note: This will conflict with [https://github.com/apache/spark/pull/9484], but I'll try to merge [https://github.com/apache/spark/pull/9484] first and then rebase this PR.

CC: hhbyyh feynmanliang  If you have a chance to make a pass, that'd be really helpful--thanks!  Now that I'm done traveling & this PR is almost ready, I'll see about reviewing other PRs critical for 1.6.

CC: mengxr

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9513 from jkbradley/lda-pipelines.
---
 .../org/apache/spark/ml/clustering/LDA.scala  | 701 ++++++++++++++++++
 .../spark/mllib/clustering/LDAModel.scala     |  29 +-
 .../apache/spark/ml/clustering/LDASuite.scala | 221 ++++++
 3 files changed, 946 insertions(+), 5 deletions(-)
 create mode 100644 mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
 create mode 100644 mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala

diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
new file mode 100644
index 0000000000000..f66233ed3d0f0
--- /dev/null
+++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
@@ -0,0 +1,701 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.clustering
+
+import org.apache.spark.Logging
+import org.apache.spark.annotation.{Experimental, Since}
+import org.apache.spark.ml.util.{SchemaUtils, Identifiable}
+import org.apache.spark.ml.{Estimator, Model}
+import org.apache.spark.ml.param.shared.{HasCheckpointInterval, HasFeaturesCol, HasSeed, HasMaxIter}
+import org.apache.spark.ml.param._
+import org.apache.spark.mllib.clustering.{DistributedLDAModel => OldDistributedLDAModel,
+    EMLDAOptimizer => OldEMLDAOptimizer, LDA => OldLDA, LDAModel => OldLDAModel,
+    LDAOptimizer => OldLDAOptimizer, LocalLDAModel => OldLocalLDAModel,
+    OnlineLDAOptimizer => OldOnlineLDAOptimizer}
+import org.apache.spark.mllib.linalg.{VectorUDT, Vectors, Matrix, Vector}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.{SQLContext, DataFrame, Row}
+import org.apache.spark.sql.functions.{col, monotonicallyIncreasingId, udf}
+import org.apache.spark.sql.types.StructType
+
+
+private[clustering] trait LDAParams extends Params with HasFeaturesCol with HasMaxIter
+  with HasSeed with HasCheckpointInterval {
+
+  /**
+   * Param for the number of topics (clusters) to infer. Must be > 1. Default: 10.
+   * @group param
+   */
+  @Since("1.6.0")
+  final val k = new IntParam(this, "k", "number of topics (clusters) to infer",
+    ParamValidators.gt(1))
+
+  /** @group getParam */
+  @Since("1.6.0")
+  def getK: Int = $(k)
+
+  /**
+   * Concentration parameter (commonly named "alpha") for the prior placed on documents'
+   * distributions over topics ("theta").
+   *
+   * This is the parameter to a Dirichlet distribution, where larger values mean more smoothing
+   * (more regularization).
+   *
+   * If not set by the user, then docConcentration is set automatically. If set to
+   * singleton vector [alpha], then alpha is replicated to a vector of length k in fitting.
+   * Otherwise, the [[docConcentration]] vector must be length k.
+   * (default = automatic)
+   *
+   * Optimizer-specific parameter settings:
+   *  - EM
+   *     - Currently only supports symmetric distributions, so all values in the vector should be
+   *       the same.
+   *     - Values should be > 1.0
+   *     - default = uniformly (50 / k) + 1, where 50/k is common in LDA libraries and +1 follows
+   *       from Asuncion et al. (2009), who recommend a +1 adjustment for EM.
+   *  - Online
+   *     - Values should be >= 0
+   *     - default = uniformly (1.0 / k), following the implementation from
+   *       [[https://github.com/Blei-Lab/onlineldavb]].
+   * @group param
+   */
+  @Since("1.6.0")
+  final val docConcentration = new DoubleArrayParam(this, "docConcentration",
+    "Concentration parameter (commonly named \"alpha\") for the prior placed on documents'" +
+      " distributions over topics (\"theta\").", (alpha: Array[Double]) => alpha.forall(_ >= 0.0))
+
+  /** @group getParam */
+  @Since("1.6.0")
+  def getDocConcentration: Array[Double] = $(docConcentration)
+
+  /** Get docConcentration used by spark.mllib LDA */
+  protected def getOldDocConcentration: Vector = {
+    if (isSet(docConcentration)) {
+      Vectors.dense(getDocConcentration)
+    } else {
+      Vectors.dense(-1.0)
+    }
+  }
+
+  /**
+   * Concentration parameter (commonly named "beta" or "eta") for the prior placed on topics'
+   * distributions over terms.
+   *
+   * This is the parameter to a symmetric Dirichlet distribution.
+   *
+   * Note: The topics' distributions over terms are called "beta" in the original LDA paper
+   * by Blei et al., but are called "phi" in many later papers such as Asuncion et al., 2009.
+   *
+   * If not set by the user, then topicConcentration is set automatically.
+   *  (default = automatic)
+   *
+   * Optimizer-specific parameter settings:
+   *  - EM
+   *     - Value should be > 1.0
+   *     - default = 0.1 + 1, where 0.1 gives a small amount of smoothing and +1 follows
+   *       Asuncion et al. (2009), who recommend a +1 adjustment for EM.
+   *  - Online
+   *     - Value should be >= 0
+   *     - default = (1.0 / k), following the implementation from
+   *       [[https://github.com/Blei-Lab/onlineldavb]].
+   * @group param
+   */
+  @Since("1.6.0")
+  final val topicConcentration = new DoubleParam(this, "topicConcentration",
+    "Concentration parameter (commonly named \"beta\" or \"eta\") for the prior placed on topic'" +
+      " distributions over terms.", ParamValidators.gtEq(0))
+
+  /** @group getParam */
+  @Since("1.6.0")
+  def getTopicConcentration: Double = $(topicConcentration)
+
+  /** Get topicConcentration used by spark.mllib LDA */
+  protected def getOldTopicConcentration: Double = {
+    if (isSet(topicConcentration)) {
+      getTopicConcentration
+    } else {
+      -1.0
+    }
+  }
+
+  /** Supported values for Param [[optimizer]]. */
+  @Since("1.6.0")
+  final val supportedOptimizers: Array[String] = Array("online", "em")
+
+  /**
+   * Optimizer or inference algorithm used to estimate the LDA model.
+   * Currently supported (case-insensitive):
+   *  - "online": Online Variational Bayes (default)
+   *  - "em": Expectation-Maximization
+   *
+   * For details, see the following papers:
+   *  - Online LDA:
+   *     Hoffman, Blei and Bach.  "Online Learning for Latent Dirichlet Allocation."
+   *     Neural Information Processing Systems, 2010.
+   *     [[http://www.cs.columbia.edu/~blei/papers/HoffmanBleiBach2010b.pdf]]
+   *  - EM:
+   *     Asuncion et al.  "On Smoothing and Inference for Topic Models."
+   *     Uncertainty in Artificial Intelligence, 2009.
+   *     [[http://arxiv.org/pdf/1205.2662.pdf]]
+   *
+   * @group param
+   */
+  @Since("1.6.0")
+  final val optimizer = new Param[String](this, "optimizer", "Optimizer or inference" +
+    " algorithm used to estimate the LDA model.  Supported: " + supportedOptimizers.mkString(", "),
+    (o: String) => ParamValidators.inArray(supportedOptimizers).apply(o.toLowerCase))
+
+  /** @group getParam */
+  @Since("1.6.0")
+  def getOptimizer: String = $(optimizer)
+
+  /**
+   * Output column with estimates of the topic mixture distribution for each document (often called
+   * "theta" in the literature).  Returns a vector of zeros for an empty document.
+   *
+   * This uses a variational approximation following Hoffman et al. (2010), where the approximate
+   * distribution is called "gamma."  Technically, this method returns this approximation "gamma"
+   * for each document.
+   * @group param
+   */
+  @Since("1.6.0")
+  final val topicDistributionCol = new Param[String](this, "topicDistribution", "Output column" +
+    " with estimates of the topic mixture distribution for each document (often called \"theta\"" +
+    " in the literature).  Returns a vector of zeros for an empty document.")
+
+  setDefault(topicDistributionCol -> "topicDistribution")
+
+  /** @group getParam */
+  @Since("1.6.0")
+  def getTopicDistributionCol: String = $(topicDistributionCol)
+
+  /**
+   * A (positive) learning parameter that downweights early iterations. Larger values make early
+   * iterations count less.
+   * This is called "tau0" in the Online LDA paper (Hoffman et al., 2010)
+   * Default: 1024, following Hoffman et al.
+   * @group expertParam
+   */
+  @Since("1.6.0")
+  final val learningOffset = new DoubleParam(this, "learningOffset", "A (positive) learning" +
+    " parameter that downweights early iterations. Larger values make early iterations count less.",
+    ParamValidators.gt(0))
+
+  /** @group expertGetParam */
+  @Since("1.6.0")
+  def getLearningOffset: Double = $(learningOffset)
+
+  /**
+   * Learning rate, set as an exponential decay rate.
+   * This should be between (0.5, 1.0] to guarantee asymptotic convergence.
+   * This is called "kappa" in the Online LDA paper (Hoffman et al., 2010).
+   * Default: 0.51, based on Hoffman et al.
+   * @group expertParam
+   */
+  @Since("1.6.0")
+  final val learningDecay = new DoubleParam(this, "learningDecay", "Learning rate, set as an" +
+    " exponential decay rate. This should be between (0.5, 1.0] to guarantee asymptotic" +
+    " convergence.", ParamValidators.gt(0))
+
+  /** @group expertGetParam */
+  @Since("1.6.0")
+  def getLearningDecay: Double = $(learningDecay)
+
+  /**
+   * Fraction of the corpus to be sampled and used in each iteration of mini-batch gradient descent,
+   * in range (0, 1].
+   *
+   * Note that this should be adjusted in synch with [[LDA.maxIter]]
+   * so the entire corpus is used.  Specifically, set both so that
+   * maxIterations * miniBatchFraction >= 1.
+   *
+   * Note: This is the same as the `miniBatchFraction` parameter in
+   *       [[org.apache.spark.mllib.clustering.OnlineLDAOptimizer]].
+   *
+   * Default: 0.05, i.e., 5% of total documents.
+   * @group param
+   */
+  @Since("1.6.0")
+  final val subsamplingRate = new DoubleParam(this, "subsamplingRate", "Fraction of the corpus" +
+    " to be sampled and used in each iteration of mini-batch gradient descent, in range (0, 1].",
+    ParamValidators.inRange(0.0, 1.0, lowerInclusive = false, upperInclusive = true))
+
+  /** @group getParam */
+  @Since("1.6.0")
+  def getSubsamplingRate: Double = $(subsamplingRate)
+
+  /**
+   * Indicates whether the docConcentration (Dirichlet parameter for
+   * document-topic distribution) will be optimized during training.
+   * Setting this to true will make the model more expressive and fit the training data better.
+   * Default: false
+   * @group expertParam
+   */
+  @Since("1.6.0")
+  final val optimizeDocConcentration = new BooleanParam(this, "optimizeDocConcentration",
+    "Indicates whether the docConcentration (Dirichlet parameter for document-topic" +
+      " distribution) will be optimized during training.")
+
+  /** @group expertGetParam */
+  @Since("1.6.0")
+  def getOptimizeDocConcentration: Boolean = $(optimizeDocConcentration)
+
+  /**
+   * Validates and transforms the input schema.
+   * @param schema input schema
+   * @return output schema
+   */
+  protected def validateAndTransformSchema(schema: StructType): StructType = {
+    SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT)
+    SchemaUtils.appendColumn(schema, $(topicDistributionCol), new VectorUDT)
+  }
+
+  @Since("1.6.0")
+  override def validateParams(): Unit = {
+    if (isSet(docConcentration)) {
+      if (getDocConcentration.length != 1) {
+        require(getDocConcentration.length == getK, s"LDA docConcentration was of length" +
+          s" ${getDocConcentration.length}, but k = $getK.  docConcentration must be an array of" +
+          s" length either 1 (scalar) or k (num topics).")
+      }
+      getOptimizer match {
+        case "online" =>
+          require(getDocConcentration.forall(_ >= 0),
+            "For Online LDA optimizer, docConcentration values must be >= 0.  Found values: " +
+              getDocConcentration.mkString(","))
+        case "em" =>
+          require(getDocConcentration.forall(_ >= 0),
+            "For EM optimizer, docConcentration values must be >= 1.  Found values: " +
+              getDocConcentration.mkString(","))
+      }
+    }
+    if (isSet(topicConcentration)) {
+      getOptimizer match {
+        case "online" =>
+          require(getTopicConcentration >= 0, s"For Online LDA optimizer, topicConcentration" +
+            s" must be >= 0.  Found value: $getTopicConcentration")
+        case "em" =>
+          require(getTopicConcentration >= 0, s"For EM optimizer, topicConcentration" +
+            s" must be >= 1.  Found value: $getTopicConcentration")
+      }
+    }
+  }
+
+  private[clustering] def getOldOptimizer: OldLDAOptimizer = getOptimizer match {
+    case "online" =>
+      new OldOnlineLDAOptimizer()
+        .setTau0($(learningOffset))
+        .setKappa($(learningDecay))
+        .setMiniBatchFraction($(subsamplingRate))
+        .setOptimizeDocConcentration($(optimizeDocConcentration))
+    case "em" =>
+      new OldEMLDAOptimizer()
+  }
+}
+
+
+/**
+ * :: Experimental ::
+ * Model fitted by [[LDA]].
+ *
+ * @param vocabSize  Vocabulary size (number of terms or terms in the vocabulary)
+ * @param oldLocalModel  Underlying spark.mllib model.
+ *                       If this model was produced by Online LDA, then this is the
+ *                       only model representation.
+ *                       If this model was produced by EM, then this local
+ *                       representation may be built lazily.
+ * @param sqlContext  Used to construct local DataFrames for returning query results
+ */
+@Since("1.6.0")
+@Experimental
+class LDAModel private[ml] (
+    @Since("1.6.0") override val uid: String,
+    @Since("1.6.0") val vocabSize: Int,
+    @Since("1.6.0") protected var oldLocalModel: Option[OldLocalLDAModel],
+    @Since("1.6.0") @transient protected val sqlContext: SQLContext)
+  extends Model[LDAModel] with LDAParams with Logging {
+
+  /** Returns underlying spark.mllib model */
+  @Since("1.6.0")
+  protected def getModel: OldLDAModel = oldLocalModel match {
+    case Some(m) => m
+    case None =>
+      // Should never happen.
+      throw new RuntimeException("LDAModel required local model format," +
+        " but the underlying model is missing.")
+  }
+
+  /**
+   * The features for LDA should be a [[Vector]] representing the word counts in a document.
+   * The vector should be of length vocabSize, with counts for each term (word).
+   * @group setParam
+   */
+  @Since("1.6.0")
+  def setFeaturesCol(value: String): this.type = set(featuresCol, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setSeed(value: Long): this.type = set(seed, value)
+
+  @Since("1.6.0")
+  override def copy(extra: ParamMap): LDAModel = {
+    val copied = new LDAModel(uid, vocabSize, oldLocalModel, sqlContext)
+    copyValues(copied, extra).setParent(parent)
+  }
+
+  @Since("1.6.0")
+  override def transform(dataset: DataFrame): DataFrame = {
+    if ($(topicDistributionCol).nonEmpty) {
+      val t = udf(oldLocalModel.get.getTopicDistributionMethod(sqlContext.sparkContext))
+      dataset.withColumn($(topicDistributionCol), t(col($(featuresCol))))
+    } else {
+      logWarning("LDAModel.transform was called without any output columns. Set an output column" +
+        " such as topicDistributionCol to produce results.")
+      dataset
+    }
+  }
+
+  @Since("1.6.0")
+  override def transformSchema(schema: StructType): StructType = {
+    validateAndTransformSchema(schema)
+  }
+
+  /**
+   * Value for [[docConcentration]] estimated from data.
+   * If Online LDA was used and [[optimizeDocConcentration]] was set to false,
+   * then this returns the fixed (given) value for the [[docConcentration]] parameter.
+   */
+  @Since("1.6.0")
+  def estimatedDocConcentration: Vector = getModel.docConcentration
+
+  /**
+   * Inferred topics, where each topic is represented by a distribution over terms.
+   * This is a matrix of size vocabSize x k, where each column is a topic.
+   * No guarantees are given about the ordering of the topics.
+   *
+   * WARNING: If this model is actually a [[DistributedLDAModel]] instance from EM,
+   *          then this method could involve collecting a large amount of data to the driver
+   *          (on the order of vocabSize x k).
+   */
+  @Since("1.6.0")
+  def topicsMatrix: Matrix = getModel.topicsMatrix
+
+  /** Indicates whether this instance is of type [[DistributedLDAModel]] */
+  @Since("1.6.0")
+  def isDistributed: Boolean = false
+
+  /**
+   * Calculates a lower bound on the log likelihood of the entire corpus.
+   *
+   * See Equation (16) in the Online LDA paper (Hoffman et al., 2010).
+   *
+   * WARNING: If this model was learned via a [[DistributedLDAModel]], this involves collecting
+   *          a large [[topicsMatrix]] to the driver.  This implementation may be changed in the
+   *          future.
+   *
+   * @param dataset  test corpus to use for calculating log likelihood
+   * @return variational lower bound on the log likelihood of the entire corpus
+   */
+  @Since("1.6.0")
+  def logLikelihood(dataset: DataFrame): Double = oldLocalModel match {
+    case Some(m) =>
+      val oldDataset = LDA.getOldDataset(dataset, $(featuresCol))
+      m.logLikelihood(oldDataset)
+    case None =>
+      // Should never happen.
+      throw new RuntimeException("LocalLDAModel.logLikelihood was called," +
+        " but the underlying model is missing.")
+  }
+
+  /**
+   * Calculate an upper bound bound on perplexity.  (Lower is better.)
+   * See Equation (16) in the Online LDA paper (Hoffman et al., 2010).
+   *
+   * @param dataset test corpus to use for calculating perplexity
+   * @return Variational upper bound on log perplexity per token.
+   */
+  @Since("1.6.0")
+  def logPerplexity(dataset: DataFrame): Double = oldLocalModel match {
+    case Some(m) =>
+      val oldDataset = LDA.getOldDataset(dataset, $(featuresCol))
+      m.logPerplexity(oldDataset)
+    case None =>
+      // Should never happen.
+      throw new RuntimeException("LocalLDAModel.logPerplexity was called," +
+        " but the underlying model is missing.")
+  }
+
+  /**
+   * Return the topics described by their top-weighted terms.
+   *
+   * @param maxTermsPerTopic  Maximum number of terms to collect for each topic.
+   *                          Default value of 10.
+   * @return  Local DataFrame with one topic per Row, with columns:
+   *           - "topic": IntegerType: topic index
+   *           - "termIndices": ArrayType(IntegerType): term indices, sorted in order of decreasing
+   *                            term importance
+   *           - "termWeights": ArrayType(DoubleType): corresponding sorted term weights
+   */
+  @Since("1.6.0")
+  def describeTopics(maxTermsPerTopic: Int): DataFrame = {
+    val topics = getModel.describeTopics(maxTermsPerTopic).zipWithIndex.map {
+      case ((termIndices, termWeights), topic) =>
+        (topic, termIndices.toSeq, termWeights.toSeq)
+    }
+    sqlContext.createDataFrame(topics).toDF("topic", "termIndices", "termWeights")
+  }
+
+  @Since("1.6.0")
+  def describeTopics(): DataFrame = describeTopics(10)
+}
+
+
+/**
+ * :: Experimental ::
+ *
+ * Distributed model fitted by [[LDA]] using Expectation-Maximization (EM).
+ *
+ * This model stores the inferred topics, the full training dataset, and the topic distribution
+ * for each training document.
+ */
+@Since("1.6.0")
+@Experimental
+class DistributedLDAModel private[ml] (
+    uid: String,
+    vocabSize: Int,
+    private val oldDistributedModel: OldDistributedLDAModel,
+    sqlContext: SQLContext)
+  extends LDAModel(uid, vocabSize, None, sqlContext) {
+
+  /**
+   * Convert this distributed model to a local representation.  This discards info about the
+   * training dataset.
+   */
+  @Since("1.6.0")
+  def toLocal: LDAModel = {
+    if (oldLocalModel.isEmpty) {
+      oldLocalModel = Some(oldDistributedModel.toLocal)
+    }
+    new LDAModel(uid, vocabSize, oldLocalModel, sqlContext)
+  }
+
+  @Since("1.6.0")
+  override protected def getModel: OldLDAModel = oldDistributedModel
+
+  @Since("1.6.0")
+  override def copy(extra: ParamMap): DistributedLDAModel = {
+    val copied = new DistributedLDAModel(uid, vocabSize, oldDistributedModel, sqlContext)
+    if (oldLocalModel.nonEmpty) copied.oldLocalModel = oldLocalModel
+    copyValues(copied, extra).setParent(parent)
+    copied
+  }
+
+  @Since("1.6.0")
+  override def topicsMatrix: Matrix = {
+    if (oldLocalModel.isEmpty) {
+      oldLocalModel = Some(oldDistributedModel.toLocal)
+    }
+    super.topicsMatrix
+  }
+
+  @Since("1.6.0")
+  override def isDistributed: Boolean = true
+
+  @Since("1.6.0")
+  override def logLikelihood(dataset: DataFrame): Double = {
+    if (oldLocalModel.isEmpty) {
+      oldLocalModel = Some(oldDistributedModel.toLocal)
+    }
+    super.logLikelihood(dataset)
+  }
+
+  @Since("1.6.0")
+  override def logPerplexity(dataset: DataFrame): Double = {
+    if (oldLocalModel.isEmpty) {
+      oldLocalModel = Some(oldDistributedModel.toLocal)
+    }
+    super.logPerplexity(dataset)
+  }
+
+  /**
+   * Log likelihood of the observed tokens in the training set,
+   * given the current parameter estimates:
+   *  log P(docs | topics, topic distributions for docs, Dirichlet hyperparameters)
+   *
+   * Notes:
+   *  - This excludes the prior; for that, use [[logPrior]].
+   *  - Even with [[logPrior]], this is NOT the same as the data log likelihood given the
+   *    hyperparameters.
+   *  - This is computed from the topic distributions computed during training. If you call
+   *    [[logLikelihood()]] on the same training dataset, the topic distributions will be computed
+   *    again, possibly giving different results.
+   */
+  @Since("1.6.0")
+  lazy val trainingLogLikelihood: Double = oldDistributedModel.logLikelihood
+
+  /**
+   * Log probability of the current parameter estimate:
+   * log P(topics, topic distributions for docs | Dirichlet hyperparameters)
+   */
+  @Since("1.6.0")
+  lazy val logPrior: Double = oldDistributedModel.logPrior
+}
+
+
+/**
+ * :: Experimental ::
+ *
+ * Latent Dirichlet Allocation (LDA), a topic model designed for text documents.
+ *
+ * Terminology:
+ *  - "term" = "word": an element of the vocabulary
+ *  - "token": instance of a term appearing in a document
+ *  - "topic": multinomial distribution over terms representing some concept
+ *  - "document": one piece of text, corresponding to one row in the input data
+ *
+ * References:
+ *  - Original LDA paper (journal version):
+ *    Blei, Ng, and Jordan.  "Latent Dirichlet Allocation."  JMLR, 2003.
+ *
+ * Input data (featuresCol):
+ *  LDA is given a collection of documents as input data, via the featuresCol parameter.
+ *  Each document is specified as a [[Vector]] of length vocabSize, where each entry is the
+ *  count for the corresponding term (word) in the document.  Feature transformers such as
+ *  [[org.apache.spark.ml.feature.Tokenizer]] and [[org.apache.spark.ml.feature.CountVectorizer]]
+ *  can be useful for converting text to word count vectors.
+ *
+ * @see [[http://en.wikipedia.org/wiki/Latent_Dirichlet_allocation Latent Dirichlet allocation
+ *       (Wikipedia)]]
+ */
+@Since("1.6.0")
+@Experimental
+class LDA @Since("1.6.0") (
+    @Since("1.6.0") override val uid: String) extends Estimator[LDAModel] with LDAParams {
+
+  @Since("1.6.0")
+  def this() = this(Identifiable.randomUID("lda"))
+
+  setDefault(maxIter -> 20, k -> 10, optimizer -> "online", checkpointInterval -> 10,
+    learningOffset -> 1024, learningDecay -> 0.51, subsamplingRate -> 0.05,
+    optimizeDocConcentration -> true)
+
+  /**
+   * The features for LDA should be a [[Vector]] representing the word counts in a document.
+   * The vector should be of length vocabSize, with counts for each term (word).
+   * @group setParam
+   */
+  @Since("1.6.0")
+  def setFeaturesCol(value: String): this.type = set(featuresCol, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setMaxIter(value: Int): this.type = set(maxIter, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setSeed(value: Long): this.type = set(seed, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setCheckpointInterval(value: Int): this.type = set(checkpointInterval, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setK(value: Int): this.type = set(k, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setDocConcentration(value: Array[Double]): this.type = set(docConcentration, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setDocConcentration(value: Double): this.type = set(docConcentration, Array(value))
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setTopicConcentration(value: Double): this.type = set(topicConcentration, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setOptimizer(value: String): this.type = set(optimizer, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setTopicDistributionCol(value: String): this.type = set(topicDistributionCol, value)
+
+  /** @group expertSetParam */
+  @Since("1.6.0")
+  def setLearningOffset(value: Double): this.type = set(learningOffset, value)
+
+  /** @group expertSetParam */
+  @Since("1.6.0")
+  def setLearningDecay(value: Double): this.type = set(learningDecay, value)
+
+  /** @group setParam */
+  @Since("1.6.0")
+  def setSubsamplingRate(value: Double): this.type = set(subsamplingRate, value)
+
+  /** @group expertSetParam */
+  @Since("1.6.0")
+  def setOptimizeDocConcentration(value: Boolean): this.type = set(optimizeDocConcentration, value)
+
+  @Since("1.6.0")
+  override def copy(extra: ParamMap): LDA = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def fit(dataset: DataFrame): LDAModel = {
+    transformSchema(dataset.schema, logging = true)
+    val oldLDA = new OldLDA()
+      .setK($(k))
+      .setDocConcentration(getOldDocConcentration)
+      .setTopicConcentration(getOldTopicConcentration)
+      .setMaxIterations($(maxIter))
+      .setSeed($(seed))
+      .setCheckpointInterval($(checkpointInterval))
+      .setOptimizer(getOldOptimizer)
+    // TODO: persist here, or in old LDA?
+    val oldData = LDA.getOldDataset(dataset, $(featuresCol))
+    val oldModel = oldLDA.run(oldData)
+    val newModel = oldModel match {
+      case m: OldLocalLDAModel =>
+        new LDAModel(uid, m.vocabSize, Some(m), dataset.sqlContext)
+      case m: OldDistributedLDAModel =>
+        new DistributedLDAModel(uid, m.vocabSize, m, dataset.sqlContext)
+    }
+    copyValues(newModel).setParent(this)
+  }
+
+  @Since("1.6.0")
+  override def transformSchema(schema: StructType): StructType = {
+    validateAndTransformSchema(schema)
+  }
+}
+
+
+private[clustering] object LDA {
+
+  /** Get dataset for spark.mllib LDA */
+  def getOldDataset(dataset: DataFrame, featuresCol: String): RDD[(Long, Vector)] = {
+    dataset
+      .withColumn("docId", monotonicallyIncreasingId())
+      .select("docId", featuresCol)
+      .map { case Row(docId: Long, features: Vector) =>
+        (docId, features)
+      }
+  }
+}
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala
index 31d8a9fdea1c6..cd520f09bd466 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala
@@ -183,8 +183,7 @@ abstract class LDAModel private[clustering] extends Saveable {
 /**
  * Local LDA model.
  * This model stores only the inferred topics.
- * It may be used for computing topics for new documents, but it may give less accurate answers
- * than the [[DistributedLDAModel]].
+ *
  * @param topics Inferred topics (vocabSize x k matrix).
  */
 @Since("1.3.0")
@@ -353,7 +352,7 @@ class LocalLDAModel private[clustering] (
 
     documents.map { case (id: Long, termCounts: Vector) =>
       if (termCounts.numNonzeros == 0) {
-         (id, Vectors.zeros(k))
+        (id, Vectors.zeros(k))
       } else {
         val (gamma, _) = OnlineLDAOptimizer.variationalTopicInference(
           termCounts,
@@ -366,6 +365,28 @@ class LocalLDAModel private[clustering] (
     }
   }
 
+  /** Get a method usable as a UDF for [[topicDistributions()]] */
+  private[spark] def getTopicDistributionMethod(sc: SparkContext): Vector => Vector = {
+    val expElogbeta = exp(LDAUtils.dirichletExpectation(topicsMatrix.toBreeze.toDenseMatrix.t).t)
+    val expElogbetaBc = sc.broadcast(expElogbeta)
+    val docConcentrationBrz = this.docConcentration.toBreeze
+    val gammaShape = this.gammaShape
+    val k = this.k
+
+    (termCounts: Vector) =>
+      if (termCounts.numNonzeros == 0) {
+        Vectors.zeros(k)
+      } else {
+        val (gamma, _) = OnlineLDAOptimizer.variationalTopicInference(
+          termCounts,
+          expElogbetaBc.value,
+          docConcentrationBrz,
+          gammaShape,
+          k)
+        Vectors.dense(normalize(gamma, 1.0).toArray)
+      }
+  }
+
   /**
    * Java-friendly version of [[topicDistributions]]
    */
@@ -477,8 +498,6 @@ object LocalLDAModel extends Loader[LocalLDAModel] {
 /**
  * Distributed LDA model.
  * This model stores the inferred topics, the full training dataset, and the topic distributions.
- * When computing topics for new documents, it may give more accurate answers
- * than the [[LocalLDAModel]].
  */
 @Since("1.3.0")
 class DistributedLDAModel private[clustering] (
diff --git a/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala b/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala
new file mode 100644
index 0000000000000..edb927495e8bf
--- /dev/null
+++ b/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala
@@ -0,0 +1,221 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.clustering
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.mllib.linalg.{Vector, Vectors}
+import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.sql.{DataFrame, Row, SQLContext}
+
+
+object LDASuite {
+  def generateLDAData(
+      sql: SQLContext,
+      rows: Int,
+      k: Int,
+      vocabSize: Int): DataFrame = {
+    val avgWC = 1  // average instances of each word in a doc
+    val sc = sql.sparkContext
+    val rng = new java.util.Random()
+    rng.setSeed(1)
+    val rdd = sc.parallelize(1 to rows).map { i =>
+      Vectors.dense(Array.fill(vocabSize)(rng.nextInt(2 * avgWC).toDouble))
+    }.map(v => new TestRow(v))
+    sql.createDataFrame(rdd)
+  }
+}
+
+
+class LDASuite extends SparkFunSuite with MLlibTestSparkContext {
+
+  val k: Int = 5
+  val vocabSize: Int = 30
+  @transient var dataset: DataFrame = _
+
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+    dataset = LDASuite.generateLDAData(sqlContext, 50, k, vocabSize)
+  }
+
+  test("default parameters") {
+    val lda = new LDA()
+
+    assert(lda.getFeaturesCol === "features")
+    assert(lda.getMaxIter === 20)
+    assert(lda.isDefined(lda.seed))
+    assert(lda.getCheckpointInterval === 10)
+    assert(lda.getK === 10)
+    assert(!lda.isSet(lda.docConcentration))
+    assert(!lda.isSet(lda.topicConcentration))
+    assert(lda.getOptimizer === "online")
+    assert(lda.getLearningDecay === 0.51)
+    assert(lda.getLearningOffset === 1024)
+    assert(lda.getSubsamplingRate === 0.05)
+    assert(lda.getOptimizeDocConcentration)
+    assert(lda.getTopicDistributionCol === "topicDistribution")
+  }
+
+  test("set parameters") {
+    val lda = new LDA()
+      .setFeaturesCol("test_feature")
+      .setMaxIter(33)
+      .setSeed(123)
+      .setCheckpointInterval(7)
+      .setK(9)
+      .setTopicConcentration(0.56)
+      .setTopicDistributionCol("myOutput")
+
+    assert(lda.getFeaturesCol === "test_feature")
+    assert(lda.getMaxIter === 33)
+    assert(lda.getSeed === 123)
+    assert(lda.getCheckpointInterval === 7)
+    assert(lda.getK === 9)
+    assert(lda.getTopicConcentration === 0.56)
+    assert(lda.getTopicDistributionCol === "myOutput")
+
+
+    // setOptimizer
+    lda.setOptimizer("em")
+    assert(lda.getOptimizer === "em")
+    lda.setOptimizer("online")
+    assert(lda.getOptimizer === "online")
+    lda.setLearningDecay(0.53)
+    assert(lda.getLearningDecay === 0.53)
+    lda.setLearningOffset(1027)
+    assert(lda.getLearningOffset === 1027)
+    lda.setSubsamplingRate(0.06)
+    assert(lda.getSubsamplingRate === 0.06)
+    lda.setOptimizeDocConcentration(false)
+    assert(!lda.getOptimizeDocConcentration)
+  }
+
+  test("parameters validation") {
+    val lda = new LDA()
+
+    // misc Params
+    intercept[IllegalArgumentException] {
+      new LDA().setK(1)
+    }
+    intercept[IllegalArgumentException] {
+      new LDA().setOptimizer("no_such_optimizer")
+    }
+    intercept[IllegalArgumentException] {
+      new LDA().setDocConcentration(-1.1)
+    }
+    intercept[IllegalArgumentException] {
+      new LDA().setTopicConcentration(-1.1)
+    }
+
+    // validateParams()
+    lda.validateParams()
+    lda.setDocConcentration(1.1)
+    lda.validateParams()
+    lda.setDocConcentration(Range(0, lda.getK).map(_ + 2.0).toArray)
+    lda.validateParams()
+    lda.setDocConcentration(Range(0, lda.getK - 1).map(_ + 2.0).toArray)
+    withClue("LDA docConcentration validity check failed for bad array length") {
+      intercept[IllegalArgumentException] {
+        lda.validateParams()
+      }
+    }
+
+    // Online LDA
+    intercept[IllegalArgumentException] {
+      new LDA().setLearningOffset(0)
+    }
+    intercept[IllegalArgumentException] {
+      new LDA().setLearningDecay(0)
+    }
+    intercept[IllegalArgumentException] {
+      new LDA().setSubsamplingRate(0)
+    }
+    intercept[IllegalArgumentException] {
+      new LDA().setSubsamplingRate(1.1)
+    }
+  }
+
+  test("fit & transform with Online LDA") {
+    val lda = new LDA().setK(k).setSeed(1).setOptimizer("online").setMaxIter(2)
+    val model = lda.fit(dataset)
+
+    MLTestingUtils.checkCopy(model)
+
+    assert(!model.isInstanceOf[DistributedLDAModel])
+    assert(model.vocabSize === vocabSize)
+    assert(model.estimatedDocConcentration.size === k)
+    assert(model.topicsMatrix.numRows === vocabSize)
+    assert(model.topicsMatrix.numCols === k)
+    assert(!model.isDistributed)
+
+    // transform()
+    val transformed = model.transform(dataset)
+    val expectedColumns = Array("features", lda.getTopicDistributionCol)
+    expectedColumns.foreach { column =>
+      assert(transformed.columns.contains(column))
+    }
+    transformed.select(lda.getTopicDistributionCol).collect().foreach { r =>
+      val topicDistribution = r.getAs[Vector](0)
+      assert(topicDistribution.size === k)
+      assert(topicDistribution.toArray.forall(w => w >= 0.0 && w <= 1.0))
+    }
+
+    // logLikelihood, logPerplexity
+    val ll = model.logLikelihood(dataset)
+    assert(ll <= 0.0 && ll != Double.NegativeInfinity)
+    val lp = model.logPerplexity(dataset)
+    assert(lp >= 0.0 && lp != Double.PositiveInfinity)
+
+    // describeTopics
+    val topics = model.describeTopics(3)
+    assert(topics.count() === k)
+    assert(topics.select("topic").map(_.getInt(0)).collect().toSet === Range(0, k).toSet)
+    topics.select("termIndices").collect().foreach { case r: Row =>
+      val termIndices = r.getAs[Seq[Int]](0)
+      assert(termIndices.length === 3 && termIndices.toSet.size === 3)
+    }
+    topics.select("termWeights").collect().foreach { case r: Row =>
+      val termWeights = r.getAs[Seq[Double]](0)
+      assert(termWeights.length === 3 && termWeights.forall(w => w >= 0.0 && w <= 1.0))
+    }
+  }
+
+  test("fit & transform with EM LDA") {
+    val lda = new LDA().setK(k).setSeed(1).setOptimizer("em").setMaxIter(2)
+    val model_ = lda.fit(dataset)
+
+    MLTestingUtils.checkCopy(model_)
+
+    assert(model_.isInstanceOf[DistributedLDAModel])
+    val model = model_.asInstanceOf[DistributedLDAModel]
+    assert(model.vocabSize === vocabSize)
+    assert(model.estimatedDocConcentration.size === k)
+    assert(model.topicsMatrix.numRows === vocabSize)
+    assert(model.topicsMatrix.numCols === k)
+    assert(model.isDistributed)
+
+    val localModel = model.toLocal
+    assert(!localModel.isInstanceOf[DistributedLDAModel])
+
+    // training logLikelihood, logPrior
+    val ll = model.trainingLogLikelihood
+    assert(ll <= 0.0 && ll != Double.NegativeInfinity)
+    val lp = model.logPrior
+    assert(lp <= 0.0 && lp != Double.NegativeInfinity)
+  }
+}

From 3121e78168808c015fb21da8b0d44bb33649fb81 Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Tue, 10 Nov 2015 16:25:22 -0800
Subject: [PATCH 170/518] [SPARK-9830][SPARK-11641][SQL][FOLLOW-UP] Remove
 AggregateExpression1 and update toString of Exchange

https://issues.apache.org/jira/browse/SPARK-9830

This is the follow-up pr for https://github.com/apache/spark/pull/9556 to address davies' comments.

Author: Yin Huai <yhuai@databricks.com>

Closes #9607 from yhuai/removeAgg1-followup.
---
 .../sql/catalyst/analysis/Analyzer.scala      |   2 +-
 .../sql/catalyst/analysis/CheckAnalysis.scala |  58 +++++---
 .../expressions/aggregate/Average.scala       |   2 +-
 .../aggregate/CentralMomentAgg.scala          |   2 +-
 .../expressions/aggregate/Stddev.scala        |   2 +-
 .../catalyst/expressions/aggregate/Sum.scala  |   2 +-
 .../analysis/AnalysisErrorSuite.scala         | 127 ++++++++++++++----
 .../scala/org/apache/spark/sql/SQLConf.scala  |   1 +
 .../apache/spark/sql/execution/Exchange.scala |   8 +-
 .../apache/spark/sql/execution/commands.scala |  10 ++
 10 files changed, 160 insertions(+), 54 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index b1e14390b7dc0..a9cd9a77038e7 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -532,7 +532,7 @@ class Analyzer(
                 case min: Min if isDistinct =>
                   AggregateExpression(min, Complete, isDistinct = false)
                 // We get an aggregate function, we need to wrap it in an AggregateExpression.
-                case agg2: AggregateFunction => AggregateExpression(agg2, Complete, isDistinct)
+                case agg: AggregateFunction => AggregateExpression(agg, Complete, isDistinct)
                 // This function is not an aggregate function, just return the resolved one.
                 case other => other
               }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
index 8322e9930cd5a..5a4b0c1e39ce1 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
@@ -110,17 +110,21 @@ trait CheckAnalysis {
           case Aggregate(groupingExprs, aggregateExprs, child) =>
             def checkValidAggregateExpression(expr: Expression): Unit = expr match {
               case aggExpr: AggregateExpression =>
-                // TODO: Is it possible that the child of a agg function is another
-                // agg function?
-                aggExpr.aggregateFunction.children.foreach {
-                  // This is just a sanity check, our analysis rule PullOutNondeterministic should
-                  // already pull out those nondeterministic expressions and evaluate them in
-                  // a Project node.
-                  case child if !child.deterministic =>
+                aggExpr.aggregateFunction.children.foreach { child =>
+                  child.foreach {
+                    case agg: AggregateExpression =>
+                      failAnalysis(
+                        s"It is not allowed to use an aggregate function in the argument of " +
+                          s"another aggregate function. Please use the inner aggregate function " +
+                          s"in a sub-query.")
+                    case other => // OK
+                  }
+
+                  if (!child.deterministic) {
                     failAnalysis(
                       s"nondeterministic expression ${expr.prettyString} should not " +
                         s"appear in the arguments of an aggregate function.")
-                  case child => // OK
+                  }
                 }
               case e: Attribute if !groupingExprs.exists(_.semanticEquals(e)) =>
                 failAnalysis(
@@ -133,19 +137,33 @@ trait CheckAnalysis {
               case e => e.children.foreach(checkValidAggregateExpression)
             }
 
+            def checkSupportedGroupingDataType(
+                expressionString: String,
+                dataType: DataType): Unit = dataType match {
+              case BinaryType =>
+                failAnalysis(s"expression $expressionString cannot be used in " +
+                  s"grouping expression because it is in binary type or its inner field is " +
+                  s"in binary type")
+              case a: ArrayType =>
+                failAnalysis(s"expression $expressionString cannot be used in " +
+                  s"grouping expression because it is in array type or its inner field is " +
+                  s"in array type")
+              case m: MapType =>
+                failAnalysis(s"expression $expressionString cannot be used in " +
+                  s"grouping expression because it is in map type or its inner field is " +
+                  s"in map type")
+              case s: StructType =>
+                s.fields.foreach { f =>
+                  checkSupportedGroupingDataType(expressionString, f.dataType)
+                }
+              case udt: UserDefinedType[_] =>
+                checkSupportedGroupingDataType(expressionString, udt.sqlType)
+              case _ => // OK
+            }
+
             def checkValidGroupingExprs(expr: Expression): Unit = {
-              expr.dataType match {
-                case BinaryType =>
-                  failAnalysis(s"binary type expression ${expr.prettyString} cannot be used " +
-                    "in grouping expression")
-                case a: ArrayType =>
-                  failAnalysis(s"array type expression ${expr.prettyString} cannot be used " +
-                    "in grouping expression")
-                case m: MapType =>
-                  failAnalysis(s"map type expression ${expr.prettyString} cannot be used " +
-                    "in grouping expression")
-                case _ => // OK
-              }
+              checkSupportedGroupingDataType(expr.prettyString, expr.dataType)
+
               if (!expr.deterministic) {
                 // This is just a sanity check, our analysis rule PullOutNondeterministic should
                 // already pull out those nondeterministic expressions and evaluate them in
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
index 7f9e5034702e9..94ac4bf09b90b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Average.scala
@@ -34,7 +34,7 @@ case class Average(child: Expression) extends DeclarativeAggregate {
   // Return data type.
   override def dataType: DataType = resultType
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType))
+  override def inputTypes: Seq[AbstractDataType] = Seq(NumericType)
 
   override def checkInputDataTypes(): TypeCheckResult =
     TypeUtils.checkForNumericExpr(child.dataType, "function average")
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
index 984ce7f24dacc..de5872ab11eb1 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
@@ -57,7 +57,7 @@ abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate w
 
   override def dataType: DataType = DoubleType
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType))
+  override def inputTypes: Seq[AbstractDataType] = Seq(NumericType)
 
   override def checkInputDataTypes(): TypeCheckResult =
     TypeUtils.checkForNumericExpr(child.dataType, s"function $prettyName")
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
index 5b9eb7ae02f25..2748009623355 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
@@ -50,7 +50,7 @@ abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
 
   override def dataType: DataType = resultType
 
-  override def inputTypes: Seq[AbstractDataType] = Seq(TypeCollection(NumericType))
+  override def inputTypes: Seq[AbstractDataType] = Seq(NumericType)
 
   override def checkInputDataTypes(): TypeCheckResult =
     TypeUtils.checkForNumericExpr(child.dataType, "function stddev")
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
index c005ec9657211..cfb042e0aa782 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Sum.scala
@@ -32,7 +32,7 @@ case class Sum(child: Expression) extends DeclarativeAggregate {
   override def dataType: DataType = resultType
 
   override def inputTypes: Seq[AbstractDataType] =
-    Seq(TypeCollection(LongType, DoubleType, DecimalType, NullType))
+    Seq(TypeCollection(LongType, DoubleType, DecimalType))
 
   override def checkInputDataTypes(): TypeCheckResult =
     TypeUtils.checkForNumericExpr(child.dataType, "function sum")
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
index 5a2368e329976..2e7c3bd67b554 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
@@ -23,8 +23,59 @@ import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.Inner
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.util.{GenericArrayData, ArrayData}
 import org.apache.spark.sql.types._
 
+import scala.beans.{BeanProperty, BeanInfo}
+
+@BeanInfo
+private[sql] case class GroupableData(@BeanProperty data: Int)
+
+private[sql] class GroupableUDT extends UserDefinedType[GroupableData] {
+
+  override def sqlType: DataType = IntegerType
+
+  override def serialize(obj: Any): Int = {
+    obj match {
+      case groupableData: GroupableData => groupableData.data
+    }
+  }
+
+  override def deserialize(datum: Any): GroupableData = {
+    datum match {
+      case data: Int => GroupableData(data)
+    }
+  }
+
+  override def userClass: Class[GroupableData] = classOf[GroupableData]
+
+  private[spark] override def asNullable: GroupableUDT = this
+}
+
+@BeanInfo
+private[sql] case class UngroupableData(@BeanProperty data: Array[Int])
+
+private[sql] class UngroupableUDT extends UserDefinedType[UngroupableData] {
+
+  override def sqlType: DataType = ArrayType(IntegerType)
+
+  override def serialize(obj: Any): ArrayData = {
+    obj match {
+      case groupableData: UngroupableData => new GenericArrayData(groupableData.data)
+    }
+  }
+
+  override def deserialize(datum: Any): UngroupableData = {
+    datum match {
+      case data: Array[Int] => UngroupableData(data)
+    }
+  }
+
+  override def userClass: Class[UngroupableData] = classOf[UngroupableData]
+
+  private[spark] override def asNullable: UngroupableUDT = this
+}
+
 case class TestFunction(
     children: Seq[Expression],
     inputTypes: Seq[AbstractDataType])
@@ -194,39 +245,65 @@ class AnalysisErrorSuite extends AnalysisTest {
     assert(error.message.contains("Conflicting attributes"))
   }
 
-  test("aggregation can't work on binary and map types") {
-    val plan =
-      Aggregate(
-        AttributeReference("a", BinaryType)(exprId = ExprId(2)) :: Nil,
-        Alias(sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
-        LocalRelation(
-          AttributeReference("a", BinaryType)(exprId = ExprId(2)),
-          AttributeReference("b", IntegerType)(exprId = ExprId(1))))
+  test("check grouping expression data types") {
+    def checkDataType(dataType: DataType, shouldSuccess: Boolean): Unit = {
+      val plan =
+        Aggregate(
+          AttributeReference("a", dataType)(exprId = ExprId(2)) :: Nil,
+          Alias(sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
+          LocalRelation(
+            AttributeReference("a", dataType)(exprId = ExprId(2)),
+            AttributeReference("b", IntegerType)(exprId = ExprId(1))))
+
+      shouldSuccess match {
+        case true =>
+          assertAnalysisSuccess(plan, true)
+        case false =>
+          assertAnalysisError(plan, "expression a cannot be used in grouping expression" :: Nil)
+      }
 
-    assertAnalysisError(plan,
-      "binary type expression a cannot be used in grouping expression" :: Nil)
+    }
 
-    val plan2 =
-      Aggregate(
-        AttributeReference("a", MapType(IntegerType, StringType))(exprId = ExprId(2)) :: Nil,
-        Alias(sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
-        LocalRelation(
-          AttributeReference("a", MapType(IntegerType, StringType))(exprId = ExprId(2)),
-          AttributeReference("b", IntegerType)(exprId = ExprId(1))))
+    val supportedDataTypes = Seq(
+      StringType,
+      NullType, BooleanType,
+      ByteType, ShortType, IntegerType, LongType,
+      FloatType, DoubleType, DecimalType(25, 5), DecimalType(6, 5),
+      DateType, TimestampType,
+      new StructType()
+        .add("f1", FloatType, nullable = true)
+        .add("f2", StringType, nullable = true),
+      new GroupableUDT())
+    supportedDataTypes.foreach { dataType =>
+      checkDataType(dataType, shouldSuccess = true)
+    }
 
-    assertAnalysisError(plan2,
-      "map type expression a cannot be used in grouping expression" :: Nil)
+    val unsupportedDataTypes = Seq(
+      BinaryType,
+      ArrayType(IntegerType),
+      MapType(StringType, LongType),
+      new StructType()
+        .add("f1", FloatType, nullable = true)
+        .add("f2", ArrayType(BooleanType, containsNull = true), nullable = true),
+      new UngroupableUDT())
+    unsupportedDataTypes.foreach { dataType =>
+      checkDataType(dataType, shouldSuccess = false)
+    }
+  }
 
-    val plan3 =
+  test("we should fail analysis when we find nested aggregate functions") {
+    val plan =
       Aggregate(
-        AttributeReference("a", ArrayType(IntegerType))(exprId = ExprId(2)) :: Nil,
-        Alias(sum(AttributeReference("b", IntegerType)(exprId = ExprId(1))), "c")() :: Nil,
+        AttributeReference("a", IntegerType)(exprId = ExprId(2)) :: Nil,
+        Alias(sum(sum(AttributeReference("b", IntegerType)(exprId = ExprId(1)))), "c")() :: Nil,
         LocalRelation(
-          AttributeReference("a", ArrayType(IntegerType))(exprId = ExprId(2)),
+          AttributeReference("a", IntegerType)(exprId = ExprId(2)),
           AttributeReference("b", IntegerType)(exprId = ExprId(1))))
 
-    assertAnalysisError(plan3,
-      "array type expression a cannot be used in grouping expression" :: Nil)
+    assertAnalysisError(
+      plan,
+      "It is not allowed to use an aggregate function in the argument of " +
+        "another aggregate function." :: Nil)
   }
 
   test("Join can't work on binary and map types") {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
index 89e196c066007..57d7d30e0eca2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
@@ -474,6 +474,7 @@ private[spark] object SQLConf {
   object Deprecated {
     val MAPRED_REDUCE_TASKS = "mapred.reduce.tasks"
     val EXTERNAL_SORT = "spark.sql.planner.externalSort"
+    val USE_SQL_AGGREGATE2 = "spark.sql.useAggregate2"
   }
 }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
index a4ce328c1a9eb..b733b26987bcb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
@@ -44,14 +44,14 @@ case class Exchange(
   override def nodeName: String = {
     val extraInfo = coordinator match {
       case Some(exchangeCoordinator) if exchangeCoordinator.isEstimated =>
-        "Shuffle"
+        s"(coordinator id: ${System.identityHashCode(coordinator)})"
       case Some(exchangeCoordinator) if !exchangeCoordinator.isEstimated =>
-        "May shuffle"
-      case None => "Shuffle without coordinator"
+        s"(coordinator id: ${System.identityHashCode(coordinator)})"
+      case None => ""
     }
 
     val simpleNodeName = if (tungstenMode) "TungstenExchange" else "Exchange"
-    s"$simpleNodeName($extraInfo)"
+    s"${simpleNodeName}${extraInfo}"
   }
 
   /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
index e5f60b15e7359..8b2755a58757c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
@@ -111,6 +111,16 @@ case class SetCommand(kv: Option[(String, Option[String])]) extends RunnableComm
       }
       (keyValueOutput, runFunc)
 
+    case Some((SQLConf.Deprecated.USE_SQL_AGGREGATE2, Some(value))) =>
+      val runFunc = (sqlContext: SQLContext) => {
+        logWarning(
+          s"Property ${SQLConf.Deprecated.USE_SQL_AGGREGATE2} is deprecated and " +
+            s"will be ignored. ${SQLConf.Deprecated.USE_SQL_AGGREGATE2} will " +
+            s"continue to be true.")
+        Seq(Row(SQLConf.Deprecated.USE_SQL_AGGREGATE2, "true"))
+      }
+      (keyValueOutput, runFunc)
+
     // Configures a single property.
     case Some((key, Some(value))) =>
       val runFunc = (sqlContext: SQLContext) => {

From 21c562fa03430365f5c2b7d6de1f8f60ab2140d4 Mon Sep 17 00:00:00 2001
From: Herman van Hovell <hvanhovell@questtec.nl>
Date: Tue, 10 Nov 2015 16:28:21 -0800
Subject: [PATCH 171/518] [SPARK-9241][SQL] Supporting multiple DISTINCT
 columns - follow-up (3)

This PR is a 2nd follow-up for [SPARK-9241](https://issues.apache.org/jira/browse/SPARK-9241). It contains the following improvements:
* Fix for a potential bug in distinct child expression and attribute alignment.
* Improved handling of duplicate distinct child expressions.
* Added test for distinct UDAF with multiple children.

cc yhuai

Author: Herman van Hovell <hvanhovell@questtec.nl>

Closes #9566 from hvanhovell/SPARK-9241-followup-2.
---
 .../DistinctAggregationRewriter.scala         |  9 ++--
 .../execution/AggregationQuerySuite.scala     | 41 +++++++++++++++++--
 2 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
index 397eff05686b6..c0c960471a61a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
@@ -151,11 +151,12 @@ case class DistinctAggregationRewriter(conf: CatalystConf) extends Rule[LogicalP
       }
 
       // Setup unique distinct aggregate children.
-      val distinctAggChildren = distinctAggGroups.keySet.flatten.toSeq
-      val distinctAggChildAttrMap = distinctAggChildren.map(expressionAttributePair).toMap
-      val distinctAggChildAttrs = distinctAggChildAttrMap.values.toSeq
+      val distinctAggChildren = distinctAggGroups.keySet.flatten.toSeq.distinct
+      val distinctAggChildAttrMap = distinctAggChildren.map(expressionAttributePair)
+      val distinctAggChildAttrs = distinctAggChildAttrMap.map(_._2)
 
       // Setup expand & aggregate operators for distinct aggregate expressions.
+      val distinctAggChildAttrLookup = distinctAggChildAttrMap.toMap
       val distinctAggOperatorMap = distinctAggGroups.toSeq.zipWithIndex.map {
         case ((group, expressions), i) =>
           val id = Literal(i + 1)
@@ -170,7 +171,7 @@ case class DistinctAggregationRewriter(conf: CatalystConf) extends Rule[LogicalP
           val operators = expressions.map { e =>
             val af = e.aggregateFunction
             val naf = patchAggregateFunctionChildren(af) { x =>
-              evalWithinGroup(id, distinctAggChildAttrMap(x))
+              evalWithinGroup(id, distinctAggChildAttrLookup(x))
             }
             (e, e.copy(aggregateFunction = naf, isDistinct = false))
           }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 6bf2c53440baf..8253921563b3a 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -66,6 +66,36 @@ class ScalaAggregateFunction(schema: StructType) extends UserDefinedAggregateFun
   }
 }
 
+class LongProductSum extends UserDefinedAggregateFunction {
+  def inputSchema: StructType = new StructType()
+    .add("a", LongType)
+    .add("b", LongType)
+
+  def bufferSchema: StructType = new StructType()
+    .add("product", LongType)
+
+  def dataType: DataType = LongType
+
+  def deterministic: Boolean = true
+
+  def initialize(buffer: MutableAggregationBuffer): Unit = {
+    buffer(0) = 0L
+  }
+
+  def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
+    if (!(input.isNullAt(0) || input.isNullAt(1))) {
+      buffer(0) = buffer.getLong(0) + input.getLong(0) * input.getLong(1)
+    }
+  }
+
+  def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
+    buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
+  }
+
+  def evaluate(buffer: Row): Any =
+    buffer.getLong(0)
+}
+
 abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
   import testImplicits._
 
@@ -110,6 +140,7 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
     // Register UDAFs
     sqlContext.udf.register("mydoublesum", new MyDoubleSum)
     sqlContext.udf.register("mydoubleavg", new MyDoubleAvg)
+    sqlContext.udf.register("longProductSum", new LongProductSum)
   }
 
   override def afterAll(): Unit = {
@@ -545,19 +576,21 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
           |  count(distinct value2),
           |  sum(distinct value2),
           |  count(distinct value1, value2),
+          |  longProductSum(distinct value1, value2),
           |  count(value1),
           |  sum(value1),
           |  count(value2),
           |  sum(value2),
+          |  longProductSum(value1, value2),
           |  count(*),
           |  count(1)
           |FROM agg2
           |GROUP BY key
         """.stripMargin),
-      Row(null, 3, 30, 3, 60, 3, 3, 30, 3, 60, 4, 4) ::
-        Row(1, 2, 40, 3, -10, 3, 3, 70, 3, -10, 3, 3) ::
-        Row(2, 2, 0, 1, 1, 1, 3, 1, 3, 3, 4, 4) ::
-        Row(3, 0, null, 1, 3, 0, 0, null, 1, 3, 2, 2) :: Nil)
+      Row(null, 3, 30, 3, 60, 3, -4700, 3, 30, 3, 60, -4700, 4, 4) ::
+        Row(1, 2, 40, 3, -10, 3, -100, 3, 70, 3, -10, -100, 3, 3) ::
+        Row(2, 2, 0, 1, 1, 1, 1, 3, 1, 3, 3, 2, 4, 4) ::
+        Row(3, 0, null, 1, 3, 0, 0, 0, null, 1, 3, 0, 2, 2) :: Nil)
   }
 
   test("test count") {

From a3989058c0938c8c59c278e7d1a766701cfa255b Mon Sep 17 00:00:00 2001
From: Bryan Cutler <bjcutler@us.ibm.com>
Date: Tue, 10 Nov 2015 16:32:32 -0800
Subject: [PATCH 172/518] [SPARK-10827][CORE] AppClient should not use
 `askWithReply` in `receiveAndReply`

Changed AppClient to be non-blocking in `receiveAndReply` by using a separate thread to wait for response and reply to the context.  The threads are managed by a thread pool.  Also added unit tests for the AppClient interface.

Author: Bryan Cutler <bjcutler@us.ibm.com>

Closes #9317 from BryanCutler/appClient-receiveAndReply-SPARK-10827.
---
 .../spark/deploy/client/AppClient.scala       |  33 ++-
 .../spark/deploy/client/AppClientSuite.scala  | 209 ++++++++++++++++++
 2 files changed, 238 insertions(+), 4 deletions(-)
 create mode 100644 core/src/test/scala/org/apache/spark/deploy/client/AppClientSuite.scala

diff --git a/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala b/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
index 25ea6925434ab..3f29da663b798 100644
--- a/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
@@ -49,8 +49,8 @@ private[spark] class AppClient(
   private val REGISTRATION_TIMEOUT_SECONDS = 20
   private val REGISTRATION_RETRIES = 3
 
-  private var endpoint: RpcEndpointRef = null
-  private var appId: String = null
+  @volatile private var endpoint: RpcEndpointRef = null
+  @volatile private var appId: String = null
   @volatile private var registered = false
 
   private class ClientEndpoint(override val rpcEnv: RpcEnv) extends ThreadSafeRpcEndpoint
@@ -77,6 +77,11 @@ private[spark] class AppClient(
     private val registrationRetryThread =
       ThreadUtils.newDaemonSingleThreadScheduledExecutor("appclient-registration-retry-thread")
 
+    // A thread pool to perform receive then reply actions in a thread so as not to block the
+    // event loop.
+    private val askAndReplyThreadPool =
+      ThreadUtils.newDaemonCachedThreadPool("appclient-receive-and-reply-threadpool")
+
     override def onStart(): Unit = {
       try {
         registerWithMaster(1)
@@ -200,7 +205,7 @@ private[spark] class AppClient(
 
       case r: RequestExecutors =>
         master match {
-          case Some(m) => context.reply(m.askWithRetry[Boolean](r))
+          case Some(m) => askAndReplyAsync(m, context, r)
           case None =>
             logWarning("Attempted to request executors before registering with Master.")
             context.reply(false)
@@ -208,13 +213,32 @@ private[spark] class AppClient(
 
       case k: KillExecutors =>
         master match {
-          case Some(m) => context.reply(m.askWithRetry[Boolean](k))
+          case Some(m) => askAndReplyAsync(m, context, k)
           case None =>
             logWarning("Attempted to kill executors before registering with Master.")
             context.reply(false)
         }
     }
 
+    private def askAndReplyAsync[T](
+        endpointRef: RpcEndpointRef,
+        context: RpcCallContext,
+        msg: T): Unit = {
+      // Create a thread to ask a message and reply with the result.  Allow thread to be
+      // interrupted during shutdown, otherwise context must be notified of NonFatal errors.
+      askAndReplyThreadPool.execute(new Runnable {
+        override def run(): Unit = {
+          try {
+            context.reply(endpointRef.askWithRetry[Boolean](msg))
+          } catch {
+            case ie: InterruptedException => // Cancelled
+            case NonFatal(t) =>
+              context.sendFailure(t)
+          }
+        }
+      })
+    }
+
     override def onDisconnected(address: RpcAddress): Unit = {
       if (master.exists(_.address == address)) {
         logWarning(s"Connection to $address failed; waiting for master to reconnect...")
@@ -252,6 +276,7 @@ private[spark] class AppClient(
       registrationRetryThread.shutdownNow()
       registerMasterFutures.foreach(_.cancel(true))
       registerMasterThreadPool.shutdownNow()
+      askAndReplyThreadPool.shutdownNow()
     }
 
   }
diff --git a/core/src/test/scala/org/apache/spark/deploy/client/AppClientSuite.scala b/core/src/test/scala/org/apache/spark/deploy/client/AppClientSuite.scala
new file mode 100644
index 0000000000000..1e5c05a73f8aa
--- /dev/null
+++ b/core/src/test/scala/org/apache/spark/deploy/client/AppClientSuite.scala
@@ -0,0 +1,209 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.deploy.client
+
+import scala.collection.mutable.{ArrayBuffer, SynchronizedBuffer}
+import scala.concurrent.duration._
+
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.concurrent.Eventually._
+
+import org.apache.spark._
+import org.apache.spark.deploy.{ApplicationDescription, Command}
+import org.apache.spark.deploy.DeployMessages.{MasterStateResponse, RequestMasterState}
+import org.apache.spark.deploy.master.{ApplicationInfo, Master}
+import org.apache.spark.deploy.worker.Worker
+import org.apache.spark.rpc.RpcEnv
+import org.apache.spark.util.Utils
+
+/**
+ * End-to-end tests for application client in standalone mode.
+ */
+class AppClientSuite extends SparkFunSuite with LocalSparkContext with BeforeAndAfterAll {
+  private val numWorkers = 2
+  private val conf = new SparkConf()
+  private val securityManager = new SecurityManager(conf)
+
+  private var masterRpcEnv: RpcEnv = null
+  private var workerRpcEnvs: Seq[RpcEnv] = null
+  private var master: Master = null
+  private var workers: Seq[Worker] = null
+
+  /**
+   * Start the local cluster.
+   * Note: local-cluster mode is insufficient because we want a reference to the Master.
+   */
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+    masterRpcEnv = RpcEnv.create(Master.SYSTEM_NAME, "localhost", 0, conf, securityManager)
+    workerRpcEnvs = (0 until numWorkers).map { i =>
+      RpcEnv.create(Worker.SYSTEM_NAME + i, "localhost", 0, conf, securityManager)
+    }
+    master = makeMaster()
+    workers = makeWorkers(10, 2048)
+    // Wait until all workers register with master successfully
+    eventually(timeout(60.seconds), interval(10.millis)) {
+      assert(getMasterState.workers.size === numWorkers)
+    }
+  }
+
+  override def afterAll(): Unit = {
+    workerRpcEnvs.foreach(_.shutdown())
+    masterRpcEnv.shutdown()
+    workers.foreach(_.stop())
+    master.stop()
+    workerRpcEnvs = null
+    masterRpcEnv = null
+    workers = null
+    master = null
+    super.afterAll()
+  }
+
+  test("interface methods of AppClient using local Master") {
+    val ci = new AppClientInst(masterRpcEnv.address.toSparkURL)
+
+    ci.client.start()
+
+    // Client should connect with one Master which registers the application
+    eventually(timeout(10.seconds), interval(10.millis)) {
+      val apps = getApplications()
+      assert(ci.listener.connectedIdList.size === 1, "client listener should have one connection")
+      assert(apps.size === 1, "master should have 1 registered app")
+    }
+
+    // Send message to Master to request Executors, verify request by change in executor limit
+    val numExecutorsRequested = 1
+    assert(ci.client.requestTotalExecutors(numExecutorsRequested))
+
+    eventually(timeout(10.seconds), interval(10.millis)) {
+      val apps = getApplications()
+      assert(apps.head.getExecutorLimit === numExecutorsRequested, s"executor request failed")
+    }
+
+    // Send request to kill executor, verify request was made
+    assert {
+      val apps = getApplications()
+      val executorId: String = apps.head.executors.head._2.fullId
+      ci.client.killExecutors(Seq(executorId))
+    }
+
+    // Issue stop command for Client to disconnect from Master
+    ci.client.stop()
+
+    // Verify Client is marked dead and unregistered from Master
+    eventually(timeout(10.seconds), interval(10.millis)) {
+      val apps = getApplications()
+      assert(ci.listener.deadReasonList.size === 1, "client should have been marked dead")
+      assert(apps.isEmpty, "master should have 0 registered apps")
+    }
+  }
+
+  test("request from AppClient before initialized with master") {
+    val ci = new AppClientInst(masterRpcEnv.address.toSparkURL)
+
+    // requests to master should fail immediately
+    assert(ci.client.requestTotalExecutors(3) === false)
+  }
+
+  // ===============================
+  // | Utility methods for testing |
+  // ===============================
+
+  /** Return a SparkConf for applications that want to talk to our Master. */
+  private def appConf: SparkConf = {
+    new SparkConf()
+      .setMaster(masterRpcEnv.address.toSparkURL)
+      .setAppName("test")
+      .set("spark.executor.memory", "256m")
+  }
+
+  /** Make a master to which our application will send executor requests. */
+  private def makeMaster(): Master = {
+    val master = new Master(masterRpcEnv, masterRpcEnv.address, 0, securityManager, conf)
+    masterRpcEnv.setupEndpoint(Master.ENDPOINT_NAME, master)
+    master
+  }
+
+  /** Make a few workers that talk to our master. */
+  private def makeWorkers(cores: Int, memory: Int): Seq[Worker] = {
+    (0 until numWorkers).map { i =>
+      val rpcEnv = workerRpcEnvs(i)
+      val worker = new Worker(rpcEnv, 0, cores, memory, Array(masterRpcEnv.address),
+        Worker.SYSTEM_NAME + i, Worker.ENDPOINT_NAME, null, conf, securityManager)
+      rpcEnv.setupEndpoint(Worker.ENDPOINT_NAME, worker)
+      worker
+    }
+  }
+
+  /** Get the Master state */
+  private def getMasterState: MasterStateResponse = {
+    master.self.askWithRetry[MasterStateResponse](RequestMasterState)
+  }
+
+  /** Get the applictions that are active from Master */
+  private def getApplications(): Seq[ApplicationInfo] = {
+    getMasterState.activeApps
+  }
+
+  /** Application Listener to collect events */
+  private class AppClientCollector extends AppClientListener with Logging {
+    val connectedIdList = new ArrayBuffer[String] with SynchronizedBuffer[String]
+    @volatile var disconnectedCount: Int = 0
+    val deadReasonList = new ArrayBuffer[String] with SynchronizedBuffer[String]
+    val execAddedList = new ArrayBuffer[String] with SynchronizedBuffer[String]
+    val execRemovedList = new ArrayBuffer[String] with SynchronizedBuffer[String]
+
+    def connected(id: String): Unit = {
+      connectedIdList += id
+    }
+
+    def disconnected(): Unit = {
+      synchronized {
+        disconnectedCount += 1
+      }
+    }
+
+    def dead(reason: String): Unit = {
+      deadReasonList += reason
+    }
+
+    def executorAdded(
+        id: String,
+        workerId: String,
+        hostPort: String,
+        cores: Int,
+        memory: Int): Unit = {
+      execAddedList += id
+    }
+
+    def executorRemoved(id: String, message: String, exitStatus: Option[Int]): Unit = {
+      execRemovedList += id
+    }
+  }
+
+  /** Create AppClient and supporting objects */
+  private class AppClientInst(masterUrl: String) {
+    val rpcEnv = RpcEnv.create("spark", Utils.localHostName(), 0, conf, securityManager)
+    private val cmd = new Command(TestExecutor.getClass.getCanonicalName.stripSuffix("$"),
+      List(), Map(), Seq(), Seq(), Seq())
+    private val desc = new ApplicationDescription("AppClientSuite", Some(1), 512, cmd, "ignored")
+    val listener = new AppClientCollector
+    val client = new AppClient(rpcEnv, Array(masterUrl), desc, listener, new SparkConf)
+  }
+
+}

From c0e48dfa611fa5d94132af7e6f6731f60ab833da Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Tue, 10 Nov 2015 16:42:28 -0800
Subject: [PATCH 173/518] [SPARK-11566] [MLLIB] [PYTHON] Refactoring
 GaussianMixtureModel.gaussians in Python

cc jkbradley

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>

Closes #9534 from yu-iskw/SPARK-11566.
---
 .../python/GaussianMixtureModelWrapper.scala  | 21 ++++++-------------
 python/pyspark/mllib/clustering.py            |  2 +-
 2 files changed, 7 insertions(+), 16 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/api/python/GaussianMixtureModelWrapper.scala b/mllib/src/main/scala/org/apache/spark/mllib/api/python/GaussianMixtureModelWrapper.scala
index 0ec88ef77d695..6a3b20c88d2d2 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/api/python/GaussianMixtureModelWrapper.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/api/python/GaussianMixtureModelWrapper.scala
@@ -17,14 +17,11 @@
 
 package org.apache.spark.mllib.api.python
 
-import java.util.{List => JList}
-
-import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.JavaConverters
 
 import org.apache.spark.SparkContext
-import org.apache.spark.mllib.linalg.{Vector, Vectors, Matrix}
 import org.apache.spark.mllib.clustering.GaussianMixtureModel
+import org.apache.spark.mllib.linalg.{Vector, Vectors}
 
 /**
   * Wrapper around GaussianMixtureModel to provide helper methods in Python
@@ -36,17 +33,11 @@ private[python] class GaussianMixtureModelWrapper(model: GaussianMixtureModel) {
   /**
     * Returns gaussians as a List of Vectors and Matrices corresponding each MultivariateGaussian
     */
-  val gaussians: JList[Object] = {
-    val modelGaussians = model.gaussians
-    var i = 0
-    var mu = ArrayBuffer.empty[Vector]
-    var sigma = ArrayBuffer.empty[Matrix]
-    while (i < k) {
-      mu += modelGaussians(i).mu
-      sigma += modelGaussians(i).sigma
-      i += 1
+  val gaussians: Array[Byte] = {
+    val modelGaussians = model.gaussians.map { gaussian =>
+      Array[Any](gaussian.mu, gaussian.sigma)
     }
-    List(mu.toArray, sigma.toArray).map(_.asInstanceOf[Object]).asJava
+    SerDe.dumps(JavaConverters.seqAsJavaListConverter(modelGaussians).asJava)
   }
 
   def save(sc: SparkContext, path: String): Unit = model.save(sc, path)
diff --git a/python/pyspark/mllib/clustering.py b/python/pyspark/mllib/clustering.py
index 1fa061dc2da99..c9e6f1dec6bf8 100644
--- a/python/pyspark/mllib/clustering.py
+++ b/python/pyspark/mllib/clustering.py
@@ -266,7 +266,7 @@ def gaussians(self):
         """
         return [
             MultivariateGaussian(gaussian[0], gaussian[1])
-            for gaussian in zip(*self.call("gaussians"))]
+            for gaussian in self.call("gaussians")]
 
     @property
     @since('1.4.0')

From 33112f9c48680c33d663978f76806ebf0ea39789 Mon Sep 17 00:00:00 2001
From: Imran Rashid <irashid@cloudera.com>
Date: Tue, 10 Nov 2015 16:50:22 -0800
Subject: [PATCH 174/518] [SPARK-10192][CORE] simple test w/ failure involving
 a shared dependency

just trying to increase test coverage in the scheduler, this already works.  It includes a regression test for SPARK-9809

copied some test utils from https://github.com/apache/spark/pull/5636, we can wait till that is merged first

Author: Imran Rashid <irashid@cloudera.com>

Closes #8402 from squito/test_retry_in_shared_shuffle_dep.
---
 .../spark/scheduler/DAGSchedulerSuite.scala   | 51 ++++++++++++++++++-
 1 file changed, 49 insertions(+), 2 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
index 3816b8c4a09aa..068b49bd5844b 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
@@ -594,11 +594,17 @@ class DAGSchedulerSuite
    * @param stageId - The current stageId
    * @param attemptIdx - The current attempt count
    */
-  private def completeNextResultStageWithSuccess(stageId: Int, attemptIdx: Int): Unit = {
+  private def completeNextResultStageWithSuccess(
+      stageId: Int,
+      attemptIdx: Int,
+      partitionToResult: Int => Int = _ => 42): Unit = {
     val stageAttempt = taskSets.last
     checkStageId(stageId, attemptIdx, stageAttempt)
     assert(scheduler.stageIdToStage(stageId).isInstanceOf[ResultStage])
-    complete(stageAttempt, stageAttempt.tasks.zipWithIndex.map(_ => (Success, 42)).toSeq)
+    val taskResults = stageAttempt.tasks.zipWithIndex.map { case (task, idx) =>
+      (Success, partitionToResult(idx))
+    }
+    complete(stageAttempt, taskResults.toSeq)
   }
 
   /**
@@ -1054,6 +1060,47 @@ class DAGSchedulerSuite
     assertDataStructuresEmpty()
   }
 
+  /**
+   * Run two jobs, with a shared dependency.  We simulate a fetch failure in the second job, which
+   * requires regenerating some outputs of the shared dependency.  One key aspect of this test is
+   * that the second job actually uses a different stage for the shared dependency (a "skipped"
+   * stage).
+   */
+  test("shuffle fetch failure in a reused shuffle dependency") {
+    // Run the first job successfully, which creates one shuffle dependency
+
+    val shuffleMapRdd = new MyRDD(sc, 2, Nil)
+    val shuffleDep = new ShuffleDependency(shuffleMapRdd, null)
+    val reduceRdd = new MyRDD(sc, 2, List(shuffleDep))
+    submit(reduceRdd, Array(0, 1))
+
+    completeShuffleMapStageSuccessfully(0, 0, 2)
+    completeNextResultStageWithSuccess(1, 0)
+    assert(results === Map(0 -> 42, 1 -> 42))
+    assertDataStructuresEmpty()
+
+    // submit another job w/ the shared dependency, and have a fetch failure
+    val reduce2 = new MyRDD(sc, 2, List(shuffleDep))
+    submit(reduce2, Array(0, 1))
+    // Note that the stage numbering here is only b/c the shared dependency produces a new, skipped
+    // stage.  If instead it reused the existing stage, then this would be stage 2
+    completeNextStageWithFetchFailure(3, 0, shuffleDep)
+    scheduler.resubmitFailedStages()
+
+    // the scheduler now creates a new task set to regenerate the missing map output, but this time
+    // using a different stage, the "skipped" one
+
+    // SPARK-9809 -- this stage is submitted without a task for each partition (because some of
+    // the shuffle map output is still available from stage 0); make sure we've still got internal
+    // accumulators setup
+    assert(scheduler.stageIdToStage(2).internalAccumulators.nonEmpty)
+    completeShuffleMapStageSuccessfully(2, 0, 2)
+    completeNextResultStageWithSuccess(3, 1, idx => idx + 1234)
+    assert(results === Map(0 -> 1234, 1 -> 1235))
+
+    assertDataStructuresEmpty()
+  }
+
   /**
    * This test runs a three stage job, with a fetch failure in stage 1.  but during the retry, we
    * have completions from both the first & second attempt of stage 1.  So all the map output is

From 3e0a6cf1e02a19b37c68d3026415d53bb57a576b Mon Sep 17 00:00:00 2001
From: tedyu <yuzhihong@gmail.com>
Date: Tue, 10 Nov 2015 16:51:25 -0800
Subject: [PATCH 175/518] [SPARK-11572] Exit AsynchronousListenerBus thread
 when stop() is called

As vonnagy reported in the following thread:
http://search-hadoop.com/m/q3RTtk982kvIow22

Attempts to join the thread in AsynchronousListenerBus resulted in lock up because AsynchronousListenerBus thread was still getting messages `SparkListenerExecutorMetricsUpdate` from the DAGScheduler

Author: tedyu <yuzhihong@gmail.com>

Closes #9546 from ted-yu/master.
---
 .../org/apache/spark/util/AsynchronousListenerBus.scala  | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala b/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
index 61b5a4cecddce..b8481eabc7618 100644
--- a/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
+++ b/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
@@ -67,15 +67,12 @@ private[spark] abstract class AsynchronousListenerBus[L <: AnyRef, E](name: Stri
           processingEvent = true
         }
         try {
-          val event = eventQueue.poll
-          if (event == null) {
+          if (stopped.get()) {
             // Get out of the while loop and shutdown the daemon thread
-            if (!stopped.get) {
-              throw new IllegalStateException("Polling `null` from eventQueue means" +
-                " the listener bus has been stopped. So `stopped` must be true")
-            }
             return
           }
+          val event = eventQueue.poll
+          assert(event != null, "event queue was empty but the listener bus was not stopped")
           postToAll(event)
         } finally {
           self.synchronized {

From 900917541651abe7125f0d205085d2ab6a00d92c Mon Sep 17 00:00:00 2001
From: tedyu <yuzhihong@gmail.com>
Date: Tue, 10 Nov 2015 16:52:26 -0800
Subject: [PATCH 176/518] [SPARK-11615] Drop @VisibleForTesting annotation

See http://search-hadoop.com/m/q3RTtjpe8r1iRbTj2 for discussion.

Summary: addition of VisibleForTesting annotation resulted in spark-shell malfunctioning.

Author: tedyu <yuzhihong@gmail.com>

Closes #9585 from tedyu/master.
---
 .../src/main/scala/org/apache/spark/rpc/netty/Inbox.scala | 8 ++++----
 .../org/apache/spark/ui/jobs/JobProgressListener.scala    | 2 --
 .../org/apache/spark/util/AsynchronousListenerBus.scala   | 5 ++---
 .../org/apache/spark/util/collection/ExternalSorter.scala | 3 +--
 scalastyle-config.xml                                     | 7 +++++++
 .../org/apache/spark/sql/execution/QueryExecution.scala   | 3 ---
 .../spark/network/shuffle/ShuffleTestAccessor.scala       | 1 -
 7 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala b/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala
index c72b588db57fe..464027f07cc88 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala
@@ -21,8 +21,6 @@ import javax.annotation.concurrent.GuardedBy
 
 import scala.util.control.NonFatal
 
-import com.google.common.annotations.VisibleForTesting
-
 import org.apache.spark.{Logging, SparkException}
 import org.apache.spark.rpc.{RpcAddress, RpcEndpoint, ThreadSafeRpcEndpoint}
 
@@ -193,8 +191,10 @@ private[netty] class Inbox(
 
   def isEmpty: Boolean = inbox.synchronized { messages.isEmpty }
 
-  /** Called when we are dropping a message. Test cases override this to test message dropping. */
-  @VisibleForTesting
+  /**
+   * Called when we are dropping a message. Test cases override this to test message dropping.
+   * Exposed for testing.
+   */
   protected def onDrop(message: InboxMessage): Unit = {
     logWarning(s"Drop $message because $endpointRef is stopped")
   }
diff --git a/core/src/main/scala/org/apache/spark/ui/jobs/JobProgressListener.scala b/core/src/main/scala/org/apache/spark/ui/jobs/JobProgressListener.scala
index 77d034fa5ba2c..ca37829216f22 100644
--- a/core/src/main/scala/org/apache/spark/ui/jobs/JobProgressListener.scala
+++ b/core/src/main/scala/org/apache/spark/ui/jobs/JobProgressListener.scala
@@ -21,8 +21,6 @@ import java.util.concurrent.TimeoutException
 
 import scala.collection.mutable.{HashMap, HashSet, ListBuffer}
 
-import com.google.common.annotations.VisibleForTesting
-
 import org.apache.spark._
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.executor.TaskMetrics
diff --git a/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala b/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
index b8481eabc7618..b3b54af972cb4 100644
--- a/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
+++ b/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
@@ -20,7 +20,6 @@ package org.apache.spark.util
 import java.util.concurrent._
 import java.util.concurrent.atomic.AtomicBoolean
 
-import com.google.common.annotations.VisibleForTesting
 import org.apache.spark.SparkContext
 
 /**
@@ -119,8 +118,8 @@ private[spark] abstract class AsynchronousListenerBus[L <: AnyRef, E](name: Stri
    * For testing only. Wait until there are no more events in the queue, or until the specified
    * time has elapsed. Throw `TimeoutException` if the specified time elapsed before the queue
    * emptied.
+   * Exposed for testing.
    */
-  @VisibleForTesting
   @throws(classOf[TimeoutException])
   def waitUntilEmpty(timeoutMillis: Long): Unit = {
     val finishTime = System.currentTimeMillis + timeoutMillis
@@ -137,8 +136,8 @@ private[spark] abstract class AsynchronousListenerBus[L <: AnyRef, E](name: Stri
 
   /**
    * For testing only. Return whether the listener daemon thread is still alive.
+   * Exposed for testing.
    */
-  @VisibleForTesting
   def listenerThreadIsAlive: Boolean = listenerThread.isAlive
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala b/core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala
index a44e72b7c16d3..bd6844d045cad 100644
--- a/core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala
+++ b/core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala
@@ -23,7 +23,6 @@ import java.util.Comparator
 import scala.collection.mutable.ArrayBuffer
 import scala.collection.mutable
 
-import com.google.common.annotations.VisibleForTesting
 import com.google.common.io.ByteStreams
 
 import org.apache.spark._
@@ -608,8 +607,8 @@ private[spark] class ExternalSorter[K, V, C](
    *
    * For now, we just merge all the spilled files in once pass, but this can be modified to
    * support hierarchical merging.
+   * Exposed for testing.
    */
-  @VisibleForTesting
   def partitionedIterator: Iterator[(Int, Iterator[Product2[K, C]])] = {
     val usingMap = aggregator.isDefined
     val collection: WritablePartitionedPairCollection[K, C] = if (usingMap) map else buffer
diff --git a/scalastyle-config.xml b/scalastyle-config.xml
index 64a0c71bbef2a..050c3f360476f 100644
--- a/scalastyle-config.xml
+++ b/scalastyle-config.xml
@@ -150,6 +150,13 @@ This file is divided into 3 sections:
       // scalastyle:on println]]></customMessage>
   </check>
 
+  <check customId="visiblefortesting" level="error" class="org.scalastyle.file.RegexChecker" enabled="true">
+    <parameters><parameter name="regex">@VisibleForTesting</parameter></parameters>
+    <customMessage><![CDATA[
+      @VisibleForTesting auses classpath issues. Please note this in the java doc instead (SPARK-11615).
+    ]]></customMessage>
+  </check>
+
   <check customId="classforname" level="error" class="org.scalastyle.file.RegexChecker" enabled="true">
     <parameters><parameter name="regex">Class\.forName</parameter></parameters>
     <customMessage><![CDATA[
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index c2142d03f422b..77843f53b9bd0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql.execution
 
-import com.google.common.annotations.VisibleForTesting
-
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.SQLContext
 import org.apache.spark.sql.catalyst.InternalRow
@@ -33,7 +31,6 @@ import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
  */
 class QueryExecution(val sqlContext: SQLContext, val logical: LogicalPlan) {
 
-  @VisibleForTesting
   def assertAnalyzed(): Unit = sqlContext.analyzer.checkAnalysis(analyzed)
 
   lazy val analyzed: LogicalPlan = sqlContext.analyzer.execute(logical)
diff --git a/yarn/src/test/scala/org/apache/spark/network/shuffle/ShuffleTestAccessor.scala b/yarn/src/test/scala/org/apache/spark/network/shuffle/ShuffleTestAccessor.scala
index aa46ec5100f0e..94bf579dc8247 100644
--- a/yarn/src/test/scala/org/apache/spark/network/shuffle/ShuffleTestAccessor.scala
+++ b/yarn/src/test/scala/org/apache/spark/network/shuffle/ShuffleTestAccessor.scala
@@ -19,7 +19,6 @@ package org.apache.spark.network.shuffle
 import java.io.{IOException, File}
 import java.util.concurrent.ConcurrentMap
 
-import com.google.common.annotations.VisibleForTesting
 import org.apache.hadoop.yarn.api.records.ApplicationId
 import org.fusesource.leveldbjni.JniDBFactory
 import org.iq80.leveldb.{DB, Options}

From 6600786dddc89cb16779ee56b9173f63a3af3f27 Mon Sep 17 00:00:00 2001
From: Tathagata Das <tathagata.das1565@gmail.com>
Date: Tue, 10 Nov 2015 16:54:06 -0800
Subject: [PATCH 177/518] [SPARK-11361][STREAMING] Show scopes of RDD
 operations inside DStream.foreachRDD and DStream.transform in DAG viz

Currently, when a DStream sets the scope for RDD generated by it, that scope is not allowed to be overridden by the RDD operations. So in case of `DStream.foreachRDD`, all the RDDs generated inside the foreachRDD get the same scope - `foreachRDD  <time>`, as set by the `ForeachDStream`. So it is hard to debug generated RDDs in the RDD DAG viz in the Spark UI.

This patch allows the RDD operations inside `DStream.transform` and `DStream.foreachRDD` to append their own scopes to the earlier DStream scope.

I have also slightly tweaked how callsites are set such that the short callsite reflects the RDD operation name and line number. This tweak is necessary as callsites are not managed through scopes (which support nesting and overriding) and I didnt want to add another local property to control nesting and overriding of callsites.

## Before:
![image](https://cloud.githubusercontent.com/assets/663212/10808548/fa71c0c4-7da9-11e5-9af0-5737793a146f.png)

## After:
![image](https://cloud.githubusercontent.com/assets/663212/10808659/37bc45b6-7dab-11e5-8041-c20be6a9bc26.png)

The code that was used to generate this is:
```
    val lines = ssc.socketTextStream(args(0), args(1).toInt, StorageLevel.MEMORY_AND_DISK_SER)
    val words = lines.flatMap(_.split(" "))
    val wordCounts = words.map(x => (x, 1)).reduceByKey(_ + _)
    wordCounts.foreachRDD { rdd =>
      val temp = rdd.map { _ -> 1 }.reduceByKey( _ + _)
      val temp2 = temp.map { _ -> 1}.reduceByKey(_ + _)
      val count = temp2.count
      println(count)
    }
```

Note
- The inner scopes of the RDD operations map/reduceByKey inside foreachRDD is visible
- The short callsites of stages refers to the line number of the RDD ops rather than the same line number of foreachRDD in all three cases.

Author: Tathagata Das <tathagata.das1565@gmail.com>

Closes #9315 from tdas/SPARK-11361.
---
 .../scala/org/apache/spark/SparkContext.scala |  9 ++-
 .../spark/streaming/TestOutputStream.scala    |  2 +-
 .../spark/streaming/dstream/DStream.scala     | 63 +++++++++++++---
 .../streaming/dstream/ForEachDStream.scala    | 14 +++-
 .../dstream/TransformedDStream.scala          | 13 ++++
 .../spark/streaming/DStreamScopeSuite.scala   | 75 ++++++++++++++++---
 .../spark/streaming/TestSuiteBase.scala       |  4 +-
 7 files changed, 147 insertions(+), 33 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 7421821e2601b..67270c38fa1b5 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -1787,10 +1787,11 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
    * has overridden the call site using `setCallSite()`, this will return the user's version.
    */
   private[spark] def getCallSite(): CallSite = {
-    Option(getLocalProperty(CallSite.SHORT_FORM)).map { case shortCallSite =>
-      val longCallSite = Option(getLocalProperty(CallSite.LONG_FORM)).getOrElse("")
-      CallSite(shortCallSite, longCallSite)
-    }.getOrElse(Utils.getCallSite())
+    val callSite = Utils.getCallSite()
+    CallSite(
+      Option(getLocalProperty(CallSite.SHORT_FORM)).getOrElse(callSite.shortForm),
+      Option(getLocalProperty(CallSite.LONG_FORM)).getOrElse(callSite.longForm)
+    )
   }
 
   /**
diff --git a/external/flume/src/test/scala/org/apache/spark/streaming/TestOutputStream.scala b/external/flume/src/test/scala/org/apache/spark/streaming/TestOutputStream.scala
index 1a900007b696b..79077e4a49e1a 100644
--- a/external/flume/src/test/scala/org/apache/spark/streaming/TestOutputStream.scala
+++ b/external/flume/src/test/scala/org/apache/spark/streaming/TestOutputStream.scala
@@ -37,7 +37,7 @@ class TestOutputStream[T: ClassTag](parent: DStream[T],
   extends ForEachDStream[T](parent, (rdd: RDD[T], t: Time) => {
     val collected = rdd.collect()
     output += collected
-  }) {
+  }, false) {
 
   // This is to clear the output buffer every it is read from a checkpoint
   @throws(classOf[IOException])
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/dstream/DStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/dstream/DStream.scala
index 1da0b0a54df07..1a6edf9473d84 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/dstream/DStream.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/dstream/DStream.scala
@@ -341,7 +341,7 @@ abstract class DStream[T: ClassTag] (
       // of RDD generation, else generate nothing.
       if (isTimeValid(time)) {
 
-        val rddOption = createRDDWithLocalProperties(time) {
+        val rddOption = createRDDWithLocalProperties(time, displayInnerRDDOps = false) {
           // Disable checks for existing output directories in jobs launched by the streaming
           // scheduler, since we may need to write output to an existing directory during checkpoint
           // recovery; see SPARK-4835 for more details. We need to have this call here because
@@ -373,27 +373,52 @@ abstract class DStream[T: ClassTag] (
   /**
    * Wrap a body of code such that the call site and operation scope
    * information are passed to the RDDs created in this body properly.
-   */
-  protected def createRDDWithLocalProperties[U](time: Time)(body: => U): U = {
+   * @param body RDD creation code to execute with certain local properties.
+   * @param time Current batch time that should be embedded in the scope names
+   * @param displayInnerRDDOps Whether the detailed callsites and scopes of the inner RDDs generated
+   *                           by `body` will be displayed in the UI; only the scope and callsite
+   *                           of the DStream operation that generated `this` will be displayed.
+   */
+  protected[streaming] def createRDDWithLocalProperties[U](
+      time: Time,
+      displayInnerRDDOps: Boolean)(body: => U): U = {
     val scopeKey = SparkContext.RDD_SCOPE_KEY
     val scopeNoOverrideKey = SparkContext.RDD_SCOPE_NO_OVERRIDE_KEY
     // Pass this DStream's operation scope and creation site information to RDDs through
     // thread-local properties in our SparkContext. Since this method may be called from another
     // DStream, we need to temporarily store any old scope and creation site information to
     // restore them later after setting our own.
-    val prevCallSite = ssc.sparkContext.getCallSite()
+    val prevCallSite = CallSite(
+      ssc.sparkContext.getLocalProperty(CallSite.SHORT_FORM),
+      ssc.sparkContext.getLocalProperty(CallSite.LONG_FORM)
+    )
     val prevScope = ssc.sparkContext.getLocalProperty(scopeKey)
     val prevScopeNoOverride = ssc.sparkContext.getLocalProperty(scopeNoOverrideKey)
 
     try {
-      ssc.sparkContext.setCallSite(creationSite)
+      if (displayInnerRDDOps) {
+        // Unset the short form call site, so that generated RDDs get their own
+        ssc.sparkContext.setLocalProperty(CallSite.SHORT_FORM, null)
+        ssc.sparkContext.setLocalProperty(CallSite.LONG_FORM, null)
+      } else {
+        // Set the callsite, so that the generated RDDs get the DStream's call site and
+        // the internal RDD call sites do not get displayed
+        ssc.sparkContext.setCallSite(creationSite)
+      }
+
       // Use the DStream's base scope for this RDD so we can (1) preserve the higher level
       // DStream operation name, and (2) share this scope with other DStreams created in the
       // same operation. Disallow nesting so that low-level Spark primitives do not show up.
       // TODO: merge callsites with scopes so we can just reuse the code there
       makeScope(time).foreach { s =>
         ssc.sparkContext.setLocalProperty(scopeKey, s.toJson)
-        ssc.sparkContext.setLocalProperty(scopeNoOverrideKey, "true")
+        if (displayInnerRDDOps) {
+          // Allow inner RDDs to add inner scopes
+          ssc.sparkContext.setLocalProperty(scopeNoOverrideKey, null)
+        } else {
+          // Do not allow inner RDDs to override the scope set by DStream
+          ssc.sparkContext.setLocalProperty(scopeNoOverrideKey, "true")
+        }
       }
 
       body
@@ -628,7 +653,7 @@ abstract class DStream[T: ClassTag] (
    */
   def foreachRDD(foreachFunc: RDD[T] => Unit): Unit = ssc.withScope {
     val cleanedF = context.sparkContext.clean(foreachFunc, false)
-    this.foreachRDD((r: RDD[T], t: Time) => cleanedF(r))
+    foreachRDD((r: RDD[T], t: Time) => cleanedF(r), displayInnerRDDOps = true)
   }
 
   /**
@@ -639,7 +664,23 @@ abstract class DStream[T: ClassTag] (
     // because the DStream is reachable from the outer object here, and because
     // DStreams can't be serialized with closures, we can't proactively check
     // it for serializability and so we pass the optional false to SparkContext.clean
-    new ForEachDStream(this, context.sparkContext.clean(foreachFunc, false)).register()
+    foreachRDD(foreachFunc, displayInnerRDDOps = true)
+  }
+
+  /**
+   * Apply a function to each RDD in this DStream. This is an output operator, so
+   * 'this' DStream will be registered as an output stream and therefore materialized.
+   * @param foreachFunc foreachRDD function
+   * @param displayInnerRDDOps Whether the detailed callsites and scopes of the RDDs generated
+   *                           in the `foreachFunc` to be displayed in the UI. If `false`, then
+   *                           only the scopes and callsites of `foreachRDD` will override those
+   *                           of the RDDs on the display.
+   */
+  private def foreachRDD(
+      foreachFunc: (RDD[T], Time) => Unit,
+      displayInnerRDDOps: Boolean): Unit = {
+    new ForEachDStream(this,
+      context.sparkContext.clean(foreachFunc, false), displayInnerRDDOps).register()
   }
 
   /**
@@ -730,7 +771,7 @@ abstract class DStream[T: ClassTag] (
         // scalastyle:on println
       }
     }
-    new ForEachDStream(this, context.sparkContext.clean(foreachFunc)).register()
+    foreachRDD(context.sparkContext.clean(foreachFunc), displayInnerRDDOps = false)
   }
 
   /**
@@ -900,7 +941,7 @@ abstract class DStream[T: ClassTag] (
       val file = rddToFileName(prefix, suffix, time)
       rdd.saveAsObjectFile(file)
     }
-    this.foreachRDD(saveFunc)
+    this.foreachRDD(saveFunc, displayInnerRDDOps = false)
   }
 
   /**
@@ -913,7 +954,7 @@ abstract class DStream[T: ClassTag] (
       val file = rddToFileName(prefix, suffix, time)
       rdd.saveAsTextFile(file)
     }
-    this.foreachRDD(saveFunc)
+    this.foreachRDD(saveFunc, displayInnerRDDOps = false)
   }
 
   /**
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/dstream/ForEachDStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/dstream/ForEachDStream.scala
index c109ceccc6989..4410a9977c87b 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/dstream/ForEachDStream.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/dstream/ForEachDStream.scala
@@ -22,10 +22,19 @@ import org.apache.spark.streaming.{Duration, Time}
 import org.apache.spark.streaming.scheduler.Job
 import scala.reflect.ClassTag
 
+/**
+ * An internal DStream used to represent output operations like DStream.foreachRDD.
+ * @param parent        Parent DStream
+ * @param foreachFunc   Function to apply on each RDD generated by the parent DStream
+ * @param displayInnerRDDOps Whether the detailed callsites and scopes of the RDDs generated
+ *                           by `foreachFunc` will be displayed in the UI; only the scope and
+ *                           callsite of `DStream.foreachRDD` will be displayed.
+ */
 private[streaming]
 class ForEachDStream[T: ClassTag] (
     parent: DStream[T],
-    foreachFunc: (RDD[T], Time) => Unit
+    foreachFunc: (RDD[T], Time) => Unit,
+    displayInnerRDDOps: Boolean
   ) extends DStream[Unit](parent.ssc) {
 
   override def dependencies: List[DStream[_]] = List(parent)
@@ -37,8 +46,7 @@ class ForEachDStream[T: ClassTag] (
   override def generateJob(time: Time): Option[Job] = {
     parent.getOrCompute(time) match {
       case Some(rdd) =>
-        val jobFunc = () => createRDDWithLocalProperties(time) {
-          ssc.sparkContext.setCallSite(creationSite)
+        val jobFunc = () => createRDDWithLocalProperties(time, displayInnerRDDOps) {
           foreachFunc(rdd, time)
         }
         Some(new Job(time, jobFunc))
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/dstream/TransformedDStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/dstream/TransformedDStream.scala
index 5eabdf63dc8d7..080bc873fa0a8 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/dstream/TransformedDStream.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/dstream/TransformedDStream.scala
@@ -51,4 +51,17 @@ class TransformedDStream[U: ClassTag] (
     }
     Some(transformedRDD)
   }
+
+  /**
+   * Wrap a body of code such that the call site and operation scope
+   * information are passed to the RDDs created in this body properly.
+   * This has been overriden to make sure that `displayInnerRDDOps` is always `true`, that is,
+   * the inner scopes and callsites of RDDs generated in `DStream.transform` are always
+   * displayed in the UI.
+   */
+  override protected[streaming] def createRDDWithLocalProperties[U](
+      time: Time,
+      displayInnerRDDOps: Boolean)(body: => U): U = {
+    super.createRDDWithLocalProperties(time, displayInnerRDDOps = true)(body)
+  }
 }
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/DStreamScopeSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/DStreamScopeSuite.scala
index 8844c9d74b933..bc223e648a417 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/DStreamScopeSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/DStreamScopeSuite.scala
@@ -17,12 +17,15 @@
 
 package org.apache.spark.streaming
 
+import scala.collection.mutable.ArrayBuffer
+
 import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll}
 
-import org.apache.spark.{SparkContext, SparkFunSuite}
-import org.apache.spark.rdd.RDDOperationScope
+import org.apache.spark.rdd.{RDD, RDDOperationScope}
 import org.apache.spark.streaming.dstream.DStream
 import org.apache.spark.streaming.ui.UIUtils
+import org.apache.spark.util.ManualClock
+import org.apache.spark.{SparkConf, SparkContext, SparkFunSuite}
 
 /**
  * Tests whether scope information is passed from DStream operations to RDDs correctly.
@@ -32,7 +35,9 @@ class DStreamScopeSuite extends SparkFunSuite with BeforeAndAfter with BeforeAnd
   private val batchDuration: Duration = Seconds(1)
 
   override def beforeAll(): Unit = {
-    ssc = new StreamingContext(new SparkContext("local", "test"), batchDuration)
+    val conf = new SparkConf().setMaster("local").setAppName("test")
+    conf.set("spark.streaming.clock", classOf[ManualClock].getName())
+    ssc = new StreamingContext(new SparkContext(conf), batchDuration)
   }
 
   override def afterAll(): Unit = {
@@ -103,6 +108,8 @@ class DStreamScopeSuite extends SparkFunSuite with BeforeAndAfter with BeforeAnd
 
   test("scoping nested operations") {
     val inputStream = new DummyInputDStream(ssc)
+    // countByKeyAndWindow internally uses reduceByKeyAndWindow, but only countByKeyAndWindow
+    // should appear in scope
     val countStream = inputStream.countByWindow(Seconds(10), Seconds(1))
     countStream.initialize(Time(0))
 
@@ -137,6 +144,57 @@ class DStreamScopeSuite extends SparkFunSuite with BeforeAndAfter with BeforeAnd
     testStream(countStream)
   }
 
+  test("transform should allow RDD operations to be captured in scopes") {
+    val inputStream = new DummyInputDStream(ssc)
+    val transformedStream = inputStream.transform { _.map { _ -> 1}.reduceByKey(_ + _) }
+    transformedStream.initialize(Time(0))
+
+    val transformScopeBase = transformedStream.baseScope.map(RDDOperationScope.fromJson)
+    val transformScope1 = transformedStream.getOrCompute(Time(1000)).get.scope
+    val transformScope2 = transformedStream.getOrCompute(Time(2000)).get.scope
+    val transformScope3 = transformedStream.getOrCompute(Time(3000)).get.scope
+
+    // Assert that all children RDDs inherit the DStream operation name correctly
+    assertDefined(transformScopeBase, transformScope1, transformScope2, transformScope3)
+    assert(transformScopeBase.get.name === "transform")
+    assertNestedScopeCorrect(transformScope1.get, 1000)
+    assertNestedScopeCorrect(transformScope2.get, 2000)
+    assertNestedScopeCorrect(transformScope3.get, 3000)
+
+    def assertNestedScopeCorrect(rddScope: RDDOperationScope, batchTime: Long): Unit = {
+      assert(rddScope.name === "reduceByKey")
+      assert(rddScope.parent.isDefined)
+      assertScopeCorrect(transformScopeBase.get, rddScope.parent.get, batchTime)
+    }
+  }
+
+  test("foreachRDD should allow RDD operations to be captured in scope") {
+    val inputStream = new DummyInputDStream(ssc)
+    val generatedRDDs = new ArrayBuffer[RDD[(Int, Int)]]
+    inputStream.foreachRDD { rdd =>
+      generatedRDDs += rdd.map { _ -> 1}.reduceByKey(_ + _)
+    }
+    val batchCounter = new BatchCounter(ssc)
+    ssc.start()
+    val clock = ssc.scheduler.clock.asInstanceOf[ManualClock]
+    clock.advance(3000)
+    batchCounter.waitUntilBatchesCompleted(3, 10000)
+    assert(generatedRDDs.size === 3)
+
+    val foreachBaseScope =
+      ssc.graph.getOutputStreams().head.baseScope.map(RDDOperationScope.fromJson)
+    assertDefined(foreachBaseScope)
+    assert(foreachBaseScope.get.name === "foreachRDD")
+
+    val rddScopes = generatedRDDs.map { _.scope }
+    assertDefined(rddScopes: _*)
+    rddScopes.zipWithIndex.foreach { case (rddScope, idx) =>
+      assert(rddScope.get.name === "reduceByKey")
+      assert(rddScope.get.parent.isDefined)
+      assertScopeCorrect(foreachBaseScope.get, rddScope.get.parent.get, (idx + 1) * 1000)
+    }
+  }
+
   /** Assert that the RDD operation scope properties are not set in our SparkContext. */
   private def assertPropertiesNotSet(): Unit = {
     assert(ssc != null)
@@ -149,19 +207,12 @@ class DStreamScopeSuite extends SparkFunSuite with BeforeAndAfter with BeforeAnd
       baseScope: RDDOperationScope,
       rddScope: RDDOperationScope,
       batchTime: Long): Unit = {
-    assertScopeCorrect(baseScope.id, baseScope.name, rddScope, batchTime)
-  }
-
-  /** Assert that the given RDD scope inherits the base name and ID correctly. */
-  private def assertScopeCorrect(
-      baseScopeId: String,
-      baseScopeName: String,
-      rddScope: RDDOperationScope,
-      batchTime: Long): Unit = {
+    val (baseScopeId, baseScopeName) = (baseScope.id, baseScope.name)
     val formattedBatchTime = UIUtils.formatBatchTime(
       batchTime, ssc.graph.batchDuration.milliseconds, showYYYYMMSS = false)
     assert(rddScope.id === s"${baseScopeId}_$batchTime")
     assert(rddScope.name.replaceAll("\\n", " ") === s"$baseScopeName @ $formattedBatchTime")
+    assert(rddScope.parent.isEmpty)  // There should not be any higher scope
   }
 
   /** Assert that all the specified options are defined. */
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/TestSuiteBase.scala b/streaming/src/test/scala/org/apache/spark/streaming/TestSuiteBase.scala
index 0d58a7b54412f..a45c92d9c7bc8 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/TestSuiteBase.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/TestSuiteBase.scala
@@ -98,7 +98,7 @@ class TestOutputStream[T: ClassTag](
   ) extends ForEachDStream[T](parent, (rdd: RDD[T], t: Time) => {
     val collected = rdd.collect()
     output += collected
-  }) {
+  }, false) {
 
   // This is to clear the output buffer every it is read from a checkpoint
   @throws(classOf[IOException])
@@ -122,7 +122,7 @@ class TestOutputStreamWithPartitions[T: ClassTag](
   extends ForEachDStream[T](parent, (rdd: RDD[T], t: Time) => {
     val collected = rdd.glom().collect().map(_.toSeq)
     output += collected
-  }) {
+  }, false) {
 
   // This is to clear the output buffer every it is read from a checkpoint
   @throws(classOf[IOException])

From 12c7635dc025239d3b69b9adef2f4eebb28edf48 Mon Sep 17 00:00:00 2001
From: Forest Fang <saurfang@users.noreply.github.com>
Date: Tue, 10 Nov 2015 16:55:58 -0800
Subject: [PATCH 178/518] [MINOR] Fix typo in AggregationQuerySuite.scala

Author: Forest Fang <saurfang@users.noreply.github.com>

Closes #9357 from saurfang/patch-1.
---
 .../spark/sql/hive/execution/AggregationQuerySuite.scala      | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 8253921563b3a..22d2aefd699b5 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -795,14 +795,14 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
       val df = sqlContext.createDataFrame(rdd, schema)
 
       val allColumns = df.schema.fields.map(f => col(f.name))
-      val expectedAnaswer =
+      val expectedAnswer =
         data
           .find(r => r.getInt(0) == 50)
           .getOrElse(fail("A row with id 50 should be the expected answer."))
       checkAnswer(
         df.groupBy().agg(udaf(allColumns: _*)),
         // udaf returns a Row as the output value.
-        Row(expectedAnaswer)
+        Row(expectedAnswer)
       )
     }
   }

From 745e45d5ff7fe251c0d5197b7e08b1f80807b005 Mon Sep 17 00:00:00 2001
From: Marc Prud'hommeaux <mwp1@cornell.edu>
Date: Tue, 10 Nov 2015 16:57:08 -0800
Subject: [PATCH 179/518] [MINOR] License header formatting fix

The header wasn't indented properly.

Author: Marc Prud'hommeaux <mwp1@cornell.edu>

Closes #9312 from mprudhom/patch-1.
---
 .../datasources/jdbc/DefaultSource.scala      | 30 +++++++++----------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/DefaultSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/DefaultSource.scala
index 6773afc794f9c..f522303be94ad 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/DefaultSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/DefaultSource.scala
@@ -1,19 +1,19 @@
 /*
-* Licensed to the Apache Software Foundation (ASF) under one or more
-* contributor license agreements.  See the NOTICE file distributed with
-* this work for additional information regarding copyright ownership.
-* The ASF licenses this file to You under the Apache License, Version 2.0
-* (the "License"); you may not use this file except in compliance with
-* the License.  You may obtain a copy of the License at
-*
-*    http://www.apache.org/licenses/LICENSE-2.0
-*
-* Unless required by applicable law or agreed to in writing, software
-* distributed under the License is distributed on an "AS IS" BASIS,
-* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-* See the License for the specific language governing permissions and
-* limitations under the License.
-*/
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 
 package org.apache.spark.sql.execution.datasources.jdbc
 

From 6e101d2e9d6e08a6a63f7065c1e87a5338f763ea Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Tue, 10 Nov 2015 18:45:48 -0800
Subject: [PATCH 180/518] [SPARK-6726][ML] Import/export for spark.ml
 LogisticRegressionModel

This PR adds model save/load for spark.ml's LogisticRegressionModel.  It also does minor refactoring of the default save/load classes to reuse code.

CC: mengxr

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9606 from jkbradley/logreg-io2.
---
 .../classification/LogisticRegression.scala   | 68 ++++++++++++++++-
 .../org/apache/spark/ml/util/ReadWrite.scala  | 74 +++++++++++++++++--
 .../LogisticRegressionSuite.scala             | 17 ++++-
 .../spark/ml/util/DefaultReadWriteTest.scala  |  4 +-
 4 files changed, 152 insertions(+), 11 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
index f5fca686df144..a88f52674102c 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
@@ -21,13 +21,14 @@ import scala.collection.mutable
 
 import breeze.linalg.{DenseVector => BDV}
 import breeze.optimize.{CachedDiffFunction, DiffFunction, LBFGS => BreezeLBFGS, OWLQN => BreezeOWLQN}
+import org.apache.hadoop.fs.Path
 
 import org.apache.spark.{Logging, SparkException}
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.ml.feature.Instance
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg._
 import org.apache.spark.mllib.linalg.BLAS._
 import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics
@@ -396,7 +397,7 @@ class LogisticRegressionModel private[ml] (
     val coefficients: Vector,
     val intercept: Double)
   extends ProbabilisticClassificationModel[Vector, LogisticRegressionModel]
-  with LogisticRegressionParams {
+  with LogisticRegressionParams with Writable {
 
   @deprecated("Use coefficients instead.", "1.6.0")
   def weights: Vector = coefficients
@@ -510,8 +511,71 @@ class LogisticRegressionModel private[ml] (
     // Note: We should use getThreshold instead of $(threshold) since getThreshold is overridden.
     if (probability(1) > getThreshold) 1 else 0
   }
+
+  /**
+   * Returns a [[Writer]] instance for this ML instance.
+   *
+   * For [[LogisticRegressionModel]], this does NOT currently save the training [[summary]].
+   * An option to save [[summary]] may be added in the future.
+   */
+  override def write: Writer = new LogisticRegressionWriter(this)
+}
+
+
+/** [[Writer]] instance for [[LogisticRegressionModel]] */
+private[classification] class LogisticRegressionWriter(instance: LogisticRegressionModel)
+  extends Writer with Logging {
+
+  private case class Data(
+      numClasses: Int,
+      numFeatures: Int,
+      intercept: Double,
+      coefficients: Vector)
+
+  override protected def saveImpl(path: String): Unit = {
+    // Save metadata and Params
+    DefaultParamsWriter.saveMetadata(instance, path, sc)
+    // Save model data: numClasses, numFeatures, intercept, coefficients
+    val data = Data(instance.numClasses, instance.numFeatures, instance.intercept,
+      instance.coefficients)
+    val dataPath = new Path(path, "data").toString
+    sqlContext.createDataFrame(Seq(data)).write.format("parquet").save(dataPath)
+  }
+}
+
+
+object LogisticRegressionModel extends Readable[LogisticRegressionModel] {
+
+  override def read: Reader[LogisticRegressionModel] = new LogisticRegressionReader
+
+  override def load(path: String): LogisticRegressionModel = read.load(path)
 }
 
+
+private[classification] class LogisticRegressionReader extends Reader[LogisticRegressionModel] {
+
+  /** Checked against metadata when loading model */
+  private val className = "org.apache.spark.ml.classification.LogisticRegressionModel"
+
+  override def load(path: String): LogisticRegressionModel = {
+    val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+
+    val dataPath = new Path(path, "data").toString
+    val data = sqlContext.read.format("parquet").load(dataPath)
+      .select("numClasses", "numFeatures", "intercept", "coefficients").head()
+    // We will need numClasses, numFeatures in the future for multinomial logreg support.
+    // val numClasses = data.getInt(0)
+    // val numFeatures = data.getInt(1)
+    val intercept = data.getDouble(2)
+    val coefficients = data.getAs[Vector](3)
+    val model = new LogisticRegressionModel(metadata.uid, coefficients, intercept)
+
+    DefaultParamsReader.getAndSetParams(model, metadata)
+    model
+  }
+}
+
+
 /**
  * MultiClassSummarizer computes the number of distinct labels and corresponding counts,
  * and validates the data to see if the labels used for k class multi-label classification
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index cbdf913ba8dfa..85f888c9f2f67 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -175,6 +175,21 @@ trait Readable[T] {
 private[ml] class DefaultParamsWriter(instance: Params) extends Writer {
 
   override protected def saveImpl(path: String): Unit = {
+    DefaultParamsWriter.saveMetadata(instance, path, sc)
+  }
+}
+
+private[ml] object DefaultParamsWriter {
+
+  /**
+   * Saves metadata + Params to: path + "/metadata"
+   *  - class
+   *  - timestamp
+   *  - sparkVersion
+   *  - uid
+   *  - paramMap
+   */
+  def saveMetadata(instance: Params, path: String, sc: SparkContext): Unit = {
     val uid = instance.uid
     val cls = instance.getClass.getName
     val params = instance.extractParamMap().toSeq.asInstanceOf[Seq[ParamPair[Any]]]
@@ -201,14 +216,61 @@ private[ml] class DefaultParamsWriter(instance: Params) extends Writer {
 private[ml] class DefaultParamsReader[T] extends Reader[T] {
 
   override def load(path: String): T = {
-    implicit val format = DefaultFormats
+    val metadata = DefaultParamsReader.loadMetadata(path, sc)
+    val cls = Utils.classForName(metadata.className)
+    val instance =
+      cls.getConstructor(classOf[String]).newInstance(metadata.uid).asInstanceOf[Params]
+    DefaultParamsReader.getAndSetParams(instance, metadata)
+    instance.asInstanceOf[T]
+  }
+}
+
+private[ml] object DefaultParamsReader {
+
+  /**
+   * All info from metadata file.
+   * @param params  paramMap, as a [[JValue]]
+   * @param metadataStr  Full metadata file String (for debugging)
+   */
+  case class Metadata(
+      className: String,
+      uid: String,
+      timestamp: Long,
+      sparkVersion: String,
+      params: JValue,
+      metadataStr: String)
+
+  /**
+   * Load metadata from file.
+   * @param expectedClassName  If non empty, this is checked against the loaded metadata.
+   * @throws IllegalArgumentException if expectedClassName is specified and does not match metadata
+   */
+  def loadMetadata(path: String, sc: SparkContext, expectedClassName: String = ""): Metadata = {
     val metadataPath = new Path(path, "metadata").toString
     val metadataStr = sc.textFile(metadataPath, 1).first()
     val metadata = parse(metadataStr)
-    val cls = Utils.classForName((metadata \ "class").extract[String])
+
+    implicit val format = DefaultFormats
+    val className = (metadata \ "class").extract[String]
     val uid = (metadata \ "uid").extract[String]
-    val instance = cls.getConstructor(classOf[String]).newInstance(uid).asInstanceOf[Params]
-    (metadata \ "paramMap") match {
+    val timestamp = (metadata \ "timestamp").extract[Long]
+    val sparkVersion = (metadata \ "sparkVersion").extract[String]
+    val params = metadata \ "paramMap"
+    if (expectedClassName.nonEmpty) {
+      require(className == expectedClassName, s"Error loading metadata: Expected class name" +
+        s" $expectedClassName but found class name $className")
+    }
+
+    Metadata(className, uid, timestamp, sparkVersion, params, metadataStr)
+  }
+
+  /**
+   * Extract Params from metadata, and set them in the instance.
+   * This works if all Params implement [[org.apache.spark.ml.param.Param.jsonDecode()]].
+   */
+  def getAndSetParams(instance: Params, metadata: Metadata): Unit = {
+    implicit val format = DefaultFormats
+    metadata.params match {
       case JObject(pairs) =>
         pairs.foreach { case (paramName, jsonValue) =>
           val param = instance.getParam(paramName)
@@ -216,8 +278,8 @@ private[ml] class DefaultParamsReader[T] extends Reader[T] {
           instance.set(param, value)
         }
       case _ =>
-        throw new IllegalArgumentException(s"Cannot recognize JSON metadata: $metadataStr.")
+        throw new IllegalArgumentException(
+          s"Cannot recognize JSON metadata: ${metadata.metadataStr}.")
     }
-    instance.asInstanceOf[T]
   }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
index 325faf37e8eea..51b06b7eb6d53 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
@@ -23,7 +23,7 @@ import scala.util.Random
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.feature.Instance
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{Identifiable, DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.classification.LogisticRegressionSuite._
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.regression.LabeledPoint
@@ -31,7 +31,8 @@ import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.{DataFrame, Row}
 
-class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
+class LogisticRegressionSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   @transient var dataset: DataFrame = _
   @transient var binaryDataset: DataFrame = _
@@ -869,6 +870,18 @@ class LogisticRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(model1a0.intercept !~= model1a1.intercept absTol 1E-3)
     assert(model1a0.coefficients ~== model1b.coefficients absTol 1E-3)
     assert(model1a0.intercept ~== model1b.intercept absTol 1E-3)
+  }
 
+  test("read/write") {
+    // Set some Params to make sure set Params are serialized.
+    val lr = new LogisticRegression()
+      .setElasticNetParam(0.1)
+      .setMaxIter(2)
+      .fit(dataset)
+    val lr2 = testDefaultReadWrite(lr)
+    assert(lr.intercept === lr2.intercept)
+    assert(lr.coefficients.toArray === lr2.coefficients.toArray)
+    assert(lr.numClasses === lr2.numClasses)
+    assert(lr.numFeatures === lr2.numFeatures)
   }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
index 4545b0f281f5a..cac4bd9aa3ab8 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
@@ -31,8 +31,9 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
    * Checks "overwrite" option and params.
    * @param instance ML instance to test saving/loading
    * @tparam T ML instance type
+   * @return  Instance loaded from file
    */
-  def testDefaultReadWrite[T <: Params with Writable](instance: T): Unit = {
+  def testDefaultReadWrite[T <: Params with Writable](instance: T): T = {
     val uid = instance.uid
     val path = new File(tempDir, uid).getPath
 
@@ -61,6 +62,7 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
     val load = instance.getClass.getMethod("load", classOf[String])
     val another = load.invoke(instance, path).asInstanceOf[T]
     assert(another.uid === instance.uid)
+    another
   }
 }
 

From fac53d8ec015e27d034dfe30ed8ce7d83f07efa6 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Tue, 10 Nov 2015 22:24:00 -0800
Subject: [PATCH 181/518] [SPARK-10192][HOTFIX] Fix NPE in test that was added
 in #8402

This fixes an NPE introduced in SPARK-10192 / #8402.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9620 from JoshRosen/SPARK-10192-hotfix.
---
 .../scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
index 068b49bd5844b..4d6b25455226f 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
@@ -1070,7 +1070,7 @@ class DAGSchedulerSuite
     // Run the first job successfully, which creates one shuffle dependency
 
     val shuffleMapRdd = new MyRDD(sc, 2, Nil)
-    val shuffleDep = new ShuffleDependency(shuffleMapRdd, null)
+    val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(2))
     val reduceRdd = new MyRDD(sc, 2, List(shuffleDep))
     submit(reduceRdd, Array(0, 1))
 

From 1a8e0468a1c07e99ad395eb0e4dc072c5cf7393a Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Tue, 10 Nov 2015 22:45:17 -0800
Subject: [PATCH 182/518] [SPARK-11468] [SPARKR] add stddev/variance agg
 functions for Column

Checked names, none of them should conflict with anything in base

shivaram davies rxin

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9489 from felixcheung/rstddev.
---
 R/pkg/NAMESPACE                  |  10 ++
 R/pkg/R/functions.R              | 186 +++++++++++++++++++++++++++++--
 R/pkg/R/generics.R               |  40 +++++++
 R/pkg/R/group.R                  |   8 +-
 R/pkg/inst/tests/test_sparkSQL.R |  83 +++++++++++---
 5 files changed, 297 insertions(+), 30 deletions(-)

diff --git a/R/pkg/NAMESPACE b/R/pkg/NAMESPACE
index 52fd6c9f76c54..2ee7d6f94f1bc 100644
--- a/R/pkg/NAMESPACE
+++ b/R/pkg/NAMESPACE
@@ -155,6 +155,7 @@ exportMethods("%in%",
               "isNaN",
               "isNotNull",
               "isNull",
+              "kurtosis",
               "lag",
               "last",
               "last_day",
@@ -207,12 +208,17 @@ exportMethods("%in%",
               "shiftLeft",
               "shiftRight",
               "shiftRightUnsigned",
+              "sd",
               "sign",
               "signum",
               "sin",
               "sinh",
               "size",
+              "skewness",
               "soundex",
+              "stddev",
+              "stddev_pop",
+              "stddev_samp",
               "sqrt",
               "startsWith",
               "substr",
@@ -231,6 +237,10 @@ exportMethods("%in%",
               "unhex",
               "unix_timestamp",
               "upper",
+              "var",
+              "variance",
+              "var_pop",
+              "var_samp",
               "weekofyear",
               "when",
               "year")
diff --git a/R/pkg/R/functions.R b/R/pkg/R/functions.R
index 0b280870295a2..3d0255a62f155 100644
--- a/R/pkg/R/functions.R
+++ b/R/pkg/R/functions.R
@@ -520,6 +520,22 @@ setMethod("isNaN",
             column(jc)
           })
 
+#' kurtosis
+#'
+#' Aggregate function: returns the kurtosis of the values in a group.
+#'
+#' @rdname kurtosis
+#' @name kurtosis
+#' @family agg_funcs
+#' @export
+#' @examples \dontrun{kurtosis(df$c)}
+setMethod("kurtosis",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "kurtosis", x@jc)
+            column(jc)
+          })
+
 #' last
 #'
 #' Aggregate function: returns the last value in a group.
@@ -861,6 +877,28 @@ setMethod("rtrim",
             column(jc)
           })
 
+#' sd
+#'
+#' Aggregate function: alias for \link{stddev_samp}
+#'
+#' @rdname sd
+#' @name sd
+#' @family agg_funcs
+#' @seealso \link{stddev_pop}, \link{stddev_samp}
+#' @export
+#' @examples
+#'\dontrun{
+#'stddev(df$c)
+#'select(df, stddev(df$age))
+#'agg(df, sd(df$age))
+#'}
+setMethod("sd",
+          signature(x = "Column"),
+          function(x, na.rm = FALSE) {
+            # In R, sample standard deviation is calculated with the sd() function.
+            stddev_samp(x)
+          })
+
 #' second
 #'
 #' Extracts the seconds as an integer from a given date/timestamp/string.
@@ -958,6 +996,22 @@ setMethod("size",
             column(jc)
           })
 
+#' skewness
+#'
+#' Aggregate function: returns the skewness of the values in a group.
+#'
+#' @rdname skewness
+#' @name skewness
+#' @family agg_funcs
+#' @export
+#' @examples \dontrun{skewness(df$c)}
+setMethod("skewness",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "skewness", x@jc)
+            column(jc)
+          })
+
 #' soundex
 #'
 #' Return the soundex code for the specified expression.
@@ -974,6 +1028,49 @@ setMethod("soundex",
             column(jc)
           })
 
+#' @rdname sd
+#' @name stddev
+setMethod("stddev",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "stddev", x@jc)
+            column(jc)
+          })
+
+#' stddev_pop
+#'
+#' Aggregate function: returns the population standard deviation of the expression in a group.
+#'
+#' @rdname stddev_pop
+#' @name stddev_pop
+#' @family agg_funcs
+#' @seealso \link{sd}, \link{stddev_samp}
+#' @export
+#' @examples \dontrun{stddev_pop(df$c)}
+setMethod("stddev_pop",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "stddev_pop", x@jc)
+            column(jc)
+          })
+
+#' stddev_samp
+#'
+#' Aggregate function: returns the unbiased sample standard deviation of the expression in a group.
+#'
+#' @rdname stddev_samp
+#' @name stddev_samp
+#' @family agg_funcs
+#' @seealso \link{stddev_pop}, \link{sd}
+#' @export
+#' @examples \dontrun{stddev_samp(df$c)}
+setMethod("stddev_samp",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "stddev_samp", x@jc)
+            column(jc)
+          })
+
 #' sqrt
 #'
 #' Computes the square root of the specified float value.
@@ -1168,6 +1265,71 @@ setMethod("upper",
             column(jc)
           })
 
+#' var
+#'
+#' Aggregate function: alias for \link{var_samp}.
+#'
+#' @rdname var
+#' @name var
+#' @family agg_funcs
+#' @seealso \link{var_pop}, \link{var_samp}
+#' @export
+#' @examples
+#'\dontrun{
+#'variance(df$c)
+#'select(df, var_pop(df$age))
+#'agg(df, var(df$age))
+#'}
+setMethod("var",
+          signature(x = "Column"),
+          function(x, y = NULL, na.rm = FALSE, use) {
+            # In R, sample variance is calculated with the var() function.
+            var_samp(x)
+          })
+
+#' @rdname var
+#' @name variance
+setMethod("variance",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "variance", x@jc)
+            column(jc)
+          })
+
+#' var_pop
+#'
+#' Aggregate function: returns the population variance of the values in a group.
+#'
+#' @rdname var_pop
+#' @name var_pop
+#' @family agg_funcs
+#' @seealso \link{var}, \link{var_samp}
+#' @export
+#' @examples \dontrun{var_pop(df$c)}
+setMethod("var_pop",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "var_pop", x@jc)
+            column(jc)
+          })
+
+#' var_samp
+#'
+#' Aggregate function: returns the unbiased variance of the values in a group.
+#'
+#' @rdname var_samp
+#' @name var_samp
+#' @family agg_funcs
+#' @seealso \link{var_pop}, \link{var}
+#' @export
+#' @examples \dontrun{var_samp(df$c)}
+setMethod("var_samp",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "var_samp", x@jc)
+            column(jc)
+          })
+
 #' weekofyear
 #'
 #' Extracts the week number as an integer from a given date/timestamp/string.
@@ -2020,10 +2182,10 @@ setMethod("ifelse",
 #'
 #' Window function: returns the cumulative distribution of values within a window partition,
 #' i.e. the fraction of rows that are below the current row.
-#' 
+#'
 #'   N = total number of rows in the partition
 #'   cumeDist(x) = number of values before (and including) x / N
-#'   
+#'
 #' This is equivalent to the CUME_DIST function in SQL.
 #'
 #' @rdname cumeDist
@@ -2039,13 +2201,13 @@ setMethod("cumeDist",
           })
 
 #' denseRank
-#' 
+#'
 #' Window function: returns the rank of rows within a window partition, without any gaps.
 #' The difference between rank and denseRank is that denseRank leaves no gaps in ranking
 #' sequence when there are ties. That is, if you were ranking a competition using denseRank
 #' and had three people tie for second place, you would say that all three were in second
 #' place and that the next person came in third.
-#' 
+#'
 #' This is equivalent to the DENSE_RANK function in SQL.
 #'
 #' @rdname denseRank
@@ -2065,7 +2227,7 @@ setMethod("denseRank",
 #' Window function: returns the value that is `offset` rows before the current row, and
 #' `defaultValue` if there is less than `offset` rows before the current row. For example,
 #' an `offset` of one will return the previous row at any given point in the window partition.
-#' 
+#'
 #' This is equivalent to the LAG function in SQL.
 #'
 #' @rdname lag
@@ -2092,7 +2254,7 @@ setMethod("lag",
 #' Window function: returns the value that is `offset` rows after the current row, and
 #' `null` if there is less than `offset` rows after the current row. For example,
 #' an `offset` of one will return the next row at any given point in the window partition.
-#' 
+#'
 #' This is equivalent to the LEAD function in SQL.
 #'
 #' @rdname lead
@@ -2119,7 +2281,7 @@ setMethod("lead",
 #' Window function: returns the ntile group id (from 1 to `n` inclusive) in an ordered window
 #' partition. Fow example, if `n` is 4, the first quarter of the rows will get value 1, the second
 #' quarter will get 2, the third quarter will get 3, and the last quarter will get 4.
-#' 
+#'
 #' This is equivalent to the NTILE function in SQL.
 #'
 #' @rdname ntile
@@ -2137,9 +2299,9 @@ setMethod("ntile",
 #' percentRank
 #'
 #' Window function: returns the relative rank (i.e. percentile) of rows within a window partition.
-#' 
+#'
 #' This is computed by:
-#' 
+#'
 #'   (rank of row in its partition - 1) / (number of rows in the partition - 1)
 #'
 #' This is equivalent to the PERCENT_RANK function in SQL.
@@ -2159,12 +2321,12 @@ setMethod("percentRank",
 #' rank
 #'
 #' Window function: returns the rank of rows within a window partition.
-#' 
+#'
 #' The difference between rank and denseRank is that denseRank leaves no gaps in ranking
 #' sequence when there are ties. That is, if you were ranking a competition using denseRank
 #' and had three people tie for second place, you would say that all three were in second
 #' place and that the next person came in third.
-#' 
+#'
 #' This is equivalent to the RANK function in SQL.
 #'
 #' @rdname rank
@@ -2189,7 +2351,7 @@ setMethod("rank",
 #' rowNumber
 #'
 #' Window function: returns a sequential number starting at 1 within a window partition.
-#' 
+#'
 #' This is equivalent to the ROW_NUMBER function in SQL.
 #'
 #' @rdname rowNumber
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 89731affeb898..92ad4ee8685ee 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -798,6 +798,10 @@ setGeneric("instr", function(y, x) { standardGeneric("instr") })
 #' @export
 setGeneric("isNaN", function(x) { standardGeneric("isNaN") })
 
+#' @rdname kurtosis
+#' @export
+setGeneric("kurtosis", function(x) { standardGeneric("kurtosis") })
+
 #' @rdname lag
 #' @export
 setGeneric("lag", function(x, offset, defaultValue = NULL) { standardGeneric("lag") })
@@ -935,6 +939,10 @@ setGeneric("rpad", function(x, len, pad) { standardGeneric("rpad") })
 #' @export
 setGeneric("rtrim", function(x) { standardGeneric("rtrim") })
 
+#' @rdname sd
+#' @export
+setGeneric("sd", function(x, na.rm = FALSE) { standardGeneric("sd") })
+
 #' @rdname second
 #' @export
 setGeneric("second", function(x) { standardGeneric("second") })
@@ -967,10 +975,26 @@ setGeneric("signum", function(x) { standardGeneric("signum") })
 #' @export
 setGeneric("size", function(x) { standardGeneric("size") })
 
+#' @rdname skewness
+#' @export
+setGeneric("skewness", function(x) { standardGeneric("skewness") })
+
 #' @rdname soundex
 #' @export
 setGeneric("soundex", function(x) { standardGeneric("soundex") })
 
+#' @rdname sd
+#' @export
+setGeneric("stddev", function(x) { standardGeneric("stddev") })
+
+#' @rdname stddev_pop
+#' @export
+setGeneric("stddev_pop", function(x) { standardGeneric("stddev_pop") })
+
+#' @rdname stddev_samp
+#' @export
+setGeneric("stddev_samp", function(x) { standardGeneric("stddev_samp") })
+
 #' @rdname substring_index
 #' @export
 setGeneric("substring_index", function(x, delim, count) { standardGeneric("substring_index") })
@@ -1019,6 +1043,22 @@ setGeneric("unix_timestamp", function(x, format) { standardGeneric("unix_timesta
 #' @export
 setGeneric("upper", function(x) { standardGeneric("upper") })
 
+#' @rdname var
+#' @export
+setGeneric("var", function(x, y = NULL, na.rm = FALSE, use) { standardGeneric("var") })
+
+#' @rdname var
+#' @export
+setGeneric("variance", function(x) { standardGeneric("variance") })
+
+#' @rdname var_pop
+#' @export
+setGeneric("var_pop", function(x) { standardGeneric("var_pop") })
+
+#' @rdname var_samp
+#' @export
+setGeneric("var_samp", function(x) { standardGeneric("var_samp") })
+
 #' @rdname weekofyear
 #' @export
 setGeneric("weekofyear", function(x) { standardGeneric("weekofyear") })
diff --git a/R/pkg/R/group.R b/R/pkg/R/group.R
index 4cab1a69f601a..e5f702faee65d 100644
--- a/R/pkg/R/group.R
+++ b/R/pkg/R/group.R
@@ -79,6 +79,7 @@ setMethod("count",
 #' @param x a GroupedData
 #' @return a DataFrame
 #' @rdname agg
+#' @family agg_funcs
 #' @examples
 #' \dontrun{
 #'  df2 <- agg(df, age = "sum")  # new column name will be created as 'SUM(age#0)'
@@ -117,8 +118,11 @@ setMethod("summarize",
             agg(x, ...)
           })
 
-# sum/mean/avg/min/max
-methods <- c("sum", "mean", "avg", "min", "max")
+# Aggregate Functions by name
+methods <- c("avg", "max", "mean", "min", "sum")
+
+# These are not exposed on GroupedData: "kurtosis", "skewness", "stddev", "stddev_samp", "stddev_pop",
+# "variance", "var_samp", "var_pop"
 
 createMethod <- function(name) {
   setMethod(name,
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 06f52d021cff8..9e453a1e7c2f4 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -826,12 +826,13 @@ test_that("column functions", {
   c6 <- log(c) + (c) + log1p(c) + log2(c) + lower(c) + ltrim(c) + max(c) + md5(c)
   c7 <- mean(c) + min(c) + month(c) + negate(c) + quarter(c)
   c8 <- reverse(c) + rint(c) + round(c) + rtrim(c) + sha1(c)
-  c9 <- signum(c) + sin(c) + sinh(c) + size(c) + soundex(c) + sqrt(c) + sum(c)
+  c9 <- signum(c) + sin(c) + sinh(c) + size(c) + stddev(c) + soundex(c) + sqrt(c) + sum(c)
   c10 <- sumDistinct(c) + tan(c) + tanh(c) + toDegrees(c) + toRadians(c)
   c11 <- to_date(c) + trim(c) + unbase64(c) + unhex(c) + upper(c)
-  c12 <- lead("col", 1) + lead(c, 1) + lag("col", 1) + lag(c, 1)
-  c13 <- cumeDist() + ntile(1)
-  c14 <- denseRank() + percentRank() + rank() + rowNumber()
+  c12 <- variance(c)
+  c13 <- lead("col", 1) + lead(c, 1) + lag("col", 1) + lag(c, 1)
+  c14 <- cumeDist() + ntile(1)
+  c15 <- denseRank() + percentRank() + rank() + rowNumber()
 
   # Test if base::rank() is exposed
   expect_equal(class(rank())[[1]], "Column")
@@ -849,6 +850,12 @@ test_that("column functions", {
   expect_equal(collect(df3)[[2, 1]], FALSE)
   expect_equal(collect(df3)[[3, 1]], TRUE)
 
+  expect_equal(collect(select(df, sum(df$age)))[1, 1], 49)
+
+  expect_true(abs(collect(select(df, stddev(df$age)))[1, 1] - 7.778175) < 1e-6)
+
+  expect_equal(collect(select(df, var_pop(df$age)))[1, 1], 30.25)
+
   df4 <- createDataFrame(sqlContext, list(list(a = "010101")))
   expect_equal(collect(select(df4, conv(df4$a, 2, 16)))[1, 1], "15")
 })
@@ -976,7 +983,7 @@ test_that("when(), otherwise() and ifelse() on a DataFrame", {
   expect_equal(collect(select(df, ifelse(df$a > 1 & df$b > 2, 0, 1)))[, 1], c(1, 0))
 })
 
-test_that("group by", {
+test_that("group by, agg functions", {
   df <- jsonFile(sqlContext, jsonPath)
   df1 <- agg(df, name = "max", age = "sum")
   expect_equal(1, count(df1))
@@ -997,20 +1004,64 @@ test_that("group by", {
   expect_is(df_summarized, "DataFrame")
   expect_equal(3, count(df_summarized))
 
-  df3 <- agg(gd, age = "sum")
-  expect_is(df3, "DataFrame")
-  expect_equal(3, count(df3))
-
-  df3 <- agg(gd, age = sum(df$age))
+  df3 <- agg(gd, age = "stddev")
   expect_is(df3, "DataFrame")
-  expect_equal(3, count(df3))
-  expect_equal(columns(df3), c("name", "age"))
+  df3_local <- collect(df3)
+  expect_equal(0, df3_local[df3_local$name == "Andy",][1, 2])
 
-  df4 <- sum(gd, "age")
+  df4 <- agg(gd, sumAge = sum(df$age))
   expect_is(df4, "DataFrame")
   expect_equal(3, count(df4))
-  expect_equal(3, count(mean(gd, "age")))
-  expect_equal(3, count(max(gd, "age")))
+  expect_equal(columns(df4), c("name", "sumAge"))
+
+  df5 <- sum(gd, "age")
+  expect_is(df5, "DataFrame")
+  expect_equal(3, count(df5))
+
+  expect_equal(3, count(mean(gd)))
+  expect_equal(3, count(max(gd)))
+  expect_equal(30, collect(max(gd))[1, 2])
+  expect_equal(1, collect(count(gd))[1, 2])
+
+  mockLines2 <- c("{\"name\":\"ID1\", \"value\": \"10\"}",
+                  "{\"name\":\"ID1\", \"value\": \"10\"}",
+                  "{\"name\":\"ID1\", \"value\": \"22\"}",
+                  "{\"name\":\"ID2\", \"value\": \"-3\"}")
+  jsonPath2 <- tempfile(pattern="sparkr-test", fileext=".tmp")
+  writeLines(mockLines2, jsonPath2)
+  gd2 <- groupBy(jsonFile(sqlContext, jsonPath2), "name")
+  df6 <- agg(gd2, value = "sum")
+  df6_local <- collect(df6)
+  expect_equal(42, df6_local[df6_local$name == "ID1",][1, 2])
+  expect_equal(-3, df6_local[df6_local$name == "ID2",][1, 2])
+
+  df7 <- agg(gd2, value = "stddev")
+  df7_local <- collect(df7)
+  expect_true(abs(df7_local[df7_local$name == "ID1",][1, 2] - 6.928203) < 1e-6)
+  expect_equal(0, df7_local[df7_local$name == "ID2",][1, 2])
+
+  mockLines3 <- c("{\"name\":\"Andy\", \"age\":30}",
+                  "{\"name\":\"Andy\", \"age\":30}",
+                  "{\"name\":\"Justin\", \"age\":19}",
+                  "{\"name\":\"Justin\", \"age\":1}")
+  jsonPath3 <- tempfile(pattern="sparkr-test", fileext=".tmp")
+  writeLines(mockLines3, jsonPath3)
+  df8 <- jsonFile(sqlContext, jsonPath3)
+  gd3 <- groupBy(df8, "name")
+  gd3_local <- collect(sum(gd3))
+  expect_equal(60, gd3_local[gd3_local$name == "Andy",][1, 2])
+  expect_equal(20, gd3_local[gd3_local$name == "Justin",][1, 2])
+
+  expect_true(abs(collect(agg(df, sd(df$age)))[1, 1] - 7.778175) < 1e-6)
+  gd3_local <- collect(agg(gd3, var(df8$age)))
+  expect_equal(162, gd3_local[gd3_local$name == "Justin",][1, 2])
+
+  # make sure base:: or stats::sd, var are working
+  expect_true(abs(sd(1:2) - 0.7071068) < 1e-6)
+  expect_true(abs(var(1:5, 1:5) - 2.5) < 1e-6)
+
+  unlink(jsonPath2)
+  unlink(jsonPath3)
 })
 
 test_that("arrange() and orderBy() on a DataFrame", {
@@ -1238,7 +1289,7 @@ test_that("mutate(), transform(), rename() and names()", {
   expect_equal(columns(transformedDF)[4], "newAge2")
   expect_equal(first(filter(transformedDF, transformedDF$name == "Andy"))$newAge, -30)
 
-  # test if transform on local data frames works
+  # test if base::transform on local data frames works
   # ensure the proper signature is used - otherwise this will fail to run
   attach(airquality)
   result <- transform(Ozone, logOzone = log(Ozone))

From bd70244b3cda62cc447fd4cc343d4eb5ddaec893 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Tue, 10 Nov 2015 22:46:17 -0800
Subject: [PATCH 183/518] [SPARK-11463] [PYSPARK] only install signal in main
 thread

Only install signal in main thread, or it will fail to create context in not-main thread.

Author: Davies Liu <davies@databricks.com>

Closes #9574 from davies/python_signal.
---
 python/pyspark/context.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/python/pyspark/context.py b/python/pyspark/context.py
index afd74d937a413..77710a13394c6 100644
--- a/python/pyspark/context.py
+++ b/python/pyspark/context.py
@@ -21,6 +21,7 @@
 import shutil
 import signal
 import sys
+import threading
 from threading import RLock
 from tempfile import NamedTemporaryFile
 
@@ -222,7 +223,9 @@ def _do_init(self, master, appName, sparkHome, pyFiles, environment, batchSize,
         def signal_handler(signal, frame):
             self.cancelAllJobs()
 
-        signal.signal(signal.SIGINT, signal_handler)
+        # see http://stackoverflow.com/questions/23206787/
+        if isinstance(threading.current_thread(), threading._MainThread):
+            signal.signal(signal.SIGINT, signal_handler)
 
     def _initialize_context(self, jconf):
         """

From 99f5f988612b3093d73d9ce98819767e822fcbff Mon Sep 17 00:00:00 2001
From: Tathagata Das <tathagata.das1565@gmail.com>
Date: Tue, 10 Nov 2015 23:16:18 -0800
Subject: [PATCH 184/518] [SPARK-11290][STREAMING] Basic implementation of
 trackStateByKey

Current updateStateByKey provides stateful processing in Spark Streaming. It allows the user to maintain per-key state and manage that state using an updateFunction. The updateFunction is called for each key, and it uses new data and existing state of the key, to generate an updated state. However, based on community feedback, we have learnt the following lessons.
* Need for more optimized state management that does not scan every key
* Need to make it easier to implement common use cases - (a) timeout of idle data, (b) returning items other than state

The high level idea that of this PR
* Introduce a new API trackStateByKey that, allows the user to update per-key state, and emit arbitrary records. The new API is necessary as this will have significantly different semantics than the existing updateStateByKey API. This API will have direct support for timeouts.
* Internally, the system will keep the state data as a map/list within the partitions of the state RDDs. The new data RDDs will be partitioned appropriately, and for all the key-value data, it will lookup the map/list in the state RDD partition and create a new list/map of updated state data. The new state RDD partition will be created based on the update data and if necessary, with old data.
Here is the detailed design doc. Please take a look and provide feedback as comments.
https://docs.google.com/document/d/1NoALLyd83zGs1hNGMm0Pc5YOVgiPpMHugGMk6COqxxE/edit#heading=h.ph3w0clkd4em

This is still WIP. Major things left to be done.
- [x] Implement basic functionality of state tracking, with initial RDD and timeouts
- [x] Unit tests for state tracking
- [x] Unit tests for initial RDD and timeout
- [ ] Unit tests for TrackStateRDD
       - [x] state creating, updating, removing
       - [ ] emitting
       - [ ] checkpointing
- [x] Misc unit tests for State, TrackStateSpec, etc.
- [x] Update docs and experimental tags

Author: Tathagata Das <tathagata.das1565@gmail.com>

Closes #9256 from tdas/trackStateByKey.
---
 .../streaming/StatefulNetworkWordCount.scala  |  25 +-
 .../org/apache/spark/streaming/State.scala    | 193 +++++++
 .../apache/spark/streaming/StateSpec.scala    | 212 ++++++++
 .../dstream/PairDStreamFunctions.scala        |  46 +-
 .../streaming/dstream/TrackStateDStream.scala | 142 +++++
 .../spark/streaming/rdd/TrackStateRDD.scala   | 188 +++++++
 .../spark/streaming/util/StateMap.scala       | 337 ++++++++++++
 .../spark/streaming/StateMapSuite.scala       | 314 +++++++++++
 .../streaming/TrackStateByKeySuite.scala      | 494 ++++++++++++++++++
 .../streaming/rdd/TrackStateRDDSuite.scala    | 193 +++++++
 10 files changed, 2125 insertions(+), 19 deletions(-)
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/State.scala
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
 create mode 100644 streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala
 create mode 100644 streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala
 create mode 100644 streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala

diff --git a/examples/src/main/scala/org/apache/spark/examples/streaming/StatefulNetworkWordCount.scala b/examples/src/main/scala/org/apache/spark/examples/streaming/StatefulNetworkWordCount.scala
index 02ba1c2eed0f7..be2ae0b473363 100644
--- a/examples/src/main/scala/org/apache/spark/examples/streaming/StatefulNetworkWordCount.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/streaming/StatefulNetworkWordCount.scala
@@ -44,18 +44,6 @@ object StatefulNetworkWordCount {
 
     StreamingExamples.setStreamingLogLevels()
 
-    val updateFunc = (values: Seq[Int], state: Option[Int]) => {
-      val currentCount = values.sum
-
-      val previousCount = state.getOrElse(0)
-
-      Some(currentCount + previousCount)
-    }
-
-    val newUpdateFunc = (iterator: Iterator[(String, Seq[Int], Option[Int])]) => {
-      iterator.flatMap(t => updateFunc(t._2, t._3).map(s => (t._1, s)))
-    }
-
     val sparkConf = new SparkConf().setAppName("StatefulNetworkWordCount")
     // Create the context with a 1 second batch size
     val ssc = new StreamingContext(sparkConf, Seconds(1))
@@ -71,9 +59,16 @@ object StatefulNetworkWordCount {
     val wordDstream = words.map(x => (x, 1))
 
     // Update the cumulative count using updateStateByKey
-    // This will give a Dstream made of state (which is the cumulative count of the words)
-    val stateDstream = wordDstream.updateStateByKey[Int](newUpdateFunc,
-      new HashPartitioner (ssc.sparkContext.defaultParallelism), true, initialRDD)
+    // This will give a DStream made of state (which is the cumulative count of the words)
+    val trackStateFunc = (batchTime: Time, word: String, one: Option[Int], state: State[Int]) => {
+      val sum = one.getOrElse(0) + state.getOption.getOrElse(0)
+      val output = (word, sum)
+      state.update(sum)
+      Some(output)
+    }
+
+    val stateDstream = wordDstream.trackStateByKey(
+      StateSpec.function(trackStateFunc).initialState(initialRDD))
     stateDstream.print()
     ssc.start()
     ssc.awaitTermination()
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/State.scala b/streaming/src/main/scala/org/apache/spark/streaming/State.scala
new file mode 100644
index 0000000000000..7dd1b72f80499
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/State.scala
@@ -0,0 +1,193 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming
+
+import scala.language.implicitConversions
+
+import org.apache.spark.annotation.Experimental
+
+/**
+ * :: Experimental ::
+ * Abstract class for getting and updating the tracked state in the `trackStateByKey` operation of
+ * a [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
+ * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
+ *
+ * Scala example of using `State`:
+ * {{{
+ *    // A tracking function that maintains an integer state and return a String
+ *    def trackStateFunc(data: Option[Int], state: State[Int]): Option[String] = {
+ *      // Check if state exists
+ *      if (state.exists) {
+ *        val existingState = state.get  // Get the existing state
+ *        val shouldRemove = ...         // Decide whether to remove the state
+ *        if (shouldRemove) {
+ *          state.remove()     // Remove the state
+ *        } else {
+ *          val newState = ...
+ *          state.update(newState)    // Set the new state
+ *        }
+ *      } else {
+ *        val initialState = ...
+ *        state.update(initialState)  // Set the initial state
+ *      }
+ *      ... // return something
+ *    }
+ *
+ * }}}
+ *
+ * Java example:
+ * {{{
+ *      TODO(@zsxwing)
+ * }}}
+ */
+@Experimental
+sealed abstract class State[S] {
+
+  /** Whether the state already exists */
+  def exists(): Boolean
+
+  /**
+   * Get the state if it exists, otherwise it will throw `java.util.NoSuchElementException`.
+   * Check with `exists()` whether the state exists or not before calling `get()`.
+   *
+   * @throws java.util.NoSuchElementException If the state does not exist.
+   */
+  def get(): S
+
+  /**
+   * Update the state with a new value.
+   *
+   * State cannot be updated if it has been already removed (that is, `remove()` has already been
+   * called) or it is going to be removed due to timeout (that is, `isTimingOut()` is `true`).
+   *
+   * @throws java.lang.IllegalArgumentException If the state has already been removed, or is
+   *                                            going to be removed
+   */
+  def update(newState: S): Unit
+
+  /**
+   * Remove the state if it exists.
+   *
+   * State cannot be updated if it has been already removed (that is, `remove()` has already been
+   * called) or it is going to be removed due to timeout (that is, `isTimingOut()` is `true`).
+   */
+  def remove(): Unit
+
+  /**
+   * Whether the state is timing out and going to be removed by the system after the current batch.
+   * This timeout can occur if timeout duration has been specified in the
+   * [[org.apache.spark.streaming.StateSpec StatSpec]] and the key has not received any new data
+   * for that timeout duration.
+   */
+  def isTimingOut(): Boolean
+
+  /**
+   * Get the state as an [[scala.Option]]. It will be `Some(state)` if it exists, otherwise `None`.
+   */
+  @inline final def getOption(): Option[S] = if (exists) Some(get()) else None
+
+  @inline final override def toString(): String = {
+    getOption.map { _.toString }.getOrElse("<state not set>")
+  }
+}
+
+/** Internal implementation of the [[State]] interface */
+private[streaming] class StateImpl[S] extends State[S] {
+
+  private var state: S = null.asInstanceOf[S]
+  private var defined: Boolean = false
+  private var timingOut: Boolean = false
+  private var updated: Boolean = false
+  private var removed: Boolean = false
+
+  // ========= Public API =========
+  override def exists(): Boolean = {
+    defined
+  }
+
+  override def get(): S = {
+    if (defined) {
+      state
+    } else {
+      throw new NoSuchElementException("State is not set")
+    }
+  }
+
+  override def update(newState: S): Unit = {
+    require(!removed, "Cannot update the state after it has been removed")
+    require(!timingOut, "Cannot update the state that is timing out")
+    state = newState
+    defined = true
+    updated = true
+  }
+
+  override def isTimingOut(): Boolean = {
+    timingOut
+  }
+
+  override def remove(): Unit = {
+    require(!timingOut, "Cannot remove the state that is timing out")
+    require(!removed, "Cannot remove the state that has already been removed")
+    defined = false
+    updated = false
+    removed = true
+  }
+
+  // ========= Internal API =========
+
+  /** Whether the state has been marked for removing */
+  def isRemoved(): Boolean = {
+    removed
+  }
+
+  /** Whether the state has been been updated */
+  def isUpdated(): Boolean = {
+    updated
+  }
+
+  /**
+   * Update the internal data and flags in `this` to the given state option.
+   * This method allows `this` object to be reused across many state records.
+   */
+  def wrap(optionalState: Option[S]): Unit = {
+    optionalState match {
+      case Some(newState) =>
+        this.state = newState
+        defined = true
+
+      case None =>
+        this.state = null.asInstanceOf[S]
+        defined = false
+    }
+    timingOut = false
+    removed = false
+    updated = false
+  }
+
+  /**
+   * Update the internal data and flags in `this` to the given state that is going to be timed out.
+   * This method allows `this` object to be reused across many state records.
+   */
+  def wrapTiminoutState(newState: S): Unit = {
+    this.state = newState
+    defined = true
+    timingOut = true
+    removed = false
+    updated = false
+  }
+}
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala b/streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala
new file mode 100644
index 0000000000000..c9fe35e74c1c7
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala
@@ -0,0 +1,212 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming
+
+import scala.reflect.ClassTag
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.api.java.JavaPairRDD
+import org.apache.spark.rdd.RDD
+import org.apache.spark.util.ClosureCleaner
+import org.apache.spark.{HashPartitioner, Partitioner}
+
+
+/**
+ * :: Experimental ::
+ * Abstract class representing all the specifications of the DStream transformation
+ * `trackStateByKey` operation of a
+ * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
+ * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
+ * Use the [[org.apache.spark.streaming.StateSpec StateSpec.apply()]] or
+ * [[org.apache.spark.streaming.StateSpec StateSpec.create()]] to create instances of
+ * this class.
+ *
+ * Example in Scala:
+ * {{{
+ *    def trackingFunction(data: Option[ValueType], wrappedState: State[StateType]): EmittedType = {
+ *      ...
+ *    }
+ *
+ *    val spec = StateSpec.function(trackingFunction).numPartitions(10)
+ *
+ *    val emittedRecordDStream = keyValueDStream.trackStateByKey[StateType, EmittedDataType](spec)
+ * }}}
+ *
+ * Example in Java:
+ * {{{
+ *    StateStateSpec[KeyType, ValueType, StateType, EmittedDataType] spec =
+ *      StateStateSpec.function[KeyType, ValueType, StateType, EmittedDataType](trackingFunction)
+ *                    .numPartition(10);
+ *
+ *    JavaDStream[EmittedDataType] emittedRecordDStream =
+ *      javaPairDStream.trackStateByKey[StateType, EmittedDataType](spec);
+ * }}}
+ */
+@Experimental
+sealed abstract class StateSpec[KeyType, ValueType, StateType, EmittedType] extends Serializable {
+
+  /** Set the RDD containing the initial states that will be used by `trackStateByKey` */
+  def initialState(rdd: RDD[(KeyType, StateType)]): this.type
+
+  /** Set the RDD containing the initial states that will be used by `trackStateByKey` */
+  def initialState(javaPairRDD: JavaPairRDD[KeyType, StateType]): this.type
+
+  /**
+   * Set the number of partitions by which the state RDDs generated by `trackStateByKey`
+   * will be partitioned. Hash partitioning will be used.
+   */
+  def numPartitions(numPartitions: Int): this.type
+
+  /**
+   * Set the partitioner by which the state RDDs generated by `trackStateByKey` will be
+   * be partitioned.
+   */
+  def partitioner(partitioner: Partitioner): this.type
+
+  /**
+   * Set the duration after which the state of an idle key will be removed. A key and its state is
+   * considered idle if it has not received any data for at least the given duration. The state
+   * tracking function will be called one final time on the idle states that are going to be
+   * removed; [[org.apache.spark.streaming.State State.isTimingOut()]] set
+   * to `true` in that call.
+   */
+  def timeout(idleDuration: Duration): this.type
+}
+
+
+/**
+ * :: Experimental ::
+ * Builder object for creating instances of [[org.apache.spark.streaming.StateSpec StateSpec]]
+ * that is used for specifying the parameters of the DStream transformation
+ * `trackStateByKey` operation of a
+ * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
+ * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
+ *
+ * Example in Scala:
+ * {{{
+ *    def trackingFunction(data: Option[ValueType], wrappedState: State[StateType]): EmittedType = {
+ *      ...
+ *    }
+ *
+ *    val spec = StateSpec.function(trackingFunction).numPartitions(10)
+ *
+ *    val emittedRecordDStream = keyValueDStream.trackStateByKey[StateType, EmittedDataType](spec)
+ * }}}
+ *
+ * Example in Java:
+ * {{{
+ *    StateStateSpec[KeyType, ValueType, StateType, EmittedDataType] spec =
+ *      StateStateSpec.function[KeyType, ValueType, StateType, EmittedDataType](trackingFunction)
+ *                    .numPartition(10);
+ *
+ *    JavaDStream[EmittedDataType] emittedRecordDStream =
+ *      javaPairDStream.trackStateByKey[StateType, EmittedDataType](spec);
+ * }}}
+ */
+@Experimental
+object StateSpec {
+  /**
+   * Create a [[org.apache.spark.streaming.StateSpec StateSpec]] for setting all the specifications
+   * `trackStateByKey` operation on a
+   * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
+   * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
+   * @param trackingFunction The function applied on every data item to manage the associated state
+   *                         and generate the emitted data
+   * @tparam KeyType      Class of the keys
+   * @tparam ValueType    Class of the values
+   * @tparam StateType    Class of the states data
+   * @tparam EmittedType  Class of the emitted data
+   */
+  def function[KeyType, ValueType, StateType, EmittedType](
+      trackingFunction: (Time, KeyType, Option[ValueType], State[StateType]) => Option[EmittedType]
+    ): StateSpec[KeyType, ValueType, StateType, EmittedType] = {
+    ClosureCleaner.clean(trackingFunction, checkSerializable = true)
+    new StateSpecImpl(trackingFunction)
+  }
+
+  /**
+   * Create a [[org.apache.spark.streaming.StateSpec StateSpec]] for setting all the specifications
+   * `trackStateByKey` operation on a
+   * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
+   * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
+   * @param trackingFunction The function applied on every data item to manage the associated state
+   *                         and generate the emitted data
+   * @tparam ValueType    Class of the values
+   * @tparam StateType    Class of the states data
+   * @tparam EmittedType  Class of the emitted data
+   */
+  def function[KeyType, ValueType, StateType, EmittedType](
+      trackingFunction: (Option[ValueType], State[StateType]) => EmittedType
+    ): StateSpec[KeyType, ValueType, StateType, EmittedType] = {
+    ClosureCleaner.clean(trackingFunction, checkSerializable = true)
+    val wrappedFunction =
+      (time: Time, key: Any, value: Option[ValueType], state: State[StateType]) => {
+        Some(trackingFunction(value, state))
+      }
+    new StateSpecImpl(wrappedFunction)
+  }
+}
+
+
+/** Internal implementation of [[org.apache.spark.streaming.StateSpec]] interface. */
+private[streaming]
+case class StateSpecImpl[K, V, S, T](
+    function: (Time, K, Option[V], State[S]) => Option[T]) extends StateSpec[K, V, S, T] {
+
+  require(function != null)
+
+  @volatile private var partitioner: Partitioner = null
+  @volatile private var initialStateRDD: RDD[(K, S)] = null
+  @volatile private var timeoutInterval: Duration = null
+
+  override def initialState(rdd: RDD[(K, S)]): this.type = {
+    this.initialStateRDD = rdd
+    this
+  }
+
+  override def initialState(javaPairRDD: JavaPairRDD[K, S]): this.type = {
+    this.initialStateRDD = javaPairRDD.rdd
+    this
+  }
+
+
+  override def numPartitions(numPartitions: Int): this.type = {
+    this.partitioner(new HashPartitioner(numPartitions))
+    this
+  }
+
+  override def partitioner(partitioner: Partitioner): this.type = {
+    this.partitioner = partitioner
+    this
+  }
+
+  override def timeout(interval: Duration): this.type = {
+    this.timeoutInterval = interval
+    this
+  }
+
+  // ================= Private Methods =================
+
+  private[streaming] def getFunction(): (Time, K, Option[V], State[S]) => Option[T] = function
+
+  private[streaming] def getInitialStateRDD(): Option[RDD[(K, S)]] = Option(initialStateRDD)
+
+  private[streaming] def getPartitioner(): Option[Partitioner] = Option(partitioner)
+
+  private[streaming] def getTimeoutInterval(): Option[Duration] = Option(timeoutInterval)
+}
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/dstream/PairDStreamFunctions.scala b/streaming/src/main/scala/org/apache/spark/streaming/dstream/PairDStreamFunctions.scala
index 71bec96d46c8d..fb691eed27e32 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/dstream/PairDStreamFunctions.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/dstream/PairDStreamFunctions.scala
@@ -24,19 +24,19 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.mapred.{JobConf, OutputFormat}
 import org.apache.hadoop.mapreduce.{OutputFormat => NewOutputFormat}
 
-import org.apache.spark.{HashPartitioner, Partitioner}
+import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
-import org.apache.spark.streaming.{Duration, Time}
 import org.apache.spark.streaming.StreamingContext.rddToFileName
+import org.apache.spark.streaming._
 import org.apache.spark.util.{SerializableConfiguration, SerializableJobConf}
+import org.apache.spark.{HashPartitioner, Partitioner}
 
 /**
  * Extra functions available on DStream of (key, value) pairs through an implicit conversion.
  */
 class PairDStreamFunctions[K, V](self: DStream[(K, V)])
     (implicit kt: ClassTag[K], vt: ClassTag[V], ord: Ordering[K])
-  extends Serializable
-{
+  extends Serializable {
   private[streaming] def ssc = self.ssc
 
   private[streaming] def sparkContext = self.context.sparkContext
@@ -350,6 +350,44 @@ class PairDStreamFunctions[K, V](self: DStream[(K, V)])
     )
   }
 
+  /**
+   * :: Experimental ::
+   * Return a new DStream of data generated by combining the key-value data in `this` stream
+   * with a continuously updated per-key state. The user-provided state tracking function is
+   * applied on each keyed data item along with its corresponding state. The function can choose to
+   * update/remove the state and return a transformed data, which forms the
+   * [[org.apache.spark.streaming.dstream.TrackStateDStream]].
+   *
+   * The specifications of this transformation is made through the
+   * [[org.apache.spark.streaming.StateSpec StateSpec]] class. Besides the tracking function, there
+   * are a number of optional parameters - initial state data, number of partitions, timeouts, etc.
+   * See the [[org.apache.spark.streaming.StateSpec StateSpec spec docs]] for more details.
+   *
+   * Example of using `trackStateByKey`:
+   * {{{
+   *    def trackingFunction(data: Option[Int], wrappedState: State[Int]): String = {
+   *      // Check if state exists, accordingly update/remove state and return transformed data
+   *    }
+   *
+   *    val spec = StateSpec.function(trackingFunction).numPartitions(10)
+   *
+   *    val trackStateDStream = keyValueDStream.trackStateByKey[Int, String](spec)
+   * }}}
+   *
+   * @param spec          Specification of this transformation
+   * @tparam StateType    Class type of the state
+   * @tparam EmittedType  Class type of the tranformed data return by the tracking function
+   */
+  @Experimental
+  def trackStateByKey[StateType: ClassTag, EmittedType: ClassTag](
+      spec: StateSpec[K, V, StateType, EmittedType]
+    ): TrackStateDStream[K, V, StateType, EmittedType] = {
+    new TrackStateDStreamImpl[K, V, StateType, EmittedType](
+      self,
+      spec.asInstanceOf[StateSpecImpl[K, V, StateType, EmittedType]]
+    )
+  }
+
   /**
    * Return a new "state" DStream where the state for each key is updated by applying
    * the given function on the previous state of the key and the new values of each key.
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
new file mode 100644
index 0000000000000..58d89c93bcbef
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.dstream
+
+import scala.reflect.ClassTag
+
+import org.apache.spark._
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.rdd.{EmptyRDD, RDD}
+import org.apache.spark.storage.StorageLevel
+import org.apache.spark.streaming._
+import org.apache.spark.streaming.rdd.{TrackStateRDD, TrackStateRDDRecord}
+
+/**
+ * :: Experimental ::
+ * DStream representing the stream of records emitted by the tracking function in the
+ * `trackStateByKey` operation on a
+ * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]].
+ * Additionally, it also gives access to the stream of state snapshots, that is, the state data of
+ * all keys after a batch has updated them.
+ *
+ * @tparam KeyType Class of the state key
+ * @tparam StateType Class of the state data
+ * @tparam EmittedType Class of the emitted records
+ */
+@Experimental
+sealed abstract class TrackStateDStream[KeyType, ValueType, StateType, EmittedType: ClassTag](
+    ssc: StreamingContext) extends DStream[EmittedType](ssc) {
+
+  /** Return a pair DStream where each RDD is the snapshot of the state of all the keys. */
+  def stateSnapshots(): DStream[(KeyType, StateType)]
+}
+
+/** Internal implementation of the [[TrackStateDStream]] */
+private[streaming] class TrackStateDStreamImpl[
+    KeyType: ClassTag, ValueType: ClassTag, StateType: ClassTag, EmittedType: ClassTag](
+    dataStream: DStream[(KeyType, ValueType)],
+    spec: StateSpecImpl[KeyType, ValueType, StateType, EmittedType])
+  extends TrackStateDStream[KeyType, ValueType, StateType, EmittedType](dataStream.context) {
+
+  private val internalStream =
+    new InternalTrackStateDStream[KeyType, ValueType, StateType, EmittedType](dataStream, spec)
+
+  override def slideDuration: Duration = internalStream.slideDuration
+
+  override def dependencies: List[DStream[_]] = List(internalStream)
+
+  override def compute(validTime: Time): Option[RDD[EmittedType]] = {
+    internalStream.getOrCompute(validTime).map { _.flatMap[EmittedType] { _.emittedRecords } }
+  }
+
+  /**
+   * Forward the checkpoint interval to the internal DStream that computes the state maps. This
+   * to make sure that this DStream does not get checkpointed, only the internal stream.
+   */
+  override def checkpoint(checkpointInterval: Duration): DStream[EmittedType] = {
+    internalStream.checkpoint(checkpointInterval)
+    this
+  }
+
+  /** Return a pair DStream where each RDD is the snapshot of the state of all the keys. */
+  def stateSnapshots(): DStream[(KeyType, StateType)] = {
+    internalStream.flatMap {
+      _.stateMap.getAll().map { case (k, s, _) => (k, s) }.toTraversable }
+  }
+
+  def keyClass: Class[_] = implicitly[ClassTag[KeyType]].runtimeClass
+
+  def valueClass: Class[_] = implicitly[ClassTag[ValueType]].runtimeClass
+
+  def stateClass: Class[_] = implicitly[ClassTag[StateType]].runtimeClass
+
+  def emittedClass: Class[_] = implicitly[ClassTag[EmittedType]].runtimeClass
+}
+
+/**
+ * A DStream that allows per-key state to be maintains, and arbitrary records to be generated
+ * based on updates to the state. This is the main DStream that implements the `trackStateByKey`
+ * operation on DStreams.
+ *
+ * @param parent Parent (key, value) stream that is the source
+ * @param spec Specifications of the trackStateByKey operation
+ * @tparam K   Key type
+ * @tparam V   Value type
+ * @tparam S   Type of the state maintained
+ * @tparam E   Type of the emitted data
+ */
+private[streaming]
+class InternalTrackStateDStream[K: ClassTag, V: ClassTag, S: ClassTag, E: ClassTag](
+    parent: DStream[(K, V)], spec: StateSpecImpl[K, V, S, E])
+  extends DStream[TrackStateRDDRecord[K, S, E]](parent.context) {
+
+  persist(StorageLevel.MEMORY_ONLY)
+
+  private val partitioner = spec.getPartitioner().getOrElse(
+    new HashPartitioner(ssc.sc.defaultParallelism))
+
+  private val trackingFunction = spec.getFunction()
+
+  override def slideDuration: Duration = parent.slideDuration
+
+  override def dependencies: List[DStream[_]] = List(parent)
+
+  /** Enable automatic checkpointing */
+  override val mustCheckpoint = true
+
+  /** Method that generates a RDD for the given time */
+  override def compute(validTime: Time): Option[RDD[TrackStateRDDRecord[K, S, E]]] = {
+    // Get the previous state or create a new empty state RDD
+    val prevStateRDD = getOrCompute(validTime - slideDuration).getOrElse {
+      TrackStateRDD.createFromPairRDD[K, V, S, E](
+        spec.getInitialStateRDD().getOrElse(new EmptyRDD[(K, S)](ssc.sparkContext)),
+        partitioner, validTime
+      )
+    }
+
+    // Compute the new state RDD with previous state RDD and partitioned data RDD
+    parent.getOrCompute(validTime).map { dataRDD =>
+      val partitionedDataRDD = dataRDD.partitionBy(partitioner)
+      val timeoutThresholdTime = spec.getTimeoutInterval().map { interval =>
+        (validTime - interval).milliseconds
+      }
+      new TrackStateRDD(
+        prevStateRDD, partitionedDataRDD, trackingFunction, validTime, timeoutThresholdTime)
+    }
+  }
+}
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala b/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
new file mode 100644
index 0000000000000..ed7cea26d0608
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.rdd
+
+import java.io.{IOException, ObjectInputStream, ObjectOutputStream}
+
+import scala.collection.mutable.ArrayBuffer
+import scala.reflect.ClassTag
+
+import org.apache.spark.rdd.{MapPartitionsRDD, RDD}
+import org.apache.spark.streaming.{Time, StateImpl, State}
+import org.apache.spark.streaming.util.{EmptyStateMap, StateMap}
+import org.apache.spark.util.Utils
+import org.apache.spark._
+
+/**
+ * Record storing the keyed-state [[TrackStateRDD]]. Each record contains a [[StateMap]] and a
+ * sequence of records returned by the tracking function of `trackStateByKey`.
+ */
+private[streaming] case class TrackStateRDDRecord[K, S, T](
+    var stateMap: StateMap[K, S], var emittedRecords: Seq[T])
+
+/**
+ * Partition of the [[TrackStateRDD]], which depends on corresponding partitions of prev state
+ * RDD, and a partitioned keyed-data RDD
+ */
+private[streaming] class TrackStateRDDPartition(
+    idx: Int,
+    @transient private var prevStateRDD: RDD[_],
+    @transient private var partitionedDataRDD: RDD[_]) extends Partition {
+
+  private[rdd] var previousSessionRDDPartition: Partition = null
+  private[rdd] var partitionedDataRDDPartition: Partition = null
+
+  override def index: Int = idx
+  override def hashCode(): Int = idx
+
+  @throws(classOf[IOException])
+  private def writeObject(oos: ObjectOutputStream): Unit = Utils.tryOrIOException {
+    // Update the reference to parent split at the time of task serialization
+    previousSessionRDDPartition = prevStateRDD.partitions(index)
+    partitionedDataRDDPartition = partitionedDataRDD.partitions(index)
+    oos.defaultWriteObject()
+  }
+}
+
+
+/**
+ * RDD storing the keyed-state of `trackStateByKey` and corresponding emitted records.
+ * Each partition of this RDD has a single record of type [[TrackStateRDDRecord]]. This contains a
+ * [[StateMap]] (containing the keyed-states) and the sequence of records returned by the tracking
+ * function of  `trackStateByKey`.
+ * @param prevStateRDD The previous TrackStateRDD on whose StateMap data `this` RDD will be created
+ * @param partitionedDataRDD The partitioned data RDD which is used update the previous StateMaps
+ *                           in the `prevStateRDD` to create `this` RDD
+ * @param trackingFunction The function that will be used to update state and return new data
+ * @param batchTime        The time of the batch to which this RDD belongs to. Use to update
+ */
+private[streaming] class TrackStateRDD[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
+    private var prevStateRDD: RDD[TrackStateRDDRecord[K, S, T]],
+    private var partitionedDataRDD: RDD[(K, V)],
+    trackingFunction: (Time, K, Option[V], State[S]) => Option[T],
+    batchTime: Time, timeoutThresholdTime: Option[Long]
+  ) extends RDD[TrackStateRDDRecord[K, S, T]](
+    partitionedDataRDD.sparkContext,
+    List(
+      new OneToOneDependency[TrackStateRDDRecord[K, S, T]](prevStateRDD),
+      new OneToOneDependency(partitionedDataRDD))
+  ) {
+
+  @volatile private var doFullScan = false
+
+  require(prevStateRDD.partitioner.nonEmpty)
+  require(partitionedDataRDD.partitioner == prevStateRDD.partitioner)
+
+  override val partitioner = prevStateRDD.partitioner
+
+  override def checkpoint(): Unit = {
+    super.checkpoint()
+    doFullScan = true
+  }
+
+  override def compute(
+      partition: Partition, context: TaskContext): Iterator[TrackStateRDDRecord[K, S, T]] = {
+
+    val stateRDDPartition = partition.asInstanceOf[TrackStateRDDPartition]
+    val prevStateRDDIterator = prevStateRDD.iterator(
+      stateRDDPartition.previousSessionRDDPartition, context)
+    val dataIterator = partitionedDataRDD.iterator(
+      stateRDDPartition.partitionedDataRDDPartition, context)
+
+    // Create a new state map by cloning the previous one (if it exists) or by creating an empty one
+    val newStateMap = if (prevStateRDDIterator.hasNext) {
+      prevStateRDDIterator.next().stateMap.copy()
+    } else {
+      new EmptyStateMap[K, S]()
+    }
+
+    val emittedRecords = new ArrayBuffer[T]
+    val wrappedState = new StateImpl[S]()
+
+    // Call the tracking function on each record in the data RDD partition, and accordingly
+    // update the states touched, and the data returned by the tracking function.
+    dataIterator.foreach { case (key, value) =>
+      wrappedState.wrap(newStateMap.get(key))
+      val emittedRecord = trackingFunction(batchTime, key, Some(value), wrappedState)
+      if (wrappedState.isRemoved) {
+        newStateMap.remove(key)
+      } else if (wrappedState.isUpdated) {
+        newStateMap.put(key, wrappedState.get(), batchTime.milliseconds)
+      }
+      emittedRecords ++= emittedRecord
+    }
+
+    // If the RDD is expected to be doing a full scan of all the data in the StateMap,
+    // then use this opportunity to filter out those keys that have timed out.
+    // For each of them call the tracking function.
+    if (doFullScan && timeoutThresholdTime.isDefined) {
+      newStateMap.getByTime(timeoutThresholdTime.get).foreach { case (key, state, _) =>
+        wrappedState.wrapTiminoutState(state)
+        val emittedRecord = trackingFunction(batchTime, key, None, wrappedState)
+        emittedRecords ++= emittedRecord
+        newStateMap.remove(key)
+      }
+    }
+
+    Iterator(TrackStateRDDRecord(newStateMap, emittedRecords))
+  }
+
+  override protected def getPartitions: Array[Partition] = {
+    Array.tabulate(prevStateRDD.partitions.length) { i =>
+      new TrackStateRDDPartition(i, prevStateRDD, partitionedDataRDD)}
+  }
+
+  override def clearDependencies(): Unit = {
+    super.clearDependencies()
+    prevStateRDD = null
+    partitionedDataRDD = null
+  }
+
+  def setFullScan(): Unit = {
+    doFullScan = true
+  }
+}
+
+private[streaming] object TrackStateRDD {
+
+  def createFromPairRDD[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
+      pairRDD: RDD[(K, S)],
+      partitioner: Partitioner,
+      updateTime: Time): TrackStateRDD[K, V, S, T] = {
+
+    val rddOfTrackStateRecords = pairRDD.partitionBy(partitioner).mapPartitions ({ iterator =>
+      val stateMap = StateMap.create[K, S](SparkEnv.get.conf)
+      iterator.foreach { case (key, state) => stateMap.put(key, state, updateTime.milliseconds) }
+      Iterator(TrackStateRDDRecord(stateMap, Seq.empty[T]))
+    }, preservesPartitioning = true)
+
+    val emptyDataRDD = pairRDD.sparkContext.emptyRDD[(K, V)].partitionBy(partitioner)
+
+    val noOpFunc = (time: Time, key: K, value: Option[V], state: State[S]) => None
+
+    new TrackStateRDD[K, V, S, T](rddOfTrackStateRecords, emptyDataRDD, noOpFunc, updateTime, None)
+  }
+}
+
+private[streaming] class EmittedRecordsRDD[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
+    parent: TrackStateRDD[K, V, S, T]) extends RDD[T](parent) {
+  override protected def getPartitions: Array[Partition] = parent.partitions
+  override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
+    parent.compute(partition, context).flatMap { _.emittedRecords }
+  }
+}
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
new file mode 100644
index 0000000000000..ed622ef7bf700
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
@@ -0,0 +1,337 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.util
+
+import java.io.{ObjectInputStream, ObjectOutputStream}
+
+import scala.reflect.ClassTag
+
+import org.apache.spark.SparkConf
+import org.apache.spark.streaming.util.OpenHashMapBasedStateMap._
+import org.apache.spark.util.collection.OpenHashMap
+
+/** Internal interface for defining the map that keeps track of sessions. */
+private[streaming] abstract class StateMap[K: ClassTag, S: ClassTag] extends Serializable {
+
+  /** Get the state for a key if it exists */
+  def get(key: K): Option[S]
+
+  /** Get all the keys and states whose updated time is older than the given threshold time */
+  def getByTime(threshUpdatedTime: Long): Iterator[(K, S, Long)]
+
+  /** Get all the keys and states in this map. */
+  def getAll(): Iterator[(K, S, Long)]
+
+  /** Add or update state */
+  def put(key: K, state: S, updatedTime: Long): Unit
+
+  /** Remove a key */
+  def remove(key: K): Unit
+
+  /**
+   * Shallow copy `this` map to create a new state map.
+   * Updates to the new map should not mutate `this` map.
+   */
+  def copy(): StateMap[K, S]
+
+  def toDebugString(): String = toString()
+}
+
+/** Companion object for [[StateMap]], with utility methods */
+private[streaming] object StateMap {
+  def empty[K: ClassTag, S: ClassTag]: StateMap[K, S] = new EmptyStateMap[K, S]
+
+  def create[K: ClassTag, S: ClassTag](conf: SparkConf): StateMap[K, S] = {
+    val deltaChainThreshold = conf.getInt("spark.streaming.sessionByKey.deltaChainThreshold",
+      DELTA_CHAIN_LENGTH_THRESHOLD)
+    new OpenHashMapBasedStateMap[K, S](64, deltaChainThreshold)
+  }
+}
+
+/** Implementation of StateMap interface representing an empty map */
+private[streaming] class EmptyStateMap[K: ClassTag, S: ClassTag] extends StateMap[K, S] {
+  override def put(key: K, session: S, updateTime: Long): Unit = {
+    throw new NotImplementedError("put() should not be called on an EmptyStateMap")
+  }
+  override def get(key: K): Option[S] = None
+  override def getByTime(threshUpdatedTime: Long): Iterator[(K, S, Long)] = Iterator.empty
+  override def getAll(): Iterator[(K, S, Long)] = Iterator.empty
+  override def copy(): StateMap[K, S] = this
+  override def remove(key: K): Unit = { }
+  override def toDebugString(): String = ""
+}
+
+/** Implementation of StateMap based on Spark's [[org.apache.spark.util.collection.OpenHashMap]] */
+private[streaming] class OpenHashMapBasedStateMap[K: ClassTag, S: ClassTag](
+    @transient @volatile var parentStateMap: StateMap[K, S],
+    initialCapacity: Int = 64,
+    deltaChainThreshold: Int = DELTA_CHAIN_LENGTH_THRESHOLD
+  ) extends StateMap[K, S] { self =>
+
+  def this(initialCapacity: Int, deltaChainThreshold: Int) = this(
+    new EmptyStateMap[K, S],
+    initialCapacity = initialCapacity,
+    deltaChainThreshold = deltaChainThreshold)
+
+  def this(deltaChainThreshold: Int) = this(
+    initialCapacity = 64, deltaChainThreshold = deltaChainThreshold)
+
+  def this() = this(DELTA_CHAIN_LENGTH_THRESHOLD)
+
+  @transient @volatile private var deltaMap =
+    new OpenHashMap[K, StateInfo[S]](initialCapacity)
+
+  /** Get the session data if it exists */
+  override def get(key: K): Option[S] = {
+    val stateInfo = deltaMap(key)
+    if (stateInfo != null) {
+      if (!stateInfo.deleted) {
+        Some(stateInfo.data)
+      } else {
+        None
+      }
+    } else {
+      parentStateMap.get(key)
+    }
+  }
+
+  /** Get all the keys and states whose updated time is older than the give threshold time */
+  override def getByTime(threshUpdatedTime: Long): Iterator[(K, S, Long)] = {
+    val oldStates = parentStateMap.getByTime(threshUpdatedTime).filter { case (key, value, _) =>
+      !deltaMap.contains(key)
+    }
+
+    val updatedStates = deltaMap.iterator.filter { case (_, stateInfo) =>
+      !stateInfo.deleted && stateInfo.updateTime < threshUpdatedTime
+    }.map { case (key, stateInfo) =>
+      (key, stateInfo.data, stateInfo.updateTime)
+    }
+    oldStates ++ updatedStates
+  }
+
+  /** Get all the keys and states in this map. */
+  override def getAll(): Iterator[(K, S, Long)] = {
+
+    val oldStates = parentStateMap.getAll().filter { case (key, _, _) =>
+      !deltaMap.contains(key)
+    }
+
+    val updatedStates = deltaMap.iterator.filter { ! _._2.deleted }.map { case (key, stateInfo) =>
+      (key, stateInfo.data, stateInfo.updateTime)
+    }
+    oldStates ++ updatedStates
+  }
+
+  /** Add or update state */
+  override def put(key: K, state: S, updateTime: Long): Unit = {
+    val stateInfo = deltaMap(key)
+    if (stateInfo != null) {
+      stateInfo.update(state, updateTime)
+    } else {
+      deltaMap.update(key, new StateInfo(state, updateTime))
+    }
+  }
+
+  /** Remove a state */
+  override def remove(key: K): Unit = {
+    val stateInfo = deltaMap(key)
+    if (stateInfo != null) {
+      stateInfo.markDeleted()
+    } else {
+      val newInfo = new StateInfo[S](deleted = true)
+      deltaMap.update(key, newInfo)
+    }
+  }
+
+  /**
+   * Shallow copy the map to create a new session store. Updates to the new map
+   * should not mutate `this` map.
+   */
+  override def copy(): StateMap[K, S] = {
+    new OpenHashMapBasedStateMap[K, S](this, deltaChainThreshold = deltaChainThreshold)
+  }
+
+  /** Whether the delta chain lenght is long enough that it should be compacted */
+  def shouldCompact: Boolean = {
+    deltaChainLength >= deltaChainThreshold
+  }
+
+  /** Length of the delta chains of this map */
+  def deltaChainLength: Int = parentStateMap match {
+    case map: OpenHashMapBasedStateMap[_, _] => map.deltaChainLength + 1
+    case _ => 0
+  }
+
+  /**
+   * Approximate number of keys in the map. This is an overestimation that is mainly used to
+   * reserve capacity in a new map at delta compaction time.
+   */
+  def approxSize: Int = deltaMap.size + {
+    parentStateMap match {
+      case s: OpenHashMapBasedStateMap[_, _] => s.approxSize
+      case _ => 0
+    }
+  }
+
+  /** Get all the data of this map as string formatted as a tree based on the delta depth */
+  override def toDebugString(): String = {
+    val tabs = if (deltaChainLength > 0) {
+      ("    " * (deltaChainLength - 1)) + "+--- "
+    } else ""
+    parentStateMap.toDebugString() + "\n" + deltaMap.iterator.mkString(tabs, "\n" + tabs, "")
+  }
+
+  override def toString(): String = {
+    s"[${System.identityHashCode(this)}, ${System.identityHashCode(parentStateMap)}]"
+  }
+
+  /**
+   * Serialize the map data. Besides serialization, this method actually compact the deltas
+   * (if needed) in a single pass over all the data in the map.
+   */
+
+  private def writeObject(outputStream: ObjectOutputStream): Unit = {
+    // Write all the non-transient fields, especially class tags, etc.
+    outputStream.defaultWriteObject()
+
+    // Write the data in the delta of this state map
+    outputStream.writeInt(deltaMap.size)
+    val deltaMapIterator = deltaMap.iterator
+    var deltaMapCount = 0
+    while (deltaMapIterator.hasNext) {
+      deltaMapCount += 1
+      val (key, stateInfo) = deltaMapIterator.next()
+      outputStream.writeObject(key)
+      outputStream.writeObject(stateInfo)
+    }
+    assert(deltaMapCount == deltaMap.size)
+
+    // Write the data in the parent state map while copying the data into a new parent map for
+    // compaction (if needed)
+    val doCompaction = shouldCompact
+    val newParentSessionStore = if (doCompaction) {
+      val initCapacity = if (approxSize > 0) approxSize else 64
+      new OpenHashMapBasedStateMap[K, S](initialCapacity = initCapacity, deltaChainThreshold)
+    } else { null }
+
+    val iterOfActiveSessions = parentStateMap.getAll()
+
+    var parentSessionCount = 0
+
+    // First write the approximate size of the data to be written, so that readObject can
+    // allocate appropriately sized OpenHashMap.
+    outputStream.writeInt(approxSize)
+
+    while(iterOfActiveSessions.hasNext) {
+      parentSessionCount += 1
+
+      val (key, state, updateTime) = iterOfActiveSessions.next()
+      outputStream.writeObject(key)
+      outputStream.writeObject(state)
+      outputStream.writeLong(updateTime)
+
+      if (doCompaction) {
+        newParentSessionStore.deltaMap.update(
+          key, StateInfo(state, updateTime, deleted = false))
+      }
+    }
+
+    // Write the final limit marking object with the correct count of records written.
+    val limiterObj = new LimitMarker(parentSessionCount)
+    outputStream.writeObject(limiterObj)
+    if (doCompaction) {
+      parentStateMap = newParentSessionStore
+    }
+  }
+
+  /** Deserialize the map data. */
+  private def readObject(inputStream: ObjectInputStream): Unit = {
+
+    // Read the non-transient fields, especially class tags, etc.
+    inputStream.defaultReadObject()
+
+    // Read the data of the delta
+    val deltaMapSize = inputStream.readInt()
+    deltaMap = new OpenHashMap[K, StateInfo[S]]()
+    var deltaMapCount = 0
+    while (deltaMapCount < deltaMapSize) {
+      val key = inputStream.readObject().asInstanceOf[K]
+      val sessionInfo = inputStream.readObject().asInstanceOf[StateInfo[S]]
+      deltaMap.update(key, sessionInfo)
+      deltaMapCount += 1
+    }
+
+
+    // Read the data of the parent map. Keep reading records, until the limiter is reached
+    // First read the approximate number of records to expect and allocate properly size
+    // OpenHashMap
+    val parentSessionStoreSizeHint = inputStream.readInt()
+    val newParentSessionStore = new OpenHashMapBasedStateMap[K, S](
+      initialCapacity = parentSessionStoreSizeHint, deltaChainThreshold)
+
+    // Read the records until the limit marking object has been reached
+    var parentSessionLoopDone = false
+    while(!parentSessionLoopDone) {
+      val obj = inputStream.readObject()
+      if (obj.isInstanceOf[LimitMarker]) {
+        parentSessionLoopDone = true
+        val expectedCount = obj.asInstanceOf[LimitMarker].num
+        assert(expectedCount == newParentSessionStore.deltaMap.size)
+      } else {
+        val key = obj.asInstanceOf[K]
+        val state = inputStream.readObject().asInstanceOf[S]
+        val updateTime = inputStream.readLong()
+        newParentSessionStore.deltaMap.update(
+          key, StateInfo(state, updateTime, deleted = false))
+      }
+    }
+    parentStateMap = newParentSessionStore
+  }
+}
+
+/**
+ * Companion object of [[OpenHashMapBasedStateMap]] having associated helper
+ * classes and methods
+ */
+private[streaming] object OpenHashMapBasedStateMap {
+
+  /** Internal class to represent the state information */
+  case class StateInfo[S](
+      var data: S = null.asInstanceOf[S],
+      var updateTime: Long = -1,
+      var deleted: Boolean = false) {
+
+    def markDeleted(): Unit = {
+      deleted = true
+    }
+
+    def update(newData: S, newUpdateTime: Long): Unit = {
+      data = newData
+      updateTime = newUpdateTime
+      deleted = false
+    }
+  }
+
+  /**
+   * Internal class to represent a marker the demarkate the the end of all state data in the
+   * serialized bytes.
+   */
+  class LimitMarker(val num: Int) extends Serializable
+
+  val DELTA_CHAIN_LENGTH_THRESHOLD = 20
+}
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala
new file mode 100644
index 0000000000000..48d3b41b66cbf
--- /dev/null
+++ b/streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala
@@ -0,0 +1,314 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming
+
+import scala.collection.{immutable, mutable, Map}
+import scala.util.Random
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.streaming.util.{EmptyStateMap, OpenHashMapBasedStateMap, StateMap}
+import org.apache.spark.util.Utils
+
+class StateMapSuite extends SparkFunSuite {
+
+  test("EmptyStateMap") {
+    val map = new EmptyStateMap[Int, Int]
+    intercept[scala.NotImplementedError] {
+      map.put(1, 1, 1)
+    }
+    assert(map.get(1) === None)
+    assert(map.getByTime(10000).isEmpty)
+    assert(map.getAll().isEmpty)
+    map.remove(1)   // no exception
+    assert(map.copy().eq(map))
+  }
+
+  test("OpenHashMapBasedStateMap - put, get, getByTime, getAll, remove") {
+    val map = new OpenHashMapBasedStateMap[Int, Int]()
+
+    map.put(1, 100, 10)
+    assert(map.get(1) === Some(100))
+    assert(map.get(2) === None)
+    assert(map.getByTime(11).toSet === Set((1, 100, 10)))
+    assert(map.getByTime(10).toSet === Set.empty)
+    assert(map.getByTime(9).toSet === Set.empty)
+    assert(map.getAll().toSet === Set((1, 100, 10)))
+
+    map.put(2, 200, 20)
+    assert(map.getByTime(21).toSet === Set((1, 100, 10), (2, 200, 20)))
+    assert(map.getByTime(11).toSet === Set((1, 100, 10)))
+    assert(map.getByTime(10).toSet === Set.empty)
+    assert(map.getByTime(9).toSet === Set.empty)
+    assert(map.getAll().toSet === Set((1, 100, 10), (2, 200, 20)))
+
+    map.remove(1)
+    assert(map.get(1) === None)
+    assert(map.getAll().toSet === Set((2, 200, 20)))
+  }
+
+  test("OpenHashMapBasedStateMap - put, get, getByTime, getAll, remove with copy") {
+    val parentMap = new OpenHashMapBasedStateMap[Int, Int]()
+    parentMap.put(1, 100, 1)
+    parentMap.put(2, 200, 2)
+    parentMap.remove(1)
+
+    // Create child map and make changes
+    val map = parentMap.copy()
+    assert(map.get(1) === None)
+    assert(map.get(2) === Some(200))
+    assert(map.getByTime(10).toSet === Set((2, 200, 2)))
+    assert(map.getByTime(2).toSet === Set.empty)
+    assert(map.getAll().toSet === Set((2, 200, 2)))
+
+    // Add new items
+    map.put(3, 300, 3)
+    assert(map.get(3) === Some(300))
+    map.put(4, 400, 4)
+    assert(map.get(4) === Some(400))
+    assert(map.getByTime(10).toSet === Set((2, 200, 2), (3, 300, 3), (4, 400, 4)))
+    assert(map.getByTime(4).toSet === Set((2, 200, 2), (3, 300, 3)))
+    assert(map.getAll().toSet === Set((2, 200, 2), (3, 300, 3), (4, 400, 4)))
+    assert(parentMap.getAll().toSet === Set((2, 200, 2)))
+
+    // Remove items
+    map.remove(4)
+    assert(map.get(4) === None)       // item added in this map, then removed in this map
+    map.remove(2)
+    assert(map.get(2) === None)       // item removed in parent map, then added in this map
+    assert(map.getAll().toSet === Set((3, 300, 3)))
+    assert(parentMap.getAll().toSet === Set((2, 200, 2)))
+
+    // Update items
+    map.put(1, 1000, 100)
+    assert(map.get(1) === Some(1000)) // item removed in parent map, then added in this map
+    map.put(2, 2000, 200)
+    assert(map.get(2) === Some(2000)) // item added in parent map, then removed + added in this map
+    map.put(3, 3000, 300)
+    assert(map.get(3) === Some(3000)) // item added + updated in this map
+    map.put(4, 4000, 400)
+    assert(map.get(4) === Some(4000)) // item removed + updated in this map
+
+    assert(map.getAll().toSet ===
+      Set((1, 1000, 100), (2, 2000, 200), (3, 3000, 300), (4, 4000, 400)))
+    assert(parentMap.getAll().toSet === Set((2, 200, 2)))
+
+    map.remove(2)         // remove item present in parent map, so that its not visible in child map
+
+    // Create child map and see availability of items
+    val childMap = map.copy()
+    assert(childMap.getAll().toSet === map.getAll().toSet)
+    assert(childMap.get(1) === Some(1000))  // item removed in grandparent, but added in parent map
+    assert(childMap.get(2) === None)        // item added in grandparent, but removed in parent map
+    assert(childMap.get(3) === Some(3000))  // item added and updated in parent map
+
+    childMap.put(2, 20000, 200)
+    assert(childMap.get(2) === Some(20000)) // item map
+  }
+
+  test("OpenHashMapBasedStateMap - serializing and deserializing") {
+    val map1 = new OpenHashMapBasedStateMap[Int, Int]()
+    map1.put(1, 100, 1)
+    map1.put(2, 200, 2)
+
+    val map2 = map1.copy()
+    map2.put(3, 300, 3)
+    map2.put(4, 400, 4)
+
+    val map3 = map2.copy()
+    map3.put(3, 600, 3)
+    map3.remove(2)
+
+    // Do not test compaction
+    assert(map3.asInstanceOf[OpenHashMapBasedStateMap[_, _]].shouldCompact === false)
+
+    val deser_map3 = Utils.deserialize[StateMap[Int, Int]](
+      Utils.serialize(map3), Thread.currentThread().getContextClassLoader)
+    assertMap(deser_map3, map3, 1, "Deserialized map not same as original map")
+  }
+
+  test("OpenHashMapBasedStateMap - serializing and deserializing with compaction") {
+    val targetDeltaLength = 10
+    val deltaChainThreshold = 5
+
+    var map = new OpenHashMapBasedStateMap[Int, Int](
+      deltaChainThreshold = deltaChainThreshold)
+
+    // Make large delta chain with length more than deltaChainThreshold
+    for(i <- 1 to targetDeltaLength) {
+      map.put(Random.nextInt(), Random.nextInt(), 1)
+      map = map.copy().asInstanceOf[OpenHashMapBasedStateMap[Int, Int]]
+    }
+    assert(map.deltaChainLength > deltaChainThreshold)
+    assert(map.shouldCompact === true)
+
+    val deser_map = Utils.deserialize[OpenHashMapBasedStateMap[Int, Int]](
+      Utils.serialize(map), Thread.currentThread().getContextClassLoader)
+    assert(deser_map.deltaChainLength < deltaChainThreshold)
+    assert(deser_map.shouldCompact === false)
+    assertMap(deser_map, map, 1, "Deserialized + compacted map not same as original map")
+  }
+
+  test("OpenHashMapBasedStateMap - all possible sequences of operations with copies ") {
+    /*
+     * This tests the map using all permutations of sequences operations, across multiple map
+     * copies as well as between copies. It is to ensure complete coverage, though it is
+     * kind of hard to debug this. It is set up as follows.
+     *
+     * - For any key, there can be 2 types of update ops on a state map - put or remove
+     *
+     * - These operations are done on a test map in "sets". After each set, the map is "copied"
+     *   to create a new map, and the next set of operations are done on the new one. This tests
+     *   whether the map data persistes correctly across copies.
+     *
+     * - Within each set, there are a number of operations to test whether the map correctly
+     *   updates and removes data without affecting the parent state map.
+     *
+     * - Overall this creates (numSets * numOpsPerSet) operations, each of which that can 2 types
+     *   of operations. This leads to a total of [2 ^ (numSets * numOpsPerSet)] different sequence
+     *   of operations, which we will test with different keys.
+     *
+     * Example: With numSets = 2, and numOpsPerSet = 2 give numTotalOps = 4. This means that
+     * 2 ^ 4 = 16 possible permutations needs to be tested using 16 keys.
+     * _______________________________________________
+     * |         |      Set1       |     Set2        |
+     * |         |-----------------|-----------------|
+     * |         |   Op1    Op2   |c|   Op3    Op4   |
+     * |---------|----------------|o|----------------|
+     * | key 0   |   put    put   |p|   put    put   |
+     * | key 1   |   put    put   |y|   put    rem   |
+     * | key 2   |   put    put   | |   rem    put   |
+     * | key 3   |   put    put   |t|   rem    rem   |
+     * | key 4   |   put    rem   |h|   put    put   |
+     * | key 5   |   put    rem   |e|   put    rem   |
+     * | key 6   |   put    rem   | |   rem    put   |
+     * | key 7   |   put    rem   |s|   rem    rem   |
+     * | key 8   |   rem    put   |t|   put    put   |
+     * | key 9   |   rem    put   |a|   put    rem   |
+     * | key 10  |   rem    put   |t|   rem    put   |
+     * | key 11  |   rem    put   |e|   rem    rem   |
+     * | key 12  |   rem    rem   | |   put    put   |
+     * | key 13  |   rem    rem   |m|   put    rem   |
+     * | key 14  |   rem    rem   |a|   rem    put   |
+     * | key 15  |   rem    rem   |p|   rem    rem   |
+     * |_________|________________|_|________________|
+     */
+
+    val numTypeMapOps = 2   // 0 = put a new value, 1 = remove value
+    val numSets = 3
+    val numOpsPerSet = 3    // to test seq of ops like update -> remove -> update in same set
+    val numTotalOps = numOpsPerSet * numSets
+    val numKeys = math.pow(numTypeMapOps, numTotalOps).toInt  // to get all combinations of ops
+
+    val refMap = new mutable.HashMap[Int, (Int, Long)]()
+    var prevSetRefMap: immutable.Map[Int, (Int, Long)] = null
+
+    var stateMap: StateMap[Int, Int] = new OpenHashMapBasedStateMap[Int, Int]()
+    var prevSetStateMap: StateMap[Int, Int] = null
+
+    var time = 1L
+
+    for (setId <- 0 until numSets) {
+      for (opInSetId <- 0 until numOpsPerSet) {
+        val opId = setId * numOpsPerSet + opInSetId
+        for (keyId <- 0 until numKeys) {
+          time += 1
+          // Find the operation type that needs to be done
+          // This is similar to finding the nth bit value of a binary number
+          // E.g.  nth bit from the right of any binary number B is [ B / (2 ^ (n - 1)) ] % 2
+          val opCode =
+            (keyId / math.pow(numTypeMapOps, numTotalOps - opId - 1).toInt) % numTypeMapOps
+          opCode match {
+            case 0 =>
+              val value = Random.nextInt()
+              stateMap.put(keyId, value, time)
+              refMap.put(keyId, (value, time))
+            case 1 =>
+              stateMap.remove(keyId)
+              refMap.remove(keyId)
+          }
+        }
+
+        // Test whether the current state map after all key updates is correct
+        assertMap(stateMap, refMap, time, "State map does not match reference map")
+
+        // Test whether the previous map before copy has not changed
+        if (prevSetStateMap != null && prevSetRefMap != null) {
+          assertMap(prevSetStateMap, prevSetRefMap, time,
+            "Parent state map somehow got modified, does not match corresponding reference map")
+        }
+      }
+
+      // Copy the map and remember the previous maps for future tests
+      prevSetStateMap = stateMap
+      prevSetRefMap = refMap.toMap
+      stateMap = stateMap.copy()
+
+      // Assert that the copied map has the same data
+      assertMap(stateMap, prevSetRefMap, time,
+        "State map does not match reference map after copying")
+    }
+    assertMap(stateMap, refMap.toMap, time, "Final state map does not match reference map")
+  }
+
+  // Assert whether all the data and operations on a state map matches that of a reference state map
+  private def assertMap(
+      mapToTest: StateMap[Int, Int],
+      refMapToTestWith: StateMap[Int, Int],
+      time: Long,
+      msg: String): Unit = {
+    withClue(msg) {
+      // Assert all the data is same as the reference map
+      assert(mapToTest.getAll().toSet === refMapToTestWith.getAll().toSet)
+
+      // Assert that get on every key returns the right value
+      for (keyId <- refMapToTestWith.getAll().map { _._1 }) {
+        assert(mapToTest.get(keyId) === refMapToTestWith.get(keyId))
+      }
+
+      // Assert that every time threshold returns the correct data
+      for (t <- 0L to (time + 1)) {
+        assert(mapToTest.getByTime(t).toSet ===  refMapToTestWith.getByTime(t).toSet)
+      }
+    }
+  }
+
+  // Assert whether all the data and operations on a state map matches that of a reference map
+  private def assertMap(
+      mapToTest: StateMap[Int, Int],
+      refMapToTestWith: Map[Int, (Int, Long)],
+      time: Long,
+      msg: String): Unit = {
+    withClue(msg) {
+      // Assert all the data is same as the reference map
+      assert(mapToTest.getAll().toSet ===
+        refMapToTestWith.iterator.map { x => (x._1, x._2._1, x._2._2) }.toSet)
+
+      // Assert that get on every key returns the right value
+      for (keyId <- refMapToTestWith.keys) {
+        assert(mapToTest.get(keyId) === refMapToTestWith.get(keyId).map { _._1 })
+      }
+
+      // Assert that every time threshold returns the correct data
+      for (t <- 0L to (time + 1)) {
+        val expectedRecords =
+          refMapToTestWith.iterator.filter { _._2._2 < t }.map { x => (x._1, x._2._1, x._2._2) }
+        assert(mapToTest.getByTime(t).toSet ===  expectedRecords.toSet)
+      }
+    }
+  }
+}
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala
new file mode 100644
index 0000000000000..e3072b4442840
--- /dev/null
+++ b/streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala
@@ -0,0 +1,494 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming
+
+import java.io.File
+
+import scala.collection.mutable.{ArrayBuffer, SynchronizedBuffer}
+import scala.reflect.ClassTag
+
+import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll}
+
+import org.apache.spark.streaming.dstream.{TrackStateDStream, TrackStateDStreamImpl}
+import org.apache.spark.util.{ManualClock, Utils}
+import org.apache.spark.{SparkConf, SparkContext, SparkFunSuite}
+
+class TrackStateByKeySuite extends SparkFunSuite with BeforeAndAfterAll with BeforeAndAfter {
+
+  private var sc: SparkContext = null
+  private var ssc: StreamingContext = null
+  private var checkpointDir: File = null
+  private val batchDuration = Seconds(1)
+
+  before {
+    StreamingContext.getActive().foreach {
+      _.stop(stopSparkContext = false)
+    }
+    checkpointDir = Utils.createTempDir("checkpoint")
+
+    ssc = new StreamingContext(sc, batchDuration)
+    ssc.checkpoint(checkpointDir.toString)
+  }
+
+  after {
+    StreamingContext.getActive().foreach {
+      _.stop(stopSparkContext = false)
+    }
+  }
+
+  override def beforeAll(): Unit = {
+    val conf = new SparkConf().setMaster("local").setAppName("TrackStateByKeySuite")
+    conf.set("spark.streaming.clock", classOf[ManualClock].getName())
+    sc = new SparkContext(conf)
+  }
+
+  test("state - get, exists, update, remove, ") {
+    var state: StateImpl[Int] = null
+
+    def testState(
+        expectedData: Option[Int],
+        shouldBeUpdated: Boolean = false,
+        shouldBeRemoved: Boolean = false,
+        shouldBeTimingOut: Boolean = false
+      ): Unit = {
+      if (expectedData.isDefined) {
+        assert(state.exists)
+        assert(state.get() === expectedData.get)
+        assert(state.getOption() === expectedData)
+        assert(state.getOption.getOrElse(-1) === expectedData.get)
+      } else {
+        assert(!state.exists)
+        intercept[NoSuchElementException] {
+          state.get()
+        }
+        assert(state.getOption() === None)
+        assert(state.getOption.getOrElse(-1) === -1)
+      }
+
+      assert(state.isTimingOut() === shouldBeTimingOut)
+      if (shouldBeTimingOut) {
+        intercept[IllegalArgumentException] {
+          state.remove()
+        }
+        intercept[IllegalArgumentException] {
+          state.update(-1)
+        }
+      }
+
+      assert(state.isUpdated() === shouldBeUpdated)
+
+      assert(state.isRemoved() === shouldBeRemoved)
+      if (shouldBeRemoved) {
+        intercept[IllegalArgumentException] {
+          state.remove()
+        }
+        intercept[IllegalArgumentException] {
+          state.update(-1)
+        }
+      }
+    }
+
+    state = new StateImpl[Int]()
+    testState(None)
+
+    state.wrap(None)
+    testState(None)
+
+    state.wrap(Some(1))
+    testState(Some(1))
+
+    state.update(2)
+    testState(Some(2), shouldBeUpdated = true)
+
+    state = new StateImpl[Int]()
+    state.update(2)
+    testState(Some(2), shouldBeUpdated = true)
+
+    state.remove()
+    testState(None, shouldBeRemoved = true)
+
+    state.wrapTiminoutState(3)
+    testState(Some(3), shouldBeTimingOut = true)
+  }
+
+  test("trackStateByKey - basic operations with simple API") {
+    val inputData =
+      Seq(
+        Seq(),
+        Seq("a"),
+        Seq("a", "b"),
+        Seq("a", "b", "c"),
+        Seq("a", "b"),
+        Seq("a"),
+        Seq()
+      )
+
+    val outputData =
+      Seq(
+        Seq(),
+        Seq(1),
+        Seq(2, 1),
+        Seq(3, 2, 1),
+        Seq(4, 3),
+        Seq(5),
+        Seq()
+      )
+
+    val stateData =
+      Seq(
+        Seq(),
+        Seq(("a", 1)),
+        Seq(("a", 2), ("b", 1)),
+        Seq(("a", 3), ("b", 2), ("c", 1)),
+        Seq(("a", 4), ("b", 3), ("c", 1)),
+        Seq(("a", 5), ("b", 3), ("c", 1)),
+        Seq(("a", 5), ("b", 3), ("c", 1))
+      )
+
+    // state maintains running count, and updated count is returned
+    val trackStateFunc = (value: Option[Int], state: State[Int]) => {
+      val sum = value.getOrElse(0) + state.getOption.getOrElse(0)
+      state.update(sum)
+      sum
+    }
+
+    testOperation[String, Int, Int](
+      inputData, StateSpec.function(trackStateFunc), outputData, stateData)
+  }
+
+  test("trackStateByKey - basic operations with advanced API") {
+    val inputData =
+      Seq(
+        Seq(),
+        Seq("a"),
+        Seq("a", "b"),
+        Seq("a", "b", "c"),
+        Seq("a", "b"),
+        Seq("a"),
+        Seq()
+      )
+
+    val outputData =
+      Seq(
+        Seq(),
+        Seq("aa"),
+        Seq("aa", "bb"),
+        Seq("aa", "bb", "cc"),
+        Seq("aa", "bb"),
+        Seq("aa"),
+        Seq()
+      )
+
+    val stateData =
+      Seq(
+        Seq(),
+        Seq(("a", 1)),
+        Seq(("a", 2), ("b", 1)),
+        Seq(("a", 3), ("b", 2), ("c", 1)),
+        Seq(("a", 4), ("b", 3), ("c", 1)),
+        Seq(("a", 5), ("b", 3), ("c", 1)),
+        Seq(("a", 5), ("b", 3), ("c", 1))
+      )
+
+    // state maintains running count, key string doubled and returned
+    val trackStateFunc = (batchTime: Time, key: String, value: Option[Int], state: State[Int]) => {
+      val sum = value.getOrElse(0) + state.getOption.getOrElse(0)
+      state.update(sum)
+      Some(key * 2)
+    }
+
+    testOperation(inputData, StateSpec.function(trackStateFunc), outputData, stateData)
+  }
+
+  test("trackStateByKey - type inferencing and class tags") {
+
+    // Simple track state function with value as Int, state as Double and emitted type as Double
+    val simpleFunc = (value: Option[Int], state: State[Double]) => {
+      0L
+    }
+
+    // Advanced track state function with key as String, value as Int, state as Double and
+    // emitted type as Double
+    val advancedFunc = (time: Time, key: String, value: Option[Int], state: State[Double]) => {
+      Some(0L)
+    }
+
+    def testTypes(dstream: TrackStateDStream[_, _, _, _]): Unit = {
+      val dstreamImpl = dstream.asInstanceOf[TrackStateDStreamImpl[_, _, _, _]]
+      assert(dstreamImpl.keyClass === classOf[String])
+      assert(dstreamImpl.valueClass === classOf[Int])
+      assert(dstreamImpl.stateClass === classOf[Double])
+      assert(dstreamImpl.emittedClass === classOf[Long])
+    }
+
+    val inputStream = new TestInputStream[(String, Int)](ssc, Seq.empty, numPartitions = 2)
+
+    // Defining StateSpec inline with trackStateByKey and simple function implicitly gets the types
+    val simpleFunctionStateStream1 = inputStream.trackStateByKey(
+      StateSpec.function(simpleFunc).numPartitions(1))
+    testTypes(simpleFunctionStateStream1)
+
+    // Separately defining StateSpec with simple function requires explicitly specifying types
+    val simpleFuncSpec = StateSpec.function[String, Int, Double, Long](simpleFunc)
+    val simpleFunctionStateStream2 = inputStream.trackStateByKey(simpleFuncSpec)
+    testTypes(simpleFunctionStateStream2)
+
+    // Separately defining StateSpec with advanced function implicitly gets the types
+    val advFuncSpec1 = StateSpec.function(advancedFunc)
+    val advFunctionStateStream1 = inputStream.trackStateByKey(advFuncSpec1)
+    testTypes(advFunctionStateStream1)
+
+    // Defining StateSpec inline with trackStateByKey and advanced func implicitly gets the types
+    val advFunctionStateStream2 = inputStream.trackStateByKey(
+      StateSpec.function(simpleFunc).numPartitions(1))
+    testTypes(advFunctionStateStream2)
+
+    // Defining StateSpec inline with trackStateByKey and advanced func implicitly gets the types
+    val advFuncSpec2 = StateSpec.function[String, Int, Double, Long](advancedFunc)
+    val advFunctionStateStream3 = inputStream.trackStateByKey[Double, Long](advFuncSpec2)
+    testTypes(advFunctionStateStream3)
+  }
+
+  test("trackStateByKey - states as emitted records") {
+    val inputData =
+      Seq(
+        Seq(),
+        Seq("a"),
+        Seq("a", "b"),
+        Seq("a", "b", "c"),
+        Seq("a", "b"),
+        Seq("a"),
+        Seq()
+      )
+
+    val outputData =
+      Seq(
+        Seq(),
+        Seq(("a", 1)),
+        Seq(("a", 2), ("b", 1)),
+        Seq(("a", 3), ("b", 2), ("c", 1)),
+        Seq(("a", 4), ("b", 3)),
+        Seq(("a", 5)),
+        Seq()
+      )
+
+    val stateData =
+      Seq(
+        Seq(),
+        Seq(("a", 1)),
+        Seq(("a", 2), ("b", 1)),
+        Seq(("a", 3), ("b", 2), ("c", 1)),
+        Seq(("a", 4), ("b", 3), ("c", 1)),
+        Seq(("a", 5), ("b", 3), ("c", 1)),
+        Seq(("a", 5), ("b", 3), ("c", 1))
+      )
+
+    val trackStateFunc = (time: Time, key: String, value: Option[Int], state: State[Int]) => {
+      val sum = value.getOrElse(0) + state.getOption.getOrElse(0)
+      val output = (key, sum)
+      state.update(sum)
+      Some(output)
+    }
+
+    testOperation(inputData, StateSpec.function(trackStateFunc), outputData, stateData)
+  }
+
+  test("trackStateByKey - initial states, with nothing emitted") {
+
+    val initialState = Seq(("a", 5), ("b", 10), ("c", -20), ("d", 0))
+
+    val inputData =
+      Seq(
+        Seq(),
+        Seq("a"),
+        Seq("a", "b"),
+        Seq("a", "b", "c"),
+        Seq("a", "b"),
+        Seq("a"),
+        Seq()
+      )
+
+    val outputData = Seq.fill(inputData.size)(Seq.empty[Int])
+
+    val stateData =
+      Seq(
+        Seq(("a", 5), ("b", 10), ("c", -20), ("d", 0)),
+        Seq(("a", 6), ("b", 10), ("c", -20), ("d", 0)),
+        Seq(("a", 7), ("b", 11), ("c", -20), ("d", 0)),
+        Seq(("a", 8), ("b", 12), ("c", -19), ("d", 0)),
+        Seq(("a", 9), ("b", 13), ("c", -19), ("d", 0)),
+        Seq(("a", 10), ("b", 13), ("c", -19), ("d", 0)),
+        Seq(("a", 10), ("b", 13), ("c", -19), ("d", 0))
+      )
+
+    val trackStateFunc = (time: Time, key: String, value: Option[Int], state: State[Int]) => {
+      val sum = value.getOrElse(0) + state.getOption.getOrElse(0)
+      val output = (key, sum)
+      state.update(sum)
+      None.asInstanceOf[Option[Int]]
+    }
+
+    val trackStateSpec = StateSpec.function(trackStateFunc).initialState(sc.makeRDD(initialState))
+    testOperation(inputData, trackStateSpec, outputData, stateData)
+  }
+
+  test("trackStateByKey - state removing") {
+    val inputData =
+      Seq(
+        Seq(),
+        Seq("a"),
+        Seq("a", "b"), // a will be removed
+        Seq("a", "b", "c"), // b will be removed
+        Seq("a", "b", "c"), // a and c will be removed
+        Seq("a", "b"), // b will be removed
+        Seq("a"), // a will be removed
+        Seq()
+      )
+
+    // States that were removed
+    val outputData =
+      Seq(
+        Seq(),
+        Seq(),
+        Seq("a"),
+        Seq("b"),
+        Seq("a", "c"),
+        Seq("b"),
+        Seq("a"),
+        Seq()
+      )
+
+    val stateData =
+      Seq(
+        Seq(),
+        Seq(("a", 1)),
+        Seq(("b", 1)),
+        Seq(("a", 1), ("c", 1)),
+        Seq(("b", 1)),
+        Seq(("a", 1)),
+        Seq(),
+        Seq()
+      )
+
+    val trackStateFunc = (time: Time, key: String, value: Option[Int], state: State[Int]) => {
+      if (state.exists) {
+        state.remove()
+        Some(key)
+      } else {
+        state.update(value.get)
+        None
+      }
+    }
+
+    testOperation(
+      inputData, StateSpec.function(trackStateFunc).numPartitions(1), outputData, stateData)
+  }
+
+  test("trackStateByKey - state timing out") {
+    val inputData =
+      Seq(
+        Seq("a", "b", "c"),
+        Seq("a", "b"),
+        Seq("a"),
+        Seq(), // c will time out
+        Seq(), // b will time out
+        Seq("a") // a will not time out
+      ) ++ Seq.fill(20)(Seq("a")) // a will continue to stay active
+
+    val trackStateFunc = (time: Time, key: String, value: Option[Int], state: State[Int]) => {
+      if (value.isDefined) {
+        state.update(1)
+      }
+      if (state.isTimingOut) {
+        Some(key)
+      } else {
+        None
+      }
+    }
+
+    val (collectedOutputs, collectedStateSnapshots) = getOperationOutput(
+      inputData, StateSpec.function(trackStateFunc).timeout(Seconds(3)), 20)
+
+    // b and c should be emitted once each, when they were marked as expired
+    assert(collectedOutputs.flatten.sorted === Seq("b", "c"))
+
+    // States for a, b, c should be defined at one point of time
+    assert(collectedStateSnapshots.exists {
+      _.toSet == Set(("a", 1), ("b", 1), ("c", 1))
+    })
+
+    // Finally state should be defined only for a
+    assert(collectedStateSnapshots.last.toSet === Set(("a", 1)))
+  }
+
+
+  private def testOperation[K: ClassTag, S: ClassTag, T: ClassTag](
+      input: Seq[Seq[K]],
+      trackStateSpec: StateSpec[K, Int, S, T],
+      expectedOutputs: Seq[Seq[T]],
+      expectedStateSnapshots: Seq[Seq[(K, S)]]
+    ): Unit = {
+    require(expectedOutputs.size == expectedStateSnapshots.size)
+
+    val (collectedOutputs, collectedStateSnapshots) =
+      getOperationOutput(input, trackStateSpec, expectedOutputs.size)
+    assert(expectedOutputs, collectedOutputs, "outputs")
+    assert(expectedStateSnapshots, collectedStateSnapshots, "state snapshots")
+  }
+
+  private def getOperationOutput[K: ClassTag, S: ClassTag, T: ClassTag](
+      input: Seq[Seq[K]],
+      trackStateSpec: StateSpec[K, Int, S, T],
+      numBatches: Int
+    ): (Seq[Seq[T]], Seq[Seq[(K, S)]]) = {
+
+    // Setup the stream computation
+    val inputStream = new TestInputStream(ssc, input, numPartitions = 2)
+    val trackeStateStream = inputStream.map(x => (x, 1)).trackStateByKey(trackStateSpec)
+    val collectedOutputs = new ArrayBuffer[Seq[T]] with SynchronizedBuffer[Seq[T]]
+    val outputStream = new TestOutputStream(trackeStateStream, collectedOutputs)
+    val collectedStateSnapshots = new ArrayBuffer[Seq[(K, S)]] with SynchronizedBuffer[Seq[(K, S)]]
+    val stateSnapshotStream = new TestOutputStream(
+      trackeStateStream.stateSnapshots(), collectedStateSnapshots)
+    outputStream.register()
+    stateSnapshotStream.register()
+
+    val batchCounter = new BatchCounter(ssc)
+    ssc.start()
+
+    val clock = ssc.scheduler.clock.asInstanceOf[ManualClock]
+    clock.advance(batchDuration.milliseconds * numBatches)
+
+    batchCounter.waitUntilBatchesCompleted(numBatches, 10000)
+    (collectedOutputs, collectedStateSnapshots)
+  }
+
+  private def assert[U](expected: Seq[Seq[U]], collected: Seq[Seq[U]], typ: String) {
+    val debugString = "\nExpected:\n" + expected.mkString("\n") +
+      "\nCollected:\n" + collected.mkString("\n")
+    assert(expected.size === collected.size,
+      s"number of collected $typ (${collected.size}) different from expected (${expected.size})" +
+        debugString)
+    expected.zip(collected).foreach { case (c, e) =>
+      assert(c.toSet === e.toSet,
+        s"collected $typ is different from expected $debugString"
+      )
+    }
+  }
+}
+
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
new file mode 100644
index 0000000000000..fc5f26607ef98
--- /dev/null
+++ b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
@@ -0,0 +1,193 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.rdd
+
+import scala.collection.mutable.ArrayBuffer
+import scala.reflect.ClassTag
+
+import org.scalatest.BeforeAndAfterAll
+
+import org.apache.spark.rdd.RDD
+import org.apache.spark.streaming.{Time, State}
+import org.apache.spark.{HashPartitioner, SparkConf, SparkContext, SparkFunSuite}
+
+class TrackStateRDDSuite extends SparkFunSuite with BeforeAndAfterAll {
+
+  private var sc = new SparkContext(
+    new SparkConf().setMaster("local").setAppName("TrackStateRDDSuite"))
+
+  override def afterAll(): Unit = {
+    sc.stop()
+  }
+
+  test("creation from pair RDD") {
+    val data = Seq((1, "1"), (2, "2"), (3, "3"))
+    val partitioner = new HashPartitioner(10)
+    val rdd = TrackStateRDD.createFromPairRDD[Int, Int, String, Int](
+      sc.parallelize(data), partitioner, Time(123))
+    assertRDD[Int, Int, String, Int](rdd, data.map { x => (x._1, x._2, 123)}.toSet, Set.empty)
+    assert(rdd.partitions.size === partitioner.numPartitions)
+
+    assert(rdd.partitioner === Some(partitioner))
+  }
+
+  test("states generated by TrackStateRDD") {
+    val initStates = Seq(("k1", 0), ("k2", 0))
+    val initTime = 123
+    val initStateWthTime = initStates.map { x => (x._1, x._2, initTime) }.toSet
+    val partitioner = new HashPartitioner(2)
+    val initStateRDD = TrackStateRDD.createFromPairRDD[String, Int, Int, Int](
+      sc.parallelize(initStates), partitioner, Time(initTime)).persist()
+    assertRDD(initStateRDD, initStateWthTime, Set.empty)
+
+    val updateTime = 345
+
+    /**
+     * Test that the test state RDD, when operated with new data,
+     * creates a new state RDD with expected states
+     */
+    def testStateUpdates(
+        testStateRDD: TrackStateRDD[String, Int, Int, Int],
+        testData: Seq[(String, Int)],
+        expectedStates: Set[(String, Int, Int)]): TrackStateRDD[String, Int, Int, Int] = {
+
+      // Persist the test TrackStateRDD so that its not recomputed while doing the next operation.
+      // This is to make sure that we only track which state keys are being touched in the next op.
+      testStateRDD.persist().count()
+
+      // To track which keys are being touched
+      TrackStateRDDSuite.touchedStateKeys.clear()
+
+      val trackingFunc = (time: Time, key: String, data: Option[Int], state: State[Int]) => {
+
+        // Track the key that has been touched
+        TrackStateRDDSuite.touchedStateKeys += key
+
+        // If the data is 0, do not do anything with the state
+        // else if the data is 1, increment the state if it exists, or set new state to 0
+        // else if the data is 2, remove the state if it exists
+        data match {
+          case Some(1) =>
+            if (state.exists()) { state.update(state.get + 1) }
+            else state.update(0)
+          case Some(2) =>
+            state.remove()
+          case _ =>
+        }
+        None.asInstanceOf[Option[Int]]  // Do not return anything, not being tested
+      }
+      val newDataRDD = sc.makeRDD(testData).partitionBy(testStateRDD.partitioner.get)
+
+      // Assert that the new state RDD has expected state data
+      val newStateRDD = assertOperation(
+        testStateRDD, newDataRDD, trackingFunc, updateTime, expectedStates, Set.empty)
+
+      // Assert that the function was called only for the keys present in the data
+      assert(TrackStateRDDSuite.touchedStateKeys.size === testData.size,
+        "More number of keys are being touched than that is expected")
+      assert(TrackStateRDDSuite.touchedStateKeys.toSet === testData.toMap.keys,
+        "Keys not in the data are being touched unexpectedly")
+
+      // Assert that the test RDD's data has not changed
+      assertRDD(initStateRDD, initStateWthTime, Set.empty)
+      newStateRDD
+    }
+
+    // Test no-op, no state should change
+    testStateUpdates(initStateRDD, Seq(), initStateWthTime)   // should not scan any state
+    testStateUpdates(
+      initStateRDD, Seq(("k1", 0)), initStateWthTime)         // should not update existing state
+    testStateUpdates(
+      initStateRDD, Seq(("k3", 0)), initStateWthTime)         // should not create new state
+
+    // Test creation of new state
+    val rdd1 = testStateUpdates(initStateRDD, Seq(("k3", 1)), // should create k3's state as 0
+      Set(("k1", 0, initTime), ("k2", 0, initTime), ("k3", 0, updateTime)))
+
+    val rdd2 = testStateUpdates(rdd1, Seq(("k4", 1)),         // should create k4's state as 0
+      Set(("k1", 0, initTime), ("k2", 0, initTime), ("k3", 0, updateTime), ("k4", 0, updateTime)))
+
+    // Test updating of state
+    val rdd3 = testStateUpdates(
+      initStateRDD, Seq(("k1", 1)),                   // should increment k1's state 0 -> 1
+      Set(("k1", 1, updateTime), ("k2", 0, initTime)))
+
+    val rdd4 = testStateUpdates(rdd3,
+      Seq(("x", 0), ("k2", 1), ("k2", 1), ("k3", 1)),  // should update k2, 0 -> 2 and create k3, 0
+      Set(("k1", 1, updateTime), ("k2", 2, updateTime), ("k3", 0, updateTime)))
+
+    val rdd5 = testStateUpdates(
+      rdd4, Seq(("k3", 1)),                           // should update k3's state 0 -> 2
+      Set(("k1", 1, updateTime), ("k2", 2, updateTime), ("k3", 1, updateTime)))
+
+    // Test removing of state
+    val rdd6 = testStateUpdates(                      // should remove k1's state
+      initStateRDD, Seq(("k1", 2)), Set(("k2", 0, initTime)))
+
+    val rdd7 = testStateUpdates(                      // should remove k2's state
+      rdd6, Seq(("k2", 2), ("k0", 2), ("k3", 1)), Set(("k3", 0, updateTime)))
+
+    val rdd8 = testStateUpdates(
+      rdd7, Seq(("k3", 2)), Set()                     //
+    )
+  }
+
+  /** Assert whether the `trackStateByKey` operation generates expected results */
+  private def assertOperation[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
+      testStateRDD: TrackStateRDD[K, V, S, T],
+      newDataRDD: RDD[(K, V)],
+      trackStateFunc: (Time, K, Option[V], State[S]) => Option[T],
+      currentTime: Long,
+      expectedStates: Set[(K, S, Int)],
+      expectedEmittedRecords: Set[T],
+      doFullScan: Boolean = false
+    ): TrackStateRDD[K, V, S, T] = {
+
+    val partitionedNewDataRDD = if (newDataRDD.partitioner != testStateRDD.partitioner) {
+      newDataRDD.partitionBy(testStateRDD.partitioner.get)
+    } else {
+      newDataRDD
+    }
+
+    val newStateRDD = new TrackStateRDD[K, V, S, T](
+      testStateRDD, newDataRDD, trackStateFunc, Time(currentTime), None)
+    if (doFullScan) newStateRDD.setFullScan()
+
+    // Persist to make sure that it gets computed only once and we can track precisely how many
+    // state keys the computing touched
+    newStateRDD.persist()
+    assertRDD(newStateRDD, expectedStates, expectedEmittedRecords)
+    newStateRDD
+  }
+
+  /** Assert whether the [[TrackStateRDD]] has the expected state ad emitted records */
+  private def assertRDD[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
+      trackStateRDD: TrackStateRDD[K, V, S, T],
+      expectedStates: Set[(K, S, Int)],
+      expectedEmittedRecords: Set[T]): Unit = {
+    val states = trackStateRDD.flatMap { _.stateMap.getAll() }.collect().toSet
+    val emittedRecords = trackStateRDD.flatMap { _.emittedRecords }.collect().toSet
+    assert(states === expectedStates, "states after track state operation were not as expected")
+    assert(emittedRecords === expectedEmittedRecords,
+      "emitted records after track state operation were not as expected")
+  }
+}
+
+object TrackStateRDDSuite {
+  private val touchedStateKeys = new ArrayBuffer[String]()
+}

From 1bc41125ee6306e627be212969854f639969c440 Mon Sep 17 00:00:00 2001
From: hyukjinkwon <gurwls223@gmail.com>
Date: Wed, 11 Nov 2015 16:46:04 +0800
Subject: [PATCH 185/518] [SPARK-11500][SQL] Not deterministic order of columns
 when using merging schemas.

https://issues.apache.org/jira/browse/SPARK-11500

As filed in SPARK-11500, if merging schemas is enabled, the order of files to touch is a matter which might affect the ordering of the output columns.

This was mostly because of the use of `Set` and `Map` so I replaced them to `LinkedHashSet` and `LinkedHashMap` to keep the insertion order.

Also, I changed `reduceOption` to `reduceLeftOption`, and replaced the order of `filesToTouch` from `metadataStatuses ++ commonMetadataStatuses ++ needMerged` to  `needMerged ++ metadataStatuses ++ commonMetadataStatuses` in order to touch the part-files first which always have the schema in footers whereas the others might not exist.

One nit is, If merging schemas is not enabled, but when multiple files are given, there is no guarantee of the output order, since there might not be a summary file for the first file, which ends up putting ahead the columns of the other files.

However, I thought this should be okay since disabling merging schemas means (assumes) all the files have the same schemas.

In addition, in the test code for this, I only checked the names of fields.

Author: hyukjinkwon <gurwls223@gmail.com>

Closes #9517 from HyukjinKwon/SPARK-11500.
---
 .../datasources/parquet/ParquetRelation.scala | 29 +++++++++++++++----
 .../apache/spark/sql/sources/interfaces.scala | 22 +++++++-------
 .../ParquetHadoopFsRelationSuite.scala        | 21 +++++++++++++-
 3 files changed, 55 insertions(+), 17 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
index 5a7c6b95b565f..21337b2932aac 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
@@ -383,7 +383,7 @@ private[sql] class ParquetRelation(
     var schema: StructType = _
 
     // Cached leaves
-    var cachedLeaves: Set[FileStatus] = null
+    var cachedLeaves: mutable.LinkedHashSet[FileStatus] = null
 
     /**
      * Refreshes `FileStatus`es, footers, partition spec, and table schema.
@@ -396,13 +396,13 @@ private[sql] class ParquetRelation(
         !cachedLeaves.equals(currentLeafStatuses)
 
       if (leafStatusesChanged) {
-        cachedLeaves = currentLeafStatuses.toIterator.toSet
+        cachedLeaves = currentLeafStatuses
 
         // Lists `FileStatus`es of all leaf nodes (files) under all base directories.
         val leaves = currentLeafStatuses.filter { f =>
           isSummaryFile(f.getPath) ||
             !(f.getPath.getName.startsWith("_") || f.getPath.getName.startsWith("."))
-        }.toArray
+        }.toArray.sortBy(_.getPath.toString)
 
         dataStatuses = leaves.filterNot(f => isSummaryFile(f.getPath))
         metadataStatuses =
@@ -465,13 +465,30 @@ private[sql] class ParquetRelation(
           // You should enable this configuration only if you are very sure that for the parquet
           // part-files to read there are corresponding summary files containing correct schema.
 
+          // As filed in SPARK-11500, the order of files to touch is a matter, which might affect
+          // the ordering of the output columns. There are several things to mention here.
+          //
+          //  1. If mergeRespectSummaries config is false, then it merges schemas by reducing from
+          //     the first part-file so that the columns of the lexicographically first file show
+          //     first.
+          //
+          //  2. If mergeRespectSummaries config is true, then there should be, at least,
+          //     "_metadata"s for all given files, so that we can ensure the columns of
+          //     the lexicographically first file show first.
+          //
+          //  3. If shouldMergeSchemas is false, but when multiple files are given, there is
+          //     no guarantee of the output order, since there might not be a summary file for the
+          //     lexicographically first file, which ends up putting ahead the columns of
+          //     the other files. However, this should be okay since not enabling
+          //     shouldMergeSchemas means (assumes) all the files have the same schemas.
+
           val needMerged: Seq[FileStatus] =
             if (mergeRespectSummaries) {
               Seq()
             } else {
               dataStatuses
             }
-          (metadataStatuses ++ commonMetadataStatuses ++ needMerged).toSeq
+          needMerged ++ metadataStatuses ++ commonMetadataStatuses
         } else {
           // Tries any "_common_metadata" first. Parquet files written by old versions or Parquet
           // don't have this.
@@ -768,10 +785,10 @@ private[sql] object ParquetRelation extends Logging {
 
           footers.map { footer =>
             ParquetRelation.readSchemaFromFooter(footer, converter)
-          }.reduceOption(_ merge _).iterator
+          }.reduceLeftOption(_ merge _).iterator
         }.collect()
 
-    partiallyMergedSchemas.reduceOption(_ merge _)
+    partiallyMergedSchemas.reduceLeftOption(_ merge _)
   }
 
   /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index e296d631f0f30..5b8841bc154a5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -428,11 +428,11 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
   private var _partitionSpec: PartitionSpec = _
 
   private class FileStatusCache {
-    var leafFiles = mutable.Map.empty[Path, FileStatus]
+    var leafFiles = mutable.LinkedHashMap.empty[Path, FileStatus]
 
     var leafDirToChildrenFiles = mutable.Map.empty[Path, Array[FileStatus]]
 
-    private def listLeafFiles(paths: Array[String]): Set[FileStatus] = {
+    private def listLeafFiles(paths: Array[String]): mutable.LinkedHashSet[FileStatus] = {
       if (paths.length >= sqlContext.conf.parallelPartitionDiscoveryThreshold) {
         HadoopFsRelation.listLeafFilesInParallel(paths, hadoopConf, sqlContext.sparkContext)
       } else {
@@ -450,10 +450,11 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
 
         val (dirs, files) = statuses.partition(_.isDir)
 
+        // It uses [[LinkedHashSet]] since the order of files can affect the results. (SPARK-11500)
         if (dirs.isEmpty) {
-          files.toSet
+          mutable.LinkedHashSet(files: _*)
         } else {
-          files.toSet ++ listLeafFiles(dirs.map(_.getPath.toString))
+          mutable.LinkedHashSet(files: _*) ++ listLeafFiles(dirs.map(_.getPath.toString))
         }
       }
     }
@@ -464,7 +465,7 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
       leafFiles.clear()
       leafDirToChildrenFiles.clear()
 
-      leafFiles ++= files.map(f => f.getPath -> f).toMap
+      leafFiles ++= files.map(f => f.getPath -> f)
       leafDirToChildrenFiles ++= files.toArray.groupBy(_.getPath.getParent)
     }
   }
@@ -475,8 +476,8 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
     cache
   }
 
-  protected def cachedLeafStatuses(): Set[FileStatus] = {
-    fileStatusCache.leafFiles.values.toSet
+  protected def cachedLeafStatuses(): mutable.LinkedHashSet[FileStatus] = {
+    mutable.LinkedHashSet(fileStatusCache.leafFiles.values.toArray: _*)
   }
 
   final private[sql] def partitionSpec: PartitionSpec = {
@@ -834,7 +835,7 @@ private[sql] object HadoopFsRelation extends Logging {
   def listLeafFilesInParallel(
       paths: Array[String],
       hadoopConf: Configuration,
-      sparkContext: SparkContext): Set[FileStatus] = {
+      sparkContext: SparkContext): mutable.LinkedHashSet[FileStatus] = {
     logInfo(s"Listing leaf files and directories in parallel under: ${paths.mkString(", ")}")
 
     val serializableConfiguration = new SerializableConfiguration(hadoopConf)
@@ -854,9 +855,10 @@ private[sql] object HadoopFsRelation extends Logging {
         status.getAccessTime)
     }.collect()
 
-    fakeStatuses.map { f =>
+    val hadoopFakeStatuses = fakeStatuses.map { f =>
       new FileStatus(
         f.length, f.isDir, f.blockReplication, f.blockSize, f.modificationTime, new Path(f.path))
-    }.toSet
+    }
+    mutable.LinkedHashSet(hadoopFakeStatuses: _*)
   }
 }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
index e2d754e806403..e866493ee6c96 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
@@ -23,7 +23,7 @@ import com.google.common.io.Files
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.deploy.SparkHadoopUtil
-import org.apache.spark.sql.{execution, AnalysisException, SaveMode}
+import org.apache.spark.sql._
 import org.apache.spark.sql.types._
 
 
@@ -155,4 +155,23 @@ class ParquetHadoopFsRelationSuite extends HadoopFsRelationTest {
       assert(physicalPlan.collect { case p: execution.Filter => p }.length === 1)
     }
   }
+
+  test("SPARK-11500: Not deterministic order of columns when using merging schemas.") {
+    import testImplicits._
+    withSQLConf(SQLConf.PARQUET_SCHEMA_MERGING_ENABLED.key -> "true") {
+      withTempPath { dir =>
+        val pathOne = s"${dir.getCanonicalPath}/part=1"
+        Seq(1, 1).zipWithIndex.toDF("a", "b").write.parquet(pathOne)
+        val pathTwo = s"${dir.getCanonicalPath}/part=2"
+        Seq(1, 1).zipWithIndex.toDF("c", "b").write.parquet(pathTwo)
+        val pathThree = s"${dir.getCanonicalPath}/part=3"
+        Seq(1, 1).zipWithIndex.toDF("d", "b").write.parquet(pathThree)
+
+        // The schema consists of the leading columns of the first part-file
+        // in the lexicographic order.
+        assert(sqlContext.read.parquet(dir.getCanonicalPath).schema.map(_.name)
+          === Seq("a", "b", "c", "d", "part"))
+      }
+    }
+  }
 }

From 1510c527b4f5ee0953ae42313ef9e16d2f5864c4 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 11 Nov 2015 09:33:07 -0800
Subject: [PATCH 186/518] [SPARK-10371][SQL][FOLLOW-UP] fix code style

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9627 from cloud-fan/follow.
---
 .../expressions/EquivalentExpressions.scala   | 22 +++++++++----------
 .../sql/catalyst/expressions/Expression.scala | 21 +++++++++---------
 .../expressions/codegen/CodeGenerator.scala   | 20 ++++++++---------
 3 files changed, 30 insertions(+), 33 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
index e7380d21f98af..f83df494ba8a6 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
@@ -29,29 +29,27 @@ class EquivalentExpressions {
    * Wrapper around an Expression that provides semantic equality.
    */
   case class Expr(e: Expression) {
-    val hash = e.semanticHash()
     override def equals(o: Any): Boolean = o match {
       case other: Expr => e.semanticEquals(other.e)
       case _ => false
     }
-    override def hashCode: Int = hash
+    override val hashCode: Int = e.semanticHash()
   }
 
   // For each expression, the set of equivalent expressions.
-  private val equivalenceMap: mutable.HashMap[Expr, mutable.MutableList[Expression]] =
-      new mutable.HashMap[Expr, mutable.MutableList[Expression]]
+  private val equivalenceMap = mutable.HashMap.empty[Expr, mutable.MutableList[Expression]]
 
   /**
    * Adds each expression to this data structure, grouping them with existing equivalent
    * expressions. Non-recursive.
-   * Returns if there was already a matching expression.
+   * Returns true if there was already a matching expression.
    */
   def addExpr(expr: Expression): Boolean = {
     if (expr.deterministic) {
       val e: Expr = Expr(expr)
       val f = equivalenceMap.get(e)
       if (f.isDefined) {
-        f.get.+= (expr)
+        f.get += expr
         true
       } else {
         equivalenceMap.put(e, mutable.MutableList(expr))
@@ -63,19 +61,19 @@ class EquivalentExpressions {
   }
 
   /**
-   * Adds the expression to this datastructure recursively. Stops if a matching expression
+   * Adds the expression to this data structure recursively. Stops if a matching expression
    * is found. That is, if `expr` has already been added, its children are not added.
    * If ignoreLeaf is true, leaf nodes are ignored.
    */
   def addExprTree(root: Expression, ignoreLeaf: Boolean = true): Unit = {
     val skip = root.isInstanceOf[LeafExpression] && ignoreLeaf
-    if (!skip && root.deterministic && !addExpr(root)) {
-     root.children.foreach(addExprTree(_, ignoreLeaf))
+    if (!skip && !addExpr(root)) {
+      root.children.foreach(addExprTree(_, ignoreLeaf))
     }
   }
 
   /**
-   * Returns all fo the expression trees that are equivalent to `e`. Returns
+   * Returns all of the expression trees that are equivalent to `e`. Returns
    * an empty collection if there are none.
    */
   def getEquivalentExprs(e: Expression): Seq[Expression] = {
@@ -90,8 +88,8 @@ class EquivalentExpressions {
   }
 
   /**
-   * Returns the state of the datastructure as a string. If all is false, skips sets of equivalent
-   * expressions with cardinality 1.
+   * Returns the state of the data structure as a string. If `all` is false, skips sets of
+   * equivalent expressions with cardinality 1.
    */
   def debugString(all: Boolean = false): String = {
     val sb: mutable.StringBuilder = new StringBuilder()
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
index 7d5741eefcc7d..540ed3500616a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
@@ -92,17 +92,16 @@ abstract class Expression extends TreeNode[Expression] {
    * @return [[GeneratedExpressionCode]]
    */
   def gen(ctx: CodeGenContext): GeneratedExpressionCode = {
-    val subExprState = ctx.subExprEliminationExprs.get(this)
-    if (subExprState.isDefined) {
+    ctx.subExprEliminationExprs.get(this).map { subExprState =>
       // This expression is repeated meaning the code to evaluated has already been added
       // as a function, `subExprState.fnName`. Just call that.
       val code =
         s"""
            |/* $this */
-           |${subExprState.get.fnName}(${ctx.INPUT_ROW});
-           |""".stripMargin.trim
-      GeneratedExpressionCode(code, subExprState.get.code.isNull, subExprState.get.code.value)
-    } else {
+           |${subExprState.fnName}(${ctx.INPUT_ROW});
+         """.stripMargin.trim
+      GeneratedExpressionCode(code, subExprState.code.isNull, subExprState.code.value)
+    }.getOrElse {
       val isNull = ctx.freshName("isNull")
       val primitive = ctx.freshName("primitive")
       val ve = GeneratedExpressionCode("", isNull, primitive)
@@ -157,7 +156,7 @@ abstract class Expression extends TreeNode[Expression] {
         case (i1, i2) => i1 == i2
       }
     }
-    // Non-determinstic expressions cannot be equal
+    // Non-deterministic expressions cannot be semantic equal
     if (!deterministic || !other.deterministic) return false
     val elements1 = this.productIterator.toSeq
     val elements2 = other.asInstanceOf[Product].productIterator.toSeq
@@ -174,11 +173,11 @@ abstract class Expression extends TreeNode[Expression] {
       var hash: Int = 17
       e.foreach(i => {
         val h: Int = i match {
-          case (e: Expression) => e.semanticHash()
-          case (Some(e: Expression)) => e.semanticHash()
-          case (t: Traversable[_]) => computeHash(t.toSeq)
+          case e: Expression => e.semanticHash()
+          case Some(e: Expression) => e.semanticHash()
+          case t: Traversable[_] => computeHash(t.toSeq)
           case null => 0
-          case (o) => o.hashCode()
+          case other => other.hashCode()
         }
         hash = hash * 37 + h
       })
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
index 60a3d6018496c..5a4bba232b04b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
@@ -109,15 +109,15 @@ class CodeGenContext {
 
   // State used for subexpression elimination.
   case class SubExprEliminationState(
-    val isLoaded: String, code: GeneratedExpressionCode, val fnName: String)
+      isLoaded: String,
+      code: GeneratedExpressionCode,
+      fnName: String)
 
   // Foreach expression that is participating in subexpression elimination, the state to use.
-  val subExprEliminationExprs: mutable.HashMap[Expression, SubExprEliminationState] =
-    mutable.HashMap[Expression, SubExprEliminationState]()
+  val subExprEliminationExprs = mutable.HashMap.empty[Expression, SubExprEliminationState]
 
   // The collection of isLoaded variables that need to be reset on each row.
-  val subExprIsLoadedVariables: mutable.ArrayBuffer[String] =
-    mutable.ArrayBuffer.empty[String]
+  val subExprIsLoadedVariables = mutable.ArrayBuffer.empty[String]
 
   final val JAVA_BOOLEAN = "boolean"
   final val JAVA_BYTE = "byte"
@@ -361,7 +361,7 @@ class CodeGenContext {
       val expr = e.head
       val isLoaded = freshName("isLoaded")
       val isNull = freshName("isNull")
-      val primitive = freshName("primitive")
+      val value = freshName("value")
       val fnName = freshName("evalExpr")
 
       // Generate the code for this expression tree and wrap it in a function.
@@ -373,13 +373,13 @@ class CodeGenContext {
            |    ${code.code.trim}
            |    $isLoaded = true;
            |    $isNull = ${code.isNull};
-           |    $primitive = ${code.value};
+           |    $value = ${code.value};
            |  }
            |}
            """.stripMargin
       code.code = fn
       code.isNull = isNull
-      code.value = primitive
+      code.value = value
 
       addNewFunction(fnName, fn)
 
@@ -406,8 +406,8 @@ class CodeGenContext {
       // optimizations.
       addMutableState("boolean", isLoaded, s"$isLoaded = false;")
       addMutableState("boolean", isNull, s"$isNull = false;")
-      addMutableState(javaType(expr.dataType), primitive,
-        s"$primitive = ${defaultValue(expr.dataType)};")
+      addMutableState(javaType(expr.dataType), value,
+        s"$value = ${defaultValue(expr.dataType)};")
       subExprIsLoadedVariables += isLoaded
 
       val state = SubExprEliminationState(isLoaded, code, fnName)

From 27524a3a9ccee6fbe56149180ebfb3f74e0957e7 Mon Sep 17 00:00:00 2001
From: Yuming Wang <q79969786@gmail.com>
Date: Wed, 11 Nov 2015 09:43:26 -0800
Subject: [PATCH 187/518] [SPARK-11626][ML] ml.feature.Word2Vec.transform()
 function very slow

org.apache.spark.ml.feature.Word2Vec.transform() very slow. we should not read broadcast every sentence.

Author: Yuming Wang <q79969786@gmail.com>
Author: yuming.wang <q79969786@gmail.com>
Author: Xiangrui Meng <meng@databricks.com>

Closes #9592 from 979969786/master.
---
 .../apache/spark/ml/feature/Word2Vec.scala    | 34 +++++++++----------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
index 9edab3af913ca..5c64cb09d5948 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
@@ -17,18 +17,16 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.SparkContext
+import org.apache.spark.annotation.Experimental
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
 import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
 import org.apache.spark.mllib.feature
-import org.apache.spark.mllib.linalg.{VectorUDT, Vector, Vectors}
-import org.apache.spark.mllib.linalg.BLAS._
-import org.apache.spark.sql.DataFrame
+import org.apache.spark.mllib.linalg.{BLAS, Vector, VectorUDT, Vectors}
+import org.apache.spark.sql.{DataFrame, SQLContext}
 import org.apache.spark.sql.functions._
-import org.apache.spark.sql.SQLContext
 import org.apache.spark.sql.types._
 
 /**
@@ -148,10 +146,9 @@ final class Word2Vec(override val uid: String) extends Estimator[Word2VecModel]
 @Experimental
 class Word2VecModel private[ml] (
     override val uid: String,
-    wordVectors: feature.Word2VecModel)
+    @transient wordVectors: feature.Word2VecModel)
   extends Model[Word2VecModel] with Word2VecBase {
 
-
   /**
    * Returns a dataframe with two fields, "word" and "vector", with "word" being a String and
    * and the vector the DenseVector that it is mapped to.
@@ -197,22 +194,23 @@ class Word2VecModel private[ml] (
    */
   override def transform(dataset: DataFrame): DataFrame = {
     transformSchema(dataset.schema, logging = true)
-    val bWordVectors = dataset.sqlContext.sparkContext.broadcast(wordVectors)
+    val vectors = wordVectors.getVectors
+      .mapValues(vv => Vectors.dense(vv.map(_.toDouble)))
+      .map(identity) // mapValues doesn't return a serializable map (SI-7005)
+    val bVectors = dataset.sqlContext.sparkContext.broadcast(vectors)
+    val d = $(vectorSize)
     val word2Vec = udf { sentence: Seq[String] =>
       if (sentence.size == 0) {
-        Vectors.sparse($(vectorSize), Array.empty[Int], Array.empty[Double])
+        Vectors.sparse(d, Array.empty[Int], Array.empty[Double])
       } else {
-        val cum = Vectors.zeros($(vectorSize))
-        val model = bWordVectors.value.getVectors
-        for (word <- sentence) {
-          if (model.contains(word)) {
-            axpy(1.0, bWordVectors.value.transform(word), cum)
-          } else {
-            // pass words which not belong to model
+        val sum = Vectors.zeros(d)
+        sentence.foreach { word =>
+          bVectors.value.get(word).foreach { v =>
+            BLAS.axpy(1.0, v, sum)
           }
         }
-        scal(1.0 / sentence.size, cum)
-        cum
+        BLAS.scal(1.0 / sentence.size, sum)
+        sum
       }
     }
     dataset.withColumn($(outputCol), word2Vec(col($(inputCol))))

From 95daff6459fc749949c2d71a0b7ab1c5be854f70 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 11 Nov 2015 10:17:54 -0800
Subject: [PATCH 188/518] [SPARK-11646] WholeTextFileRDD should return Text
 rather than String

If it returns Text, we can reuse this in Spark SQL to provide a WholeTextFile data source and directly convert the Text into UTF8String without extra string decoding and encoding.

Author: Reynold Xin <rxin@databricks.com>

Closes #9622 from rxin/SPARK-11646.
---
 .../scala/org/apache/spark/SparkContext.scala |  6 +-
 .../input/WholeTextFileInputFormat.scala      |  6 +-
 .../input/WholeTextFileRecordReader.scala     | 12 ++--
 .../org/apache/spark/rdd/NewHadoopRDD.scala   | 33 +----------
 .../apache/spark/rdd/WholeTextFileRDD.scala   | 56 +++++++++++++++++++
 5 files changed, 69 insertions(+), 44 deletions(-)
 create mode 100644 core/src/main/scala/org/apache/spark/rdd/WholeTextFileRDD.scala

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 67270c38fa1b5..43a241686fd7d 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -863,10 +863,10 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
     new WholeTextFileRDD(
       this,
       classOf[WholeTextFileInputFormat],
-      classOf[String],
-      classOf[String],
+      classOf[Text],
+      classOf[Text],
       updateConf,
-      minPartitions).setName(path)
+      minPartitions).setName(path).map(record => (record._1.toString, record._2.toString))
   }
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/input/WholeTextFileInputFormat.scala b/core/src/main/scala/org/apache/spark/input/WholeTextFileInputFormat.scala
index 1ba34a11414a2..413408723b54d 100644
--- a/core/src/main/scala/org/apache/spark/input/WholeTextFileInputFormat.scala
+++ b/core/src/main/scala/org/apache/spark/input/WholeTextFileInputFormat.scala
@@ -20,6 +20,7 @@ package org.apache.spark.input
 import scala.collection.JavaConverters._
 
 import org.apache.hadoop.fs.Path
+import org.apache.hadoop.io.Text
 import org.apache.hadoop.mapreduce.InputSplit
 import org.apache.hadoop.mapreduce.JobContext
 import org.apache.hadoop.mapreduce.lib.input.CombineFileInputFormat
@@ -33,14 +34,13 @@ import org.apache.hadoop.mapreduce.TaskAttemptContext
  */
 
 private[spark] class WholeTextFileInputFormat
-  extends CombineFileInputFormat[String, String] with Configurable {
+  extends CombineFileInputFormat[Text, Text] with Configurable {
 
   override protected def isSplitable(context: JobContext, file: Path): Boolean = false
 
   override def createRecordReader(
       split: InputSplit,
-      context: TaskAttemptContext): RecordReader[String, String] = {
-
+      context: TaskAttemptContext): RecordReader[Text, Text] = {
     val reader =
       new ConfigurableCombineFileRecordReader(split, context, classOf[WholeTextFileRecordReader])
     reader.setConf(getConf)
diff --git a/core/src/main/scala/org/apache/spark/input/WholeTextFileRecordReader.scala b/core/src/main/scala/org/apache/spark/input/WholeTextFileRecordReader.scala
index 31bde8a78f3c6..b56b2aa88a414 100644
--- a/core/src/main/scala/org/apache/spark/input/WholeTextFileRecordReader.scala
+++ b/core/src/main/scala/org/apache/spark/input/WholeTextFileRecordReader.scala
@@ -49,7 +49,7 @@ private[spark] class WholeTextFileRecordReader(
     split: CombineFileSplit,
     context: TaskAttemptContext,
     index: Integer)
-  extends RecordReader[String, String] with Configurable {
+  extends RecordReader[Text, Text] with Configurable {
 
   private[this] val path = split.getPath(index)
   private[this] val fs = path.getFileSystem(
@@ -58,8 +58,8 @@ private[spark] class WholeTextFileRecordReader(
   // True means the current file has been processed, then skip it.
   private[this] var processed = false
 
-  private[this] val key = path.toString
-  private[this] var value: String = null
+  private[this] val key: Text = new Text(path.toString)
+  private[this] var value: Text = null
 
   override def initialize(split: InputSplit, context: TaskAttemptContext): Unit = {}
 
@@ -67,9 +67,9 @@ private[spark] class WholeTextFileRecordReader(
 
   override def getProgress: Float = if (processed) 1.0f else 0.0f
 
-  override def getCurrentKey: String = key
+  override def getCurrentKey: Text = key
 
-  override def getCurrentValue: String = value
+  override def getCurrentValue: Text = value
 
   override def nextKeyValue(): Boolean = {
     if (!processed) {
@@ -83,7 +83,7 @@ private[spark] class WholeTextFileRecordReader(
         ByteStreams.toByteArray(fileIn)
       }
 
-      value = new Text(innerBuffer).toString
+      value = new Text(innerBuffer)
       Closeables.close(fileIn, false)
       processed = true
       true
diff --git a/core/src/main/scala/org/apache/spark/rdd/NewHadoopRDD.scala b/core/src/main/scala/org/apache/spark/rdd/NewHadoopRDD.scala
index 9c4b70844bdbe..d1960990da0fe 100644
--- a/core/src/main/scala/org/apache/spark/rdd/NewHadoopRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/NewHadoopRDD.scala
@@ -28,12 +28,11 @@ import org.apache.hadoop.mapreduce._
 import org.apache.hadoop.mapreduce.lib.input.{CombineFileSplit, FileSplit}
 
 import org.apache.spark.annotation.DeveloperApi
-import org.apache.spark.input.WholeTextFileInputFormat
 import org.apache.spark._
 import org.apache.spark.executor.DataReadMethod
 import org.apache.spark.mapreduce.SparkHadoopMapReduceUtil
 import org.apache.spark.rdd.NewHadoopRDD.NewHadoopMapPartitionsWithSplitRDD
-import org.apache.spark.util.{SerializableConfiguration, ShutdownHookManager, Utils}
+import org.apache.spark.util.{SerializableConfiguration, ShutdownHookManager}
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.storage.StorageLevel
 
@@ -59,7 +58,6 @@ private[spark] class NewHadoopPartition(
  * @param inputFormatClass Storage format of the data to be read.
  * @param keyClass Class of the key associated with the inputFormatClass.
  * @param valueClass Class of the value associated with the inputFormatClass.
- * @param conf The Hadoop configuration.
  */
 @DeveloperApi
 class NewHadoopRDD[K, V](
@@ -282,32 +280,3 @@ private[spark] object NewHadoopRDD {
     }
   }
 }
-
-private[spark] class WholeTextFileRDD(
-    sc : SparkContext,
-    inputFormatClass: Class[_ <: WholeTextFileInputFormat],
-    keyClass: Class[String],
-    valueClass: Class[String],
-    conf: Configuration,
-    minPartitions: Int)
-  extends NewHadoopRDD[String, String](sc, inputFormatClass, keyClass, valueClass, conf) {
-
-  override def getPartitions: Array[Partition] = {
-    val inputFormat = inputFormatClass.newInstance
-    val conf = getConf
-    inputFormat match {
-      case configurable: Configurable =>
-        configurable.setConf(conf)
-      case _ =>
-    }
-    val jobContext = newJobContext(conf, jobId)
-    inputFormat.setMinPartitions(jobContext, minPartitions)
-    val rawSplits = inputFormat.getSplits(jobContext).toArray
-    val result = new Array[Partition](rawSplits.size)
-    for (i <- 0 until rawSplits.size) {
-      result(i) = new NewHadoopPartition(id, i, rawSplits(i).asInstanceOf[InputSplit with Writable])
-    }
-    result
-  }
-}
-
diff --git a/core/src/main/scala/org/apache/spark/rdd/WholeTextFileRDD.scala b/core/src/main/scala/org/apache/spark/rdd/WholeTextFileRDD.scala
new file mode 100644
index 0000000000000..e3f14fe7ef0f8
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/rdd/WholeTextFileRDD.scala
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.rdd
+
+import org.apache.hadoop.conf.{Configurable, Configuration}
+import org.apache.hadoop.io.{Text, Writable}
+import org.apache.hadoop.mapreduce.InputSplit
+
+import org.apache.spark.{Partition, SparkContext}
+import org.apache.spark.input.WholeTextFileInputFormat
+
+/**
+ * An RDD that reads a bunch of text files in, and each text file becomes one record.
+ */
+private[spark] class WholeTextFileRDD(
+    sc : SparkContext,
+    inputFormatClass: Class[_ <: WholeTextFileInputFormat],
+    keyClass: Class[Text],
+    valueClass: Class[Text],
+    conf: Configuration,
+    minPartitions: Int)
+  extends NewHadoopRDD[Text, Text](sc, inputFormatClass, keyClass, valueClass, conf) {
+
+  override def getPartitions: Array[Partition] = {
+    val inputFormat = inputFormatClass.newInstance
+    val conf = getConf
+    inputFormat match {
+      case configurable: Configurable =>
+        configurable.setConf(conf)
+      case _ =>
+    }
+    val jobContext = newJobContext(conf, jobId)
+    inputFormat.setMinPartitions(jobContext, minPartitions)
+    val rawSplits = inputFormat.getSplits(jobContext).toArray
+    val result = new Array[Partition](rawSplits.size)
+    for (i <- 0 until rawSplits.size) {
+      result(i) = new NewHadoopPartition(id, i, rawSplits(i).asInstanceOf[InputSplit with Writable])
+    }
+    result
+  }
+}

From c964fc101585171aee76996981fe2c9fdafc614e Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 11 Nov 2015 10:19:09 -0800
Subject: [PATCH 189/518] [SQL][MINOR] rename present to finish in Aggregator

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9617 from cloud-fan/tmp.
---
 .../sql/execution/aggregate/TypedAggregateExpression.scala  | 2 +-
 .../scala/org/apache/spark/sql/expressions/Aggregator.scala | 2 +-
 .../scala/org/apache/spark/sql/DatasetAggregatorSuite.scala | 6 +++---
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
index 0e5bc1f9abf28..b5a87c56e681b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
@@ -136,7 +136,7 @@ case class TypedAggregateExpression(
 
   override def eval(buffer: InternalRow): Any = {
     val b = boundB.shift(mutableAggBufferOffset).fromRow(buffer)
-    val result = cEncoder.toRow(aggregator.present(b))
+    val result = cEncoder.toRow(aggregator.finish(b))
     dataType match {
       case _: StructType => result
       case _ => result.get(0, dataType)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
index 3c1c457e06d5f..2aa5a7d540316 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
@@ -65,7 +65,7 @@ abstract class Aggregator[-A, B, C] {
   /**
    * Transform the output of the reduction.
    */
-  def present(reduction: B): C
+  def finish(reduction: B): C
 
   /**
    * Returns this `Aggregator` as a [[TypedColumn]] that can be used in [[Dataset]] or [[DataFrame]]
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
index 206095a519762..002d5c18f0a25 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -36,7 +36,7 @@ class SumOf[I, N : Numeric](f: I => N) extends Aggregator[I, N, N] with Serializ
 
   override def merge(b1: N, b2: N): N = numeric.plus(b1, b2)
 
-  override def present(reduction: N): N = reduction
+  override def finish(reduction: N): N = reduction
 }
 
 object TypedAverage extends Aggregator[(String, Int), (Long, Long), Double] with Serializable {
@@ -50,7 +50,7 @@ object TypedAverage extends Aggregator[(String, Int), (Long, Long), Double] with
     (b1._1 + b2._1, b1._2 + b2._2)
   }
 
-  override def present(countAndSum: (Long, Long)): Double = countAndSum._2 / countAndSum._1
+  override def finish(countAndSum: (Long, Long)): Double = countAndSum._2 / countAndSum._1
 }
 
 object ComplexResultAgg extends Aggregator[(String, Int), (Long, Long), (Long, Long)]
@@ -66,7 +66,7 @@ object ComplexResultAgg extends Aggregator[(String, Int), (Long, Long), (Long, L
     (b1._1 + b2._1, b1._2 + b2._2)
   }
 
-  override def present(reduction: (Long, Long)): (Long, Long) = reduction
+  override def finish(reduction: (Long, Long)): (Long, Long) = reduction
 }
 
 class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {

From 9c57bc0efce0ac37d8319666f5a8d3e8dce7651c Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 11 Nov 2015 10:21:53 -0800
Subject: [PATCH 190/518] [SPARK-11656][SQL] support typed aggregate in project
 list

insert `aEncoder` like we do in `agg`

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9630 from cloud-fan/select.
---
 .../scala/org/apache/spark/sql/Dataset.scala  | 20 +++++++++++++++----
 .../spark/sql/DatasetAggregatorSuite.scala    | 11 ++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index a7e5ab19bf846..87dae6b331593 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -21,14 +21,15 @@ import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.catalyst.analysis.UnresolvedAlias
 import org.apache.spark.api.java.function._
 
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAlias
 import org.apache.spark.sql.catalyst.plans.Inner
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.{Queryable, QueryExecution}
+import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
 import org.apache.spark.sql.types.StructType
 
 /**
@@ -359,7 +360,7 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def select[U1: Encoder](c1: TypedColumn[T, U1]): Dataset[U1] = {
-    new Dataset[U1](sqlContext, Project(Alias(c1.expr, "_1")() :: Nil, logicalPlan))
+    new Dataset[U1](sqlContext, Project(Alias(withEncoder(c1).expr, "_1")() :: Nil, logicalPlan))
   }
 
   /**
@@ -368,11 +369,12 @@ class Dataset[T] private[sql](
    * that cast appropriately for the user facing interface.
    */
   protected def selectUntyped(columns: TypedColumn[_, _]*): Dataset[_] = {
-    val aliases = columns.zipWithIndex.map { case (c, i) => Alias(c.expr, s"_${i + 1}")() }
+    val withEncoders = columns.map(withEncoder)
+    val aliases = withEncoders.zipWithIndex.map { case (c, i) => Alias(c.expr, s"_${i + 1}")() }
     val unresolvedPlan = Project(aliases, logicalPlan)
     val execution = new QueryExecution(sqlContext, unresolvedPlan)
     // Rebind the encoders to the nested schema that will be produced by the select.
-    val encoders = columns.map(_.encoder.asInstanceOf[ExpressionEncoder[_]]).zip(aliases).map {
+    val encoders = withEncoders.map(_.encoder.asInstanceOf[ExpressionEncoder[_]]).zip(aliases).map {
       case (e: ExpressionEncoder[_], a) if !e.flat =>
         e.nested(a.toAttribute).resolve(execution.analyzed.output)
       case (e, a) =>
@@ -381,6 +383,16 @@ class Dataset[T] private[sql](
     new Dataset(sqlContext, execution, ExpressionEncoder.tuple(encoders))
   }
 
+  private def withEncoder(c: TypedColumn[_, _]): TypedColumn[_, _] = {
+    val e = c.expr transform {
+      case ta: TypedAggregateExpression if ta.aEncoder.isEmpty =>
+        ta.copy(
+          aEncoder = Some(encoder.asInstanceOf[ExpressionEncoder[Any]]),
+          children = queryExecution.analyzed.output)
+    }
+    new TypedColumn(e, c.encoder)
+  }
+
   /**
    * Returns a new [[Dataset]] by computing the given [[Column]] expressions for each element.
    * @since 1.6.0
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
index 002d5c18f0a25..d4f0ab76cfbf2 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -114,4 +114,15 @@ class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
         ComplexResultAgg.toColumn),
       ("a", 2.0, (2L, 4L)), ("b", 3.0, (1L, 3L)))
   }
+
+  test("typed aggregation: in project list") {
+    val ds = Seq(1, 3, 2, 5).toDS()
+
+    checkAnswer(
+      ds.select(sum((i: Int) => i)),
+      11)
+    checkAnswer(
+      ds.select(sum((i: Int) => i), sum((i: Int) => i * 2)),
+      11 -> 22)
+  }
 }

From ec2b807212e568c9e98cd80746bcb61e02c7a98e Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 11 Nov 2015 10:52:23 -0800
Subject: [PATCH 191/518] [SPARK-11564][SQL][FOLLOW-UP] clean up java tuple
 encoder

We need to support custom classes like java beans and combine them into tuple, and it's very hard to do it with the  TypeTag-based approach.
We should keep only the compose-based way to create tuple encoder.

This PR also move `Encoder` to `org.apache.spark.sql`

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9567 from cloud-fan/java.
---
 .../sql/{catalyst/encoders => }/Encoder.scala | 65 ++--------------
 .../catalyst/encoders/ExpressionEncoder.scala | 10 +--
 .../spark/sql/catalyst/encoders/package.scala |  3 +-
 .../plans/logical/basicOperators.scala        |  1 +
 .../scala/org/apache/spark/sql/Column.scala   |  2 +-
 .../org/apache/spark/sql/DataFrame.scala      |  2 -
 .../org/apache/spark/sql/GroupedDataset.scala |  2 +-
 .../org/apache/spark/sql/SQLContext.scala     |  2 +-
 .../aggregate/TypedAggregateExpression.scala  |  3 +-
 .../spark/sql/expressions/Aggregator.scala    |  3 +-
 .../org/apache/spark/sql/functions.scala      |  2 +-
 .../apache/spark/sql/JavaDatasetSuite.java    | 78 ++++++++++---------
 .../spark/sql/DatasetAggregatorSuite.scala    |  4 +-
 .../org/apache/spark/sql/QueryTest.scala      |  1 -
 14 files changed, 65 insertions(+), 113 deletions(-)
 rename sql/catalyst/src/main/scala/org/apache/spark/sql/{catalyst/encoders => }/Encoder.scala (71%)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
similarity index 71%
rename from sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala
rename to sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index 6569b900fed90..1ff7340557e6c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -15,13 +15,14 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.catalyst.encoders
+package org.apache.spark.sql
 
-import scala.reflect.ClassTag
-
-import org.apache.spark.util.Utils
-import org.apache.spark.sql.types.{ObjectType, StructField, StructType}
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types.{ObjectType, StructField, StructType}
+import org.apache.spark.util.Utils
+
+import scala.reflect.ClassTag
 
 /**
  * Used to convert a JVM object of type `T` to and from the internal Spark SQL representation.
@@ -38,9 +39,7 @@ trait Encoder[T] extends Serializable {
   def clsTag: ClassTag[T]
 }
 
-object Encoder {
-  import scala.reflect.runtime.universe._
-
+object Encoders {
   def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder(flat = true)
   def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder(flat = true)
   def SHORT: Encoder[java.lang.Short] = ExpressionEncoder(flat = true)
@@ -129,54 +128,4 @@ object Encoder {
       constructExpression,
       ClassTag.apply(cls))
   }
-
-  def typeTagOfTuple2[T1 : TypeTag, T2 : TypeTag]: TypeTag[(T1, T2)] = typeTag[(T1, T2)]
-
-  private def getTypeTag[T](c: Class[T]): TypeTag[T] = {
-    import scala.reflect.api
-
-    // val mirror = runtimeMirror(c.getClassLoader)
-    val mirror = rootMirror
-    val sym = mirror.staticClass(c.getName)
-    val tpe = sym.selfType
-    TypeTag(mirror, new api.TypeCreator {
-      def apply[U <: api.Universe with Singleton](m: api.Mirror[U]) =
-        if (m eq mirror) tpe.asInstanceOf[U # Type]
-        else throw new IllegalArgumentException(
-          s"Type tag defined in $mirror cannot be migrated to other mirrors.")
-    })
-  }
-
-  def forTuple[T1, T2](c1: Class[T1], c2: Class[T2]): Encoder[(T1, T2)] = {
-    implicit val typeTag1 = getTypeTag(c1)
-    implicit val typeTag2 = getTypeTag(c2)
-    ExpressionEncoder[(T1, T2)]()
-  }
-
-  def forTuple[T1, T2, T3](c1: Class[T1], c2: Class[T2], c3: Class[T3]): Encoder[(T1, T2, T3)] = {
-    implicit val typeTag1 = getTypeTag(c1)
-    implicit val typeTag2 = getTypeTag(c2)
-    implicit val typeTag3 = getTypeTag(c3)
-    ExpressionEncoder[(T1, T2, T3)]()
-  }
-
-  def forTuple[T1, T2, T3, T4](
-      c1: Class[T1], c2: Class[T2], c3: Class[T3], c4: Class[T4]): Encoder[(T1, T2, T3, T4)] = {
-    implicit val typeTag1 = getTypeTag(c1)
-    implicit val typeTag2 = getTypeTag(c2)
-    implicit val typeTag3 = getTypeTag(c3)
-    implicit val typeTag4 = getTypeTag(c4)
-    ExpressionEncoder[(T1, T2, T3, T4)]()
-  }
-
-  def forTuple[T1, T2, T3, T4, T5](
-      c1: Class[T1], c2: Class[T2], c3: Class[T3], c4: Class[T4], c5: Class[T5])
-    : Encoder[(T1, T2, T3, T4, T5)] = {
-    implicit val typeTag1 = getTypeTag(c1)
-    implicit val typeTag2 = getTypeTag(c2)
-    implicit val typeTag3 = getTypeTag(c3)
-    implicit val typeTag4 = getTypeTag(c4)
-    implicit val typeTag5 = getTypeTag(c5)
-    ExpressionEncoder[(T1, T2, T3, T4, T5)]()
-  }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index 005c0627f56b8..294afde5347e2 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -17,18 +17,18 @@
 
 package org.apache.spark.sql.catalyst.encoders
 
-import org.apache.spark.sql.catalyst.analysis.{SimpleAnalyzer, UnresolvedExtractValue, UnresolvedAttribute}
-import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project}
-import org.apache.spark.util.Utils
-
 import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.{typeTag, TypeTag}
 
+import org.apache.spark.util.Utils
+import org.apache.spark.sql.Encoder
+import org.apache.spark.sql.catalyst.analysis.{SimpleAnalyzer, UnresolvedExtractValue, UnresolvedAttribute}
+import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project}
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateSafeProjection, GenerateUnsafeProjection}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.ScalaReflection
-import org.apache.spark.sql.types.{StructField, DataType, ObjectType, StructType}
+import org.apache.spark.sql.types.{StructField, ObjectType, StructType}
 
 /**
  * A factory for constructing encoders that convert objects and primitves to and from the
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/package.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/package.scala
index d4642a500672e..2c35adca9c925 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/package.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/package.scala
@@ -17,10 +17,11 @@
 
 package org.apache.spark.sql.catalyst
 
+import org.apache.spark.sql.Encoder
+
 package object encoders {
   private[sql] def encoderFor[A : Encoder]: ExpressionEncoder[A] = implicitly[Encoder[A]] match {
     case e: ExpressionEncoder[A] => e
     case _ => sys.error(s"Only expression encoders are supported today")
   }
 }
-
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index 764f8aaebddf1..597f03e752707 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.catalyst.plans.logical
 
+import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index d26b6c3579205..f0f275e91f1a3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -23,7 +23,7 @@ import org.apache.spark.annotation.Experimental
 import org.apache.spark.Logging
 import org.apache.spark.sql.functions.lit
 import org.apache.spark.sql.catalyst.analysis._
-import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
+import org.apache.spark.sql.catalyst.encoders.encoderFor
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.util.DataTypeParser
 import org.apache.spark.sql.types._
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 691b476fff8d6..a492099b9392b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -23,7 +23,6 @@ import java.util.Properties
 import scala.language.implicitConversions
 import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.TypeTag
-import scala.util.control.NonFatal
 
 import com.fasterxml.jackson.core.JsonFactory
 import org.apache.commons.lang3.StringUtils
@@ -35,7 +34,6 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
-import org.apache.spark.sql.catalyst.encoders.Encoder
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.{Inner, JoinType}
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, ScalaReflection, SqlParser}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index db61499229284..61e2a9545069b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -24,7 +24,7 @@ import scala.collection.JavaConverters._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.function.{Function2 => JFunction2, Function3 => JFunction3, _}
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttribute}
-import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor, Encoder}
+import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor}
 import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression, Alias, Attribute}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index 1cf1e30f967cd..cd1fdc4edb39d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -33,7 +33,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.scheduler.{SparkListener, SparkListenerApplicationEnd}
 import org.apache.spark.sql.SQLConf.SQLConfEntry
 import org.apache.spark.sql.catalyst.analysis._
-import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
+import org.apache.spark.sql.catalyst.encoders.encoderFor
 import org.apache.spark.sql.catalyst.errors.DialectException
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.optimizer.{DefaultOptimizer, Optimizer}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
index b5a87c56e681b..dfcbac8687b3e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
@@ -20,9 +20,10 @@ package org.apache.spark.sql.execution.aggregate
 import scala.language.existentials
 
 import org.apache.spark.Logging
+import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.expressions.Aggregator
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
+import org.apache.spark.sql.catalyst.encoders.encoderFor
 import org.apache.spark.sql.catalyst.expressions.aggregate.ImperativeAggregate
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
index 2aa5a7d540316..360c9a5bc15e7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
@@ -17,7 +17,8 @@
 
 package org.apache.spark.sql.expressions
 
-import org.apache.spark.sql.catalyst.encoders.{encoderFor, Encoder}
+import org.apache.spark.sql.Encoder
+import org.apache.spark.sql.catalyst.encoders.encoderFor
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete}
 import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
 import org.apache.spark.sql.{Dataset, DataFrame, TypedColumn}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index a59d738010f7c..ab49ed4b5d4ec 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -26,7 +26,7 @@ import scala.util.Try
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.sql.catalyst.{SqlParser, ScalaReflection}
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunction, Star}
-import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, Encoder}
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.BroadcastHint
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index 2da63d1b96706..33d8388f615ae 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -30,8 +30,8 @@
 import org.apache.spark.SparkContext;
 import org.apache.spark.api.java.function.*;
 import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.sql.catalyst.encoders.Encoder;
-import org.apache.spark.sql.catalyst.encoders.Encoder$;
+import org.apache.spark.sql.Encoder;
+import org.apache.spark.sql.Encoders;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.GroupedDataset;
 import org.apache.spark.sql.test.TestSQLContext;
@@ -41,7 +41,6 @@
 public class JavaDatasetSuite implements Serializable {
   private transient JavaSparkContext jsc;
   private transient TestSQLContext context;
-  private transient Encoder$ e = Encoder$.MODULE$;
 
   @Before
   public void setUp() {
@@ -66,7 +65,7 @@ private <T1, T2> Tuple2<T1, T2> tuple2(T1 t1, T2 t2) {
   @Test
   public void testCollect() {
     List<String> data = Arrays.asList("hello", "world");
-    Dataset<String> ds = context.createDataset(data, e.STRING());
+    Dataset<String> ds = context.createDataset(data, Encoders.STRING());
     List<String> collected = ds.collectAsList();
     Assert.assertEquals(Arrays.asList("hello", "world"), collected);
   }
@@ -74,7 +73,7 @@ public void testCollect() {
   @Test
   public void testTake() {
     List<String> data = Arrays.asList("hello", "world");
-    Dataset<String> ds = context.createDataset(data, e.STRING());
+    Dataset<String> ds = context.createDataset(data, Encoders.STRING());
     List<String> collected = ds.takeAsList(1);
     Assert.assertEquals(Arrays.asList("hello"), collected);
   }
@@ -82,7 +81,7 @@ public void testTake() {
   @Test
   public void testCommonOperation() {
     List<String> data = Arrays.asList("hello", "world");
-    Dataset<String> ds = context.createDataset(data, e.STRING());
+    Dataset<String> ds = context.createDataset(data, Encoders.STRING());
     Assert.assertEquals("hello", ds.first());
 
     Dataset<String> filtered = ds.filter(new FilterFunction<String>() {
@@ -99,7 +98,7 @@ public boolean call(String v) throws Exception {
       public Integer call(String v) throws Exception {
         return v.length();
       }
-    }, e.INT());
+    }, Encoders.INT());
     Assert.assertEquals(Arrays.asList(5, 5), mapped.collectAsList());
 
     Dataset<String> parMapped = ds.mapPartitions(new MapPartitionsFunction<String, String>() {
@@ -111,7 +110,7 @@ public Iterable<String> call(Iterator<String> it) throws Exception {
         }
         return ls;
       }
-    }, e.STRING());
+    }, Encoders.STRING());
     Assert.assertEquals(Arrays.asList("HELLO", "WORLD"), parMapped.collectAsList());
 
     Dataset<String> flatMapped = ds.flatMap(new FlatMapFunction<String, String>() {
@@ -123,7 +122,7 @@ public Iterable<String> call(String s) throws Exception {
         }
         return ls;
       }
-    }, e.STRING());
+    }, Encoders.STRING());
     Assert.assertEquals(
       Arrays.asList("h", "e", "l", "l", "o", "w", "o", "r", "l", "d"),
       flatMapped.collectAsList());
@@ -133,7 +132,7 @@ public Iterable<String> call(String s) throws Exception {
   public void testForeach() {
     final Accumulator<Integer> accum = jsc.accumulator(0);
     List<String> data = Arrays.asList("a", "b", "c");
-    Dataset<String> ds = context.createDataset(data, e.STRING());
+    Dataset<String> ds = context.createDataset(data, Encoders.STRING());
 
     ds.foreach(new ForeachFunction<String>() {
       @Override
@@ -147,7 +146,7 @@ public void call(String s) throws Exception {
   @Test
   public void testReduce() {
     List<Integer> data = Arrays.asList(1, 2, 3);
-    Dataset<Integer> ds = context.createDataset(data, e.INT());
+    Dataset<Integer> ds = context.createDataset(data, Encoders.INT());
 
     int reduced = ds.reduce(new ReduceFunction<Integer>() {
       @Override
@@ -161,13 +160,13 @@ public Integer call(Integer v1, Integer v2) throws Exception {
   @Test
   public void testGroupBy() {
     List<String> data = Arrays.asList("a", "foo", "bar");
-    Dataset<String> ds = context.createDataset(data, e.STRING());
+    Dataset<String> ds = context.createDataset(data, Encoders.STRING());
     GroupedDataset<Integer, String> grouped = ds.groupBy(new MapFunction<String, Integer>() {
       @Override
       public Integer call(String v) throws Exception {
         return v.length();
       }
-    }, e.INT());
+    }, Encoders.INT());
 
     Dataset<String> mapped = grouped.map(new MapGroupFunction<Integer, String, String>() {
       @Override
@@ -178,7 +177,7 @@ public String call(Integer key, Iterator<String> values) throws Exception {
         }
         return sb.toString();
       }
-    }, e.STRING());
+    }, Encoders.STRING());
 
     Assert.assertEquals(Arrays.asList("1a", "3foobar"), mapped.collectAsList());
 
@@ -193,27 +192,27 @@ public Iterable<String> call(Integer key, Iterator<String> values) throws Except
           return Collections.singletonList(sb.toString());
         }
       },
-      e.STRING());
+      Encoders.STRING());
 
     Assert.assertEquals(Arrays.asList("1a", "3foobar"), flatMapped.collectAsList());
 
     List<Integer> data2 = Arrays.asList(2, 6, 10);
-    Dataset<Integer> ds2 = context.createDataset(data2, e.INT());
+    Dataset<Integer> ds2 = context.createDataset(data2, Encoders.INT());
     GroupedDataset<Integer, Integer> grouped2 = ds2.groupBy(new MapFunction<Integer, Integer>() {
       @Override
       public Integer call(Integer v) throws Exception {
         return v / 2;
       }
-    }, e.INT());
+    }, Encoders.INT());
 
     Dataset<String> cogrouped = grouped.cogroup(
       grouped2,
       new CoGroupFunction<Integer, String, Integer, String>() {
         @Override
         public Iterable<String> call(
-            Integer key,
-            Iterator<String> left,
-            Iterator<Integer> right) throws Exception {
+          Integer key,
+          Iterator<String> left,
+          Iterator<Integer> right) throws Exception {
           StringBuilder sb = new StringBuilder(key.toString());
           while (left.hasNext()) {
             sb.append(left.next());
@@ -225,7 +224,7 @@ public Iterable<String> call(
           return Collections.singletonList(sb.toString());
         }
       },
-      e.STRING());
+      Encoders.STRING());
 
     Assert.assertEquals(Arrays.asList("1a#2", "3foobar#6", "5#10"), cogrouped.collectAsList());
   }
@@ -233,8 +232,9 @@ public Iterable<String> call(
   @Test
   public void testGroupByColumn() {
     List<String> data = Arrays.asList("a", "foo", "bar");
-    Dataset<String> ds = context.createDataset(data, e.STRING());
-    GroupedDataset<Integer, String> grouped = ds.groupBy(length(col("value"))).asKey(e.INT());
+    Dataset<String> ds = context.createDataset(data, Encoders.STRING());
+    GroupedDataset<Integer, String> grouped =
+      ds.groupBy(length(col("value"))).asKey(Encoders.INT());
 
     Dataset<String> mapped = grouped.map(
       new MapGroupFunction<Integer, String, String>() {
@@ -247,7 +247,7 @@ public String call(Integer key, Iterator<String> data) throws Exception {
           return sb.toString();
         }
       },
-      e.STRING());
+      Encoders.STRING());
 
     Assert.assertEquals(Arrays.asList("1a", "3foobar"), mapped.collectAsList());
   }
@@ -255,11 +255,11 @@ public String call(Integer key, Iterator<String> data) throws Exception {
   @Test
   public void testSelect() {
     List<Integer> data = Arrays.asList(2, 6);
-    Dataset<Integer> ds = context.createDataset(data, e.INT());
+    Dataset<Integer> ds = context.createDataset(data, Encoders.INT());
 
     Dataset<Tuple2<Integer, String>> selected = ds.select(
       expr("value + 1"),
-      col("value").cast("string")).as(e.tuple(e.INT(), e.STRING()));
+      col("value").cast("string")).as(Encoders.tuple(Encoders.INT(), Encoders.STRING()));
 
     Assert.assertEquals(
       Arrays.asList(tuple2(3, "2"), tuple2(7, "6")),
@@ -269,14 +269,14 @@ public void testSelect() {
   @Test
   public void testSetOperation() {
     List<String> data = Arrays.asList("abc", "abc", "xyz");
-    Dataset<String> ds = context.createDataset(data, e.STRING());
+    Dataset<String> ds = context.createDataset(data, Encoders.STRING());
 
     Assert.assertEquals(
       Arrays.asList("abc", "xyz"),
       sort(ds.distinct().collectAsList().toArray(new String[0])));
 
     List<String> data2 = Arrays.asList("xyz", "foo", "foo");
-    Dataset<String> ds2 = context.createDataset(data2, e.STRING());
+    Dataset<String> ds2 = context.createDataset(data2, Encoders.STRING());
 
     Dataset<String> intersected = ds.intersect(ds2);
     Assert.assertEquals(Arrays.asList("xyz"), intersected.collectAsList());
@@ -298,9 +298,9 @@ private <T extends Comparable<T>> List<T> sort(T[] data) {
   @Test
   public void testJoin() {
     List<Integer> data = Arrays.asList(1, 2, 3);
-    Dataset<Integer> ds = context.createDataset(data, e.INT()).as("a");
+    Dataset<Integer> ds = context.createDataset(data, Encoders.INT()).as("a");
     List<Integer> data2 = Arrays.asList(2, 3, 4);
-    Dataset<Integer> ds2 = context.createDataset(data2, e.INT()).as("b");
+    Dataset<Integer> ds2 = context.createDataset(data2, Encoders.INT()).as("b");
 
     Dataset<Tuple2<Integer, Integer>> joined =
       ds.joinWith(ds2, col("a.value").equalTo(col("b.value")));
@@ -311,26 +311,28 @@ public void testJoin() {
 
   @Test
   public void testTupleEncoder() {
-    Encoder<Tuple2<Integer, String>> encoder2 = e.tuple(e.INT(), e.STRING());
+    Encoder<Tuple2<Integer, String>> encoder2 = Encoders.tuple(Encoders.INT(), Encoders.STRING());
     List<Tuple2<Integer, String>> data2 = Arrays.asList(tuple2(1, "a"), tuple2(2, "b"));
     Dataset<Tuple2<Integer, String>> ds2 = context.createDataset(data2, encoder2);
     Assert.assertEquals(data2, ds2.collectAsList());
 
-    Encoder<Tuple3<Integer, Long, String>> encoder3 = e.tuple(e.INT(), e.LONG(), e.STRING());
+    Encoder<Tuple3<Integer, Long, String>> encoder3 =
+      Encoders.tuple(Encoders.INT(), Encoders.LONG(), Encoders.STRING());
     List<Tuple3<Integer, Long, String>> data3 =
       Arrays.asList(new Tuple3<Integer, Long, String>(1, 2L, "a"));
     Dataset<Tuple3<Integer, Long, String>> ds3 = context.createDataset(data3, encoder3);
     Assert.assertEquals(data3, ds3.collectAsList());
 
     Encoder<Tuple4<Integer, String, Long, String>> encoder4 =
-      e.tuple(e.INT(), e.STRING(), e.LONG(), e.STRING());
+      Encoders.tuple(Encoders.INT(), Encoders.STRING(), Encoders.LONG(), Encoders.STRING());
     List<Tuple4<Integer, String, Long, String>> data4 =
       Arrays.asList(new Tuple4<Integer, String, Long, String>(1, "b", 2L, "a"));
     Dataset<Tuple4<Integer, String, Long, String>> ds4 = context.createDataset(data4, encoder4);
     Assert.assertEquals(data4, ds4.collectAsList());
 
     Encoder<Tuple5<Integer, String, Long, String, Boolean>> encoder5 =
-      e.tuple(e.INT(), e.STRING(), e.LONG(), e.STRING(), e.BOOLEAN());
+      Encoders.tuple(Encoders.INT(), Encoders.STRING(), Encoders.LONG(), Encoders.STRING(),
+        Encoders.BOOLEAN());
     List<Tuple5<Integer, String, Long, String, Boolean>> data5 =
       Arrays.asList(new Tuple5<Integer, String, Long, String, Boolean>(1, "b", 2L, "a", true));
     Dataset<Tuple5<Integer, String, Long, String, Boolean>> ds5 =
@@ -342,7 +344,7 @@ public void testTupleEncoder() {
   public void testNestedTupleEncoder() {
     // test ((int, string), string)
     Encoder<Tuple2<Tuple2<Integer, String>, String>> encoder =
-      e.tuple(e.tuple(e.INT(), e.STRING()), e.STRING());
+      Encoders.tuple(Encoders.tuple(Encoders.INT(), Encoders.STRING()), Encoders.STRING());
     List<Tuple2<Tuple2<Integer, String>, String>> data =
       Arrays.asList(tuple2(tuple2(1, "a"), "a"), tuple2(tuple2(2, "b"), "b"));
     Dataset<Tuple2<Tuple2<Integer, String>, String>> ds = context.createDataset(data, encoder);
@@ -350,7 +352,8 @@ public void testNestedTupleEncoder() {
 
     // test (int, (string, string, long))
     Encoder<Tuple2<Integer, Tuple3<String, String, Long>>> encoder2 =
-      e.tuple(e.INT(), e.tuple(e.STRING(), e.STRING(), e.LONG()));
+      Encoders.tuple(Encoders.INT(),
+        Encoders.tuple(Encoders.STRING(), Encoders.STRING(), Encoders.LONG()));
     List<Tuple2<Integer, Tuple3<String, String, Long>>> data2 =
       Arrays.asList(tuple2(1, new Tuple3<String, String, Long>("a", "b", 3L)));
     Dataset<Tuple2<Integer, Tuple3<String, String, Long>>> ds2 =
@@ -359,7 +362,8 @@ public void testNestedTupleEncoder() {
 
     // test (int, ((string, long), string))
     Encoder<Tuple2<Integer, Tuple2<Tuple2<String, Long>, String>>> encoder3 =
-      e.tuple(e.INT(), e.tuple(e.tuple(e.STRING(), e.LONG()), e.STRING()));
+      Encoders.tuple(Encoders.INT(),
+        Encoders.tuple(Encoders.tuple(Encoders.STRING(), Encoders.LONG()), Encoders.STRING()));
     List<Tuple2<Integer, Tuple2<Tuple2<String, Long>, String>>> data3 =
       Arrays.asList(tuple2(1, tuple2(tuple2("a", 2L), "b")));
     Dataset<Tuple2<Integer, Tuple2<Tuple2<String, Long>, String>>> ds3 =
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
index d4f0ab76cfbf2..378cd365276b3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -17,13 +17,11 @@
 
 package org.apache.spark.sql
 
-import org.apache.spark.sql.catalyst.encoders.Encoder
-import org.apache.spark.sql.functions._
 
 import scala.language.postfixOps
 
 import org.apache.spark.sql.test.SharedSQLContext
-
+import org.apache.spark.sql.functions._
 import org.apache.spark.sql.expressions.Aggregator
 
 /** An `Aggregator` that adds up any numeric type returned by the given function. */
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
index 3c174efe73ffe..7a8b7ae5bf265 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
@@ -24,7 +24,6 @@ import scala.collection.JavaConverters._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.columnar.InMemoryRelation
-import org.apache.spark.sql.catalyst.encoders.Encoder
 
 abstract class QueryTest extends PlanTest {
 

From e71ba56586ba64da18f412bc0e0263777c46ac4a Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 11 Nov 2015 11:04:04 -0800
Subject: [PATCH 192/518] [SQL][MINOR] remove newLongEncoder in functions

it may shadows the one from implicits in some case.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9629 from cloud-fan/minor.
---
 .../src/main/scala/org/apache/spark/sql/functions.scala     | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index ab49ed4b5d4ec..b6330e230afef 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -83,9 +83,6 @@ object functions extends LegacyFunctions {
     Column(func.toAggregateExpression(isDistinct))
   }
 
-  private implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder[Long](flat = true)
-
-
   /**
    * Returns a [[Column]] based on the given column name.
    *
@@ -269,7 +266,8 @@ object functions extends LegacyFunctions {
    * @group agg_funcs
    * @since 1.3.0
    */
-  def count(columnName: String): TypedColumn[Any, Long] = count(Column(columnName)).as[Long]
+  def count(columnName: String): TypedColumn[Any, Long] =
+    count(Column(columnName)).as(ExpressionEncoder[Long](flat = true))
 
   /**
    * Aggregate function: returns the number of distinct items in a group.

From 529a1d3380c4c23fed068ad05a6376162c4b76d6 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Wed, 11 Nov 2015 11:16:39 -0800
Subject: [PATCH 193/518] [SPARK-6152] Use shaded ASM5 to support closure
 cleaning of Java 8 compiled classes

This patch modifies Spark's closure cleaner (and a few other places) to use ASM 5, which is necessary in order to support cleaning of closures that were compiled by Java 8.

In order to avoid ASM dependency conflicts, Spark excludes ASM from all of its dependencies and uses a shaded version of ASM 4 that comes from `reflectasm` (see [SPARK-782](https://issues.apache.org/jira/browse/SPARK-782) and #232). This patch updates Spark to use a shaded version of ASM 5.0.4 that was published by the Apache XBean project; the POM used to create the shaded artifact can be found at https://github.com/apache/geronimo-xbean/blob/xbean-4.4/xbean-asm5-shaded/pom.xml.

http://movingfulcrum.tumblr.com/post/80826553604/asm-framework-50-the-missing-migration-guide was a useful resource while upgrading the code to use the new ASM5 opcodes.

I also added a new regression tests in the `java8-tests` subproject; the existing tests were insufficient to catch this bug, which only affected Scala 2.11 user code which was compiled targeting Java 8.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9512 from JoshRosen/SPARK-6152.
---
 core/pom.xml                                  |  4 ++
 .../apache/spark/util/ClosureCleaner.scala    | 25 +++++-----
 docs/building-spark.md                        |  4 ++
 .../org/apache/spark/JDK8ScalaSuite.scala     | 27 ++++++++++
 graphx/pom.xml                                |  4 ++
 .../spark/graphx/util/BytecodeUtils.scala     | 16 +++---
 pom.xml                                       |  8 +++
 project/SparkBuild.scala                      | 26 +++++++++-
 repl/pom.xml                                  |  4 ++
 .../spark/repl/ExecutorClassLoader.scala      |  9 ++--
 sql/core/pom.xml                              |  5 ++
 .../execution/metric/SQLMetricsSuite.scala    | 50 ++++++++-----------
 12 files changed, 125 insertions(+), 57 deletions(-)
 create mode 100644 extras/java8-tests/src/test/scala/org/apache/spark/JDK8ScalaSuite.scala

diff --git a/core/pom.xml b/core/pom.xml
index 570a25cf325a2..7e1205a076f26 100644
--- a/core/pom.xml
+++ b/core/pom.xml
@@ -51,6 +51,10 @@
       <groupId>com.twitter</groupId>
       <artifactId>chill-java</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.xbean</groupId>
+      <artifactId>xbean-asm5-shaded</artifactId>
+    </dependency>
     <dependency>
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-client</artifactId>
diff --git a/core/src/main/scala/org/apache/spark/util/ClosureCleaner.scala b/core/src/main/scala/org/apache/spark/util/ClosureCleaner.scala
index 1b49dca9dc78b..e27d2e6c94f7b 100644
--- a/core/src/main/scala/org/apache/spark/util/ClosureCleaner.scala
+++ b/core/src/main/scala/org/apache/spark/util/ClosureCleaner.scala
@@ -21,8 +21,8 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
 
 import scala.collection.mutable.{Map, Set}
 
-import com.esotericsoftware.reflectasm.shaded.org.objectweb.asm.{ClassReader, ClassVisitor, MethodVisitor, Type}
-import com.esotericsoftware.reflectasm.shaded.org.objectweb.asm.Opcodes._
+import org.apache.xbean.asm5.{ClassReader, ClassVisitor, MethodVisitor, Type}
+import org.apache.xbean.asm5.Opcodes._
 
 import org.apache.spark.{Logging, SparkEnv, SparkException}
 
@@ -325,11 +325,11 @@ private[spark] object ClosureCleaner extends Logging {
 private[spark] class ReturnStatementInClosureException
   extends SparkException("Return statements aren't allowed in Spark closures")
 
-private class ReturnStatementFinder extends ClassVisitor(ASM4) {
+private class ReturnStatementFinder extends ClassVisitor(ASM5) {
   override def visitMethod(access: Int, name: String, desc: String,
       sig: String, exceptions: Array[String]): MethodVisitor = {
     if (name.contains("apply")) {
-      new MethodVisitor(ASM4) {
+      new MethodVisitor(ASM5) {
         override def visitTypeInsn(op: Int, tp: String) {
           if (op == NEW && tp.contains("scala/runtime/NonLocalReturnControl")) {
             throw new ReturnStatementInClosureException
@@ -337,7 +337,7 @@ private class ReturnStatementFinder extends ClassVisitor(ASM4) {
         }
       }
     } else {
-      new MethodVisitor(ASM4) {}
+      new MethodVisitor(ASM5) {}
     }
   }
 }
@@ -361,7 +361,7 @@ private[util] class FieldAccessFinder(
     findTransitively: Boolean,
     specificMethod: Option[MethodIdentifier[_]] = None,
     visitedMethods: Set[MethodIdentifier[_]] = Set.empty)
-  extends ClassVisitor(ASM4) {
+  extends ClassVisitor(ASM5) {
 
   override def visitMethod(
       access: Int,
@@ -376,7 +376,7 @@ private[util] class FieldAccessFinder(
       return null
     }
 
-    new MethodVisitor(ASM4) {
+    new MethodVisitor(ASM5) {
       override def visitFieldInsn(op: Int, owner: String, name: String, desc: String) {
         if (op == GETFIELD) {
           for (cl <- fields.keys if cl.getName == owner.replace('/', '.')) {
@@ -385,7 +385,8 @@ private[util] class FieldAccessFinder(
         }
       }
 
-      override def visitMethodInsn(op: Int, owner: String, name: String, desc: String) {
+      override def visitMethodInsn(
+          op: Int, owner: String, name: String, desc: String, itf: Boolean) {
         for (cl <- fields.keys if cl.getName == owner.replace('/', '.')) {
           // Check for calls a getter method for a variable in an interpreter wrapper object.
           // This means that the corresponding field will be accessed, so we should save it.
@@ -408,7 +409,7 @@ private[util] class FieldAccessFinder(
   }
 }
 
-private class InnerClosureFinder(output: Set[Class[_]]) extends ClassVisitor(ASM4) {
+private class InnerClosureFinder(output: Set[Class[_]]) extends ClassVisitor(ASM5) {
   var myName: String = null
 
   // TODO: Recursively find inner closures that we indirectly reference, e.g.
@@ -423,9 +424,9 @@ private class InnerClosureFinder(output: Set[Class[_]]) extends ClassVisitor(ASM
 
   override def visitMethod(access: Int, name: String, desc: String,
       sig: String, exceptions: Array[String]): MethodVisitor = {
-    new MethodVisitor(ASM4) {
-      override def visitMethodInsn(op: Int, owner: String, name: String,
-          desc: String) {
+    new MethodVisitor(ASM5) {
+      override def visitMethodInsn(
+          op: Int, owner: String, name: String, desc: String, itf: Boolean) {
         val argTypes = Type.getArgumentTypes(desc)
         if (op == INVOKESPECIAL && name == "<init>" && argTypes.length > 0
             && argTypes(0).toString.startsWith("L") // is it an object?
diff --git a/docs/building-spark.md b/docs/building-spark.md
index 4f73adb85446c..3d38edbdad4bc 100644
--- a/docs/building-spark.md
+++ b/docs/building-spark.md
@@ -190,6 +190,10 @@ Running only Java 8 tests and nothing else.
 
     mvn install -DskipTests -Pjava8-tests
 
+or
+
+    sbt -Pjava8-tests java8-tests/test
+
 Java 8 tests are run when `-Pjava8-tests` profile is enabled, they will run in spite of `-DskipTests`.
 For these tests to run your system must have a JDK 8 installation.
 If you have JDK 8 installed but it is not the system default, you can set JAVA_HOME to point to JDK 8 before running the tests.
diff --git a/extras/java8-tests/src/test/scala/org/apache/spark/JDK8ScalaSuite.scala b/extras/java8-tests/src/test/scala/org/apache/spark/JDK8ScalaSuite.scala
new file mode 100644
index 0000000000000..fa0681db41088
--- /dev/null
+++ b/extras/java8-tests/src/test/scala/org/apache/spark/JDK8ScalaSuite.scala
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark
+
+/**
+ * Test cases where JDK8-compiled Scala user code is used with Spark.
+ */
+class JDK8ScalaSuite extends SparkFunSuite with SharedSparkContext {
+  test("basic RDD closure test (SPARK-6152)") {
+    sc.parallelize(1 to 1000).map(x => x * x).count()
+  }
+}
diff --git a/graphx/pom.xml b/graphx/pom.xml
index 987b831021a54..8cd66c5b2e826 100644
--- a/graphx/pom.xml
+++ b/graphx/pom.xml
@@ -47,6 +47,10 @@
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.xbean</groupId>
+      <artifactId>xbean-asm5-shaded</artifactId>
+    </dependency>
     <dependency>
       <groupId>com.google.guava</groupId>
       <artifactId>guava</artifactId>
diff --git a/graphx/src/main/scala/org/apache/spark/graphx/util/BytecodeUtils.scala b/graphx/src/main/scala/org/apache/spark/graphx/util/BytecodeUtils.scala
index 74a7de18d4161..a6d0cb6409664 100644
--- a/graphx/src/main/scala/org/apache/spark/graphx/util/BytecodeUtils.scala
+++ b/graphx/src/main/scala/org/apache/spark/graphx/util/BytecodeUtils.scala
@@ -22,11 +22,10 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
 import scala.collection.mutable.HashSet
 import scala.language.existentials
 
-import org.apache.spark.util.Utils
-
-import com.esotericsoftware.reflectasm.shaded.org.objectweb.asm.{ClassReader, ClassVisitor, MethodVisitor}
-import com.esotericsoftware.reflectasm.shaded.org.objectweb.asm.Opcodes._
+import org.apache.xbean.asm5.{ClassReader, ClassVisitor, MethodVisitor}
+import org.apache.xbean.asm5.Opcodes._
 
+import org.apache.spark.util.Utils
 
 /**
  * Includes an utility function to test whether a function accesses a specific attribute
@@ -107,18 +106,19 @@ private[graphx] object BytecodeUtils {
    *   MethodInvocationFinder("spark/graph/Foo", "test")
    * its methodsInvoked variable will contain the set of methods invoked directly by
    * Foo.test(). Interface invocations are not returned as part of the result set because we cannot
-   * determine the actual metod invoked by inspecting the bytecode.
+   * determine the actual method invoked by inspecting the bytecode.
    */
   private class MethodInvocationFinder(className: String, methodName: String)
-    extends ClassVisitor(ASM4) {
+    extends ClassVisitor(ASM5) {
 
     val methodsInvoked = new HashSet[(Class[_], String)]
 
     override def visitMethod(access: Int, name: String, desc: String,
                              sig: String, exceptions: Array[String]): MethodVisitor = {
       if (name == methodName) {
-        new MethodVisitor(ASM4) {
-          override def visitMethodInsn(op: Int, owner: String, name: String, desc: String) {
+        new MethodVisitor(ASM5) {
+          override def visitMethodInsn(
+              op: Int, owner: String, name: String, desc: String, itf: Boolean) {
             if (op == INVOKEVIRTUAL || op == INVOKESPECIAL || op == INVOKESTATIC) {
               if (!skipClass(owner)) {
                 methodsInvoked.add((Utils.classForName(owner.replace("/", ".")), name))
diff --git a/pom.xml b/pom.xml
index c499a80aa0f43..01afa80617891 100644
--- a/pom.xml
+++ b/pom.xml
@@ -393,6 +393,14 @@
           </exclusion>
         </exclusions>
       </dependency>
+      <!-- This artifact is a shaded version of ASM 5.0.4. The POM that was used to produce this
+           is at https://github.com/apache/geronimo-xbean/tree/xbean-4.4/xbean-asm5-shaded
+           For context on why we shade ASM, see SPARK-782 and SPARK-6152. -->
+      <dependency>
+        <groupId>org.apache.xbean</groupId>
+        <artifactId>xbean-asm5-shaded</artifactId>
+        <version>4.4</version>
+      </dependency>
 
       <!-- Shaded deps marked as provided. These are promoted to compile scope
            in the modules where we want the shaded classes to appear in the
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index b7c619224329f..570c9e50ed1f4 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -57,6 +57,9 @@ object BuildCommons {
   val sparkHome = buildLocation
 
   val testTempDir = s"$sparkHome/target/tmp"
+
+  val javacJVMVersion = settingKey[String]("source and target JVM version for javac")
+  val scalacJVMVersion = settingKey[String]("source and target JVM version for scalac")
 }
 
 object SparkBuild extends PomBuild {
@@ -154,9 +157,17 @@ object SparkBuild extends PomBuild {
       if (major.toInt >= 1 && minor.toInt >= 8) Seq("-Xdoclint:all", "-Xdoclint:-missing") else Seq.empty
     },
 
-    javacOptions in Compile ++= Seq("-encoding", "UTF-8"),
+    javacJVMVersion := "1.7",
+    scalacJVMVersion := "1.7",
+
+    javacOptions in Compile ++= Seq(
+      "-encoding", "UTF-8",
+      "-source", javacJVMVersion.value,
+      "-target", javacJVMVersion.value
+    ),
 
     scalacOptions in Compile ++= Seq(
+      s"-target:jvm-${scalacJVMVersion.value}",
       "-sourcepath", (baseDirectory in ThisBuild).value.getAbsolutePath  // Required for relative source links in scaladoc
     ),
 
@@ -241,8 +252,9 @@ object SparkBuild extends PomBuild {
 
   enable(Flume.settings)(streamingFlumeSink)
 
-  enable(DockerIntegrationTests.settings)(dockerIntegrationTests)
+  enable(Java8TestSettings.settings)(java8Tests)
 
+  enable(DockerIntegrationTests.settings)(dockerIntegrationTests)
 
   /**
    * Adds the ability to run the spark shell directly from SBT without building an assembly
@@ -591,6 +603,16 @@ object Unidoc {
   )
 }
 
+object Java8TestSettings {
+  import BuildCommons._
+
+  lazy val settings = Seq(
+    javacJVMVersion := "1.8",
+    // Targeting Java 8 bytecode is only supported in Scala 2.11.4 and higher:
+    scalacJVMVersion := (if (System.getProperty("scala-2.11") == "true") "1.8" else "1.7")
+  )
+}
+
 object TestSettings {
   import BuildCommons._
 
diff --git a/repl/pom.xml b/repl/pom.xml
index fb0a0e1286c80..154c99d23c7f4 100644
--- a/repl/pom.xml
+++ b/repl/pom.xml
@@ -95,6 +95,10 @@
       <groupId>org.apache.spark</groupId>
       <artifactId>spark-test-tags_${scala.binary.version}</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.xbean</groupId>
+      <artifactId>xbean-asm5-shaded</artifactId>
+    </dependency>
 
     <!-- Explicit listing of transitive deps that are shaded. Otherwise, odd compiler crashes. -->
     <dependency>
diff --git a/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala b/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
index 004941d5f50ae..3d2d235a00c93 100644
--- a/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
+++ b/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
@@ -23,15 +23,14 @@ import java.net.{HttpURLConnection, URI, URL, URLEncoder}
 import scala.util.control.NonFatal
 
 import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.xbean.asm5._
+import org.apache.xbean.asm5.Opcodes._
 
 import org.apache.spark.{SparkConf, SparkEnv, Logging}
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.util.Utils
 import org.apache.spark.util.ParentClassLoader
 
-import com.esotericsoftware.reflectasm.shaded.org.objectweb.asm._
-import com.esotericsoftware.reflectasm.shaded.org.objectweb.asm.Opcodes._
-
 /**
  * A ClassLoader that reads classes from a Hadoop FileSystem or HTTP URI,
  * used to load classes defined by the interpreter when the REPL is used.
@@ -192,7 +191,7 @@ class ExecutorClassLoader(conf: SparkConf, classUri: String, parent: ClassLoader
 }
 
 class ConstructorCleaner(className: String, cv: ClassVisitor)
-extends ClassVisitor(ASM4, cv) {
+extends ClassVisitor(ASM5, cv) {
   override def visitMethod(access: Int, name: String, desc: String,
       sig: String, exceptions: Array[String]): MethodVisitor = {
     val mv = cv.visitMethod(access, name, desc, sig, exceptions)
@@ -202,7 +201,7 @@ extends ClassVisitor(ASM4, cv) {
       // field in the class to point to it, but do nothing otherwise.
       mv.visitCode()
       mv.visitVarInsn(ALOAD, 0) // load this
-      mv.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "<init>", "()V")
+      mv.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "<init>", "()V", false)
       mv.visitVarInsn(ALOAD, 0) // load this
       // val classType = className.replace('.', '/')
       // mv.visitFieldInsn(PUTSTATIC, classType, "MODULE$", "L" + classType + ";")
diff --git a/sql/core/pom.xml b/sql/core/pom.xml
index c96855e261ee8..9fd6b5a07ec86 100644
--- a/sql/core/pom.xml
+++ b/sql/core/pom.xml
@@ -110,6 +110,11 @@
       <artifactId>mockito-core</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.xbean</groupId>
+      <artifactId>xbean-asm5-shaded</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <build>
     <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index 4b4f5c6c45c7a..97162249d9951 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -21,8 +21,8 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
 
 import scala.collection.mutable
 
-import com.esotericsoftware.reflectasm.shaded.org.objectweb.asm._
-import com.esotericsoftware.reflectasm.shaded.org.objectweb.asm.Opcodes._
+import org.apache.xbean.asm5._
+import org.apache.xbean.asm5.Opcodes._
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql._
@@ -41,22 +41,20 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
       l += 1L
       l.add(1L)
     }
-    BoxingFinder.getClassReader(f.getClass).foreach { cl =>
-      val boxingFinder = new BoxingFinder()
-      cl.accept(boxingFinder, 0)
-      assert(boxingFinder.boxingInvokes.isEmpty, s"Found boxing: ${boxingFinder.boxingInvokes}")
-    }
+    val cl = BoxingFinder.getClassReader(f.getClass)
+    val boxingFinder = new BoxingFinder()
+    cl.accept(boxingFinder, 0)
+    assert(boxingFinder.boxingInvokes.isEmpty, s"Found boxing: ${boxingFinder.boxingInvokes}")
   }
 
   test("Normal accumulator should do boxing") {
     // We need this test to make sure BoxingFinder works.
     val l = sparkContext.accumulator(0L)
     val f = () => { l += 1L }
-    BoxingFinder.getClassReader(f.getClass).foreach { cl =>
-      val boxingFinder = new BoxingFinder()
-      cl.accept(boxingFinder, 0)
-      assert(boxingFinder.boxingInvokes.nonEmpty, "Found find boxing in this test")
-    }
+    val cl = BoxingFinder.getClassReader(f.getClass)
+    val boxingFinder = new BoxingFinder()
+    cl.accept(boxingFinder, 0)
+    assert(boxingFinder.boxingInvokes.nonEmpty, "Found find boxing in this test")
   }
 
   /**
@@ -486,7 +484,7 @@ private class BoxingFinder(
     method: MethodIdentifier[_] = null,
     val boxingInvokes: mutable.Set[String] = mutable.Set.empty,
     visitedMethods: mutable.Set[MethodIdentifier[_]] = mutable.Set.empty)
-  extends ClassVisitor(ASM4) {
+  extends ClassVisitor(ASM5) {
 
   private val primitiveBoxingClassName =
     Set("java/lang/Long",
@@ -503,11 +501,12 @@ private class BoxingFinder(
     MethodVisitor = {
     if (method != null && (method.name != name || method.desc != desc)) {
       // If method is specified, skip other methods.
-      return new MethodVisitor(ASM4) {}
+      return new MethodVisitor(ASM5) {}
     }
 
-    new MethodVisitor(ASM4) {
-      override def visitMethodInsn(op: Int, owner: String, name: String, desc: String) {
+    new MethodVisitor(ASM5) {
+      override def visitMethodInsn(
+          op: Int, owner: String, name: String, desc: String, itf: Boolean) {
         if (op == INVOKESPECIAL && name == "<init>" || op == INVOKESTATIC && name == "valueOf") {
           if (primitiveBoxingClassName.contains(owner)) {
             // Find boxing methods, e.g, new java.lang.Long(l) or java.lang.Long.valueOf(l)
@@ -522,10 +521,9 @@ private class BoxingFinder(
           if (!visitedMethods.contains(m)) {
             // Keep track of visited methods to avoid potential infinite cycles
             visitedMethods += m
-            BoxingFinder.getClassReader(classOfMethodOwner).foreach { cl =>
-              visitedMethods += m
-              cl.accept(new BoxingFinder(m, boxingInvokes, visitedMethods), 0)
-            }
+            val cl = BoxingFinder.getClassReader(classOfMethodOwner)
+            visitedMethods += m
+            cl.accept(new BoxingFinder(m, boxingInvokes, visitedMethods), 0)
           }
         }
       }
@@ -535,22 +533,14 @@ private class BoxingFinder(
 
 private object BoxingFinder {
 
-  def getClassReader(cls: Class[_]): Option[ClassReader] = {
+  def getClassReader(cls: Class[_]): ClassReader = {
     val className = cls.getName.replaceFirst("^.*\\.", "") + ".class"
     val resourceStream = cls.getResourceAsStream(className)
     val baos = new ByteArrayOutputStream(128)
     // Copy data over, before delegating to ClassReader -
     // else we can run out of open file handles.
     Utils.copyStream(resourceStream, baos, true)
-    // ASM4 doesn't support Java 8 classes, which requires ASM5.
-    // So if the class is ASM5 (E.g., java.lang.Long when using JDK8 runtime to run these codes),
-    // then ClassReader will throw IllegalArgumentException,
-    // However, since this is only for testing, it's safe to skip these classes.
-    try {
-      Some(new ClassReader(new ByteArrayInputStream(baos.toByteArray)))
-    } catch {
-      case _: IllegalArgumentException => None
-    }
+    new ClassReader(new ByteArrayInputStream(baos.toByteArray))
   }
 
 }

From 27029bc8f6246514bd0947500c94cf38dc8616c3 Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Wed, 11 Nov 2015 11:24:55 -0800
Subject: [PATCH 194/518] [SPARK-11639][STREAMING][FLAKY-TEST] Implement
 BlockingWriteAheadLog for testing the BatchedWriteAheadLog

Several elements could be drained if the main thread is not fast enough. zsxwing warned me about a similar problem, but missed it here :( Submitting the fix using a waiter.

cc tdas

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9605 from brkyvz/fix-flaky-test.
---
 .../streaming/util/BatchedWriteAheadLog.scala |   3 +
 .../streaming/util/WriteAheadLogSuite.scala   | 124 +++++++++++-------
 2 files changed, 80 insertions(+), 47 deletions(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/BatchedWriteAheadLog.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/BatchedWriteAheadLog.scala
index 9727ed2ba1445..6e6ed8d819721 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/BatchedWriteAheadLog.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/BatchedWriteAheadLog.scala
@@ -182,6 +182,9 @@ private[util] class BatchedWriteAheadLog(val wrappedLog: WriteAheadLog, conf: Sp
       buffer.clear()
     }
   }
+
+  /** Method for querying the queue length. Should only be used in tests. */
+  private def getQueueLength(): Int = walWriteQueue.size()
 }
 
 /** Static methods for aggregating and de-aggregating records. */
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
index e96f4c2a29347..9e13f25c2efea 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
@@ -18,15 +18,14 @@ package org.apache.spark.streaming.util
 
 import java.io._
 import java.nio.ByteBuffer
-import java.util.concurrent.{ExecutionException, ThreadPoolExecutor}
-import java.util.concurrent.atomic.AtomicInteger
+import java.util.{Iterator => JIterator}
+import java.util.concurrent.ThreadPoolExecutor
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent._
 import scala.concurrent.duration._
 import scala.language.{implicitConversions, postfixOps}
-import scala.util.{Failure, Success}
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
@@ -37,12 +36,12 @@ import org.mockito.invocation.InvocationOnMock
 import org.mockito.stubbing.Answer
 import org.scalatest.concurrent.Eventually
 import org.scalatest.concurrent.Eventually._
-import org.scalatest.{BeforeAndAfterEach, BeforeAndAfter}
+import org.scalatest.{PrivateMethodTester, BeforeAndAfterEach, BeforeAndAfter}
 import org.scalatest.mock.MockitoSugar
 
 import org.apache.spark.streaming.scheduler._
 import org.apache.spark.util.{ThreadUtils, ManualClock, Utils}
-import org.apache.spark.{SparkException, SparkConf, SparkFunSuite}
+import org.apache.spark.{SparkConf, SparkFunSuite}
 
 /** Common tests for WriteAheadLogs that we would like to test with different configurations. */
 abstract class CommonWriteAheadLogTests(
@@ -315,7 +314,11 @@ class FileBasedWriteAheadLogWithFileCloseAfterWriteSuite
 class BatchedWriteAheadLogSuite extends CommonWriteAheadLogTests(
     allowBatching = true,
     closeFileAfterWrite = false,
-    "BatchedWriteAheadLog") with MockitoSugar with BeforeAndAfterEach with Eventually {
+    "BatchedWriteAheadLog")
+  with MockitoSugar
+  with BeforeAndAfterEach
+  with Eventually
+  with PrivateMethodTester {
 
   import BatchedWriteAheadLog._
   import WriteAheadLogSuite._
@@ -326,6 +329,8 @@ class BatchedWriteAheadLogSuite extends CommonWriteAheadLogTests(
   private var walBatchingExecutionContext: ExecutionContextExecutorService = _
   private val sparkConf = new SparkConf()
 
+  private val queueLength = PrivateMethod[Int]('getQueueLength)
+
   override def beforeEach(): Unit = {
     wal = mock[WriteAheadLog]
     walHandle = mock[WriteAheadLogRecordHandle]
@@ -366,7 +371,7 @@ class BatchedWriteAheadLogSuite extends CommonWriteAheadLogTests(
   }
 
   // we make the write requests in separate threads so that we don't block the test thread
-  private def promiseWriteEvent(wal: WriteAheadLog, event: String, time: Long): Promise[Unit] = {
+  private def writeAsync(wal: WriteAheadLog, event: String, time: Long): Promise[Unit] = {
     val p = Promise[Unit]()
     p.completeWith(Future {
       val v = wal.write(event, time)
@@ -375,28 +380,9 @@ class BatchedWriteAheadLogSuite extends CommonWriteAheadLogTests(
     p
   }
 
-  /**
-   * In order to block the writes on the writer thread, we mock the write method, and block it
-   * for some time with a promise.
-   */
-  private def writeBlockingPromise(wal: WriteAheadLog): Promise[Any] = {
-    // we would like to block the write so that we can queue requests
-    val promise = Promise[Any]()
-    when(wal.write(any[ByteBuffer], any[Long])).thenAnswer(
-      new Answer[WriteAheadLogRecordHandle] {
-        override def answer(invocation: InvocationOnMock): WriteAheadLogRecordHandle = {
-          Await.ready(promise.future, 4.seconds)
-          walHandle
-        }
-      }
-    )
-    promise
-  }
-
   test("BatchedWriteAheadLog - name log with aggregated entries with the timestamp of last entry") {
-    val batchedWal = new BatchedWriteAheadLog(wal, sparkConf)
-    // block the write so that we can batch some records
-    val promise = writeBlockingPromise(wal)
+    val blockingWal = new BlockingWriteAheadLog(wal, walHandle)
+    val batchedWal = new BatchedWriteAheadLog(blockingWal, sparkConf)
 
     val event1 = "hello"
     val event2 = "world"
@@ -406,21 +392,27 @@ class BatchedWriteAheadLogSuite extends CommonWriteAheadLogTests(
 
     // The queue.take() immediately takes the 3, and there is nothing left in the queue at that
     // moment. Then the promise blocks the writing of 3. The rest get queued.
-    promiseWriteEvent(batchedWal, event1, 3L)
-    // rest of the records will be batched while it takes 3 to get written
-    promiseWriteEvent(batchedWal, event2, 5L)
-    promiseWriteEvent(batchedWal, event3, 8L)
-    promiseWriteEvent(batchedWal, event4, 12L)
-    promiseWriteEvent(batchedWal, event5, 10L)
+    writeAsync(batchedWal, event1, 3L)
+    eventually(timeout(1 second)) {
+      assert(blockingWal.isBlocked)
+      assert(batchedWal.invokePrivate(queueLength()) === 0)
+    }
+    // rest of the records will be batched while it takes time for 3 to get written
+    writeAsync(batchedWal, event2, 5L)
+    writeAsync(batchedWal, event3, 8L)
+    writeAsync(batchedWal, event4, 12L)
+    writeAsync(batchedWal, event5, 10L)
     eventually(timeout(1 second)) {
       assert(walBatchingThreadPool.getActiveCount === 5)
+      assert(batchedWal.invokePrivate(queueLength()) === 4)
     }
-    promise.success(true)
+    blockingWal.allowWrite()
 
     val buffer1 = wrapArrayArrayByte(Array(event1))
     val buffer2 = wrapArrayArrayByte(Array(event2, event3, event4, event5))
 
     eventually(timeout(1 second)) {
+      assert(batchedWal.invokePrivate(queueLength()) === 0)
       verify(wal, times(1)).write(meq(buffer1), meq(3L))
       // the file name should be the timestamp of the last record, as events should be naturally
       // in order of timestamp, and we need the last element.
@@ -437,27 +429,32 @@ class BatchedWriteAheadLogSuite extends CommonWriteAheadLogTests(
   }
 
   test("BatchedWriteAheadLog - fail everything in queue during shutdown") {
-    val batchedWal = new BatchedWriteAheadLog(wal, sparkConf)
+    val blockingWal = new BlockingWriteAheadLog(wal, walHandle)
+    val batchedWal = new BatchedWriteAheadLog(blockingWal, sparkConf)
 
-    // block the write so that we can batch some records
-    writeBlockingPromise(wal)
-
-    val event1 = ("hello", 3L)
-    val event2 = ("world", 5L)
-    val event3 = ("this", 8L)
-    val event4 = ("is", 9L)
-    val event5 = ("doge", 10L)
+    val event1 = "hello"
+    val event2 = "world"
+    val event3 = "this"
 
     // The queue.take() immediately takes the 3, and there is nothing left in the queue at that
     // moment. Then the promise blocks the writing of 3. The rest get queued.
-    val writePromises = Seq(event1, event2, event3, event4, event5).map { event =>
-      promiseWriteEvent(batchedWal, event._1, event._2)
+    val promise1 = writeAsync(batchedWal, event1, 3L)
+    eventually(timeout(1 second)) {
+      assert(blockingWal.isBlocked)
+      assert(batchedWal.invokePrivate(queueLength()) === 0)
     }
+    // rest of the records will be batched while it takes time for 3 to get written
+    val promise2 = writeAsync(batchedWal, event2, 5L)
+    val promise3 = writeAsync(batchedWal, event3, 8L)
 
     eventually(timeout(1 second)) {
-      assert(walBatchingThreadPool.getActiveCount === 5)
+      assert(walBatchingThreadPool.getActiveCount === 3)
+      assert(blockingWal.isBlocked)
+      assert(batchedWal.invokePrivate(queueLength()) === 2) // event1 is being written
     }
 
+    val writePromises = Seq(promise1, promise2, promise3)
+
     batchedWal.close()
     eventually(timeout(1 second)) {
       assert(writePromises.forall(_.isCompleted))
@@ -641,4 +638,37 @@ object WriteAheadLogSuite {
   def wrapArrayArrayByte[T](records: Array[T]): ByteBuffer = {
     ByteBuffer.wrap(Utils.serialize[Array[Array[Byte]]](records.map(Utils.serialize[T])))
   }
+
+  /**
+   * A wrapper WriteAheadLog that blocks the write function to allow batching with the
+   * BatchedWriteAheadLog.
+   */
+  class BlockingWriteAheadLog(
+      wal: WriteAheadLog,
+      handle: WriteAheadLogRecordHandle) extends WriteAheadLog {
+    @volatile private var isWriteCalled: Boolean = false
+    @volatile private var blockWrite: Boolean = true
+
+    override def write(record: ByteBuffer, time: Long): WriteAheadLogRecordHandle = {
+      isWriteCalled = true
+      eventually(Eventually.timeout(2 second)) {
+        assert(!blockWrite)
+      }
+      wal.write(record, time)
+      isWriteCalled = false
+      handle
+    }
+    override def read(segment: WriteAheadLogRecordHandle): ByteBuffer = wal.read(segment)
+    override def readAll(): JIterator[ByteBuffer] = wal.readAll()
+    override def clean(threshTime: Long, waitForCompletion: Boolean): Unit = {
+      wal.clean(threshTime, waitForCompletion)
+    }
+    override def close(): Unit = wal.close()
+
+    def allowWrite(): Unit = {
+      blockWrite = false
+    }
+
+    def isBlocked: Boolean = isWriteCalled
+  }
 }

From df97df2b39194f60051f78cce23f0ba6cfe4b1df Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 11 Nov 2015 12:47:02 -0800
Subject: [PATCH 195/518] [SPARK-11644][SQL] Remove the option to turn off
 unsafe and codegen.

Author: Reynold Xin <rxin@databricks.com>

Closes #9618 from rxin/SPARK-11644.
---
 .../scala/org/apache/spark/sql/SQLConf.scala  |  27 +---
 .../apache/spark/sql/execution/Exchange.scala |   2 +-
 .../spark/sql/execution/QueryExecution.scala  |   1 -
 .../spark/sql/execution/SparkPlan.scala       | 120 +++++++----------
 .../spark/sql/execution/SparkPlanner.scala    |   4 -
 .../spark/sql/execution/SparkStrategies.scala |   6 +-
 .../spark/sql/execution/aggregate/utils.scala |  13 +-
 .../apache/spark/sql/execution/commands.scala |  27 ++++
 .../spark/sql/execution/joins/HashJoin.scala  |   4 +-
 .../sql/execution/joins/HashOuterJoin.scala   |   6 +-
 .../sql/execution/joins/HashSemiJoin.scala    |   9 +-
 .../sql/execution/joins/SortMergeJoin.scala   |   7 +-
 .../execution/joins/SortMergeOuterJoin.scala  |   7 +-
 .../sql/execution/local/HashJoinNode.scala    |   5 +-
 .../spark/sql/execution/local/LocalNode.scala |  80 +++++-------
 .../apache/spark/sql/sources/interfaces.scala |   8 +-
 .../org/apache/spark/sql/DataFrameSuite.scala |  31 +----
 .../spark/sql/DataFrameTungstenSuite.scala    |  68 +++++-----
 .../org/apache/spark/sql/SQLQuerySuite.scala  |  23 +---
 .../sql/execution/TungstenSortSuite.scala     |  13 --
 .../execution/joins/BroadcastJoinSuite.scala  |   4 +-
 .../execution/local/HashJoinNodeSuite.scala   |  23 +---
 .../local/NestedLoopJoinNodeSuite.scala       |  21 +--
 .../execution/metric/SQLMetricsSuite.scala    | 123 +++++-------------
 .../execution/AggregationQuerySuite.scala     |  44 +------
 .../sql/hive/execution/HiveExplainSuite.scala |   3 +-
 .../sql/hive/execution/HiveUDFSuite.scala     |  72 +++++-----
 27 files changed, 257 insertions(+), 494 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
index 57d7d30e0eca2..e02b502b7b4d5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
@@ -252,24 +252,8 @@ private[spark] object SQLConf {
         "not be provided to ExchangeCoordinator.",
       isPublic = false)
 
-  val TUNGSTEN_ENABLED = booleanConf("spark.sql.tungsten.enabled",
-    defaultValue = Some(true),
-    doc = "When true, use the optimized Tungsten physical execution backend which explicitly " +
-          "manages memory and dynamically generates bytecode for expression evaluation.")
-
-  val CODEGEN_ENABLED = booleanConf("spark.sql.codegen",
-    defaultValue = Some(true),  // use TUNGSTEN_ENABLED as default
-    doc = "When true, code will be dynamically generated at runtime for expression evaluation in" +
-      " a specific query.",
-    isPublic = false)
-
-  val UNSAFE_ENABLED = booleanConf("spark.sql.unsafe.enabled",
-    defaultValue = Some(true),  // use TUNGSTEN_ENABLED as default
-    doc = "When true, use the new optimized Tungsten physical execution backend.",
-    isPublic = false)
-
   val SUBEXPRESSION_ELIMINATION_ENABLED = booleanConf("spark.sql.subexpressionElimination.enabled",
-    defaultValue = Some(true),  // use CODEGEN_ENABLED as default
+    defaultValue = Some(true),
     doc = "When true, common subexpressions will be eliminated.",
     isPublic = false)
 
@@ -475,6 +459,9 @@ private[spark] object SQLConf {
     val MAPRED_REDUCE_TASKS = "mapred.reduce.tasks"
     val EXTERNAL_SORT = "spark.sql.planner.externalSort"
     val USE_SQL_AGGREGATE2 = "spark.sql.useAggregate2"
+    val TUNGSTEN_ENABLED = "spark.sql.tungsten.enabled"
+    val CODEGEN_ENABLED = "spark.sql.codegen"
+    val UNSAFE_ENABLED = "spark.sql.unsafe.enabled"
   }
 }
 
@@ -541,14 +528,10 @@ private[sql] class SQLConf extends Serializable with CatalystConf {
 
   private[spark] def sortMergeJoinEnabled: Boolean = getConf(SORTMERGE_JOIN)
 
-  private[spark] def codegenEnabled: Boolean = getConf(CODEGEN_ENABLED, getConf(TUNGSTEN_ENABLED))
-
   def caseSensitiveAnalysis: Boolean = getConf(SQLConf.CASE_SENSITIVE)
 
-  private[spark] def unsafeEnabled: Boolean = getConf(UNSAFE_ENABLED, getConf(TUNGSTEN_ENABLED))
-
   private[spark] def subexpressionEliminationEnabled: Boolean =
-    getConf(SUBEXPRESSION_ELIMINATION_ENABLED, codegenEnabled)
+    getConf(SUBEXPRESSION_ELIMINATION_ENABLED)
 
   private[spark] def autoBroadcastJoinThreshold: Int = getConf(AUTO_BROADCASTJOIN_THRESHOLD)
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
index b733b26987bcb..d0e4e068092f9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
@@ -58,7 +58,7 @@ case class Exchange(
    * Returns true iff we can support the data type, and we are not doing range partitioning.
    */
   private lazy val tungstenMode: Boolean = {
-    unsafeEnabled && codegenEnabled && GenerateUnsafeProjection.canSupport(child.schema) &&
+    GenerateUnsafeProjection.canSupport(child.schema) &&
       !newPartitioning.isInstanceOf[RangePartitioning]
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index 77843f53b9bd0..5da5aea17e25b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
@@ -77,7 +77,6 @@ class QueryExecution(val sqlContext: SQLContext, val logical: LogicalPlan) {
        |${stringOrError(optimizedPlan)}
        |== Physical Plan ==
        |${stringOrError(executedPlan)}
-       |Code Generation: ${stringOrError(executedPlan.codegenEnabled)}
     """.stripMargin.trim
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
index 8650ac500b652..1b833002f434c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
@@ -54,18 +54,8 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
   protected def sparkContext = sqlContext.sparkContext
 
   // sqlContext will be null when we are being deserialized on the slaves.  In this instance
-  // the value of codegenEnabled/unsafeEnabled will be set by the desserializer after the
+  // the value of subexpressionEliminationEnabled will be set by the desserializer after the
   // constructor has run.
-  val codegenEnabled: Boolean = if (sqlContext != null) {
-    sqlContext.conf.codegenEnabled
-  } else {
-    false
-  }
-  val unsafeEnabled: Boolean = if (sqlContext != null) {
-    sqlContext.conf.unsafeEnabled
-  } else {
-    false
-  }
   val subexpressionEliminationEnabled: Boolean = if (sqlContext != null) {
     sqlContext.conf.subexpressionEliminationEnabled
   } else {
@@ -233,83 +223,63 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
 
   protected def newProjection(
       expressions: Seq[Expression], inputSchema: Seq[Attribute]): Projection = {
-    log.debug(
-      s"Creating Projection: $expressions, inputSchema: $inputSchema, codegen:$codegenEnabled")
-    if (codegenEnabled) {
-      try {
-        GenerateProjection.generate(expressions, inputSchema)
-      } catch {
-        case e: Exception =>
-          if (isTesting) {
-            throw e
-          } else {
-            log.error("Failed to generate projection, fallback to interpret", e)
-            new InterpretedProjection(expressions, inputSchema)
-          }
-      }
-    } else {
-      new InterpretedProjection(expressions, inputSchema)
+    log.debug(s"Creating Projection: $expressions, inputSchema: $inputSchema")
+    try {
+      GenerateProjection.generate(expressions, inputSchema)
+    } catch {
+      case e: Exception =>
+        if (isTesting) {
+          throw e
+        } else {
+          log.error("Failed to generate projection, fallback to interpret", e)
+          new InterpretedProjection(expressions, inputSchema)
+        }
     }
   }
 
   protected def newMutableProjection(
-      expressions: Seq[Expression],
-      inputSchema: Seq[Attribute]): () => MutableProjection = {
-    log.debug(
-      s"Creating MutableProj: $expressions, inputSchema: $inputSchema, codegen:$codegenEnabled")
-    if(codegenEnabled) {
-      try {
-        GenerateMutableProjection.generate(expressions, inputSchema)
-      } catch {
-        case e: Exception =>
-          if (isTesting) {
-            throw e
-          } else {
-            log.error("Failed to generate mutable projection, fallback to interpreted", e)
-            () => new InterpretedMutableProjection(expressions, inputSchema)
-          }
-      }
-    } else {
-      () => new InterpretedMutableProjection(expressions, inputSchema)
+      expressions: Seq[Expression], inputSchema: Seq[Attribute]): () => MutableProjection = {
+    log.debug(s"Creating MutableProj: $expressions, inputSchema: $inputSchema")
+    try {
+      GenerateMutableProjection.generate(expressions, inputSchema)
+    } catch {
+      case e: Exception =>
+        if (isTesting) {
+          throw e
+        } else {
+          log.error("Failed to generate mutable projection, fallback to interpreted", e)
+          () => new InterpretedMutableProjection(expressions, inputSchema)
+        }
     }
   }
 
   protected def newPredicate(
       expression: Expression, inputSchema: Seq[Attribute]): (InternalRow) => Boolean = {
-    if (codegenEnabled) {
-      try {
-        GeneratePredicate.generate(expression, inputSchema)
-      } catch {
-        case e: Exception =>
-          if (isTesting) {
-            throw e
-          } else {
-            log.error("Failed to generate predicate, fallback to interpreted", e)
-            InterpretedPredicate.create(expression, inputSchema)
-          }
-      }
-    } else {
-      InterpretedPredicate.create(expression, inputSchema)
+    try {
+      GeneratePredicate.generate(expression, inputSchema)
+    } catch {
+      case e: Exception =>
+        if (isTesting) {
+          throw e
+        } else {
+          log.error("Failed to generate predicate, fallback to interpreted", e)
+          InterpretedPredicate.create(expression, inputSchema)
+        }
     }
   }
 
   protected def newOrdering(
-      order: Seq[SortOrder],
-      inputSchema: Seq[Attribute]): Ordering[InternalRow] = {
-    if (codegenEnabled) {
-      try {
-        GenerateOrdering.generate(order, inputSchema)
-      } catch {
-        case e: Exception =>
-          if (isTesting) {
-            throw e
-          } else {
-            log.error("Failed to generate ordering, fallback to interpreted", e)
-            new InterpretedOrdering(order, inputSchema)
-          }
-      }
-    } else {
-      new InterpretedOrdering(order, inputSchema)
+      order: Seq[SortOrder], inputSchema: Seq[Attribute]): Ordering[InternalRow] = {
+    try {
+      GenerateOrdering.generate(order, inputSchema)
+    } catch {
+      case e: Exception =>
+        if (isTesting) {
+          throw e
+        } else {
+          log.error("Failed to generate ordering, fallback to interpreted", e)
+          new InterpretedOrdering(order, inputSchema)
+        }
     }
   }
   /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
index a10d1edcc91aa..cf482ae4a05ee 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
@@ -27,10 +27,6 @@ import org.apache.spark.sql.execution.datasources.DataSourceStrategy
 class SparkPlanner(val sqlContext: SQLContext) extends SparkStrategies {
   val sparkContext: SparkContext = sqlContext.sparkContext
 
-  def codegenEnabled: Boolean = sqlContext.conf.codegenEnabled
-
-  def unsafeEnabled: Boolean = sqlContext.conf.unsafeEnabled
-
   def numPartitions: Int = sqlContext.conf.numShufflePartitions
 
   def strategies: Seq[Strategy] =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index d65cb1bae7fb5..96242f160aa51 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -327,8 +327,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
      *               if necessary.
      */
     def getSortOperator(sortExprs: Seq[SortOrder], global: Boolean, child: SparkPlan): SparkPlan = {
-      if (sqlContext.conf.unsafeEnabled && sqlContext.conf.codegenEnabled &&
-        TungstenSort.supportsSchema(child.schema)) {
+      if (TungstenSort.supportsSchema(child.schema)) {
         execution.TungstenSort(sortExprs, global, child)
       } else {
         execution.Sort(sortExprs, global, child)
@@ -368,8 +367,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
       case logical.Project(projectList, child) =>
         // If unsafe mode is enabled and we support these data types in Unsafe, use the
         // Tungsten project. Otherwise, use the normal project.
-        if (sqlContext.conf.unsafeEnabled &&
-          UnsafeProjection.canSupport(projectList) && UnsafeProjection.canSupport(child.schema)) {
+        if (UnsafeProjection.canSupport(projectList) && UnsafeProjection.canSupport(child.schema)) {
           execution.TungstenProject(projectList, planLater(child)) :: Nil
         } else {
           execution.Project(projectList, planLater(child)) :: Nil
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
index 79abf2d5929be..a70e41436c7aa 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
@@ -59,13 +59,10 @@ object Utils {
       resultExpressions: Seq[NamedExpression],
       child: SparkPlan): Seq[SparkPlan] = {
     // Check if we can use TungstenAggregate.
-    val usesTungstenAggregate =
-      child.sqlContext.conf.unsafeEnabled &&
-      TungstenAggregate.supportsAggregate(
+    val usesTungstenAggregate = TungstenAggregate.supportsAggregate(
         groupingExpressions,
         aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
 
-
     // 1. Create an Aggregate Operator for partial aggregations.
 
     val groupingAttributes = groupingExpressions.map(_.toAttribute)
@@ -144,11 +141,9 @@ object Utils {
       child: SparkPlan): Seq[SparkPlan] = {
 
     val aggregateExpressions = functionsWithDistinct ++ functionsWithoutDistinct
-    val usesTungstenAggregate =
-      child.sqlContext.conf.unsafeEnabled &&
-        TungstenAggregate.supportsAggregate(
-          groupingExpressions,
-          aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
+    val usesTungstenAggregate = TungstenAggregate.supportsAggregate(
+      groupingExpressions,
+      aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
 
     // functionsWithDistinct is guaranteed to be non-empty. Even though it may contain more than one
     // DISTINCT aggregate function, all of those functions will have the same column expression.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
index 8b2755a58757c..e29c281b951f2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
@@ -121,6 +121,33 @@ case class SetCommand(kv: Option[(String, Option[String])]) extends RunnableComm
       }
       (keyValueOutput, runFunc)
 
+    case Some((SQLConf.Deprecated.TUNGSTEN_ENABLED, Some(value))) =>
+      val runFunc = (sqlContext: SQLContext) => {
+        logWarning(
+          s"Property ${SQLConf.Deprecated.TUNGSTEN_ENABLED} is deprecated and " +
+            s"will be ignored. Tungsten will continue to be used.")
+        Seq(Row(SQLConf.Deprecated.TUNGSTEN_ENABLED, "true"))
+      }
+      (keyValueOutput, runFunc)
+
+    case Some((SQLConf.Deprecated.CODEGEN_ENABLED, Some(value))) =>
+      val runFunc = (sqlContext: SQLContext) => {
+        logWarning(
+          s"Property ${SQLConf.Deprecated.CODEGEN_ENABLED} is deprecated and " +
+            s"will be ignored. Codegen will continue to be used.")
+        Seq(Row(SQLConf.Deprecated.CODEGEN_ENABLED, "true"))
+      }
+      (keyValueOutput, runFunc)
+
+    case Some((SQLConf.Deprecated.UNSAFE_ENABLED, Some(value))) =>
+      val runFunc = (sqlContext: SQLContext) => {
+        logWarning(
+          s"Property ${SQLConf.Deprecated.UNSAFE_ENABLED} is deprecated and " +
+            s"will be ignored. Unsafe mode will continue to be used.")
+        Seq(Row(SQLConf.Deprecated.UNSAFE_ENABLED, "true"))
+      }
+      (keyValueOutput, runFunc)
+
     // Configures a single property.
     case Some((key, Some(value))) =>
       val runFunc = (sqlContext: SQLContext) => {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala
index 7ce4a517838cb..997f7f494f4a3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala
@@ -45,9 +45,7 @@ trait HashJoin {
   override def output: Seq[Attribute] = left.output ++ right.output
 
   protected[this] def isUnsafeMode: Boolean = {
-    (self.codegenEnabled && self.unsafeEnabled
-      && UnsafeProjection.canSupport(buildKeys)
-      && UnsafeProjection.canSupport(self.schema))
+    UnsafeProjection.canSupport(buildKeys) && UnsafeProjection.canSupport(self.schema)
   }
 
   override def outputsUnsafeRows: Boolean = isUnsafeMode
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashOuterJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashOuterJoin.scala
index 15b06b1537f8c..3633f356b014b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashOuterJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashOuterJoin.scala
@@ -65,9 +65,9 @@ trait HashOuterJoin {
   }
 
   protected[this] def isUnsafeMode: Boolean = {
-    (self.codegenEnabled && self.unsafeEnabled && joinType != FullOuter
-      && UnsafeProjection.canSupport(buildKeys)
-      && UnsafeProjection.canSupport(self.schema))
+    joinType != FullOuter &&
+      UnsafeProjection.canSupport(buildKeys) &&
+      UnsafeProjection.canSupport(self.schema)
   }
 
   override def outputsUnsafeRows: Boolean = isUnsafeMode
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashSemiJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashSemiJoin.scala
index beb141ade616d..c7d13e0a72a87 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashSemiJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashSemiJoin.scala
@@ -34,11 +34,10 @@ trait HashSemiJoin {
   override def output: Seq[Attribute] = left.output
 
   protected[this] def supportUnsafe: Boolean = {
-    (self.codegenEnabled && self.unsafeEnabled
-      && UnsafeProjection.canSupport(leftKeys)
-      && UnsafeProjection.canSupport(rightKeys)
-      && UnsafeProjection.canSupport(left.schema)
-      && UnsafeProjection.canSupport(right.schema))
+    UnsafeProjection.canSupport(leftKeys) &&
+      UnsafeProjection.canSupport(rightKeys) &&
+      UnsafeProjection.canSupport(left.schema) &&
+      UnsafeProjection.canSupport(right.schema)
   }
 
   override def outputsUnsafeRows: Boolean = supportUnsafe
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoin.scala
index 17030947b7bbc..7aee8e3dd3fce 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoin.scala
@@ -54,10 +54,9 @@ case class SortMergeJoin(
     requiredOrders(leftKeys) :: requiredOrders(rightKeys) :: Nil
 
   protected[this] def isUnsafeMode: Boolean = {
-    (codegenEnabled && unsafeEnabled
-      && UnsafeProjection.canSupport(leftKeys)
-      && UnsafeProjection.canSupport(rightKeys)
-      && UnsafeProjection.canSupport(schema))
+    UnsafeProjection.canSupport(leftKeys) &&
+      UnsafeProjection.canSupport(rightKeys) &&
+      UnsafeProjection.canSupport(schema)
   }
 
   override def outputsUnsafeRows: Boolean = isUnsafeMode
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
index 7e854e6702f77..5f1590c463836 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
@@ -90,10 +90,9 @@ case class SortMergeOuterJoin(
   }
 
   private def isUnsafeMode: Boolean = {
-    (codegenEnabled && unsafeEnabled
-      && UnsafeProjection.canSupport(leftKeys)
-      && UnsafeProjection.canSupport(rightKeys)
-      && UnsafeProjection.canSupport(schema))
+    UnsafeProjection.canSupport(leftKeys) &&
+      UnsafeProjection.canSupport(rightKeys) &&
+      UnsafeProjection.canSupport(schema)
   }
 
   override def outputsUnsafeRows: Boolean = isUnsafeMode
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/local/HashJoinNode.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/local/HashJoinNode.scala
index b1dc719ca8508..aef655727fbbb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/local/HashJoinNode.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/local/HashJoinNode.scala
@@ -46,10 +46,7 @@ trait HashJoinNode {
   private[this] var joinKeys: Projection = _
 
   protected def isUnsafeMode: Boolean = {
-    (codegenEnabled &&
-      unsafeEnabled &&
-      UnsafeProjection.canSupport(schema) &&
-      UnsafeProjection.canSupport(streamedKeys))
+    UnsafeProjection.canSupport(schema) && UnsafeProjection.canSupport(streamedKeys)
   }
 
   private def streamSideKeyGenerator: Projection = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/local/LocalNode.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/local/LocalNode.scala
index f96b62a67a254..d3381eac91d43 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/local/LocalNode.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/local/LocalNode.scala
@@ -35,10 +35,6 @@ import org.apache.spark.sql.types.StructType
  */
 abstract class LocalNode(conf: SQLConf) extends QueryPlan[LocalNode] with Logging {
 
-  protected val codegenEnabled: Boolean = conf.codegenEnabled
-
-  protected val unsafeEnabled: Boolean = conf.unsafeEnabled
-
   private[this] lazy val isTesting: Boolean = sys.props.contains("spark.testing")
 
   /**
@@ -111,21 +107,17 @@ abstract class LocalNode(conf: SQLConf) extends QueryPlan[LocalNode] with Loggin
       expressions: Seq[Expression],
       inputSchema: Seq[Attribute]): Projection = {
     log.debug(
-      s"Creating Projection: $expressions, inputSchema: $inputSchema, codegen:$codegenEnabled")
-    if (codegenEnabled) {
-      try {
-        GenerateProjection.generate(expressions, inputSchema)
-      } catch {
-        case NonFatal(e) =>
-          if (isTesting) {
-            throw e
-          } else {
-            log.error("Failed to generate projection, fallback to interpret", e)
-            new InterpretedProjection(expressions, inputSchema)
-          }
-      }
-    } else {
-      new InterpretedProjection(expressions, inputSchema)
+      s"Creating Projection: $expressions, inputSchema: $inputSchema")
+    try {
+      GenerateProjection.generate(expressions, inputSchema)
+    } catch {
+      case NonFatal(e) =>
+        if (isTesting) {
+          throw e
+        } else {
+          log.error("Failed to generate projection, fallback to interpret", e)
+          new InterpretedProjection(expressions, inputSchema)
+        }
     }
   }
 
@@ -133,41 +125,33 @@ abstract class LocalNode(conf: SQLConf) extends QueryPlan[LocalNode] with Loggin
       expressions: Seq[Expression],
       inputSchema: Seq[Attribute]): () => MutableProjection = {
     log.debug(
-      s"Creating MutableProj: $expressions, inputSchema: $inputSchema, codegen:$codegenEnabled")
-    if (codegenEnabled) {
-      try {
-        GenerateMutableProjection.generate(expressions, inputSchema)
-      } catch {
-        case NonFatal(e) =>
-          if (isTesting) {
-            throw e
-          } else {
-            log.error("Failed to generate mutable projection, fallback to interpreted", e)
-            () => new InterpretedMutableProjection(expressions, inputSchema)
-          }
-      }
-    } else {
-      () => new InterpretedMutableProjection(expressions, inputSchema)
+      s"Creating MutableProj: $expressions, inputSchema: $inputSchema")
+    try {
+      GenerateMutableProjection.generate(expressions, inputSchema)
+    } catch {
+      case NonFatal(e) =>
+        if (isTesting) {
+          throw e
+        } else {
+          log.error("Failed to generate mutable projection, fallback to interpreted", e)
+          () => new InterpretedMutableProjection(expressions, inputSchema)
+        }
     }
   }
 
   protected def newPredicate(
       expression: Expression,
       inputSchema: Seq[Attribute]): (InternalRow) => Boolean = {
-    if (codegenEnabled) {
-      try {
-        GeneratePredicate.generate(expression, inputSchema)
-      } catch {
-        case NonFatal(e) =>
-          if (isTesting) {
-            throw e
-          } else {
-            log.error("Failed to generate predicate, fallback to interpreted", e)
-            InterpretedPredicate.create(expression, inputSchema)
-          }
-      }
-    } else {
-      InterpretedPredicate.create(expression, inputSchema)
+    try {
+      GeneratePredicate.generate(expression, inputSchema)
+    } catch {
+      case NonFatal(e) =>
+        if (isTesting) {
+          throw e
+        } else {
+          log.error("Failed to generate predicate, fallback to interpreted", e)
+          InterpretedPredicate.create(expression, inputSchema)
+        }
     }
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index 5b8841bc154a5..48de693a999d8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -423,8 +423,6 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
 
   private val hadoopConf = new Configuration(sqlContext.sparkContext.hadoopConfiguration)
 
-  private val codegenEnabled = sqlContext.conf.codegenEnabled
-
   private var _partitionSpec: PartitionSpec = _
 
   private class FileStatusCache {
@@ -661,7 +659,6 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
   def buildScan(requiredColumns: Array[String], inputFiles: Array[FileStatus]): RDD[Row] = {
     // Yeah, to workaround serialization...
     val dataSchema = this.dataSchema
-    val codegenEnabled = this.codegenEnabled
     val needConversion = this.needConversion
 
     val requiredOutput = requiredColumns.map { col =>
@@ -678,11 +675,8 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
       }
 
     converted.mapPartitions { rows =>
-      val buildProjection = if (codegenEnabled) {
+      val buildProjection =
         GenerateMutableProjection.generate(requiredOutput, dataSchema.toAttributes)
-      } else {
-        () => new InterpretedMutableProjection(requiredOutput, dataSchema.toAttributes)
-      }
 
       val projectedRows = {
         val mutableProjection = buildProjection()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index f3a7aa280367a..e4f23fe17b757 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -621,11 +621,7 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
   }
 
   test("SPARK-6899: type should match when using codegen") {
-    withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "true") {
-      checkAnswer(
-        decimalData.agg(avg('a)),
-        Row(new java.math.BigDecimal(2.0)))
-    }
+    checkAnswer(decimalData.agg(avg('a)), Row(new java.math.BigDecimal(2.0)))
   }
 
   test("SPARK-7133: Implement struct, array, and map field accessor") {
@@ -844,31 +840,16 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
   }
 
   test("SPARK-8608: call `show` on local DataFrame with random columns should return same value") {
-    // Make sure we can pass this test for both codegen mode and interpreted mode.
-    withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "true") {
-      val df = testData.select(rand(33))
-      assert(df.showString(5) == df.showString(5))
-    }
-
-    withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "false") {
-      val df = testData.select(rand(33))
-      assert(df.showString(5) == df.showString(5))
-    }
+    val df = testData.select(rand(33))
+    assert(df.showString(5) == df.showString(5))
 
     // We will reuse the same Expression object for LocalRelation.
-    val df = (1 to 10).map(Tuple1.apply).toDF().select(rand(33))
-    assert(df.showString(5) == df.showString(5))
+    val df1 = (1 to 10).map(Tuple1.apply).toDF().select(rand(33))
+    assert(df1.showString(5) == df1.showString(5))
   }
 
   test("SPARK-8609: local DataFrame with random columns should return same value after sort") {
-    // Make sure we can pass this test for both codegen mode and interpreted mode.
-    withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "true") {
-      checkAnswer(testData.sort(rand(33)), testData.sort(rand(33)))
-    }
-
-    withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "false") {
-      checkAnswer(testData.sort(rand(33)), testData.sort(rand(33)))
-    }
+    checkAnswer(testData.sort(rand(33)), testData.sort(rand(33)))
 
     // We will reuse the same Expression object for LocalRelation.
     val df = (1 to 10).map(Tuple1.apply).toDF()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameTungstenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameTungstenSuite.scala
index 7ae12a7895f7e..68e99d6a6b816 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameTungstenSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameTungstenSuite.scala
@@ -31,52 +31,46 @@ class DataFrameTungstenSuite extends QueryTest with SharedSQLContext {
   import testImplicits._
 
   test("test simple types") {
-    withSQLConf(SQLConf.UNSAFE_ENABLED.key -> "true") {
-      val df = sparkContext.parallelize(Seq((1, 2))).toDF("a", "b")
-      assert(df.select(struct("a", "b")).first().getStruct(0) === Row(1, 2))
-    }
+    val df = sparkContext.parallelize(Seq((1, 2))).toDF("a", "b")
+    assert(df.select(struct("a", "b")).first().getStruct(0) === Row(1, 2))
   }
 
   test("test struct type") {
-    withSQLConf(SQLConf.UNSAFE_ENABLED.key -> "true") {
-      val struct = Row(1, 2L, 3.0F, 3.0)
-      val data = sparkContext.parallelize(Seq(Row(1, struct)))
+    val struct = Row(1, 2L, 3.0F, 3.0)
+    val data = sparkContext.parallelize(Seq(Row(1, struct)))
 
-      val schema = new StructType()
-        .add("a", IntegerType)
-        .add("b",
-          new StructType()
-            .add("b1", IntegerType)
-            .add("b2", LongType)
-            .add("b3", FloatType)
-            .add("b4", DoubleType))
+    val schema = new StructType()
+      .add("a", IntegerType)
+      .add("b",
+        new StructType()
+          .add("b1", IntegerType)
+          .add("b2", LongType)
+          .add("b3", FloatType)
+          .add("b4", DoubleType))
 
-      val df = sqlContext.createDataFrame(data, schema)
-      assert(df.select("b").first() === Row(struct))
-    }
+    val df = sqlContext.createDataFrame(data, schema)
+    assert(df.select("b").first() === Row(struct))
   }
 
   test("test nested struct type") {
-    withSQLConf(SQLConf.UNSAFE_ENABLED.key -> "true") {
-      val innerStruct = Row(1, "abcd")
-      val outerStruct = Row(1, 2L, 3.0F, 3.0, innerStruct, "efg")
-      val data = sparkContext.parallelize(Seq(Row(1, outerStruct)))
+    val innerStruct = Row(1, "abcd")
+    val outerStruct = Row(1, 2L, 3.0F, 3.0, innerStruct, "efg")
+    val data = sparkContext.parallelize(Seq(Row(1, outerStruct)))
 
-      val schema = new StructType()
-        .add("a", IntegerType)
-        .add("b",
-          new StructType()
-            .add("b1", IntegerType)
-            .add("b2", LongType)
-            .add("b3", FloatType)
-            .add("b4", DoubleType)
-            .add("b5", new StructType()
-            .add("b5a", IntegerType)
-            .add("b5b", StringType))
-            .add("b6", StringType))
+    val schema = new StructType()
+      .add("a", IntegerType)
+      .add("b",
+        new StructType()
+          .add("b1", IntegerType)
+          .add("b2", LongType)
+          .add("b3", FloatType)
+          .add("b4", DoubleType)
+          .add("b5", new StructType()
+          .add("b5a", IntegerType)
+          .add("b5b", StringType))
+          .add("b6", StringType))
 
-      val df = sqlContext.createDataFrame(data, schema)
-      assert(df.select("b").first() === Row(outerStruct))
-    }
+    val df = sqlContext.createDataFrame(data, schema)
+    assert(df.select("b").first() === Row(outerStruct))
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index 19e850a46fdfc..acabe32c67bc0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -261,8 +261,6 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
   }
 
   test("aggregation with codegen") {
-    val originalValue = sqlContext.conf.codegenEnabled
-    sqlContext.setConf(SQLConf.CODEGEN_ENABLED, true)
     // Prepare a table that we can group some rows.
     sqlContext.table("testData")
       .unionAll(sqlContext.table("testData"))
@@ -347,7 +345,6 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
         Row(null, null, null, 0) :: Nil)
     } finally {
       sqlContext.dropTempTable("testData3x")
-      sqlContext.setConf(SQLConf.CODEGEN_ENABLED, originalValue)
     }
   }
 
@@ -567,12 +564,6 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
     sortTest()
   }
 
-  test("SPARK-6927 external sorting with codegen on") {
-    withSQLConf(SQLConf.CODEGEN_ENABLED.key -> "true") {
-      sortTest()
-    }
-  }
-
   test("limit") {
     checkAnswer(
       sql("SELECT * FROM testData LIMIT 10"),
@@ -1624,12 +1615,10 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
   }
 
   test("aggregation with codegen updates peak execution memory") {
-    withSQLConf((SQLConf.CODEGEN_ENABLED.key, "true")) {
-      AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "aggregation with codegen") {
-        testCodeGen(
-          "SELECT key, count(value) FROM testData GROUP BY key",
-          (1 to 100).map(i => Row(i, 1)))
-      }
+    AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "aggregation with codegen") {
+      testCodeGen(
+        "SELECT key, count(value) FROM testData GROUP BY key",
+        (1 to 100).map(i => Row(i, 1)))
     }
   }
 
@@ -1783,9 +1772,7 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
     // This bug will be triggered when Tungsten is enabled and there are multiple
     // SortMergeJoin operators executed in the same task.
     val confs =
-      SQLConf.SORTMERGE_JOIN.key -> "true" ::
-        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1" ::
-        SQLConf.TUNGSTEN_ENABLED.key -> "true" :: Nil
+      SQLConf.SORTMERGE_JOIN.key -> "true" :: SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1" :: Nil
     withSQLConf(confs: _*) {
       val df1 = (1 to 50).map(i => (s"str_$i", i)).toDF("i", "j")
       val df2 =
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala
index 7a0f0dfd2b7f1..85486c08894c9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala
@@ -31,19 +31,6 @@ import org.apache.spark.sql.types._
 class TungstenSortSuite extends SparkPlanTest with SharedSQLContext {
   import testImplicits.localSeqToDataFrameHolder
 
-  override def beforeAll(): Unit = {
-    super.beforeAll()
-    sqlContext.conf.setConf(SQLConf.CODEGEN_ENABLED, true)
-  }
-
-  override def afterAll(): Unit = {
-    try {
-      sqlContext.conf.unsetConf(SQLConf.CODEGEN_ENABLED)
-    } finally {
-      super.afterAll()
-    }
-  }
-
   test("sort followed by limit") {
     checkThatPlansAgree(
       (1 to 100).map(v => Tuple1(v)).toDF("a"),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
index dcbfdca71acb6..5b2998c3c76d3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/BroadcastJoinSuite.scala
@@ -26,7 +26,7 @@ import org.apache.spark.sql.functions._
 import org.apache.spark.sql.{SQLConf, SQLContext, QueryTest}
 
 /**
- * Test various broadcast join operators with unsafe enabled.
+ * Test various broadcast join operators.
  *
  * Tests in this suite we need to run Spark in local-cluster mode. In particular, the use of
  * unsafe map in [[org.apache.spark.sql.execution.joins.UnsafeHashedRelation]] is not triggered
@@ -45,8 +45,6 @@ class BroadcastJoinSuite extends QueryTest with BeforeAndAfterAll {
       .setAppName("testing")
     val sc = new SparkContext(conf)
     sqlContext = new SQLContext(sc)
-    sqlContext.setConf(SQLConf.UNSAFE_ENABLED, true)
-    sqlContext.setConf(SQLConf.CODEGEN_ENABLED, true)
   }
 
   override def afterAll(): Unit = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/local/HashJoinNodeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/local/HashJoinNodeSuite.scala
index 8c2e78b2a9db7..44b0d9d4102a1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/local/HashJoinNodeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/local/HashJoinNodeSuite.scala
@@ -28,12 +28,9 @@ import org.apache.spark.sql.execution.joins.{HashedRelation, BuildLeft, BuildRig
 class HashJoinNodeSuite extends LocalNodeTest {
 
   // Test all combinations of the two dimensions: with/out unsafe and build sides
-  private val maybeUnsafeAndCodegen = Seq(false, true)
   private val buildSides = Seq(BuildLeft, BuildRight)
-  maybeUnsafeAndCodegen.foreach { unsafeAndCodegen =>
-    buildSides.foreach { buildSide =>
-      testJoin(unsafeAndCodegen, buildSide)
-    }
+  buildSides.foreach { buildSide =>
+    testJoin(buildSide)
   }
 
   /**
@@ -45,10 +42,7 @@ class HashJoinNodeSuite extends LocalNodeTest {
       buildKeys: Seq[Expression],
       buildNode: LocalNode): HashedRelation = {
 
-    val isUnsafeMode =
-      conf.codegenEnabled &&
-        conf.unsafeEnabled &&
-        UnsafeProjection.canSupport(buildKeys)
+    val isUnsafeMode = UnsafeProjection.canSupport(buildKeys)
 
     val buildSideKeyGenerator =
       if (isUnsafeMode) {
@@ -68,15 +62,10 @@ class HashJoinNodeSuite extends LocalNodeTest {
   /**
    * Test inner hash join with varying degrees of matches.
    */
-  private def testJoin(
-      unsafeAndCodegen: Boolean,
-      buildSide: BuildSide): Unit = {
-    val simpleOrUnsafe = if (!unsafeAndCodegen) "simple" else "unsafe"
-    val testNamePrefix = s"$simpleOrUnsafe / $buildSide"
+  private def testJoin(buildSide: BuildSide): Unit = {
+    val testNamePrefix = buildSide
     val someData = (1 to 100).map { i => (i, "burger" + i) }.toArray
     val conf = new SQLConf
-    conf.setConf(SQLConf.UNSAFE_ENABLED, unsafeAndCodegen)
-    conf.setConf(SQLConf.CODEGEN_ENABLED, unsafeAndCodegen)
 
     // Actual test body
     def runTest(leftInput: Array[(Int, String)], rightInput: Array[(Int, String)]): Unit = {
@@ -119,7 +108,7 @@ class HashJoinNodeSuite extends LocalNodeTest {
         .map { case (k, v) => (k, v, k, rightInputMap(k)) }
 
       Seq(makeBinaryHashJoinNode, makeBroadcastJoinNode).foreach { makeNode =>
-        val makeUnsafeNode = if (unsafeAndCodegen) wrapForUnsafe(makeNode) else makeNode
+        val makeUnsafeNode = wrapForUnsafe(makeNode)
         val hashJoinNode = makeUnsafeNode(leftNode, rightNode)
 
         val actualOutput = hashJoinNode.collect().map { row =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/local/NestedLoopJoinNodeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/local/NestedLoopJoinNodeSuite.scala
index 40299d9d5ee37..252f7cc8971f4 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/local/NestedLoopJoinNodeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/local/NestedLoopJoinNodeSuite.scala
@@ -26,30 +26,21 @@ import org.apache.spark.sql.execution.joins.{BuildLeft, BuildRight, BuildSide}
 class NestedLoopJoinNodeSuite extends LocalNodeTest {
 
   // Test all combinations of the three dimensions: with/out unsafe, build sides, and join types
-  private val maybeUnsafeAndCodegen = Seq(false, true)
   private val buildSides = Seq(BuildLeft, BuildRight)
   private val joinTypes = Seq(LeftOuter, RightOuter, FullOuter)
-  maybeUnsafeAndCodegen.foreach { unsafeAndCodegen =>
-    buildSides.foreach { buildSide =>
-      joinTypes.foreach { joinType =>
-        testJoin(unsafeAndCodegen, buildSide, joinType)
-      }
+  buildSides.foreach { buildSide =>
+    joinTypes.foreach { joinType =>
+      testJoin(buildSide, joinType)
     }
   }
 
   /**
    * Test outer nested loop joins with varying degrees of matches.
    */
-  private def testJoin(
-      unsafeAndCodegen: Boolean,
-      buildSide: BuildSide,
-      joinType: JoinType): Unit = {
-    val simpleOrUnsafe = if (!unsafeAndCodegen) "simple" else "unsafe"
-    val testNamePrefix = s"$simpleOrUnsafe / $buildSide / $joinType"
+  private def testJoin(buildSide: BuildSide, joinType: JoinType): Unit = {
+    val testNamePrefix = s"$buildSide / $joinType"
     val someData = (1 to 100).map { i => (i, "burger" + i) }.toArray
     val conf = new SQLConf
-    conf.setConf(SQLConf.UNSAFE_ENABLED, unsafeAndCodegen)
-    conf.setConf(SQLConf.CODEGEN_ENABLED, unsafeAndCodegen)
 
     // Actual test body
     def runTest(
@@ -63,7 +54,7 @@ class NestedLoopJoinNodeSuite extends LocalNodeTest {
         resolveExpressions(
           new NestedLoopJoinNode(conf, node1, node2, buildSide, joinType, Some(cond)))
       }
-      val makeUnsafeNode = if (unsafeAndCodegen) wrapForUnsafe(makeNode) else makeNode
+      val makeUnsafeNode = wrapForUnsafe(makeNode)
       val hashJoinNode = makeUnsafeNode(leftNode, rightNode)
       val expectedOutput = generateExpectedOutput(leftInput, rightInput, joinType)
       val actualOutput = hashJoinNode.collect().map { row =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index 97162249d9951..544c1ef303ae9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -110,33 +110,23 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
   }
 
   test("Project metrics") {
-    withSQLConf(
-      SQLConf.UNSAFE_ENABLED.key -> "false",
-      SQLConf.CODEGEN_ENABLED.key -> "false",
-      SQLConf.TUNGSTEN_ENABLED.key -> "false") {
-      // Assume the execution plan is
-      // PhysicalRDD(nodeId = 1) -> Project(nodeId = 0)
-      val df = person.select('name)
-      testSparkPlanMetrics(df, 1, Map(
-        0L ->("Project", Map(
-          "number of rows" -> 2L)))
-      )
-    }
+    // Assume the execution plan is
+    // PhysicalRDD(nodeId = 1) -> Project(nodeId = 0)
+    val df = person.select('name)
+    testSparkPlanMetrics(df, 1, Map(
+      0L ->("TungstenProject", Map(
+        "number of rows" -> 2L)))
+    )
   }
 
   test("TungstenProject metrics") {
-    withSQLConf(
-      SQLConf.UNSAFE_ENABLED.key -> "true",
-      SQLConf.CODEGEN_ENABLED.key -> "true",
-      SQLConf.TUNGSTEN_ENABLED.key -> "true") {
-      // Assume the execution plan is
-      // PhysicalRDD(nodeId = 1) -> TungstenProject(nodeId = 0)
-      val df = person.select('name)
-      testSparkPlanMetrics(df, 1, Map(
-        0L ->("TungstenProject", Map(
-          "number of rows" -> 2L)))
-      )
-    }
+    // Assume the execution plan is
+    // PhysicalRDD(nodeId = 1) -> TungstenProject(nodeId = 0)
+    val df = person.select('name)
+    testSparkPlanMetrics(df, 1, Map(
+      0L ->("TungstenProject", Map(
+        "number of rows" -> 2L)))
+    )
   }
 
   test("Filter metrics") {
@@ -150,71 +140,30 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
     )
   }
 
-  test("SortBasedAggregate metrics") {
-    // Because SortBasedAggregate may skip different rows if the number of partitions is different,
-    // this test should use the deterministic number of partitions.
-    withSQLConf(
-      SQLConf.UNSAFE_ENABLED.key -> "false",
-      SQLConf.CODEGEN_ENABLED.key -> "true",
-      SQLConf.TUNGSTEN_ENABLED.key -> "true") {
-      // Assume the execution plan is
-      // ... -> SortBasedAggregate(nodeId = 2) -> TungstenExchange(nodeId = 1) ->
-      // SortBasedAggregate(nodeId = 0)
-      val df = testData2.groupBy().count() // 2 partitions
-      testSparkPlanMetrics(df, 1, Map(
-        2L -> ("SortBasedAggregate", Map(
-          "number of input rows" -> 6L,
-          "number of output rows" -> 2L)),
-        0L -> ("SortBasedAggregate", Map(
-          "number of input rows" -> 2L,
-          "number of output rows" -> 1L)))
-      )
-
-      // Assume the execution plan is
-      // ... -> SortBasedAggregate(nodeId = 3) -> TungstenExchange(nodeId = 2)
-      // -> ExternalSort(nodeId = 1)-> SortBasedAggregate(nodeId = 0)
-      // 2 partitions and each partition contains 2 keys
-      val df2 = testData2.groupBy('a).count()
-      testSparkPlanMetrics(df2, 1, Map(
-        3L -> ("SortBasedAggregate", Map(
-          "number of input rows" -> 6L,
-          "number of output rows" -> 4L)),
-        0L -> ("SortBasedAggregate", Map(
-          "number of input rows" -> 4L,
-          "number of output rows" -> 3L)))
-      )
-    }
-  }
-
   test("TungstenAggregate metrics") {
-    withSQLConf(
-      SQLConf.UNSAFE_ENABLED.key -> "true",
-      SQLConf.CODEGEN_ENABLED.key -> "true",
-      SQLConf.TUNGSTEN_ENABLED.key -> "true") {
-      // Assume the execution plan is
-      // ... -> TungstenAggregate(nodeId = 2) -> Exchange(nodeId = 1)
-      // -> TungstenAggregate(nodeId = 0)
-      val df = testData2.groupBy().count() // 2 partitions
-      testSparkPlanMetrics(df, 1, Map(
-        2L -> ("TungstenAggregate", Map(
-          "number of input rows" -> 6L,
-          "number of output rows" -> 2L)),
-        0L -> ("TungstenAggregate", Map(
-          "number of input rows" -> 2L,
-          "number of output rows" -> 1L)))
-      )
+    // Assume the execution plan is
+    // ... -> TungstenAggregate(nodeId = 2) -> Exchange(nodeId = 1)
+    // -> TungstenAggregate(nodeId = 0)
+    val df = testData2.groupBy().count() // 2 partitions
+    testSparkPlanMetrics(df, 1, Map(
+      2L -> ("TungstenAggregate", Map(
+        "number of input rows" -> 6L,
+        "number of output rows" -> 2L)),
+      0L -> ("TungstenAggregate", Map(
+        "number of input rows" -> 2L,
+        "number of output rows" -> 1L)))
+    )
 
-      // 2 partitions and each partition contains 2 keys
-      val df2 = testData2.groupBy('a).count()
-      testSparkPlanMetrics(df2, 1, Map(
-        2L -> ("TungstenAggregate", Map(
-          "number of input rows" -> 6L,
-          "number of output rows" -> 4L)),
-        0L -> ("TungstenAggregate", Map(
-          "number of input rows" -> 4L,
-          "number of output rows" -> 3L)))
-      )
-    }
+    // 2 partitions and each partition contains 2 keys
+    val df2 = testData2.groupBy('a).count()
+    testSparkPlanMetrics(df2, 1, Map(
+      2L -> ("TungstenAggregate", Map(
+        "number of input rows" -> 6L,
+        "number of output rows" -> 4L)),
+      0L -> ("TungstenAggregate", Map(
+        "number of input rows" -> 4L,
+        "number of output rows" -> 3L)))
+    )
   }
 
   test("SortMergeJoin metrics") {
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 22d2aefd699b5..61e3e913c23ea 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -808,54 +808,12 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
   }
 }
 
-class SortBasedAggregationQuerySuite extends AggregationQuerySuite {
 
-  var originalUnsafeEnabled: Boolean = _
+class TungstenAggregationQuerySuite extends AggregationQuerySuite
 
-  override def beforeAll(): Unit = {
-    originalUnsafeEnabled = sqlContext.conf.unsafeEnabled
-    sqlContext.setConf(SQLConf.UNSAFE_ENABLED.key, "false")
-    super.beforeAll()
-  }
-
-  override def afterAll(): Unit = {
-    super.afterAll()
-    sqlContext.setConf(SQLConf.UNSAFE_ENABLED.key, originalUnsafeEnabled.toString)
-  }
-}
-
-class TungstenAggregationQuerySuite extends AggregationQuerySuite {
-
-  var originalUnsafeEnabled: Boolean = _
-
-  override def beforeAll(): Unit = {
-    originalUnsafeEnabled = sqlContext.conf.unsafeEnabled
-    sqlContext.setConf(SQLConf.UNSAFE_ENABLED.key, "true")
-    super.beforeAll()
-  }
-
-  override def afterAll(): Unit = {
-    super.afterAll()
-    sqlContext.setConf(SQLConf.UNSAFE_ENABLED.key, originalUnsafeEnabled.toString)
-  }
-}
 
 class TungstenAggregationQueryWithControlledFallbackSuite extends AggregationQuerySuite {
 
-  var originalUnsafeEnabled: Boolean = _
-
-  override def beforeAll(): Unit = {
-    originalUnsafeEnabled = sqlContext.conf.unsafeEnabled
-    sqlContext.setConf(SQLConf.UNSAFE_ENABLED.key, "true")
-    super.beforeAll()
-  }
-
-  override def afterAll(): Unit = {
-    super.afterAll()
-    sqlContext.setConf(SQLConf.UNSAFE_ENABLED.key, originalUnsafeEnabled.toString)
-    sqlContext.conf.unsetConf("spark.sql.TungstenAggregate.testFallbackStartsAt")
-  }
-
   override protected def checkAnswer(actual: => DataFrame, expectedAnswer: Seq[Row]): Unit = {
     (0 to 2).foreach { fallbackStartsAt =>
       sqlContext.setConf(
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveExplainSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveExplainSuite.scala
index 94162da4eae1a..a7b7ad0093915 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveExplainSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveExplainSuite.scala
@@ -37,8 +37,7 @@ class HiveExplainSuite extends QueryTest with SQLTestUtils with TestHiveSingleto
                    "== Parsed Logical Plan ==",
                    "== Analyzed Logical Plan ==",
                    "== Optimized Logical Plan ==",
-                   "== Physical Plan ==",
-                   "Code Generation")
+                   "== Physical Plan ==")
   }
 
   test("explain create table command") {
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala
index 5f9a447759b48..5ab477efc4ee0 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala
@@ -28,11 +28,11 @@ import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectIn
 import org.apache.hadoop.hive.serde2.objectinspector.{ObjectInspector, ObjectInspectorFactory}
 import org.apache.hadoop.hive.serde2.{AbstractSerDe, SerDeStats}
 import org.apache.hadoop.io.Writable
-import org.apache.spark.sql.{AnalysisException, QueryTest, Row, SQLConf}
+import org.apache.spark.sql.{AnalysisException, QueryTest, Row}
 import org.apache.spark.sql.hive.test.TestHiveSingleton
-
 import org.apache.spark.util.Utils
 
+
 case class Fields(f1: Int, f2: Int, f3: Int, f4: Int, f5: Int)
 
 // Case classes for the custom UDF's.
@@ -92,44 +92,36 @@ class HiveUDFSuite extends QueryTest with TestHiveSingleton {
   }
 
   test("Max/Min on named_struct") {
-    def testOrderInStruct(): Unit = {
-      checkAnswer(sql(
-        """
-          |SELECT max(named_struct(
-          |           "key", key,
-          |           "value", value)).value FROM src
-        """.stripMargin), Seq(Row("val_498")))
-      checkAnswer(sql(
-        """
-          |SELECT min(named_struct(
-          |           "key", key,
-          |           "value", value)).value FROM src
-        """.stripMargin), Seq(Row("val_0")))
-
-      // nested struct cases
-      checkAnswer(sql(
-        """
-          |SELECT max(named_struct(
-          |           "key", named_struct(
-                              "key", key,
-                              "value", value),
-          |           "value", value)).value FROM src
-        """.stripMargin), Seq(Row("val_498")))
-      checkAnswer(sql(
-        """
-          |SELECT min(named_struct(
-          |           "key", named_struct(
-                             "key", key,
-                             "value", value),
-          |           "value", value)).value FROM src
-        """.stripMargin), Seq(Row("val_0")))
-    }
-    val codegenDefault = hiveContext.getConf(SQLConf.CODEGEN_ENABLED)
-    hiveContext.setConf(SQLConf.CODEGEN_ENABLED, true)
-    testOrderInStruct()
-    hiveContext.setConf(SQLConf.CODEGEN_ENABLED, false)
-    testOrderInStruct()
-    hiveContext.setConf(SQLConf.CODEGEN_ENABLED, codegenDefault)
+    checkAnswer(sql(
+      """
+        |SELECT max(named_struct(
+        |           "key", key,
+        |           "value", value)).value FROM src
+      """.stripMargin), Seq(Row("val_498")))
+    checkAnswer(sql(
+      """
+        |SELECT min(named_struct(
+        |           "key", key,
+        |           "value", value)).value FROM src
+      """.stripMargin), Seq(Row("val_0")))
+
+    // nested struct cases
+    checkAnswer(sql(
+      """
+        |SELECT max(named_struct(
+        |           "key", named_struct(
+                            "key", key,
+                            "value", value),
+        |           "value", value)).value FROM src
+      """.stripMargin), Seq(Row("val_498")))
+    checkAnswer(sql(
+      """
+        |SELECT min(named_struct(
+        |           "key", named_struct(
+                           "key", key,
+                           "value", value),
+        |           "value", value)).value FROM src
+      """.stripMargin), Seq(Row("val_0")))
   }
 
   test("SPARK-6409 UDAF Average test") {

From a9a6b80c718008aac7c411dfe46355efe58dee2e Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 11 Nov 2015 12:48:51 -0800
Subject: [PATCH 196/518] [SPARK-11645][SQL] Remove OpenHashSet for the old
 aggregate.

Author: Reynold Xin <rxin@databricks.com>

Closes #9621 from rxin/SPARK-11645.
---
 .../expressions/codegen/CodeGenerator.scala   |   6 -
 .../codegen/GenerateUnsafeProjection.scala    |   7 +-
 .../spark/sql/catalyst/expressions/sets.scala | 194 ------------------
 .../sql/execution/SparkSqlSerializer.scala    | 103 +---------
 .../spark/sql/UserDefinedTypeSuite.scala      |  11 -
 5 files changed, 5 insertions(+), 316 deletions(-)
 delete mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/sets.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
index 5a4bba232b04b..ccd91d3549b53 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
@@ -33,10 +33,6 @@ import org.apache.spark.unsafe.Platform
 import org.apache.spark.unsafe.types._
 
 
-// These classes are here to avoid issues with serialization and integration with quasiquotes.
-class IntegerHashSet extends org.apache.spark.util.collection.OpenHashSet[Int]
-class LongHashSet extends org.apache.spark.util.collection.OpenHashSet[Long]
-
 /**
  * Java source for evaluating an [[Expression]] given a [[InternalRow]] of input.
  *
@@ -205,8 +201,6 @@ class CodeGenContext {
     case _: StructType => "InternalRow"
     case _: ArrayType => "ArrayData"
     case _: MapType => "MapData"
-    case dt: OpenHashSetUDT if dt.elementType == IntegerType => classOf[IntegerHashSet].getName
-    case dt: OpenHashSetUDT if dt.elementType == LongType => classOf[LongHashSet].getName
     case udt: UserDefinedType[_] => javaType(udt.sqlType)
     case ObjectType(cls) if cls.isArray => s"${javaType(ObjectType(cls.getComponentType))}[]"
     case ObjectType(cls) => cls.getName
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
index 9ef226141421b..4c17d02a23725 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
@@ -39,7 +39,6 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     case t: StructType => t.toSeq.forall(field => canSupport(field.dataType))
     case t: ArrayType if canSupport(t.elementType) => true
     case MapType(kt, vt, _) if canSupport(kt) && canSupport(vt) => true
-    case dt: OpenHashSetUDT => false  // it's not a standard UDT
     case udt: UserDefinedType[_] => canSupport(udt.sqlType)
     case _ => false
   }
@@ -309,13 +308,13 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     in.map(BindReferences.bindReference(_, inputSchema))
 
   def generate(
-    expressions: Seq[Expression],
-    subexpressionEliminationEnabled: Boolean): UnsafeProjection = {
+      expressions: Seq[Expression],
+      subexpressionEliminationEnabled: Boolean): UnsafeProjection = {
     create(canonicalize(expressions), subexpressionEliminationEnabled)
   }
 
   protected def create(expressions: Seq[Expression]): UnsafeProjection = {
-    create(expressions, false)
+    create(expressions, subexpressionEliminationEnabled = false)
   }
 
   private def create(
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/sets.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/sets.scala
deleted file mode 100644
index d124d29d534b8..0000000000000
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/sets.scala
+++ /dev/null
@@ -1,194 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.catalyst.expressions
-
-import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.codegen._
-import org.apache.spark.sql.types._
-import org.apache.spark.util.collection.OpenHashSet
-
-/** The data type for expressions returning an OpenHashSet as the result. */
-private[sql] class OpenHashSetUDT(
-    val elementType: DataType) extends UserDefinedType[OpenHashSet[Any]] {
-
-  override def sqlType: DataType = ArrayType(elementType)
-
-  /** Since we are using OpenHashSet internally, usually it will not be called. */
-  override def serialize(obj: Any): Seq[Any] = {
-    obj.asInstanceOf[OpenHashSet[Any]].iterator.toSeq
-  }
-
-  /** Since we are using OpenHashSet internally, usually it will not be called. */
-  override def deserialize(datum: Any): OpenHashSet[Any] = {
-    val iterator = datum.asInstanceOf[Seq[Any]].iterator
-    val set = new OpenHashSet[Any]
-    while(iterator.hasNext) {
-      set.add(iterator.next())
-    }
-
-    set
-  }
-
-  override def userClass: Class[OpenHashSet[Any]] = classOf[OpenHashSet[Any]]
-
-  private[spark] override def asNullable: OpenHashSetUDT = this
-}
-
-/**
- * Creates a new set of the specified type
- */
-case class NewSet(elementType: DataType) extends LeafExpression with CodegenFallback {
-
-  override def nullable: Boolean = false
-
-  override def dataType: OpenHashSetUDT = new OpenHashSetUDT(elementType)
-
-  override def eval(input: InternalRow): Any = {
-    new OpenHashSet[Any]()
-  }
-
-  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    elementType match {
-      case IntegerType | LongType =>
-        ev.isNull = "false"
-        s"""
-          ${ctx.javaType(dataType)} ${ev.value} = new ${ctx.javaType(dataType)}();
-        """
-      case _ => super.genCode(ctx, ev)
-    }
-  }
-
-  override def toString: String = s"new Set($dataType)"
-}
-
-/**
- * Adds an item to a set.
- * For performance, this expression mutates its input during evaluation.
- * Note: this expression is internal and created only by the GeneratedAggregate,
- * we don't need to do type check for it.
- */
-case class AddItemToSet(item: Expression, set: Expression)
-  extends Expression with CodegenFallback {
-
-  override def children: Seq[Expression] = item :: set :: Nil
-
-  override def nullable: Boolean = set.nullable
-
-  override def dataType: DataType = set.dataType
-
-  override def eval(input: InternalRow): Any = {
-    val itemEval = item.eval(input)
-    val setEval = set.eval(input).asInstanceOf[OpenHashSet[Any]]
-
-    if (itemEval != null) {
-      if (setEval != null) {
-        setEval.add(itemEval)
-        setEval
-      } else {
-        null
-      }
-    } else {
-      setEval
-    }
-  }
-
-  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    val elementType = set.dataType.asInstanceOf[OpenHashSetUDT].elementType
-    elementType match {
-      case IntegerType | LongType =>
-        val itemEval = item.gen(ctx)
-        val setEval = set.gen(ctx)
-        val htype = ctx.javaType(dataType)
-
-        ev.isNull = "false"
-        ev.value = setEval.value
-        itemEval.code + setEval.code +  s"""
-          if (!${itemEval.isNull} && !${setEval.isNull}) {
-           (($htype)${setEval.value}).add(${itemEval.value});
-          }
-         """
-      case _ => super.genCode(ctx, ev)
-    }
-  }
-
-  override def toString: String = s"$set += $item"
-}
-
-/**
- * Combines the elements of two sets.
- * For performance, this expression mutates its left input set during evaluation.
- * Note: this expression is internal and created only by the GeneratedAggregate,
- * we don't need to do type check for it.
- */
-case class CombineSets(left: Expression, right: Expression)
-  extends BinaryExpression with CodegenFallback {
-
-  override def nullable: Boolean = left.nullable
-  override def dataType: DataType = left.dataType
-
-  override def eval(input: InternalRow): Any = {
-    val leftEval = left.eval(input).asInstanceOf[OpenHashSet[Any]]
-    if(leftEval != null) {
-      val rightEval = right.eval(input).asInstanceOf[OpenHashSet[Any]]
-      if (rightEval != null) {
-        val iterator = rightEval.iterator
-        while(iterator.hasNext) {
-          val rightValue = iterator.next()
-          leftEval.add(rightValue)
-        }
-      }
-      leftEval
-    } else {
-      null
-    }
-  }
-
-  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    val elementType = left.dataType.asInstanceOf[OpenHashSetUDT].elementType
-    elementType match {
-      case IntegerType | LongType =>
-        val leftEval = left.gen(ctx)
-        val rightEval = right.gen(ctx)
-        val htype = ctx.javaType(dataType)
-
-        ev.isNull = leftEval.isNull
-        ev.value = leftEval.value
-        leftEval.code + rightEval.code + s"""
-          if (!${leftEval.isNull} && !${rightEval.isNull}) {
-            ${leftEval.value}.union((${htype})${rightEval.value});
-          }
-        """
-      case _ => super.genCode(ctx, ev)
-    }
-  }
-}
-
-/**
- * Returns the number of elements in the input set.
- * Note: this expression is internal and created only by the GeneratedAggregate,
- * we don't need to do type check for it.
- */
-case class CountSet(child: Expression) extends UnaryExpression with CodegenFallback {
-
-  override def dataType: DataType = LongType
-
-  protected override def nullSafeEval(input: Any): Any =
-    input.asInstanceOf[OpenHashSet[Any]].size.toLong
-
-  override def toString: String = s"$child.count()"
-}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlSerializer.scala
index b19ad4f1c563e..8317f648ccb4e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlSerializer.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlSerializer.scala
@@ -22,19 +22,16 @@ import java.util.{HashMap => JavaHashMap}
 
 import scala.reflect.ClassTag
 
-import com.clearspring.analytics.stream.cardinality.HyperLogLog
 import com.esotericsoftware.kryo.io.{Input, Output}
 import com.esotericsoftware.kryo.{Kryo, Serializer}
 import com.twitter.chill.ResourcePool
 
 import org.apache.spark.serializer.{KryoSerializer, SerializerInstance}
-import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
-import org.apache.spark.sql.catalyst.expressions.codegen.{IntegerHashSet, LongHashSet}
 import org.apache.spark.sql.types.Decimal
 import org.apache.spark.util.MutablePair
-import org.apache.spark.util.collection.OpenHashSet
 import org.apache.spark.{SparkConf, SparkEnv}
 
+
 private[sql] class SparkSqlSerializer(conf: SparkConf) extends KryoSerializer(conf) {
   override def newKryo(): Kryo = {
     val kryo = super.newKryo()
@@ -43,16 +40,9 @@ private[sql] class SparkSqlSerializer(conf: SparkConf) extends KryoSerializer(co
     kryo.register(classOf[org.apache.spark.sql.catalyst.expressions.GenericRow])
     kryo.register(classOf[org.apache.spark.sql.catalyst.expressions.GenericInternalRow])
     kryo.register(classOf[org.apache.spark.sql.catalyst.expressions.GenericMutableRow])
-    kryo.register(classOf[com.clearspring.analytics.stream.cardinality.HyperLogLog],
-                  new HyperLogLogSerializer)
     kryo.register(classOf[java.math.BigDecimal], new JavaBigDecimalSerializer)
     kryo.register(classOf[BigDecimal], new ScalaBigDecimalSerializer)
 
-    // Specific hashsets must come first TODO: Move to core.
-    kryo.register(classOf[IntegerHashSet], new IntegerHashSetSerializer)
-    kryo.register(classOf[LongHashSet], new LongHashSetSerializer)
-    kryo.register(classOf[org.apache.spark.util.collection.OpenHashSet[_]],
-                  new OpenHashSetSerializer)
     kryo.register(classOf[Decimal])
     kryo.register(classOf[JavaHashMap[_, _]])
 
@@ -62,7 +52,7 @@ private[sql] class SparkSqlSerializer(conf: SparkConf) extends KryoSerializer(co
 }
 
 private[execution] class KryoResourcePool(size: Int)
-    extends ResourcePool[SerializerInstance](size) {
+  extends ResourcePool[SerializerInstance](size) {
 
   val ser: SparkSqlSerializer = {
     val sparkConf = Option(SparkEnv.get).map(_.conf).getOrElse(new SparkConf())
@@ -116,92 +106,3 @@ private[sql] class ScalaBigDecimalSerializer extends Serializer[BigDecimal] {
     new java.math.BigDecimal(input.readString())
   }
 }
-
-private[sql] class HyperLogLogSerializer extends Serializer[HyperLogLog] {
-  def write(kryo: Kryo, output: Output, hyperLogLog: HyperLogLog) {
-    val bytes = hyperLogLog.getBytes()
-    output.writeInt(bytes.length)
-    output.writeBytes(bytes)
-  }
-
-  def read(kryo: Kryo, input: Input, tpe: Class[HyperLogLog]): HyperLogLog = {
-    val length = input.readInt()
-    val bytes = input.readBytes(length)
-    HyperLogLog.Builder.build(bytes)
-  }
-}
-
-private[sql] class OpenHashSetSerializer extends Serializer[OpenHashSet[_]] {
-  def write(kryo: Kryo, output: Output, hs: OpenHashSet[_]) {
-    val rowSerializer = kryo.getDefaultSerializer(classOf[Array[Any]]).asInstanceOf[Serializer[Any]]
-    output.writeInt(hs.size)
-    val iterator = hs.iterator
-    while(iterator.hasNext) {
-      val row = iterator.next()
-      rowSerializer.write(kryo, output, row.asInstanceOf[GenericInternalRow].values)
-    }
-  }
-
-  def read(kryo: Kryo, input: Input, tpe: Class[OpenHashSet[_]]): OpenHashSet[_] = {
-    val rowSerializer = kryo.getDefaultSerializer(classOf[Array[Any]]).asInstanceOf[Serializer[Any]]
-    val numItems = input.readInt()
-    val set = new OpenHashSet[Any](numItems + 1)
-    var i = 0
-    while (i < numItems) {
-      val row =
-        new GenericInternalRow(rowSerializer.read(
-          kryo,
-          input,
-          classOf[Array[Any]].asInstanceOf[Class[Any]]).asInstanceOf[Array[Any]])
-      set.add(row)
-      i += 1
-    }
-    set
-  }
-}
-
-private[sql] class IntegerHashSetSerializer extends Serializer[IntegerHashSet] {
-  def write(kryo: Kryo, output: Output, hs: IntegerHashSet) {
-    output.writeInt(hs.size)
-    val iterator = hs.iterator
-    while(iterator.hasNext) {
-      val value: Int = iterator.next()
-      output.writeInt(value)
-    }
-  }
-
-  def read(kryo: Kryo, input: Input, tpe: Class[IntegerHashSet]): IntegerHashSet = {
-    val numItems = input.readInt()
-    val set = new IntegerHashSet
-    var i = 0
-    while (i < numItems) {
-      val value = input.readInt()
-      set.add(value)
-      i += 1
-    }
-    set
-  }
-}
-
-private[sql] class LongHashSetSerializer extends Serializer[LongHashSet] {
-  def write(kryo: Kryo, output: Output, hs: LongHashSet) {
-    output.writeInt(hs.size)
-    val iterator = hs.iterator
-    while(iterator.hasNext) {
-      val value = iterator.next()
-      output.writeLong(value)
-    }
-  }
-
-  def read(kryo: Kryo, input: Input, tpe: Class[LongHashSet]): LongHashSet = {
-    val numItems = input.readInt()
-    val set = new LongHashSet
-    var i = 0
-    while (i < numItems) {
-      val value = input.readLong()
-      set.add(value)
-      i += 1
-    }
-    set
-  }
-}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
index e31c528f3a633..f602f2fb89ca5 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala
@@ -23,7 +23,6 @@ import scala.beans.{BeanInfo, BeanProperty}
 
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.CatalystTypeConverters
-import org.apache.spark.sql.catalyst.expressions.OpenHashSetUDT
 import org.apache.spark.sql.execution.datasources.parquet.ParquetTest
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
@@ -131,15 +130,6 @@ class UserDefinedTypeSuite extends QueryTest with SharedSQLContext with ParquetT
     df.orderBy('int).limit(1).groupBy('int).agg(first('vec)).collect()(0).getAs[MyDenseVector](0)
   }
 
-  test("OpenHashSetUDT") {
-    val openHashSetUDT = new OpenHashSetUDT(IntegerType)
-    val set = new OpenHashSet[Int]
-    (1 to 10).foreach(i => set.add(i))
-
-    val actual = openHashSetUDT.deserialize(openHashSetUDT.serialize(set))
-    assert(actual.iterator.toSet === set.iterator.toSet)
-  }
-
   test("UDTs with JSON") {
     val data = Seq(
       "{\"id\":1,\"vec\":[1.1,2.2,3.3,4.4]}",
@@ -163,7 +153,6 @@ class UserDefinedTypeSuite extends QueryTest with SharedSQLContext with ParquetT
   test("SPARK-10472 UserDefinedType.typeName") {
     assert(IntegerType.typeName === "integer")
     assert(new MyDenseVectorUDT().typeName === "mydensevector")
-    assert(new OpenHashSetUDT(IntegerType).typeName === "openhashset")
   }
 
   test("Catalyst type converter null handling for UDTs") {

From dd77e278b99e45c20fdefb1c795f3c5148d577db Mon Sep 17 00:00:00 2001
From: Nick Evans <me@nicolasevans.org>
Date: Wed, 11 Nov 2015 13:29:30 -0800
Subject: [PATCH 197/518] [SPARK-11335][STREAMING] update kafka direct python
 docs on how to get the offset ranges for a KafkaRDD

tdas koeninger

This updates the Spark Streaming + Kafka Integration Guide doc with a working method to access the offsets of a `KafkaRDD` through Python.

Author: Nick Evans <me@nicolasevans.org>

Closes #9289 from manygrams/update_kafka_direct_python_docs.
---
 docs/streaming-kafka-integration.md | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/docs/streaming-kafka-integration.md b/docs/streaming-kafka-integration.md
index ab7f0117c0b7f..b00351b2fbcc0 100644
--- a/docs/streaming-kafka-integration.md
+++ b/docs/streaming-kafka-integration.md
@@ -181,7 +181,20 @@ Next, we discuss how to use this approach in your streaming application.
 		);
 	</div>
 	<div data-lang="python" markdown="1">
-		Not supported yet
+		offsetRanges = []
+
+		def storeOffsetRanges(rdd):
+		    global offsetRanges
+		    offsetRanges = rdd.offsetRanges()
+		    return rdd
+
+		def printOffsetRanges(rdd):
+		    for o in offsetRanges:
+		        print "%s %s %s %s" % (o.topic, o.partition, o.fromOffset, o.untilOffset)
+
+		directKafkaStream\
+		    .transform(storeOffsetRanges)\
+		    .foreachRDD(printOffsetRanges)
 	</div>
    	</div>
 

From 2d76e44b1a88e08047806972b2d241a89e499bab Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Wed, 11 Nov 2015 14:30:38 -0800
Subject: [PATCH 198/518] [SPARK-11647] Attempt to reduce time/flakiness of
 Thriftserver CLI and SparkSubmit tests

This patch aims to reduce the test time and flakiness of HiveSparkSubmitSuite, SparkSubmitSuite, and CliSuite.

Key changes:

- Disable IO synchronization calls for Derby writes, since durability doesn't matter for tests. This was done for HiveCompatibilitySuite in #6651 and resulted in huge test speedups.
- Add a few missing `--conf`s to disable various Spark UIs. The CliSuite, in particular, never disabled these UIs, leaving it prone to port-contention-related flakiness.
- Fix two instances where tests defined `beforeAll()` methods which were never called because the appropriate traits were not mixed in. I updated these tests suites to extend `BeforeAndAfterEach` so that they play nicely with our `ResetSystemProperties` trait.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9623 from JoshRosen/SPARK-11647.
---
 .../spark/deploy/RPackageUtilsSuite.scala     | 12 ++++++-----
 .../spark/deploy/SparkSubmitSuite.scala       |  6 ++++--
 .../sql/hive/thriftserver/CliSuite.scala      | 21 ++++++++++++-------
 .../spark/sql/hive/HiveSparkSubmitSuite.scala | 17 +++++++++++----
 4 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/deploy/RPackageUtilsSuite.scala b/core/src/test/scala/org/apache/spark/deploy/RPackageUtilsSuite.scala
index 1ed4bae3ca21e..cc30ba223e1c3 100644
--- a/core/src/test/scala/org/apache/spark/deploy/RPackageUtilsSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/RPackageUtilsSuite.scala
@@ -33,8 +33,12 @@ import org.scalatest.BeforeAndAfterEach
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.api.r.RUtils
 import org.apache.spark.deploy.SparkSubmitUtils.MavenCoordinate
+import org.apache.spark.util.ResetSystemProperties
 
-class RPackageUtilsSuite extends SparkFunSuite with BeforeAndAfterEach {
+class RPackageUtilsSuite
+  extends SparkFunSuite
+  with BeforeAndAfterEach
+  with ResetSystemProperties {
 
   private val main = MavenCoordinate("a", "b", "c")
   private val dep1 = MavenCoordinate("a", "dep1", "c")
@@ -60,11 +64,9 @@ class RPackageUtilsSuite extends SparkFunSuite with BeforeAndAfterEach {
     }
   }
 
-  def beforeAll() {
-    System.setProperty("spark.testing", "true")
-  }
-
   override def beforeEach(): Unit = {
+    super.beforeEach()
+    System.setProperty("spark.testing", "true")
     lineBuffer.clear()
   }
 
diff --git a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
index 1fd470cd3b01d..66a50512003dc 100644
--- a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
@@ -23,7 +23,7 @@ import scala.collection.mutable.ArrayBuffer
 
 import com.google.common.base.Charsets.UTF_8
 import com.google.common.io.ByteStreams
-import org.scalatest.Matchers
+import org.scalatest.{BeforeAndAfterEach, Matchers}
 import org.scalatest.concurrent.Timeouts
 import org.scalatest.time.SpanSugar._
 
@@ -37,10 +37,12 @@ import org.apache.spark.util.{ResetSystemProperties, Utils}
 class SparkSubmitSuite
   extends SparkFunSuite
   with Matchers
+  with BeforeAndAfterEach
   with ResetSystemProperties
   with Timeouts {
 
-  def beforeAll() {
+  override def beforeEach() {
+    super.beforeEach()
     System.setProperty("spark.testing", "true")
   }
 
diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/CliSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/CliSuite.scala
index 3fa5c8528b602..fcf039916913a 100644
--- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/CliSuite.scala
+++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/CliSuite.scala
@@ -27,7 +27,7 @@ import scala.concurrent.{Await, Promise}
 import org.apache.spark.sql.test.ProcessTestUtils.ProcessOutputCapturer
 
 import org.apache.hadoop.hive.conf.HiveConf.ConfVars
-import org.scalatest.BeforeAndAfter
+import org.scalatest.BeforeAndAfterAll
 
 import org.apache.spark.util.Utils
 import org.apache.spark.{Logging, SparkFunSuite}
@@ -36,21 +36,26 @@ import org.apache.spark.{Logging, SparkFunSuite}
  * A test suite for the `spark-sql` CLI tool.  Note that all test cases share the same temporary
  * Hive metastore and warehouse.
  */
-class CliSuite extends SparkFunSuite with BeforeAndAfter with Logging {
+class CliSuite extends SparkFunSuite with BeforeAndAfterAll with Logging {
   val warehousePath = Utils.createTempDir()
   val metastorePath = Utils.createTempDir()
   val scratchDirPath = Utils.createTempDir()
 
-  before {
+  override def beforeAll(): Unit = {
+    super.beforeAll()
     warehousePath.delete()
     metastorePath.delete()
     scratchDirPath.delete()
   }
 
-  after {
-    warehousePath.delete()
-    metastorePath.delete()
-    scratchDirPath.delete()
+  override def afterAll(): Unit = {
+    try {
+      warehousePath.delete()
+      metastorePath.delete()
+      scratchDirPath.delete()
+    } finally {
+      super.afterAll()
+    }
   }
 
   /**
@@ -79,6 +84,8 @@ class CliSuite extends SparkFunSuite with BeforeAndAfter with Logging {
       val jdbcUrl = s"jdbc:derby:;databaseName=$metastorePath;create=true"
       s"""$cliScript
          |  --master local
+         |  --driver-java-options -Dderby.system.durability=test
+         |  --conf spark.ui.enabled=false
          |  --hiveconf ${ConfVars.METASTORECONNECTURLKEY}=$jdbcUrl
          |  --hiveconf ${ConfVars.METASTOREWAREHOUSE}=$warehousePath
          |  --hiveconf ${ConfVars.SCRATCHDIR}=$scratchDirPath
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala
index 10e4ae2c50308..24a3afee148c5 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala
@@ -23,7 +23,7 @@ import java.util.Date
 
 import scala.collection.mutable.ArrayBuffer
 
-import org.scalatest.Matchers
+import org.scalatest.{BeforeAndAfterEach, Matchers}
 import org.scalatest.concurrent.Timeouts
 import org.scalatest.exceptions.TestFailedDueToTimeoutException
 import org.scalatest.time.SpanSugar._
@@ -42,14 +42,14 @@ import org.apache.spark.util.{ResetSystemProperties, Utils}
 class HiveSparkSubmitSuite
   extends SparkFunSuite
   with Matchers
-  // This test suite sometimes gets extremely slow out of unknown reason on Jenkins.  Here we
-  // add a timestamp to provide more diagnosis information.
+  with BeforeAndAfterEach
   with ResetSystemProperties
   with Timeouts {
 
   // TODO: rewrite these or mark them as slow tests to be run sparingly
 
-  def beforeAll() {
+  override def beforeEach() {
+    super.beforeEach()
     System.setProperty("spark.testing", "true")
   }
 
@@ -66,6 +66,7 @@ class HiveSparkSubmitSuite
       "--master", "local-cluster[2,1,1024]",
       "--conf", "spark.ui.enabled=false",
       "--conf", "spark.master.rest.enabled=false",
+      "--driver-java-options", "-Dderby.system.durability=test",
       "--jars", jarsString,
       unusedJar.toString, "SparkSubmitClassA", "SparkSubmitClassB")
     runSparkSubmit(args)
@@ -79,6 +80,7 @@ class HiveSparkSubmitSuite
       "--master", "local-cluster[2,1,1024]",
       "--conf", "spark.ui.enabled=false",
       "--conf", "spark.master.rest.enabled=false",
+      "--driver-java-options", "-Dderby.system.durability=test",
       unusedJar.toString)
     runSparkSubmit(args)
   }
@@ -93,6 +95,7 @@ class HiveSparkSubmitSuite
     val args = Seq(
       "--conf", "spark.ui.enabled=false",
       "--conf", "spark.master.rest.enabled=false",
+      "--driver-java-options", "-Dderby.system.durability=test",
       "--class", "Main",
       testJar)
     runSparkSubmit(args)
@@ -104,6 +107,9 @@ class HiveSparkSubmitSuite
       "--class", SPARK_9757.getClass.getName.stripSuffix("$"),
       "--name", "SparkSQLConfTest",
       "--master", "local-cluster[2,1,1024]",
+      "--conf", "spark.ui.enabled=false",
+      "--conf", "spark.master.rest.enabled=false",
+      "--driver-java-options", "-Dderby.system.durability=test",
       unusedJar.toString)
     runSparkSubmit(args)
   }
@@ -114,6 +120,9 @@ class HiveSparkSubmitSuite
       "--class", SPARK_11009.getClass.getName.stripSuffix("$"),
       "--name", "SparkSQLConfTest",
       "--master", "local-cluster[2,1,1024]",
+      "--conf", "spark.ui.enabled=false",
+      "--conf", "spark.master.rest.enabled=false",
+      "--driver-java-options", "-Dderby.system.durability=test",
       unusedJar.toString)
     runSparkSubmit(args)
   }

From e1bcf6af9ba4f131f84d71660d0ab5598c0b7b67 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 11 Nov 2015 15:30:21 -0800
Subject: [PATCH 199/518] [SPARK-10827] replace volatile with Atomic* in
 AppClient.scala.

This is a followup for #9317 to replace volatile fields with AtomicBoolean and AtomicReference.

Author: Reynold Xin <rxin@databricks.com>

Closes #9611 from rxin/SPARK-10827.
---
 .../spark/deploy/client/AppClient.scala       | 68 ++++++++++---------
 1 file changed, 35 insertions(+), 33 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala b/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
index 3f29da663b798..afab362e213b5 100644
--- a/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.deploy.client
 
 import java.util.concurrent._
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
 import java.util.concurrent.{Future => JFuture, ScheduledFuture => JScheduledFuture}
 
 import scala.util.control.NonFatal
@@ -49,9 +50,9 @@ private[spark] class AppClient(
   private val REGISTRATION_TIMEOUT_SECONDS = 20
   private val REGISTRATION_RETRIES = 3
 
-  @volatile private var endpoint: RpcEndpointRef = null
-  @volatile private var appId: String = null
-  @volatile private var registered = false
+  private val endpoint = new AtomicReference[RpcEndpointRef]
+  private val appId = new AtomicReference[String]
+  private val registered = new AtomicBoolean(false)
 
   private class ClientEndpoint(override val rpcEnv: RpcEnv) extends ThreadSafeRpcEndpoint
     with Logging {
@@ -59,16 +60,17 @@ private[spark] class AppClient(
     private var master: Option[RpcEndpointRef] = None
     // To avoid calling listener.disconnected() multiple times
     private var alreadyDisconnected = false
-    @volatile private var alreadyDead = false // To avoid calling listener.dead() multiple times
-    @volatile private var registerMasterFutures: Array[JFuture[_]] = null
-    @volatile private var registrationRetryTimer: JScheduledFuture[_] = null
+    // To avoid calling listener.dead() multiple times
+    private val alreadyDead = new AtomicBoolean(false)
+    private val registerMasterFutures = new AtomicReference[Array[JFuture[_]]]
+    private val registrationRetryTimer = new AtomicReference[JScheduledFuture[_]]
 
     // A thread pool for registering with masters. Because registering with a master is a blocking
     // action, this thread pool must be able to create "masterRpcAddresses.size" threads at the same
     // time so that we can register with all masters.
     private val registerMasterThreadPool = new ThreadPoolExecutor(
       0,
-      masterRpcAddresses.size, // Make sure we can register with all masters at the same time
+      masterRpcAddresses.length, // Make sure we can register with all masters at the same time
       60L, TimeUnit.SECONDS,
       new SynchronousQueue[Runnable](),
       ThreadUtils.namedThreadFactory("appclient-register-master-threadpool"))
@@ -100,7 +102,7 @@ private[spark] class AppClient(
       for (masterAddress <- masterRpcAddresses) yield {
         registerMasterThreadPool.submit(new Runnable {
           override def run(): Unit = try {
-            if (registered) {
+            if (registered.get) {
               return
             }
             logInfo("Connecting to master " + masterAddress.toSparkURL + "...")
@@ -123,22 +125,22 @@ private[spark] class AppClient(
      * nthRetry means this is the nth attempt to register with master.
      */
     private def registerWithMaster(nthRetry: Int) {
-      registerMasterFutures = tryRegisterAllMasters()
-      registrationRetryTimer = registrationRetryThread.scheduleAtFixedRate(new Runnable {
+      registerMasterFutures.set(tryRegisterAllMasters())
+      registrationRetryTimer.set(registrationRetryThread.scheduleAtFixedRate(new Runnable {
         override def run(): Unit = {
           Utils.tryOrExit {
-            if (registered) {
-              registerMasterFutures.foreach(_.cancel(true))
+            if (registered.get) {
+              registerMasterFutures.get.foreach(_.cancel(true))
               registerMasterThreadPool.shutdownNow()
             } else if (nthRetry >= REGISTRATION_RETRIES) {
               markDead("All masters are unresponsive! Giving up.")
             } else {
-              registerMasterFutures.foreach(_.cancel(true))
+              registerMasterFutures.get.foreach(_.cancel(true))
               registerWithMaster(nthRetry + 1)
             }
           }
         }
-      }, REGISTRATION_TIMEOUT_SECONDS, REGISTRATION_TIMEOUT_SECONDS, TimeUnit.SECONDS)
+      }, REGISTRATION_TIMEOUT_SECONDS, REGISTRATION_TIMEOUT_SECONDS, TimeUnit.SECONDS))
     }
 
     /**
@@ -163,10 +165,10 @@ private[spark] class AppClient(
         // RegisteredApplications due to an unstable network.
         // 2. Receive multiple RegisteredApplication from different masters because the master is
         // changing.
-        appId = appId_
-        registered = true
+        appId.set(appId_)
+        registered.set(true)
         master = Some(masterRef)
-        listener.connected(appId)
+        listener.connected(appId.get)
 
       case ApplicationRemoved(message) =>
         markDead("Master removed our application: %s".format(message))
@@ -178,7 +180,7 @@ private[spark] class AppClient(
           cores))
         // FIXME if changing master and `ExecutorAdded` happen at the same time (the order is not
         // guaranteed), `ExecutorStateChanged` may be sent to a dead master.
-        sendToMaster(ExecutorStateChanged(appId, id, ExecutorState.RUNNING, None, None))
+        sendToMaster(ExecutorStateChanged(appId.get, id, ExecutorState.RUNNING, None, None))
         listener.executorAdded(fullId, workerId, hostPort, cores, memory)
 
       case ExecutorUpdated(id, state, message, exitStatus) =>
@@ -193,13 +195,13 @@ private[spark] class AppClient(
         logInfo("Master has changed, new master is at " + masterRef.address.toSparkURL)
         master = Some(masterRef)
         alreadyDisconnected = false
-        masterRef.send(MasterChangeAcknowledged(appId))
+        masterRef.send(MasterChangeAcknowledged(appId.get))
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
       case StopAppClient =>
         markDead("Application has been stopped.")
-        sendToMaster(UnregisterApplication(appId))
+        sendToMaster(UnregisterApplication(appId.get))
         context.reply(true)
         stop()
 
@@ -263,18 +265,18 @@ private[spark] class AppClient(
     }
 
     def markDead(reason: String) {
-      if (!alreadyDead) {
+      if (!alreadyDead.get) {
         listener.dead(reason)
-        alreadyDead = true
+        alreadyDead.set(true)
       }
     }
 
     override def onStop(): Unit = {
-      if (registrationRetryTimer != null) {
-        registrationRetryTimer.cancel(true)
+      if (registrationRetryTimer.get != null) {
+        registrationRetryTimer.get.cancel(true)
       }
       registrationRetryThread.shutdownNow()
-      registerMasterFutures.foreach(_.cancel(true))
+      registerMasterFutures.get.foreach(_.cancel(true))
       registerMasterThreadPool.shutdownNow()
       askAndReplyThreadPool.shutdownNow()
     }
@@ -283,19 +285,19 @@ private[spark] class AppClient(
 
   def start() {
     // Just launch an rpcEndpoint; it will call back into the listener.
-    endpoint = rpcEnv.setupEndpoint("AppClient", new ClientEndpoint(rpcEnv))
+    endpoint.set(rpcEnv.setupEndpoint("AppClient", new ClientEndpoint(rpcEnv)))
   }
 
   def stop() {
-    if (endpoint != null) {
+    if (endpoint.get != null) {
       try {
         val timeout = RpcUtils.askRpcTimeout(conf)
-        timeout.awaitResult(endpoint.ask[Boolean](StopAppClient))
+        timeout.awaitResult(endpoint.get.ask[Boolean](StopAppClient))
       } catch {
         case e: TimeoutException =>
           logInfo("Stop request to Master timed out; it may already be shut down.")
       }
-      endpoint = null
+      endpoint.set(null)
     }
   }
 
@@ -306,8 +308,8 @@ private[spark] class AppClient(
    * @return whether the request is acknowledged.
    */
   def requestTotalExecutors(requestedTotal: Int): Boolean = {
-    if (endpoint != null && appId != null) {
-      endpoint.askWithRetry[Boolean](RequestExecutors(appId, requestedTotal))
+    if (endpoint.get != null && appId.get != null) {
+      endpoint.get.askWithRetry[Boolean](RequestExecutors(appId.get, requestedTotal))
     } else {
       logWarning("Attempted to request executors before driver fully initialized.")
       false
@@ -319,8 +321,8 @@ private[spark] class AppClient(
    * @return whether the kill request is acknowledged.
    */
   def killExecutors(executorIds: Seq[String]): Boolean = {
-    if (endpoint != null && appId != null) {
-      endpoint.askWithRetry[Boolean](KillExecutors(appId, executorIds))
+    if (endpoint.get != null && appId.get != null) {
+      endpoint.get.askWithRetry[Boolean](KillExecutors(appId.get, executorIds))
     } else {
       logWarning("Attempted to kill executors before driver fully initialized.")
       false

From 1a21be15f655b9696ddac80aac629445a465f621 Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Wed, 11 Nov 2015 15:41:36 -0800
Subject: [PATCH 200/518] [SPARK-11672][ML] disable spark.ml read/write tests

Saw several failures on Jenkins, e.g., https://amplab.cs.berkeley.edu/jenkins/job/NewSparkPullRequestBuilder/2040/testReport/org.apache.spark.ml.util/JavaDefaultReadWriteSuite/testDefaultReadWrite/. This is the first failure in master build:

https://amplab.cs.berkeley.edu/jenkins/job/Spark-Master-SBT/3982/

I cannot reproduce it on local. So temporarily disable the tests and I will look into the issue under the same JIRA. I'm going to merge the PR after Jenkins passes compile.

Author: Xiangrui Meng <meng@databricks.com>

Closes #9641 from mengxr/SPARK-11672.
---
 .../org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java   | 4 ++--
 .../spark/ml/classification/LogisticRegressionSuite.scala     | 2 +-
 .../scala/org/apache/spark/ml/feature/BinarizerSuite.scala    | 2 +-
 .../scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java b/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
index c39538014be81..4f7aeac1ec54c 100644
--- a/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
+++ b/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
@@ -23,7 +23,7 @@
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
-import org.junit.Test;
+import org.junit.Ignore;
 
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.SQLContext;
@@ -50,7 +50,7 @@ public void tearDown() {
     Utils.deleteRecursively(tempDir);
   }
 
-  @Test
+  @Ignore // SPARK-11672
   public void testDefaultReadWrite() throws IOException {
     String uid = "my_params";
     MyParams instance = new MyParams(uid);
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
index 51b06b7eb6d53..e4c2f1baa4fa1 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
@@ -872,7 +872,7 @@ class LogisticRegressionSuite
     assert(model1a0.intercept ~== model1b.intercept absTol 1E-3)
   }
 
-  test("read/write") {
+  ignore("read/write") { // SPARK-11672
     // Set some Params to make sure set Params are serialized.
     val lr = new LogisticRegression()
       .setElasticNetParam(0.1)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
index 9dfa1439cc303..a66fe03281935 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
@@ -68,7 +68,7 @@ class BinarizerSuite extends SparkFunSuite with MLlibTestSparkContext with Defau
     }
   }
 
-  test("read/write") {
+  ignore("read/write") { // SPARK-11672
     val binarizer = new Binarizer()
       .setInputCol("feature")
       .setOutputCol("binarized_feature")
diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
index cac4bd9aa3ab8..44e09c38f9375 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
@@ -105,7 +105,7 @@ object MyParams extends Readable[MyParams] {
 class DefaultReadWriteSuite extends SparkFunSuite with MLlibTestSparkContext
   with DefaultReadWriteTest {
 
-  test("default read/write") {
+  ignore("default read/write") { // SPARK-11672
     val myParams = new MyParams("my_params")
     testDefaultReadWrite(myParams)
   }

From b8ff6888e76b437287d7d6bf2d4b9c759710a195 Mon Sep 17 00:00:00 2001
From: Andrew Ray <ray.andrew@gmail.com>
Date: Wed, 11 Nov 2015 16:23:24 -0800
Subject: [PATCH 201/518] [SPARK-8992][SQL] Add pivot to dataframe api

This adds a pivot method to the dataframe api.

Following the lead of cube and rollup this adds a Pivot operator that is translated into an Aggregate by the analyzer.

Currently the syntax is like:
~~courseSales.pivot(Seq($"year"), $"course", Seq("dotNET", "Java"), sum($"earnings"))~~

~~Would we be interested in the following syntax also/alternatively? and~~

    courseSales.groupBy($"year").pivot($"course", "dotNET", "Java").agg(sum($"earnings"))
    //or
    courseSales.groupBy($"year").pivot($"course").agg(sum($"earnings"))

Later we can add it to `SQLParser`, but as Hive doesn't support it we cant add it there, right?

~~Also what would be the suggested Java friendly method signature for this?~~

Author: Andrew Ray <ray.andrew@gmail.com>

Closes #7841 from aray/sql-pivot.
---
 .../sql/catalyst/analysis/Analyzer.scala      |  42 +++++++
 .../plans/logical/basicOperators.scala        |  14 +++
 .../org/apache/spark/sql/GroupedData.scala    | 103 ++++++++++++++++--
 .../scala/org/apache/spark/sql/SQLConf.scala  |   7 ++
 .../spark/sql/DataFramePivotSuite.scala       |  87 +++++++++++++++
 .../apache/spark/sql/test/SQLTestData.scala   |  12 ++
 6 files changed, 255 insertions(+), 10 deletions(-)
 create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index a9cd9a77038e7..2f4670b55bdba 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -72,6 +72,7 @@ class Analyzer(
       ResolveRelations ::
       ResolveReferences ::
       ResolveGroupingAnalytics ::
+      ResolvePivot ::
       ResolveSortReferences ::
       ResolveGenerate ::
       ResolveFunctions ::
@@ -166,6 +167,10 @@ class Analyzer(
       case g: GroupingAnalytics if g.child.resolved && hasUnresolvedAlias(g.aggregations) =>
         g.withNewAggs(assignAliases(g.aggregations))
 
+      case Pivot(groupByExprs, pivotColumn, pivotValues, aggregates, child)
+        if child.resolved && hasUnresolvedAlias(groupByExprs) =>
+        Pivot(assignAliases(groupByExprs), pivotColumn, pivotValues, aggregates, child)
+
       case Project(projectList, child) if child.resolved && hasUnresolvedAlias(projectList) =>
         Project(assignAliases(projectList), child)
     }
@@ -248,6 +253,43 @@ class Analyzer(
     }
   }
 
+  object ResolvePivot extends Rule[LogicalPlan] {
+    def apply(plan: LogicalPlan): LogicalPlan = plan transform {
+      case p: Pivot if !p.childrenResolved => p
+      case Pivot(groupByExprs, pivotColumn, pivotValues, aggregates, child) =>
+        val singleAgg = aggregates.size == 1
+        val pivotAggregates: Seq[NamedExpression] = pivotValues.flatMap { value =>
+          def ifExpr(expr: Expression) = {
+            If(EqualTo(pivotColumn, value), expr, Literal(null))
+          }
+          aggregates.map { aggregate =>
+            val filteredAggregate = aggregate.transformDown {
+              // Assumption is the aggregate function ignores nulls. This is true for all current
+              // AggregateFunction's with the exception of First and Last in their default mode
+              // (which we handle) and possibly some Hive UDAF's.
+              case First(expr, _) =>
+                First(ifExpr(expr), Literal(true))
+              case Last(expr, _) =>
+                Last(ifExpr(expr), Literal(true))
+              case a: AggregateFunction =>
+                a.withNewChildren(a.children.map(ifExpr))
+            }
+            if (filteredAggregate.fastEquals(aggregate)) {
+              throw new AnalysisException(
+                s"Aggregate expression required for pivot, found '$aggregate'")
+            }
+            val name = if (singleAgg) value.toString else value + "_" + aggregate.prettyString
+            Alias(filteredAggregate, name)()
+          }
+        }
+        val newGroupByExprs = groupByExprs.map {
+          case UnresolvedAlias(e) => e
+          case e => e
+        }
+        Aggregate(newGroupByExprs, groupByExprs ++ pivotAggregates, child)
+    }
+  }
+
   /**
    * Replaces [[UnresolvedRelation]]s with concrete relations from the catalog.
    */
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index 597f03e752707..32b09b59af436 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -386,6 +386,20 @@ case class Rollup(
     this.copy(aggregations = aggs)
 }
 
+case class Pivot(
+    groupByExprs: Seq[NamedExpression],
+    pivotColumn: Expression,
+    pivotValues: Seq[Literal],
+    aggregates: Seq[Expression],
+    child: LogicalPlan) extends UnaryNode {
+  override def output: Seq[Attribute] = groupByExprs.map(_.toAttribute) ++ aggregates match {
+    case agg :: Nil => pivotValues.map(value => AttributeReference(value.toString, agg.dataType)())
+    case _ => pivotValues.flatMap{ value =>
+      aggregates.map(agg => AttributeReference(value + "_" + agg.prettyString, agg.dataType)())
+    }
+  }
+}
+
 case class Limit(limitExpr: Expression, child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
index 5babf2cc0ca25..63dd7fbcbe9e4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
@@ -24,8 +24,8 @@ import org.apache.spark.annotation.Experimental
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunction, UnresolvedAlias, UnresolvedAttribute, Star}
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
-import org.apache.spark.sql.catalyst.plans.logical.{Rollup, Cube, Aggregate}
-import org.apache.spark.sql.types.NumericType
+import org.apache.spark.sql.catalyst.plans.logical.{Pivot, Rollup, Cube, Aggregate}
+import org.apache.spark.sql.types.{StringType, NumericType}
 
 
 /**
@@ -50,14 +50,8 @@ class GroupedData protected[sql](
       aggExprs
     }
 
-    val aliasedAgg = aggregates.map {
-      // Wrap UnresolvedAttribute with UnresolvedAlias, as when we resolve UnresolvedAttribute, we
-      // will remove intermediate Alias for ExtractValue chain, and we need to alias it again to
-      // make it a NamedExpression.
-      case u: UnresolvedAttribute => UnresolvedAlias(u)
-      case expr: NamedExpression => expr
-      case expr: Expression => Alias(expr, expr.prettyString)()
-    }
+    val aliasedAgg = aggregates.map(alias)
+
     groupType match {
       case GroupedData.GroupByType =>
         DataFrame(
@@ -68,9 +62,22 @@ class GroupedData protected[sql](
       case GroupedData.CubeType =>
         DataFrame(
           df.sqlContext, Cube(groupingExprs, df.logicalPlan, aliasedAgg))
+      case GroupedData.PivotType(pivotCol, values) =>
+        val aliasedGrps = groupingExprs.map(alias)
+        DataFrame(
+          df.sqlContext, Pivot(aliasedGrps, pivotCol, values, aggExprs, df.logicalPlan))
     }
   }
 
+  // Wrap UnresolvedAttribute with UnresolvedAlias, as when we resolve UnresolvedAttribute, we
+  // will remove intermediate Alias for ExtractValue chain, and we need to alias it again to
+  // make it a NamedExpression.
+  private[this] def alias(expr: Expression): NamedExpression = expr match {
+    case u: UnresolvedAttribute => UnresolvedAlias(u)
+    case expr: NamedExpression => expr
+    case expr: Expression => Alias(expr, expr.prettyString)()
+  }
+
   private[this] def aggregateNumericColumns(colNames: String*)(f: Expression => AggregateFunction)
     : DataFrame = {
 
@@ -273,6 +280,77 @@ class GroupedData protected[sql](
   def sum(colNames: String*): DataFrame = {
     aggregateNumericColumns(colNames : _*)(Sum)
   }
+
+  /**
+    * (Scala-specific) Pivots a column of the current [[DataFrame]] and preform the specified
+    * aggregation.
+    * {{{
+    *   // Compute the sum of earnings for each year by course with each course as a separate column
+    *   df.groupBy($"year").pivot($"course", "dotNET", "Java").agg(sum($"earnings"))
+    *   // Or without specifying column values
+    *   df.groupBy($"year").pivot($"course").agg(sum($"earnings"))
+    * }}}
+    * @param pivotColumn Column to pivot
+    * @param values Optional list of values of pivotColumn that will be translated to columns in the
+    *               output data frame. If values are not provided the method with do an immediate
+    *               call to .distinct() on the pivot column.
+    * @since 1.6.0
+    */
+  @scala.annotation.varargs
+  def pivot(pivotColumn: Column, values: Column*): GroupedData = groupType match {
+    case _: GroupedData.PivotType =>
+      throw new UnsupportedOperationException("repeated pivots are not supported")
+    case GroupedData.GroupByType =>
+      val pivotValues = if (values.nonEmpty) {
+        values.map {
+          case Column(literal: Literal) => literal
+          case other =>
+            throw new UnsupportedOperationException(
+              s"The values of a pivot must be literals, found $other")
+        }
+      } else {
+        // This is to prevent unintended OOM errors when the number of distinct values is large
+        val maxValues = df.sqlContext.conf.getConf(SQLConf.DATAFRAME_PIVOT_MAX_VALUES)
+        // Get the distinct values of the column and sort them so its consistent
+        val values = df.select(pivotColumn)
+          .distinct()
+          .sort(pivotColumn)
+          .map(_.get(0))
+          .take(maxValues + 1)
+          .map(Literal(_)).toSeq
+        if (values.length > maxValues) {
+          throw new RuntimeException(
+            s"The pivot column $pivotColumn has more than $maxValues distinct values, " +
+              "this could indicate an error. " +
+              "If this was intended, set \"" + SQLConf.DATAFRAME_PIVOT_MAX_VALUES.key + "\" " +
+              s"to at least the number of distinct values of the pivot column.")
+        }
+        values
+      }
+      new GroupedData(df, groupingExprs, GroupedData.PivotType(pivotColumn.expr, pivotValues))
+    case _ =>
+      throw new UnsupportedOperationException("pivot is only supported after a groupBy")
+  }
+
+  /**
+    * Pivots a column of the current [[DataFrame]] and preform the specified aggregation.
+    * {{{
+    *   // Compute the sum of earnings for each year by course with each course as a separate column
+    *   df.groupBy("year").pivot("course", "dotNET", "Java").sum("earnings")
+    *   // Or without specifying column values
+    *   df.groupBy("year").pivot("course").sum("earnings")
+    * }}}
+    * @param pivotColumn Column to pivot
+    * @param values Optional list of values of pivotColumn that will be translated to columns in the
+    *               output data frame. If values are not provided the method with do an immediate
+    *               call to .distinct() on the pivot column.
+    * @since 1.6.0
+    */
+  @scala.annotation.varargs
+  def pivot(pivotColumn: String, values: Any*): GroupedData = {
+    val resolvedPivotColumn = Column(df.resolve(pivotColumn))
+    pivot(resolvedPivotColumn, values.map(functions.lit): _*)
+  }
 }
 
 
@@ -307,4 +385,9 @@ private[sql] object GroupedData {
    * To indicate it's the ROLLUP
    */
   private[sql] object RollupType extends GroupType
+
+  /**
+    * To indicate it's the PIVOT
+    */
+  private[sql] case class PivotType(pivotCol: Expression, values: Seq[Literal]) extends GroupType
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
index e02b502b7b4d5..41d28d448ccc9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
@@ -437,6 +437,13 @@ private[spark] object SQLConf {
     defaultValue = Some(true),
     isPublic = false)
 
+  val DATAFRAME_PIVOT_MAX_VALUES = intConf(
+    "spark.sql.pivotMaxValues",
+    defaultValue = Some(10000),
+    doc = "When doing a pivot without specifying values for the pivot column this is the maximum " +
+      "number of (distinct) values that will be collected without error."
+  )
+
   val RUN_SQL_ON_FILES = booleanConf("spark.sql.runSQLOnFiles",
     defaultValue = Some(true),
     isPublic = false,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala
new file mode 100644
index 0000000000000..0c23d142670c1
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala
@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.test.SharedSQLContext
+
+class DataFramePivotSuite extends QueryTest with SharedSQLContext{
+  import testImplicits._
+
+  test("pivot courses with literals") {
+    checkAnswer(
+      courseSales.groupBy($"year").pivot($"course", lit("dotNET"), lit("Java"))
+        .agg(sum($"earnings")),
+      Row(2012, 15000.0, 20000.0) :: Row(2013, 48000.0, 30000.0) :: Nil
+    )
+  }
+
+  test("pivot year with literals") {
+    checkAnswer(
+      courseSales.groupBy($"course").pivot($"year", lit(2012), lit(2013)).agg(sum($"earnings")),
+      Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
+    )
+  }
+
+  test("pivot courses with literals and multiple aggregations") {
+    checkAnswer(
+      courseSales.groupBy($"year").pivot($"course", lit("dotNET"), lit("Java"))
+        .agg(sum($"earnings"), avg($"earnings")),
+      Row(2012, 15000.0, 7500.0, 20000.0, 20000.0) ::
+        Row(2013, 48000.0, 48000.0, 30000.0, 30000.0) :: Nil
+    )
+  }
+
+  test("pivot year with string values (cast)") {
+    checkAnswer(
+      courseSales.groupBy("course").pivot("year", "2012", "2013").sum("earnings"),
+      Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
+    )
+  }
+
+  test("pivot year with int values") {
+    checkAnswer(
+      courseSales.groupBy("course").pivot("year", 2012, 2013).sum("earnings"),
+      Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
+    )
+  }
+
+  test("pivot courses with no values") {
+    // Note Java comes before dotNet in sorted order
+    checkAnswer(
+      courseSales.groupBy($"year").pivot($"course").agg(sum($"earnings")),
+      Row(2012, 20000.0, 15000.0) :: Row(2013, 30000.0, 48000.0) :: Nil
+    )
+  }
+
+  test("pivot year with no values") {
+    checkAnswer(
+      courseSales.groupBy($"course").pivot($"year").agg(sum($"earnings")),
+      Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
+    )
+  }
+
+  test("pivot max values inforced") {
+    sqlContext.conf.setConf(SQLConf.DATAFRAME_PIVOT_MAX_VALUES, 1)
+    intercept[RuntimeException](
+      courseSales.groupBy($"year").pivot($"course")
+    )
+    sqlContext.conf.setConf(SQLConf.DATAFRAME_PIVOT_MAX_VALUES,
+      SQLConf.DATAFRAME_PIVOT_MAX_VALUES.defaultValue.get)
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala
index 520dea7f7dd92..abad0d7eaaedf 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala
@@ -242,6 +242,17 @@ private[sql] trait SQLTestData { self =>
     df
   }
 
+  protected lazy val courseSales: DataFrame = {
+    val df = sqlContext.sparkContext.parallelize(
+      CourseSales("dotNET", 2012, 10000) ::
+        CourseSales("Java", 2012, 20000) ::
+        CourseSales("dotNET", 2012, 5000) ::
+        CourseSales("dotNET", 2013, 48000) ::
+        CourseSales("Java", 2013, 30000) :: Nil).toDF()
+    df.registerTempTable("courseSales")
+    df
+  }
+
   /**
    * Initialize all test data such that all temp tables are properly registered.
    */
@@ -295,4 +306,5 @@ private[sql] object SQLTestData {
   case class Person(id: Int, name: String, age: Int)
   case class Salary(personId: Int, salary: Double)
   case class ComplexData(m: Map[String, Int], s: TestData, a: Seq[Int], b: Boolean)
+  case class CourseSales(course: String, year: Int, earnings: Double)
 }

From e49e723392b8a64d30bd90944a748eb6f5ef3a8a Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 11 Nov 2015 19:32:52 -0800
Subject: [PATCH 202/518] [SPARK-11675][SQL] Remove shuffle hash joins.

Author: Reynold Xin <rxin@databricks.com>

Closes #9645 from rxin/SPARK-11675.
---
 .../scala/org/apache/spark/sql/SQLConf.scala  |   9 +-
 .../spark/sql/execution/SparkStrategies.scala |  24 +-
 .../apache/spark/sql/execution/commands.scala |  11 +
 .../execution/joins/ShuffledHashJoin.scala    |  62 ---
 .../joins/ShuffledHashOuterJoin.scala         | 109 ----
 .../org/apache/spark/sql/JoinSuite.scala      | 523 ++++++++----------
 .../org/apache/spark/sql/SQLQuerySuite.scala  |   3 +-
 .../spark/sql/execution/PlannerSuite.scala    |  10 +-
 .../sql/execution/joins/InnerJoinSuite.scala  |  38 --
 .../sql/execution/joins/OuterJoinSuite.scala  |  12 -
 .../execution/metric/SQLMetricsSuite.scala    | 271 ++++-----
 .../spark/sql/hive/StatisticsSuite.scala      |   2 +-
 12 files changed, 357 insertions(+), 717 deletions(-)
 delete mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashJoin.scala
 delete mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashOuterJoin.scala

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
index 41d28d448ccc9..f40e603cd1939 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
@@ -353,12 +353,6 @@ private[spark] object SQLConf {
     defaultValue = Some(5 * 60),
     doc = "Timeout in seconds for the broadcast wait time in broadcast joins.")
 
-  // Options that control which operators can be chosen by the query planner.  These should be
-  // considered hints and may be ignored by future versions of Spark SQL.
-  val SORTMERGE_JOIN = booleanConf("spark.sql.planner.sortMergeJoin",
-    defaultValue = Some(true),
-    doc = "When true, use sort merge join (as opposed to hash join) by default for large joins.")
-
   // This is only used for the thriftserver
   val THRIFTSERVER_POOL = stringConf("spark.sql.thriftserver.scheduler.pool",
     doc = "Set a Fair Scheduler pool for a JDBC client session")
@@ -469,6 +463,7 @@ private[spark] object SQLConf {
     val TUNGSTEN_ENABLED = "spark.sql.tungsten.enabled"
     val CODEGEN_ENABLED = "spark.sql.codegen"
     val UNSAFE_ENABLED = "spark.sql.unsafe.enabled"
+    val SORTMERGE_JOIN = "spark.sql.planner.sortMergeJoin"
   }
 }
 
@@ -533,8 +528,6 @@ private[sql] class SQLConf extends Serializable with CatalystConf {
 
   private[spark] def nativeView: Boolean = getConf(NATIVE_VIEW)
 
-  private[spark] def sortMergeJoinEnabled: Boolean = getConf(SORTMERGE_JOIN)
-
   def caseSensitiveAnalysis: Boolean = getConf(SQLConf.CASE_SENSITIVE)
 
   private[spark] def subexpressionEliminationEnabled: Boolean =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index 96242f160aa51..90989f2cee9a6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -73,10 +73,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    *     [[org.apache.spark.sql.functions.broadcast()]] function to a DataFrame), then that side
    *     of the join will be broadcasted and the other side will be streamed, with no shuffling
    *     performed. If both sides of the join are eligible to be broadcasted then the
-   * - Sort merge: if the matching join keys are sortable and
-   *     [[org.apache.spark.sql.SQLConf.SORTMERGE_JOIN]] is enabled (default), then sort merge join
-   *     will be used.
-   * - Hash: will be chosen if neither of the above optimizations apply to this join.
+   * - Sort merge: if the matching join keys are sortable.
    */
   object EquiJoinSelection extends Strategy with PredicateHelper {
 
@@ -103,22 +100,11 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         makeBroadcastHashJoin(leftKeys, rightKeys, left, right, condition, joins.BuildLeft)
 
       case ExtractEquiJoinKeys(Inner, leftKeys, rightKeys, condition, left, right)
-        if sqlContext.conf.sortMergeJoinEnabled && RowOrdering.isOrderable(leftKeys) =>
+        if RowOrdering.isOrderable(leftKeys) =>
         val mergeJoin =
           joins.SortMergeJoin(leftKeys, rightKeys, planLater(left), planLater(right))
         condition.map(Filter(_, mergeJoin)).getOrElse(mergeJoin) :: Nil
 
-      case ExtractEquiJoinKeys(Inner, leftKeys, rightKeys, condition, left, right) =>
-        val buildSide =
-          if (right.statistics.sizeInBytes <= left.statistics.sizeInBytes) {
-            joins.BuildRight
-          } else {
-            joins.BuildLeft
-          }
-        val hashJoin = joins.ShuffledHashJoin(
-          leftKeys, rightKeys, buildSide, planLater(left), planLater(right))
-        condition.map(Filter(_, hashJoin)).getOrElse(hashJoin) :: Nil
-
       // --- Outer joins --------------------------------------------------------------------------
 
       case ExtractEquiJoinKeys(
@@ -132,14 +118,10 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
           leftKeys, rightKeys, RightOuter, condition, planLater(left), planLater(right)) :: Nil
 
       case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, condition, left, right)
-        if sqlContext.conf.sortMergeJoinEnabled && RowOrdering.isOrderable(leftKeys) =>
+        if RowOrdering.isOrderable(leftKeys) =>
         joins.SortMergeOuterJoin(
           leftKeys, rightKeys, joinType, condition, planLater(left), planLater(right)) :: Nil
 
-      case ExtractEquiJoinKeys(joinType, leftKeys, rightKeys, condition, left, right) =>
-        joins.ShuffledHashOuterJoin(
-          leftKeys, rightKeys, joinType, condition, planLater(left), planLater(right)) :: Nil
-
       // --- Cases where this strategy does not apply ---------------------------------------------
 
       case _ => Nil
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
index e29c281b951f2..24a79f289aa81 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/commands.scala
@@ -148,6 +148,17 @@ case class SetCommand(kv: Option[(String, Option[String])]) extends RunnableComm
       }
       (keyValueOutput, runFunc)
 
+      (keyValueOutput, runFunc)
+
+    case Some((SQLConf.Deprecated.SORTMERGE_JOIN, Some(value))) =>
+      val runFunc = (sqlContext: SQLContext) => {
+        logWarning(
+          s"Property ${SQLConf.Deprecated.SORTMERGE_JOIN} is deprecated and " +
+            s"will be ignored. Sort merge join will continue to be used.")
+        Seq(Row(SQLConf.Deprecated.SORTMERGE_JOIN, "true"))
+      }
+      (keyValueOutput, runFunc)
+
     // Configures a single property.
     case Some((key, Some(value))) =>
       val runFunc = (sqlContext: SQLContext) => {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashJoin.scala
deleted file mode 100644
index 755986af8b95e..0000000000000
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashJoin.scala
+++ /dev/null
@@ -1,62 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.execution.joins
-
-import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.Expression
-import org.apache.spark.sql.catalyst.plans.physical._
-import org.apache.spark.sql.execution.{BinaryNode, SparkPlan}
-import org.apache.spark.sql.execution.metric.SQLMetrics
-
-/**
- * Performs an inner hash join of two child relations by first shuffling the data using the join
- * keys.
- */
-case class ShuffledHashJoin(
-    leftKeys: Seq[Expression],
-    rightKeys: Seq[Expression],
-    buildSide: BuildSide,
-    left: SparkPlan,
-    right: SparkPlan)
-  extends BinaryNode with HashJoin {
-
-  override private[sql] lazy val metrics = Map(
-    "numLeftRows" -> SQLMetrics.createLongMetric(sparkContext, "number of left rows"),
-    "numRightRows" -> SQLMetrics.createLongMetric(sparkContext, "number of right rows"),
-    "numOutputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of output rows"))
-
-  override def outputPartitioning: Partitioning =
-    PartitioningCollection(Seq(left.outputPartitioning, right.outputPartitioning))
-
-  override def requiredChildDistribution: Seq[Distribution] =
-    ClusteredDistribution(leftKeys) :: ClusteredDistribution(rightKeys) :: Nil
-
-  protected override def doExecute(): RDD[InternalRow] = {
-    val (numBuildRows, numStreamedRows) = buildSide match {
-      case BuildLeft => (longMetric("numLeftRows"), longMetric("numRightRows"))
-      case BuildRight => (longMetric("numRightRows"), longMetric("numLeftRows"))
-    }
-    val numOutputRows = longMetric("numOutputRows")
-
-    buildPlan.execute().zipPartitions(streamedPlan.execute()) { (buildIter, streamIter) =>
-      val hashed = HashedRelation(buildIter, numBuildRows, buildSideKeyGenerator)
-      hashJoin(streamIter, numStreamedRows, hashed, numOutputRows)
-    }
-  }
-}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashOuterJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashOuterJoin.scala
deleted file mode 100644
index 6b2cb9d8f6893..0000000000000
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledHashOuterJoin.scala
+++ /dev/null
@@ -1,109 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.execution.joins
-
-import scala.collection.JavaConverters._
-
-import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.plans.physical._
-import org.apache.spark.sql.catalyst.plans.{FullOuter, JoinType, LeftOuter, RightOuter}
-import org.apache.spark.sql.execution.{BinaryNode, SparkPlan}
-import org.apache.spark.sql.execution.metric.SQLMetrics
-
-/**
- * Performs a hash based outer join for two child relations by shuffling the data using
- * the join keys. This operator requires loading the associated partition in both side into memory.
- */
-case class ShuffledHashOuterJoin(
-    leftKeys: Seq[Expression],
-    rightKeys: Seq[Expression],
-    joinType: JoinType,
-    condition: Option[Expression],
-    left: SparkPlan,
-    right: SparkPlan) extends BinaryNode with HashOuterJoin {
-
-  override private[sql] lazy val metrics = Map(
-    "numLeftRows" -> SQLMetrics.createLongMetric(sparkContext, "number of left rows"),
-    "numRightRows" -> SQLMetrics.createLongMetric(sparkContext, "number of right rows"),
-    "numOutputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of output rows"))
-
-  override def requiredChildDistribution: Seq[Distribution] =
-    ClusteredDistribution(leftKeys) :: ClusteredDistribution(rightKeys) :: Nil
-
-  override def outputPartitioning: Partitioning = joinType match {
-    case LeftOuter => left.outputPartitioning
-    case RightOuter => right.outputPartitioning
-    case FullOuter => UnknownPartitioning(left.outputPartitioning.numPartitions)
-    case x =>
-      throw new IllegalArgumentException(s"HashOuterJoin should not take $x as the JoinType")
-  }
-
-  protected override def doExecute(): RDD[InternalRow] = {
-    val numLeftRows = longMetric("numLeftRows")
-    val numRightRows = longMetric("numRightRows")
-    val numOutputRows = longMetric("numOutputRows")
-
-    val joinedRow = new JoinedRow()
-    left.execute().zipPartitions(right.execute()) { (leftIter, rightIter) =>
-      // TODO this probably can be replaced by external sort (sort merged join?)
-      joinType match {
-        case LeftOuter =>
-          val hashed = HashedRelation(rightIter, numRightRows, buildKeyGenerator)
-          val keyGenerator = streamedKeyGenerator
-          val resultProj = resultProjection
-          leftIter.flatMap( currentRow => {
-            numLeftRows += 1
-            val rowKey = keyGenerator(currentRow)
-            joinedRow.withLeft(currentRow)
-            leftOuterIterator(rowKey, joinedRow, hashed.get(rowKey), resultProj, numOutputRows)
-          })
-
-        case RightOuter =>
-          val hashed = HashedRelation(leftIter, numLeftRows, buildKeyGenerator)
-          val keyGenerator = streamedKeyGenerator
-          val resultProj = resultProjection
-          rightIter.flatMap ( currentRow => {
-            numRightRows += 1
-            val rowKey = keyGenerator(currentRow)
-            joinedRow.withRight(currentRow)
-            rightOuterIterator(rowKey, hashed.get(rowKey), joinedRow, resultProj, numOutputRows)
-          })
-
-        case FullOuter =>
-          // TODO(davies): use UnsafeRow
-          val leftHashTable =
-            buildHashTable(leftIter, numLeftRows, newProjection(leftKeys, left.output)).asScala
-          val rightHashTable =
-            buildHashTable(rightIter, numRightRows, newProjection(rightKeys, right.output)).asScala
-          (leftHashTable.keySet ++ rightHashTable.keySet).iterator.flatMap { key =>
-            fullOuterIterator(key,
-              leftHashTable.getOrElse(key, EMPTY_LIST),
-              rightHashTable.getOrElse(key, EMPTY_LIST),
-              joinedRow,
-              numOutputRows)
-          }
-
-        case x =>
-          throw new IllegalArgumentException(
-            s"ShuffledHashOuterJoin should not take $x as the JoinType")
-      }
-    }
-  }
-}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
index 3f3b837f7581c..9a3c262e9485d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala
@@ -44,8 +44,6 @@ class JoinSuite extends QueryTest with SharedSQLContext {
     val df = sql(sqlString)
     val physical = df.queryExecution.sparkPlan
     val operators = physical.collect {
-      case j: ShuffledHashJoin => j
-      case j: ShuffledHashOuterJoin => j
       case j: LeftSemiJoinHash => j
       case j: BroadcastHashJoin => j
       case j: BroadcastHashOuterJoin => j
@@ -96,75 +94,39 @@ class JoinSuite extends QueryTest with SharedSQLContext {
       ("SELECT * FROM testData full JOIN testData2 ON (key * a != key + a)",
         classOf[BroadcastNestedLoopJoin])
     ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "false") {
-      Seq(
-        ("SELECT * FROM testData JOIN testData2 ON key = a", classOf[ShuffledHashJoin]),
-        ("SELECT * FROM testData JOIN testData2 ON key = a and key = 2",
-          classOf[ShuffledHashJoin]),
-        ("SELECT * FROM testData JOIN testData2 ON key = a where key = 2",
-          classOf[ShuffledHashJoin]),
-        ("SELECT * FROM testData LEFT JOIN testData2 ON key = a", classOf[ShuffledHashOuterJoin]),
-        ("SELECT * FROM testData RIGHT JOIN testData2 ON key = a where key = 2",
-          classOf[ShuffledHashOuterJoin]),
-        ("SELECT * FROM testData right join testData2 ON key = a and key = 2",
-          classOf[ShuffledHashOuterJoin]),
-        ("SELECT * FROM testData full outer join testData2 ON key = a",
-          classOf[ShuffledHashOuterJoin])
-      ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
-    }
   }
 
-  test("SortMergeJoin shouldn't work on unsortable columns") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "true") {
-      Seq(
-        ("SELECT * FROM arrayData JOIN complexData ON data = a", classOf[ShuffledHashJoin])
-      ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
-    }
-  }
+//  ignore("SortMergeJoin shouldn't work on unsortable columns") {
+//    Seq(
+//      ("SELECT * FROM arrayData JOIN complexData ON data = a", classOf[ShuffledHashJoin])
+//    ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
+//  }
 
   test("broadcasted hash join operator selection") {
     sqlContext.cacheManager.clearCache()
     sql("CACHE TABLE testData")
-    for (sortMergeJoinEnabled <- Seq(true, false)) {
-      withClue(s"sortMergeJoinEnabled=$sortMergeJoinEnabled") {
-        withSQLConf(SQLConf.SORTMERGE_JOIN.key -> s"$sortMergeJoinEnabled") {
-          Seq(
-            ("SELECT * FROM testData join testData2 ON key = a",
-              classOf[BroadcastHashJoin]),
-            ("SELECT * FROM testData join testData2 ON key = a and key = 2",
-              classOf[BroadcastHashJoin]),
-            ("SELECT * FROM testData join testData2 ON key = a where key = 2",
-              classOf[BroadcastHashJoin])
-          ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
-        }
-      }
-    }
+    Seq(
+      ("SELECT * FROM testData join testData2 ON key = a",
+        classOf[BroadcastHashJoin]),
+      ("SELECT * FROM testData join testData2 ON key = a and key = 2",
+        classOf[BroadcastHashJoin]),
+      ("SELECT * FROM testData join testData2 ON key = a where key = 2",
+        classOf[BroadcastHashJoin])
+    ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
     sql("UNCACHE TABLE testData")
   }
 
   test("broadcasted hash outer join operator selection") {
     sqlContext.cacheManager.clearCache()
     sql("CACHE TABLE testData")
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "true") {
-      Seq(
-        ("SELECT * FROM testData LEFT JOIN testData2 ON key = a",
-          classOf[SortMergeOuterJoin]),
-        ("SELECT * FROM testData RIGHT JOIN testData2 ON key = a where key = 2",
-          classOf[BroadcastHashOuterJoin]),
-        ("SELECT * FROM testData right join testData2 ON key = a and key = 2",
-          classOf[BroadcastHashOuterJoin])
-      ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
-    }
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "false") {
-      Seq(
-        ("SELECT * FROM testData LEFT JOIN testData2 ON key = a",
-          classOf[ShuffledHashOuterJoin]),
-        ("SELECT * FROM testData RIGHT JOIN testData2 ON key = a where key = 2",
-          classOf[BroadcastHashOuterJoin]),
-        ("SELECT * FROM testData right join testData2 ON key = a and key = 2",
-          classOf[BroadcastHashOuterJoin])
-      ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
-    }
+    Seq(
+      ("SELECT * FROM testData LEFT JOIN testData2 ON key = a",
+        classOf[SortMergeOuterJoin]),
+      ("SELECT * FROM testData RIGHT JOIN testData2 ON key = a where key = 2",
+        classOf[BroadcastHashOuterJoin]),
+      ("SELECT * FROM testData right join testData2 ON key = a and key = 2",
+        classOf[BroadcastHashOuterJoin])
+    ).foreach { case (query, joinClass) => assertJoin(query, joinClass) }
     sql("UNCACHE TABLE testData")
   }
 
@@ -237,241 +199,222 @@ class JoinSuite extends QueryTest with SharedSQLContext {
         Row(2, 2, 2, 2) :: Nil)
   }
 
-  def test_outer_join(useSMJ: Boolean): Unit = {
-
-    val algo = if (useSMJ) "SortMergeOuterJoin" else "ShuffledHashOuterJoin"
-
-    test("left outer join: " + algo) {
-      withSQLConf(SQLConf.SORTMERGE_JOIN.key -> useSMJ.toString) {
-
-        checkAnswer(
-          upperCaseData.join(lowerCaseData, $"n" === $"N", "left"),
-          Row(1, "A", 1, "a") ::
-            Row(2, "B", 2, "b") ::
-            Row(3, "C", 3, "c") ::
-            Row(4, "D", 4, "d") ::
-            Row(5, "E", null, null) ::
-            Row(6, "F", null, null) :: Nil)
-
-        checkAnswer(
-          upperCaseData.join(lowerCaseData, $"n" === $"N" && $"n" > 1, "left"),
-          Row(1, "A", null, null) ::
-            Row(2, "B", 2, "b") ::
-            Row(3, "C", 3, "c") ::
-            Row(4, "D", 4, "d") ::
-            Row(5, "E", null, null) ::
-            Row(6, "F", null, null) :: Nil)
-
-        checkAnswer(
-          upperCaseData.join(lowerCaseData, $"n" === $"N" && $"N" > 1, "left"),
-          Row(1, "A", null, null) ::
-            Row(2, "B", 2, "b") ::
-            Row(3, "C", 3, "c") ::
-            Row(4, "D", 4, "d") ::
-            Row(5, "E", null, null) ::
-            Row(6, "F", null, null) :: Nil)
-
-        checkAnswer(
-          upperCaseData.join(lowerCaseData, $"n" === $"N" && $"l" > $"L", "left"),
-          Row(1, "A", 1, "a") ::
-            Row(2, "B", 2, "b") ::
-            Row(3, "C", 3, "c") ::
-            Row(4, "D", 4, "d") ::
-            Row(5, "E", null, null) ::
-            Row(6, "F", null, null) :: Nil)
-
-        // Make sure we are choosing left.outputPartitioning as the
-        // outputPartitioning for the outer join operator.
-        checkAnswer(
-          sql(
-            """
-            |SELECT l.N, count(*)
-            |FROM upperCaseData l LEFT OUTER JOIN allNulls r ON (l.N = r.a)
-            |GROUP BY l.N
-          """.
-              stripMargin),
-        Row(1, 1) ::
-          Row(2, 1) ::
-          Row(3, 1) ::
-          Row(4, 1) ::
-          Row(5, 1) ::
-          Row(6, 1) :: Nil)
-
-        checkAnswer(
-          sql(
-            """
-              |SELECT r.a, count(*)
-              |FROM upperCaseData l LEFT OUTER JOIN allNulls r ON (l.N = r.a)
-              |GROUP BY r.a
-            """.stripMargin),
-          Row(null, 6) :: Nil)
-      }
-    }
+  test("left outer join") {
+    checkAnswer(
+      upperCaseData.join(lowerCaseData, $"n" === $"N", "left"),
+      Row(1, "A", 1, "a") ::
+        Row(2, "B", 2, "b") ::
+        Row(3, "C", 3, "c") ::
+        Row(4, "D", 4, "d") ::
+        Row(5, "E", null, null) ::
+        Row(6, "F", null, null) :: Nil)
 
-    test("right outer join: " + algo) {
-      withSQLConf(SQLConf.SORTMERGE_JOIN.key -> useSMJ.toString) {
-        checkAnswer(
-          lowerCaseData.join(upperCaseData, $"n" === $"N", "right"),
-          Row(1, "a", 1, "A") ::
-            Row(2, "b", 2, "B") ::
-            Row(3, "c", 3, "C") ::
-            Row(4, "d", 4, "D") ::
-            Row(null, null, 5, "E") ::
-            Row(null, null, 6, "F") :: Nil)
-        checkAnswer(
-          lowerCaseData.join(upperCaseData, $"n" === $"N" && $"n" > 1, "right"),
-          Row(null, null, 1, "A") ::
-            Row(2, "b", 2, "B") ::
-            Row(3, "c", 3, "C") ::
-            Row(4, "d", 4, "D") ::
-            Row(null, null, 5, "E") ::
-            Row(null, null, 6, "F") :: Nil)
-        checkAnswer(
-          lowerCaseData.join(upperCaseData, $"n" === $"N" && $"N" > 1, "right"),
-          Row(null, null, 1, "A") ::
-            Row(2, "b", 2, "B") ::
-            Row(3, "c", 3, "C") ::
-            Row(4, "d", 4, "D") ::
-            Row(null, null, 5, "E") ::
-            Row(null, null, 6, "F") :: Nil)
-        checkAnswer(
-          lowerCaseData.join(upperCaseData, $"n" === $"N" && $"l" > $"L", "right"),
-          Row(1, "a", 1, "A") ::
-            Row(2, "b", 2, "B") ::
-            Row(3, "c", 3, "C") ::
-            Row(4, "d", 4, "D") ::
-            Row(null, null, 5, "E") ::
-            Row(null, null, 6, "F") :: Nil)
-
-        // Make sure we are choosing right.outputPartitioning as the
-        // outputPartitioning for the outer join operator.
-        checkAnswer(
-          sql(
-            """
-              |SELECT l.a, count(*)
-              |FROM allNulls l RIGHT OUTER JOIN upperCaseData r ON (l.a = r.N)
-              |GROUP BY l.a
-            """.stripMargin),
-          Row(null,
-            6))
-
-        checkAnswer(
-          sql(
-            """
-              |SELECT r.N, count(*)
-              |FROM allNulls l RIGHT OUTER JOIN upperCaseData r ON (l.a = r.N)
-              |GROUP BY r.N
-            """.stripMargin),
-          Row(1
-            , 1) ::
-            Row(2, 1) ::
-            Row(3, 1) ::
-            Row(4, 1) ::
-            Row(5, 1) ::
-            Row(6, 1) :: Nil)
-        }
-      }
+    checkAnswer(
+      upperCaseData.join(lowerCaseData, $"n" === $"N" && $"n" > 1, "left"),
+      Row(1, "A", null, null) ::
+        Row(2, "B", 2, "b") ::
+        Row(3, "C", 3, "c") ::
+        Row(4, "D", 4, "d") ::
+        Row(5, "E", null, null) ::
+        Row(6, "F", null, null) :: Nil)
 
-    test("full outer join: " + algo) {
-      withSQLConf(SQLConf.SORTMERGE_JOIN.key -> useSMJ.toString) {
-
-        upperCaseData.where('N <= 4).registerTempTable("left")
-        upperCaseData.where('N >= 3).registerTempTable("right")
-
-        val left = UnresolvedRelation(TableIdentifier("left"), None)
-        val right = UnresolvedRelation(TableIdentifier("right"), None)
-
-        checkAnswer(
-          left.join(right, $"left.N" === $"right.N", "full"),
-          Row(1, "A", null, null) ::
-            Row(2, "B", null, null) ::
-            Row(3, "C", 3, "C") ::
-            Row(4, "D", 4, "D") ::
-            Row(null, null, 5, "E") ::
-            Row(null, null, 6, "F") :: Nil)
-
-        checkAnswer(
-          left.join(right, ($"left.N" === $"right.N") && ($"left.N" !== 3), "full"),
-          Row(1, "A", null, null) ::
-            Row(2, "B", null, null) ::
-            Row(3, "C", null, null) ::
-            Row(null, null, 3, "C") ::
-            Row(4, "D", 4, "D") ::
-            Row(null, null, 5, "E") ::
-            Row(null, null, 6, "F") :: Nil)
-
-        checkAnswer(
-          left.join(right, ($"left.N" === $"right.N") && ($"right.N" !== 3), "full"),
-          Row(1, "A", null, null) ::
-            Row(2, "B", null, null) ::
-            Row(3, "C", null, null) ::
-            Row(null, null, 3, "C") ::
-            Row(4, "D", 4, "D") ::
-            Row(null, null, 5, "E") ::
-            Row(null, null, 6, "F") :: Nil)
-
-        // Make sure we are UnknownPartitioning as the outputPartitioning for the outer join
-        // operator.
-        checkAnswer(
-          sql(
-            """
-            |SELECT l.a, count(*)
-            |FROM allNulls l FULL OUTER JOIN upperCaseData r ON (l.a = r.N)
-            |GROUP BY l.a
-          """.
-              stripMargin),
-        Row(
-          null, 10))
-
-        checkAnswer(
-          sql(
-            """
-              |SELECT r.N, count(*)
-              |FROM allNulls l FULL OUTER JOIN upperCaseData r ON (l.a = r.N)
-              |GROUP BY r.N
-            """.stripMargin),
-          Row
-            (1, 1) ::
-            Row(2, 1) ::
-            Row(3, 1) ::
-            Row(4, 1) ::
-            Row(5, 1) ::
-            Row(6, 1) ::
-            Row(null, 4) :: Nil)
-
-        checkAnswer(
-          sql(
-            """
-              |SELECT l.N, count(*)
-              |FROM upperCaseData l FULL OUTER JOIN allNulls r ON (l.N = r.a)
-              |GROUP BY l.N
-            """.stripMargin),
-          Row(1
-            , 1) ::
-            Row(2, 1) ::
-            Row(3, 1) ::
-            Row(4, 1) ::
-            Row(5, 1) ::
-            Row(6, 1) ::
-            Row(null, 4) :: Nil)
-
-        checkAnswer(
-          sql(
-            """
-            |SELECT r.a, count(*)
-            |FROM upperCaseData l FULL OUTER JOIN allNulls r ON (l.N = r.a)
-            |GROUP BY r.a
-          """.
-              stripMargin),
-          Row(null, 10))
-      }
-    }
+    checkAnswer(
+      upperCaseData.join(lowerCaseData, $"n" === $"N" && $"N" > 1, "left"),
+      Row(1, "A", null, null) ::
+        Row(2, "B", 2, "b") ::
+        Row(3, "C", 3, "c") ::
+        Row(4, "D", 4, "d") ::
+        Row(5, "E", null, null) ::
+        Row(6, "F", null, null) :: Nil)
+
+    checkAnswer(
+      upperCaseData.join(lowerCaseData, $"n" === $"N" && $"l" > $"L", "left"),
+      Row(1, "A", 1, "a") ::
+        Row(2, "B", 2, "b") ::
+        Row(3, "C", 3, "c") ::
+        Row(4, "D", 4, "d") ::
+        Row(5, "E", null, null) ::
+        Row(6, "F", null, null) :: Nil)
+
+    // Make sure we are choosing left.outputPartitioning as the
+    // outputPartitioning for the outer join operator.
+    checkAnswer(
+      sql(
+        """
+        |SELECT l.N, count(*)
+        |FROM upperCaseData l LEFT OUTER JOIN allNulls r ON (l.N = r.a)
+        |GROUP BY l.N
+      """.
+          stripMargin),
+    Row(1, 1) ::
+      Row(2, 1) ::
+      Row(3, 1) ::
+      Row(4, 1) ::
+      Row(5, 1) ::
+      Row(6, 1) :: Nil)
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT r.a, count(*)
+          |FROM upperCaseData l LEFT OUTER JOIN allNulls r ON (l.N = r.a)
+          |GROUP BY r.a
+        """.stripMargin),
+      Row(null, 6) :: Nil)
+  }
+
+  test("right outer join") {
+    checkAnswer(
+      lowerCaseData.join(upperCaseData, $"n" === $"N", "right"),
+      Row(1, "a", 1, "A") ::
+        Row(2, "b", 2, "B") ::
+        Row(3, "c", 3, "C") ::
+        Row(4, "d", 4, "D") ::
+        Row(null, null, 5, "E") ::
+        Row(null, null, 6, "F") :: Nil)
+    checkAnswer(
+      lowerCaseData.join(upperCaseData, $"n" === $"N" && $"n" > 1, "right"),
+      Row(null, null, 1, "A") ::
+        Row(2, "b", 2, "B") ::
+        Row(3, "c", 3, "C") ::
+        Row(4, "d", 4, "D") ::
+        Row(null, null, 5, "E") ::
+        Row(null, null, 6, "F") :: Nil)
+    checkAnswer(
+      lowerCaseData.join(upperCaseData, $"n" === $"N" && $"N" > 1, "right"),
+      Row(null, null, 1, "A") ::
+        Row(2, "b", 2, "B") ::
+        Row(3, "c", 3, "C") ::
+        Row(4, "d", 4, "D") ::
+        Row(null, null, 5, "E") ::
+        Row(null, null, 6, "F") :: Nil)
+    checkAnswer(
+      lowerCaseData.join(upperCaseData, $"n" === $"N" && $"l" > $"L", "right"),
+      Row(1, "a", 1, "A") ::
+        Row(2, "b", 2, "B") ::
+        Row(3, "c", 3, "C") ::
+        Row(4, "d", 4, "D") ::
+        Row(null, null, 5, "E") ::
+        Row(null, null, 6, "F") :: Nil)
+
+    // Make sure we are choosing right.outputPartitioning as the
+    // outputPartitioning for the outer join operator.
+    checkAnswer(
+      sql(
+        """
+          |SELECT l.a, count(*)
+          |FROM allNulls l RIGHT OUTER JOIN upperCaseData r ON (l.a = r.N)
+          |GROUP BY l.a
+        """.stripMargin),
+      Row(null,
+        6))
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT r.N, count(*)
+          |FROM allNulls l RIGHT OUTER JOIN upperCaseData r ON (l.a = r.N)
+          |GROUP BY r.N
+        """.stripMargin),
+      Row(1
+        , 1) ::
+        Row(2, 1) ::
+        Row(3, 1) ::
+        Row(4, 1) ::
+        Row(5, 1) ::
+        Row(6, 1) :: Nil)
   }
 
-  // test SortMergeOuterJoin
-  test_outer_join(true)
-  // test ShuffledHashOuterJoin
-  test_outer_join(false)
+  test("full outer join") {
+    upperCaseData.where('N <= 4).registerTempTable("left")
+    upperCaseData.where('N >= 3).registerTempTable("right")
+
+    val left = UnresolvedRelation(TableIdentifier("left"), None)
+    val right = UnresolvedRelation(TableIdentifier("right"), None)
+
+    checkAnswer(
+      left.join(right, $"left.N" === $"right.N", "full"),
+      Row(1, "A", null, null) ::
+        Row(2, "B", null, null) ::
+        Row(3, "C", 3, "C") ::
+        Row(4, "D", 4, "D") ::
+        Row(null, null, 5, "E") ::
+        Row(null, null, 6, "F") :: Nil)
+
+    checkAnswer(
+      left.join(right, ($"left.N" === $"right.N") && ($"left.N" !== 3), "full"),
+      Row(1, "A", null, null) ::
+        Row(2, "B", null, null) ::
+        Row(3, "C", null, null) ::
+        Row(null, null, 3, "C") ::
+        Row(4, "D", 4, "D") ::
+        Row(null, null, 5, "E") ::
+        Row(null, null, 6, "F") :: Nil)
+
+    checkAnswer(
+      left.join(right, ($"left.N" === $"right.N") && ($"right.N" !== 3), "full"),
+      Row(1, "A", null, null) ::
+        Row(2, "B", null, null) ::
+        Row(3, "C", null, null) ::
+        Row(null, null, 3, "C") ::
+        Row(4, "D", 4, "D") ::
+        Row(null, null, 5, "E") ::
+        Row(null, null, 6, "F") :: Nil)
+
+    // Make sure we are UnknownPartitioning as the outputPartitioning for the outer join
+    // operator.
+    checkAnswer(
+      sql(
+        """
+        |SELECT l.a, count(*)
+        |FROM allNulls l FULL OUTER JOIN upperCaseData r ON (l.a = r.N)
+        |GROUP BY l.a
+      """.
+          stripMargin),
+      Row(null, 10))
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT r.N, count(*)
+          |FROM allNulls l FULL OUTER JOIN upperCaseData r ON (l.a = r.N)
+          |GROUP BY r.N
+        """.stripMargin),
+      Row
+        (1, 1) ::
+        Row(2, 1) ::
+        Row(3, 1) ::
+        Row(4, 1) ::
+        Row(5, 1) ::
+        Row(6, 1) ::
+        Row(null, 4) :: Nil)
+
+    checkAnswer(
+      sql(
+        """
+          |SELECT l.N, count(*)
+          |FROM upperCaseData l FULL OUTER JOIN allNulls r ON (l.N = r.a)
+          |GROUP BY l.N
+        """.stripMargin),
+      Row(1
+        , 1) ::
+        Row(2, 1) ::
+        Row(3, 1) ::
+        Row(4, 1) ::
+        Row(5, 1) ::
+        Row(6, 1) ::
+        Row(null, 4) :: Nil)
+
+    checkAnswer(
+      sql(
+        """
+        |SELECT r.a, count(*)
+        |FROM upperCaseData l FULL OUTER JOIN allNulls r ON (l.N = r.a)
+        |GROUP BY r.a
+      """.
+          stripMargin),
+      Row(null, 10))
+  }
 
   test("broadcasted left semi join operator selection") {
     sqlContext.cacheManager.clearCache()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index acabe32c67bc0..52a561d2e5452 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -1771,8 +1771,7 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
     // This test is for the fix of https://issues.apache.org/jira/browse/SPARK-10737.
     // This bug will be triggered when Tungsten is enabled and there are multiple
     // SortMergeJoin operators executed in the same task.
-    val confs =
-      SQLConf.SORTMERGE_JOIN.key -> "true" :: SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1" :: Nil
+    val confs = SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1" :: Nil
     withSQLConf(confs: _*) {
       val df1 = (1 to 50).map(i => (s"str_$i", i)).toDF("i", "j")
       val df2 =
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
index 44634dacbde68..8c41d79dae817 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, Literal,
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.plans.physical._
-import org.apache.spark.sql.execution.joins.{BroadcastHashJoin, ShuffledHashJoin}
+import org.apache.spark.sql.execution.joins.{SortMergeJoin, BroadcastHashJoin}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
@@ -97,10 +97,10 @@ class PlannerSuite extends SharedSQLContext {
           """.stripMargin).queryExecution.executedPlan
 
         val broadcastHashJoins = planned.collect { case join: BroadcastHashJoin => join }
-        val shuffledHashJoins = planned.collect { case join: ShuffledHashJoin => join }
+        val sortMergeJoins = planned.collect { case join: SortMergeJoin => join }
 
         assert(broadcastHashJoins.size === 1, "Should use broadcast hash join")
-        assert(shuffledHashJoins.isEmpty, "Should not use shuffled hash join")
+        assert(sortMergeJoins.isEmpty, "Should not use sort merge join")
       }
     }
 
@@ -150,10 +150,10 @@ class PlannerSuite extends SharedSQLContext {
         val planned = a.join(b, $"a.key" === $"b.key").queryExecution.executedPlan
 
         val broadcastHashJoins = planned.collect { case join: BroadcastHashJoin => join }
-        val shuffledHashJoins = planned.collect { case join: ShuffledHashJoin => join }
+        val sortMergeJoins = planned.collect { case join: SortMergeJoin => join }
 
         assert(broadcastHashJoins.size === 1, "Should use broadcast hash join")
-        assert(shuffledHashJoins.isEmpty, "Should not use shuffled hash join")
+        assert(sortMergeJoins.isEmpty, "Should not use sort merge join")
 
         sqlContext.clearCache()
       }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/InnerJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/InnerJoinSuite.scala
index 066c16e535c76..2ec17146476fe 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/InnerJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/InnerJoinSuite.scala
@@ -93,20 +93,6 @@ class InnerJoinSuite extends SparkPlanTest with SharedSQLContext {
       boundCondition.map(Filter(_, broadcastHashJoin)).getOrElse(broadcastHashJoin)
     }
 
-    def makeShuffledHashJoin(
-        leftKeys: Seq[Expression],
-        rightKeys: Seq[Expression],
-        boundCondition: Option[Expression],
-        leftPlan: SparkPlan,
-        rightPlan: SparkPlan,
-        side: BuildSide) = {
-      val shuffledHashJoin =
-        execution.joins.ShuffledHashJoin(leftKeys, rightKeys, side, leftPlan, rightPlan)
-      val filteredJoin =
-        boundCondition.map(Filter(_, shuffledHashJoin)).getOrElse(shuffledHashJoin)
-      EnsureRequirements(sqlContext).apply(filteredJoin)
-    }
-
     def makeSortMergeJoin(
         leftKeys: Seq[Expression],
         rightKeys: Seq[Expression],
@@ -143,30 +129,6 @@ class InnerJoinSuite extends SparkPlanTest with SharedSQLContext {
       }
     }
 
-    test(s"$testName using ShuffledHashJoin (build=left)") {
-      extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _) =>
-        withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
-          checkAnswer2(leftRows, rightRows, (leftPlan: SparkPlan, rightPlan: SparkPlan) =>
-            makeShuffledHashJoin(
-              leftKeys, rightKeys, boundCondition, leftPlan, rightPlan, joins.BuildLeft),
-            expectedAnswer.map(Row.fromTuple),
-            sortAnswers = true)
-        }
-      }
-    }
-
-    test(s"$testName using ShuffledHashJoin (build=right)") {
-      extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _) =>
-        withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
-          checkAnswer2(leftRows, rightRows, (leftPlan: SparkPlan, rightPlan: SparkPlan) =>
-            makeShuffledHashJoin(
-              leftKeys, rightKeys, boundCondition, leftPlan, rightPlan, joins.BuildRight),
-            expectedAnswer.map(Row.fromTuple),
-            sortAnswers = true)
-        }
-      }
-    }
-
     test(s"$testName using SortMergeJoin") {
       extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _) =>
         withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/OuterJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/OuterJoinSuite.scala
index 09e0237a7cc50..9c80714a9af4a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/OuterJoinSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/OuterJoinSuite.scala
@@ -74,18 +74,6 @@ class OuterJoinSuite extends SparkPlanTest with SharedSQLContext {
       ExtractEquiJoinKeys.unapply(join)
     }
 
-    test(s"$testName using ShuffledHashOuterJoin") {
-      extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _) =>
-        withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") {
-          checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) =>
-            EnsureRequirements(sqlContext).apply(
-              ShuffledHashOuterJoin(leftKeys, rightKeys, joinType, boundCondition, left, right)),
-            expectedAnswer.map(Row.fromTuple),
-            sortAnswers = true)
-        }
-      }
-    }
-
     if (joinType != FullOuter) {
       test(s"$testName using BroadcastHashOuterJoin") {
         extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _) =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index 544c1ef303ae9..486bfbbd70887 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -169,188 +169,123 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
   test("SortMergeJoin metrics") {
     // Because SortMergeJoin may skip different rows if the number of partitions is different, this
     // test should use the deterministic number of partitions.
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "true") {
-      val testDataForJoin = testData2.filter('a < 2) // TestData2(1, 1) :: TestData2(1, 2)
-      testDataForJoin.registerTempTable("testDataForJoin")
-      withTempTable("testDataForJoin") {
-        // Assume the execution plan is
-        // ... -> SortMergeJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
-        val df = sqlContext.sql(
-          "SELECT * FROM testData2 JOIN testDataForJoin ON testData2.a = testDataForJoin.a")
-        testSparkPlanMetrics(df, 1, Map(
-          1L -> ("SortMergeJoin", Map(
-            // It's 4 because we only read 3 rows in the first partition and 1 row in the second one
-            "number of left rows" -> 4L,
-            "number of right rows" -> 2L,
-            "number of output rows" -> 4L)))
-        )
-      }
+    val testDataForJoin = testData2.filter('a < 2) // TestData2(1, 1) :: TestData2(1, 2)
+    testDataForJoin.registerTempTable("testDataForJoin")
+    withTempTable("testDataForJoin") {
+      // Assume the execution plan is
+      // ... -> SortMergeJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
+      val df = sqlContext.sql(
+        "SELECT * FROM testData2 JOIN testDataForJoin ON testData2.a = testDataForJoin.a")
+      testSparkPlanMetrics(df, 1, Map(
+        1L -> ("SortMergeJoin", Map(
+          // It's 4 because we only read 3 rows in the first partition and 1 row in the second one
+          "number of left rows" -> 4L,
+          "number of right rows" -> 2L,
+          "number of output rows" -> 4L)))
+      )
     }
   }
 
   test("SortMergeOuterJoin metrics") {
     // Because SortMergeOuterJoin may skip different rows if the number of partitions is different,
     // this test should use the deterministic number of partitions.
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "true") {
-      val testDataForJoin = testData2.filter('a < 2) // TestData2(1, 1) :: TestData2(1, 2)
-      testDataForJoin.registerTempTable("testDataForJoin")
-      withTempTable("testDataForJoin") {
-        // Assume the execution plan is
-        // ... -> SortMergeOuterJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
-        val df = sqlContext.sql(
-          "SELECT * FROM testData2 left JOIN testDataForJoin ON testData2.a = testDataForJoin.a")
-        testSparkPlanMetrics(df, 1, Map(
-          1L -> ("SortMergeOuterJoin", Map(
-            // It's 4 because we only read 3 rows in the first partition and 1 row in the second one
-            "number of left rows" -> 6L,
-            "number of right rows" -> 2L,
-            "number of output rows" -> 8L)))
-        )
-
-        val df2 = sqlContext.sql(
-          "SELECT * FROM testDataForJoin right JOIN testData2 ON testData2.a = testDataForJoin.a")
-        testSparkPlanMetrics(df2, 1, Map(
-          1L -> ("SortMergeOuterJoin", Map(
-            // It's 4 because we only read 3 rows in the first partition and 1 row in the second one
-            "number of left rows" -> 2L,
-            "number of right rows" -> 6L,
-            "number of output rows" -> 8L)))
-        )
-      }
-    }
-  }
-
-  test("BroadcastHashJoin metrics") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "false") {
-      val df1 = Seq((1, "1"), (2, "2")).toDF("key", "value")
-      val df2 = Seq((1, "1"), (2, "2"), (3, "3"), (4, "4")).toDF("key", "value")
-      // Assume the execution plan is
-      // ... -> BroadcastHashJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
-      val df = df1.join(broadcast(df2), "key")
-      testSparkPlanMetrics(df, 2, Map(
-        1L -> ("BroadcastHashJoin", Map(
-          "number of left rows" -> 2L,
-          "number of right rows" -> 4L,
-          "number of output rows" -> 2L)))
-      )
-    }
-  }
-
-  test("ShuffledHashJoin metrics") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "false") {
-      val testDataForJoin = testData2.filter('a < 2) // TestData2(1, 1) :: TestData2(1, 2)
-      testDataForJoin.registerTempTable("testDataForJoin")
-      withTempTable("testDataForJoin") {
-        // Assume the execution plan is
-        // ... -> ShuffledHashJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
-        val df = sqlContext.sql(
-          "SELECT * FROM testData2 JOIN testDataForJoin ON testData2.a = testDataForJoin.a")
-        testSparkPlanMetrics(df, 1, Map(
-          1L -> ("ShuffledHashJoin", Map(
-            "number of left rows" -> 6L,
-            "number of right rows" -> 2L,
-            "number of output rows" -> 4L)))
-        )
-      }
-    }
-  }
-
-  test("ShuffledHashOuterJoin metrics") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "false",
-      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "0") {
-      val df1 = Seq((1, "a"), (1, "b"), (4, "c")).toDF("key", "value")
-      val df2 = Seq((1, "a"), (1, "b"), (2, "c"), (3, "d")).toDF("key2", "value")
+    val testDataForJoin = testData2.filter('a < 2) // TestData2(1, 1) :: TestData2(1, 2)
+    testDataForJoin.registerTempTable("testDataForJoin")
+    withTempTable("testDataForJoin") {
       // Assume the execution plan is
-      // ... -> ShuffledHashOuterJoin(nodeId = 0)
-      val df = df1.join(df2, $"key" === $"key2", "left_outer")
+      // ... -> SortMergeOuterJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
+      val df = sqlContext.sql(
+        "SELECT * FROM testData2 left JOIN testDataForJoin ON testData2.a = testDataForJoin.a")
       testSparkPlanMetrics(df, 1, Map(
-        0L -> ("ShuffledHashOuterJoin", Map(
-          "number of left rows" -> 3L,
-          "number of right rows" -> 4L,
-          "number of output rows" -> 5L)))
+        1L -> ("SortMergeOuterJoin", Map(
+          // It's 4 because we only read 3 rows in the first partition and 1 row in the second one
+          "number of left rows" -> 6L,
+          "number of right rows" -> 2L,
+          "number of output rows" -> 8L)))
       )
 
-      val df3 = df1.join(df2, $"key" === $"key2", "right_outer")
-      testSparkPlanMetrics(df3, 1, Map(
-        0L -> ("ShuffledHashOuterJoin", Map(
-          "number of left rows" -> 3L,
-          "number of right rows" -> 4L,
-          "number of output rows" -> 6L)))
-      )
-
-      val df4 = df1.join(df2, $"key" === $"key2", "outer")
-      testSparkPlanMetrics(df4, 1, Map(
-        0L -> ("ShuffledHashOuterJoin", Map(
-          "number of left rows" -> 3L,
-          "number of right rows" -> 4L,
-          "number of output rows" -> 7L)))
+      val df2 = sqlContext.sql(
+        "SELECT * FROM testDataForJoin right JOIN testData2 ON testData2.a = testDataForJoin.a")
+      testSparkPlanMetrics(df2, 1, Map(
+        1L -> ("SortMergeOuterJoin", Map(
+          // It's 4 because we only read 3 rows in the first partition and 1 row in the second one
+          "number of left rows" -> 2L,
+          "number of right rows" -> 6L,
+          "number of output rows" -> 8L)))
       )
     }
   }
 
+  test("BroadcastHashJoin metrics") {
+    val df1 = Seq((1, "1"), (2, "2")).toDF("key", "value")
+    val df2 = Seq((1, "1"), (2, "2"), (3, "3"), (4, "4")).toDF("key", "value")
+    // Assume the execution plan is
+    // ... -> BroadcastHashJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
+    val df = df1.join(broadcast(df2), "key")
+    testSparkPlanMetrics(df, 2, Map(
+      1L -> ("BroadcastHashJoin", Map(
+        "number of left rows" -> 2L,
+        "number of right rows" -> 4L,
+        "number of output rows" -> 2L)))
+    )
+  }
+
   test("BroadcastHashOuterJoin metrics") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "false") {
-      val df1 = Seq((1, "a"), (1, "b"), (4, "c")).toDF("key", "value")
-      val df2 = Seq((1, "a"), (1, "b"), (2, "c"), (3, "d")).toDF("key2", "value")
-      // Assume the execution plan is
-      // ... -> BroadcastHashOuterJoin(nodeId = 0)
-      val df = df1.join(broadcast(df2), $"key" === $"key2", "left_outer")
-      testSparkPlanMetrics(df, 2, Map(
-        0L -> ("BroadcastHashOuterJoin", Map(
-          "number of left rows" -> 3L,
-          "number of right rows" -> 4L,
-          "number of output rows" -> 5L)))
-      )
+    val df1 = Seq((1, "a"), (1, "b"), (4, "c")).toDF("key", "value")
+    val df2 = Seq((1, "a"), (1, "b"), (2, "c"), (3, "d")).toDF("key2", "value")
+    // Assume the execution plan is
+    // ... -> BroadcastHashOuterJoin(nodeId = 0)
+    val df = df1.join(broadcast(df2), $"key" === $"key2", "left_outer")
+    testSparkPlanMetrics(df, 2, Map(
+      0L -> ("BroadcastHashOuterJoin", Map(
+        "number of left rows" -> 3L,
+        "number of right rows" -> 4L,
+        "number of output rows" -> 5L)))
+    )
 
-      val df3 = df1.join(broadcast(df2), $"key" === $"key2", "right_outer")
-      testSparkPlanMetrics(df3, 2, Map(
-        0L -> ("BroadcastHashOuterJoin", Map(
-          "number of left rows" -> 3L,
-          "number of right rows" -> 4L,
-          "number of output rows" -> 6L)))
-      )
-    }
+    val df3 = df1.join(broadcast(df2), $"key" === $"key2", "right_outer")
+    testSparkPlanMetrics(df3, 2, Map(
+      0L -> ("BroadcastHashOuterJoin", Map(
+        "number of left rows" -> 3L,
+        "number of right rows" -> 4L,
+        "number of output rows" -> 6L)))
+    )
   }
 
   test("BroadcastNestedLoopJoin metrics") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "true") {
-      val testDataForJoin = testData2.filter('a < 2) // TestData2(1, 1) :: TestData2(1, 2)
-      testDataForJoin.registerTempTable("testDataForJoin")
-      withTempTable("testDataForJoin") {
-        // Assume the execution plan is
-        // ... -> BroadcastNestedLoopJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
-        val df = sqlContext.sql(
-          "SELECT * FROM testData2 left JOIN testDataForJoin ON " +
-            "testData2.a * testDataForJoin.a != testData2.a + testDataForJoin.a")
-        testSparkPlanMetrics(df, 3, Map(
-          1L -> ("BroadcastNestedLoopJoin", Map(
-            "number of left rows" -> 12L, // left needs to be scanned twice
-            "number of right rows" -> 2L,
-            "number of output rows" -> 12L)))
-        )
-      }
+    val testDataForJoin = testData2.filter('a < 2) // TestData2(1, 1) :: TestData2(1, 2)
+    testDataForJoin.registerTempTable("testDataForJoin")
+    withTempTable("testDataForJoin") {
+      // Assume the execution plan is
+      // ... -> BroadcastNestedLoopJoin(nodeId = 1) -> TungstenProject(nodeId = 0)
+      val df = sqlContext.sql(
+        "SELECT * FROM testData2 left JOIN testDataForJoin ON " +
+          "testData2.a * testDataForJoin.a != testData2.a + testDataForJoin.a")
+      testSparkPlanMetrics(df, 3, Map(
+        1L -> ("BroadcastNestedLoopJoin", Map(
+          "number of left rows" -> 12L, // left needs to be scanned twice
+          "number of right rows" -> 2L,
+          "number of output rows" -> 12L)))
+      )
     }
   }
 
   test("BroadcastLeftSemiJoinHash metrics") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "false") {
-      val df1 = Seq((1, "1"), (2, "2")).toDF("key", "value")
-      val df2 = Seq((1, "1"), (2, "2"), (3, "3"), (4, "4")).toDF("key2", "value")
-      // Assume the execution plan is
-      // ... -> BroadcastLeftSemiJoinHash(nodeId = 0)
-      val df = df1.join(broadcast(df2), $"key" === $"key2", "leftsemi")
-      testSparkPlanMetrics(df, 2, Map(
-        0L -> ("BroadcastLeftSemiJoinHash", Map(
-          "number of left rows" -> 2L,
-          "number of right rows" -> 4L,
-          "number of output rows" -> 2L)))
-      )
-    }
+    val df1 = Seq((1, "1"), (2, "2")).toDF("key", "value")
+    val df2 = Seq((1, "1"), (2, "2"), (3, "3"), (4, "4")).toDF("key2", "value")
+    // Assume the execution plan is
+    // ... -> BroadcastLeftSemiJoinHash(nodeId = 0)
+    val df = df1.join(broadcast(df2), $"key" === $"key2", "leftsemi")
+    testSparkPlanMetrics(df, 2, Map(
+      0L -> ("BroadcastLeftSemiJoinHash", Map(
+        "number of left rows" -> 2L,
+        "number of right rows" -> 4L,
+        "number of output rows" -> 2L)))
+    )
   }
 
   test("LeftSemiJoinHash metrics") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "true",
-      SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "0") {
+    withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "0") {
       val df1 = Seq((1, "1"), (2, "2")).toDF("key", "value")
       val df2 = Seq((1, "1"), (2, "2"), (3, "3"), (4, "4")).toDF("key2", "value")
       // Assume the execution plan is
@@ -366,19 +301,17 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
   }
 
   test("LeftSemiJoinBNL metrics") {
-    withSQLConf(SQLConf.SORTMERGE_JOIN.key -> "false") {
-      val df1 = Seq((1, "1"), (2, "2")).toDF("key", "value")
-      val df2 = Seq((1, "1"), (2, "2"), (3, "3"), (4, "4")).toDF("key2", "value")
-      // Assume the execution plan is
-      // ... -> LeftSemiJoinBNL(nodeId = 0)
-      val df = df1.join(df2, $"key" < $"key2", "leftsemi")
-      testSparkPlanMetrics(df, 2, Map(
-        0L -> ("LeftSemiJoinBNL", Map(
-          "number of left rows" -> 2L,
-          "number of right rows" -> 4L,
-          "number of output rows" -> 2L)))
-      )
-    }
+    val df1 = Seq((1, "1"), (2, "2")).toDF("key", "value")
+    val df2 = Seq((1, "1"), (2, "2"), (3, "3"), (4, "4")).toDF("key2", "value")
+    // Assume the execution plan is
+    // ... -> LeftSemiJoinBNL(nodeId = 0)
+    val df = df1.join(df2, $"key" < $"key2", "leftsemi")
+    testSparkPlanMetrics(df, 2, Map(
+      0L -> ("LeftSemiJoinBNL", Map(
+        "number of left rows" -> 2L,
+        "number of right rows" -> 4L,
+        "number of output rows" -> 2L)))
+    )
   }
 
   test("CartesianProduct metrics") {
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/StatisticsSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/StatisticsSuite.scala
index 9bb32f11b76bd..f775f1e955876 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/StatisticsSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/StatisticsSuite.scala
@@ -166,7 +166,7 @@ class StatisticsSuite extends QueryTest with TestHiveSingleton {
 
         val shj = df.queryExecution.sparkPlan.collect { case j: SortMergeJoin => j }
         assert(shj.size === 1,
-          "ShuffledHashJoin should be planned when BroadcastHashJoin is turned off")
+          "SortMergeJoin should be planned when BroadcastHashJoin is turned off")
 
         sql(s"""SET ${SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key}=$tmp""")
       }

From 39b1e36fbc284ba999e1c00b20d1b0b5de6b40b2 Mon Sep 17 00:00:00 2001
From: Daoyuan Wang <daoyuan.wang@intel.com>
Date: Wed, 11 Nov 2015 20:36:21 -0800
Subject: [PATCH 203/518] [SPARK-11396] [SQL] add native implementation of
 datetime function to_unix_timestamp

`to_unix_timestamp` is the deterministic version of `unix_timestamp`, as it accepts at least one parameters.

Since the behavior here is quite similar to `unix_timestamp`, I think the dataframe API is not necessary here.

Author: Daoyuan Wang <daoyuan.wang@intel.com>

Closes #9347 from adrian-wang/to_unix_timestamp.
---
 .../catalyst/analysis/FunctionRegistry.scala  |  1 +
 .../expressions/datetimeExpressions.scala     | 24 ++++++++++---
 .../expressions/DateExpressionsSuite.scala    | 36 +++++++++++++++++++
 .../apache/spark/sql/DateFunctionsSuite.scala | 21 +++++++++++
 4 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index dfa749d1afa5b..870808aa560e5 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -262,6 +262,7 @@ object FunctionRegistry {
     expression[Quarter]("quarter"),
     expression[Second]("second"),
     expression[ToDate]("to_date"),
+    expression[ToUnixTimestamp]("to_unix_timestamp"),
     expression[ToUTCTimestamp]("to_utc_timestamp"),
     expression[TruncDate]("trunc"),
     expression[UnixTimestamp]("unix_timestamp"),
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala
index 13cc6bb6f27b8..03c39f8404e78 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala
@@ -299,7 +299,20 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
 }
 
 /**
- * Converts time string with given pattern
+ * Converts time string with given pattern.
+ * Deterministic version of [[UnixTimestamp]], must have at least one parameter.
+ */
+case class ToUnixTimestamp(timeExp: Expression, format: Expression) extends UnixTime {
+  override def left: Expression = timeExp
+  override def right: Expression = format
+
+  def this(time: Expression) = {
+    this(time, Literal("yyyy-MM-dd HH:mm:ss"))
+  }
+}
+
+/**
+ * Converts time string with given pattern.
  * (see [http://docs.oracle.com/javase/tutorial/i18n/format/simpleDateFormat.html])
  * to Unix time stamp (in seconds), returns null if fail.
  * Note that hive Language Manual says it returns 0 if fail, but in fact it returns null.
@@ -308,9 +321,7 @@ case class DateFormatClass(left: Expression, right: Expression) extends BinaryEx
  * If the first parameter is a Date or Timestamp instead of String, we will ignore the
  * second parameter.
  */
-case class UnixTimestamp(timeExp: Expression, format: Expression)
-  extends BinaryExpression with ExpectsInputTypes {
-
+case class UnixTimestamp(timeExp: Expression, format: Expression) extends UnixTime {
   override def left: Expression = timeExp
   override def right: Expression = format
 
@@ -321,6 +332,9 @@ case class UnixTimestamp(timeExp: Expression, format: Expression)
   def this() = {
     this(CurrentTimestamp())
   }
+}
+
+abstract class UnixTime extends BinaryExpression with ExpectsInputTypes {
 
   override def inputTypes: Seq[AbstractDataType] =
     Seq(TypeCollection(StringType, DateType, TimestampType), StringType)
@@ -347,7 +361,7 @@ case class UnixTimestamp(timeExp: Expression, format: Expression)
             null
           }
         case StringType =>
-          val f = format.eval(input)
+          val f = right.eval(input)
           if (f == null) {
             null
           } else {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala
index 610d39e8493cd..53c66d8a754ed 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala
@@ -465,6 +465,42 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
       UnixTimestamp(Literal("2015-07-24"), Literal("not a valid format")), null)
   }
 
+  test("to_unix_timestamp") {
+    val sdf1 = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss")
+    val fmt2 = "yyyy-MM-dd HH:mm:ss.SSS"
+    val sdf2 = new SimpleDateFormat(fmt2)
+    val fmt3 = "yy-MM-dd"
+    val sdf3 = new SimpleDateFormat(fmt3)
+    val date1 = Date.valueOf("2015-07-24")
+    checkEvaluation(
+      ToUnixTimestamp(Literal(sdf1.format(new Timestamp(0))), Literal("yyyy-MM-dd HH:mm:ss")), 0L)
+    checkEvaluation(ToUnixTimestamp(
+      Literal(sdf1.format(new Timestamp(1000000))), Literal("yyyy-MM-dd HH:mm:ss")), 1000L)
+    checkEvaluation(
+      ToUnixTimestamp(Literal(new Timestamp(1000000)), Literal("yyyy-MM-dd HH:mm:ss")), 1000L)
+    checkEvaluation(
+      ToUnixTimestamp(Literal(date1), Literal("yyyy-MM-dd HH:mm:ss")),
+      DateTimeUtils.daysToMillis(DateTimeUtils.fromJavaDate(date1)) / 1000L)
+    checkEvaluation(
+      ToUnixTimestamp(Literal(sdf2.format(new Timestamp(-1000000))), Literal(fmt2)), -1000L)
+    checkEvaluation(ToUnixTimestamp(
+      Literal(sdf3.format(Date.valueOf("2015-07-24"))), Literal(fmt3)),
+      DateTimeUtils.daysToMillis(DateTimeUtils.fromJavaDate(Date.valueOf("2015-07-24"))) / 1000L)
+    val t1 = ToUnixTimestamp(
+      CurrentTimestamp(), Literal("yyyy-MM-dd HH:mm:ss")).eval().asInstanceOf[Long]
+    val t2 = ToUnixTimestamp(
+      CurrentTimestamp(), Literal("yyyy-MM-dd HH:mm:ss")).eval().asInstanceOf[Long]
+    assert(t2 - t1 <= 1)
+    checkEvaluation(
+      ToUnixTimestamp(Literal.create(null, DateType), Literal.create(null, StringType)), null)
+    checkEvaluation(
+      ToUnixTimestamp(Literal.create(null, DateType), Literal("yyyy-MM-dd HH:mm:ss")), null)
+    checkEvaluation(ToUnixTimestamp(
+      Literal(date1), Literal.create(null, StringType)), date1.getTime / 1000L)
+    checkEvaluation(
+      ToUnixTimestamp(Literal("2015-07-24"), Literal("not a valid format")), null)
+  }
+
   test("datediff") {
     checkEvaluation(
       DateDiff(Literal(Date.valueOf("2015-07-24")), Literal(Date.valueOf("2015-07-21"))), 3)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
index 9080c53c491ac..1266d534cc5b3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
@@ -444,6 +444,27 @@ class DateFunctionsSuite extends QueryTest with SharedSQLContext {
       Row(ts1.getTime / 1000L), Row(ts2.getTime / 1000L)))
   }
 
+  test("to_unix_timestamp") {
+    val date1 = Date.valueOf("2015-07-24")
+    val date2 = Date.valueOf("2015-07-25")
+    val ts1 = Timestamp.valueOf("2015-07-24 10:00:00.3")
+    val ts2 = Timestamp.valueOf("2015-07-25 02:02:02.2")
+    val s1 = "2015/07/24 10:00:00.5"
+    val s2 = "2015/07/25 02:02:02.6"
+    val ss1 = "2015-07-24 10:00:00"
+    val ss2 = "2015-07-25 02:02:02"
+    val fmt = "yyyy/MM/dd HH:mm:ss.S"
+    val df = Seq((date1, ts1, s1, ss1), (date2, ts2, s2, ss2)).toDF("d", "ts", "s", "ss")
+    checkAnswer(df.selectExpr("to_unix_timestamp(ts)"), Seq(
+      Row(ts1.getTime / 1000L), Row(ts2.getTime / 1000L)))
+    checkAnswer(df.selectExpr("to_unix_timestamp(ss)"), Seq(
+      Row(ts1.getTime / 1000L), Row(ts2.getTime / 1000L)))
+    checkAnswer(df.selectExpr(s"to_unix_timestamp(d, '$fmt')"), Seq(
+      Row(date1.getTime / 1000L), Row(date2.getTime / 1000L)))
+    checkAnswer(df.selectExpr(s"to_unix_timestamp(s, '$fmt')"), Seq(
+      Row(ts1.getTime / 1000L), Row(ts2.getTime / 1000L)))
+  }
+
   test("datediff") {
     val df = Seq(
       (Date.valueOf("2015-07-24"), Timestamp.valueOf("2015-07-24 01:00:00"),

From e2957bc085d39d59c09e4b33c26a05f0263200a3 Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Wed, 11 Nov 2015 21:01:14 -0800
Subject: [PATCH 204/518] [SPARK-11674][ML] add private val after @transient in
 Word2VecModel

This causes compile failure with Scala 2.11. See https://issues.scala-lang.org/browse/SI-8813. (Jenkins won't test Scala 2.11. I tested compile locally.) JoshRosen

Author: Xiangrui Meng <meng@databricks.com>

Closes #9644 from mengxr/SPARK-11674.
---
 mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
index 5c64cb09d5948..708dbeef84db4 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
@@ -146,7 +146,7 @@ final class Word2Vec(override val uid: String) extends Estimator[Word2VecModel]
 @Experimental
 class Word2VecModel private[ml] (
     override val uid: String,
-    @transient wordVectors: feature.Word2VecModel)
+    @transient private val wordVectors: feature.Word2VecModel)
   extends Model[Word2VecModel] with Word2VecBase {
 
   /**

From 14cf753704ea60f358cb870b018cbcf73654f198 Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Thu, 12 Nov 2015 16:47:00 +0800
Subject: [PATCH 205/518] [SPARK-11661][SQL] Still pushdown filters returned by
 unhandledFilters.

https://issues.apache.org/jira/browse/SPARK-11661

Author: Yin Huai <yhuai@databricks.com>

Closes #9634 from yhuai/unhandledFilters.
---
 .../datasources/DataSourceStrategy.scala      | 15 +++++--
 .../apache/spark/sql/sources/interfaces.scala |  8 ++--
 .../parquet/ParquetFilterSuite.scala          | 25 ++++++++++++
 .../spark/sql/sources/FilteredScanSuite.scala | 39 ++++++++++++-------
 .../SimpleTextHadoopFsRelationSuite.scala     |  8 ++--
 5 files changed, 71 insertions(+), 24 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
index 7265d6a4de2e6..d7c01b6e6f07e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
@@ -453,8 +453,8 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
    *
    * @return A pair of `Seq[Expression]` and `Seq[Filter]`. The first element contains all Catalyst
    *         predicate [[Expression]]s that are either not convertible or cannot be handled by
-   *         `relation`. The second element contains all converted data source [[Filter]]s that can
-   *        be handled by `relation`.
+   *         `relation`. The second element contains all converted data source [[Filter]]s that
+   *         will be pushed down to the data source.
    */
   protected[sql] def selectFilters(
     relation: BaseRelation,
@@ -476,7 +476,9 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
     // Catalyst predicate expressions that cannot be translated to data source filters.
     val unrecognizedPredicates = predicates.filterNot(translatedMap.contains)
 
-    // Data source filters that cannot be handled by `relation`
+    // Data source filters that cannot be handled by `relation`. The semantic of a unhandled filter
+    // at here is that a data source may not be able to apply this filter to every row
+    // of the underlying dataset.
     val unhandledFilters = relation.unhandledFilters(translatedMap.values.toArray).toSet
 
     val (unhandled, handled) = translated.partition {
@@ -491,6 +493,11 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
     // Translated data source filters that can be handled by `relation`
     val (_, handledFilters) = handled.unzip
 
-    (unrecognizedPredicates ++ unhandledPredicates, handledFilters)
+    // translated contains all filters that have been converted to the public Filter interface.
+    // We should always push them to the data source no matter whether the data source can apply
+    // a filter to every row or not.
+    val (_, translatedFilters) = translated.unzip
+
+    (unrecognizedPredicates ++ unhandledPredicates, translatedFilters)
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index 48de693a999d8..2be6cd45337f5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -235,9 +235,11 @@ abstract class BaseRelation {
   def needConversion: Boolean = true
 
   /**
-   * Given an array of [[Filter]]s, returns an array of [[Filter]]s that this data source relation
-   * cannot handle.  Spark SQL will apply all returned [[Filter]]s against rows returned by this
-   * data source relation.
+   * Returns the list of [[Filter]]s that this datasource may not be able to handle.
+   * These returned [[Filter]]s will be evaluated by Spark SQL after data is output by a scan.
+   * By default, this function will return all filters, as it is always safe to
+   * double evaluate a [[Filter]]. However, specific implementations can override this function to
+   * avoid double filtering when they are capable of processing a filter internally.
    *
    * @since 1.6.0
    */
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
index 579dabf73318b..2ac87ad6cd035 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
@@ -336,4 +336,29 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
       }
     }
   }
+
+  test("SPARK-11661 Still pushdown filters returned by unhandledFilters") {
+    import testImplicits._
+    withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true") {
+      withTempPath { dir =>
+        val path = s"${dir.getCanonicalPath}/part=1"
+        (1 to 3).map(i => (i, i.toString)).toDF("a", "b").write.parquet(path)
+        val df = sqlContext.read.parquet(path).filter("a = 2")
+
+        // This is the source RDD without Spark-side filtering.
+        val childRDD =
+          df
+            .queryExecution
+            .executedPlan.asInstanceOf[org.apache.spark.sql.execution.Filter]
+            .child
+            .execute()
+
+        // The result should be single row.
+        // When a filter is pushed to Parquet, Parquet can apply it to every row.
+        // So, we can check the number of rows returned from the Parquet
+        // to make sure our filter pushdown work.
+        assert(childRDD.count == 1)
+      }
+    }
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
index 2cad964e55b2b..398b8a1a661c6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala
@@ -254,7 +254,11 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext with Predic
   testPushDown("SELECT * FROM oneToTenFiltered WHERE a IN (1,3,5)", 3, Set("a", "b", "c"))
 
   testPushDown("SELECT * FROM oneToTenFiltered WHERE a = 20", 0, Set("a", "b", "c"))
-  testPushDown("SELECT * FROM oneToTenFiltered WHERE b = 1", 10, Set("a", "b", "c"))
+  testPushDown(
+    "SELECT * FROM oneToTenFiltered WHERE b = 1",
+    10,
+    Set("a", "b", "c"),
+    Set(EqualTo("b", 1)))
 
   testPushDown("SELECT * FROM oneToTenFiltered WHERE a < 5 AND a > 1", 3, Set("a", "b", "c"))
   testPushDown("SELECT * FROM oneToTenFiltered WHERE a < 3 OR a > 8", 4, Set("a", "b", "c"))
@@ -283,12 +287,23 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext with Predic
       | WHERE a + b > 9
       |   AND b < 16
       |   AND c IN ('bbbbbBBBBB', 'cccccCCCCC', 'dddddDDDDD', 'foo')
-    """.stripMargin.split("\n").map(_.trim).mkString(" "), 3, Set("a", "b"))
+    """.stripMargin.split("\n").map(_.trim).mkString(" "),
+    3,
+    Set("a", "b"),
+    Set(LessThan("b", 16)))
 
   def testPushDown(
-      sqlString: String,
-      expectedCount: Int,
-      requiredColumnNames: Set[String]): Unit = {
+    sqlString: String,
+    expectedCount: Int,
+    requiredColumnNames: Set[String]): Unit = {
+    testPushDown(sqlString, expectedCount, requiredColumnNames, Set.empty[Filter])
+  }
+
+  def testPushDown(
+    sqlString: String,
+    expectedCount: Int,
+    requiredColumnNames: Set[String],
+    expectedUnhandledFilters: Set[Filter]): Unit = {
     test(s"PushDown Returns $expectedCount: $sqlString") {
       val queryExecution = sql(sqlString).queryExecution
       val rawPlan = queryExecution.executedPlan.collect {
@@ -300,15 +315,13 @@ class FilteredScanSuite extends DataSourceTest with SharedSQLContext with Predic
       val rawCount = rawPlan.execute().count()
       assert(ColumnsRequired.set === requiredColumnNames)
 
-      assert {
-        val table = caseInsensitiveContext.table("oneToTenFiltered")
-        val relation = table.queryExecution.logical.collectFirst {
-          case LogicalRelation(r, _) => r
-        }.get
+      val table = caseInsensitiveContext.table("oneToTenFiltered")
+      val relation = table.queryExecution.logical.collectFirst {
+        case LogicalRelation(r, _) => r
+      }.get
 
-        // `relation` should be able to handle all pushed filters
-        relation.unhandledFilters(FiltersPushed.list.toArray).isEmpty
-      }
+      assert(
+        relation.unhandledFilters(FiltersPushed.list.toArray).toSet === expectedUnhandledFilters)
 
       if (rawCount != expectedCount) {
         fail(
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
index 9251a69f31a47..81af684ba0bf1 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
@@ -248,7 +248,7 @@ class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest with Predicat
     projections = Seq('c, 'p),
     filter = 'a < 3 && 'p > 0,
     requiredColumns = Seq("c", "a"),
-    pushedFilters = Nil,
+    pushedFilters = Seq(LessThan("a", 3)),
     inconvertibleFilters = Nil,
     unhandledFilters = Seq('a < 3),
     partitioningFilters = Seq('p > 0)
@@ -327,7 +327,7 @@ class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest with Predicat
     projections = Seq('b, 'p),
     filter = 'c > "val_7" && 'b < 18 && 'p > 0,
     requiredColumns = Seq("b"),
-    pushedFilters = Seq(GreaterThan("c", "val_7")),
+    pushedFilters = Seq(GreaterThan("c", "val_7"), LessThan("b", 18)),
     inconvertibleFilters = Nil,
     unhandledFilters = Seq('b < 18),
     partitioningFilters = Seq('p > 0)
@@ -344,7 +344,7 @@ class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest with Predicat
     projections = Seq('b, 'p),
     filter = 'a % 2 === 0 && 'c > "val_7" && 'b < 18 && 'p > 0,
     requiredColumns = Seq("b", "a"),
-    pushedFilters = Seq(GreaterThan("c", "val_7")),
+    pushedFilters = Seq(GreaterThan("c", "val_7"), LessThan("b", 18)),
     inconvertibleFilters = Seq('a % 2 === 0),
     unhandledFilters = Seq('b < 18),
     partitioningFilters = Seq('p > 0)
@@ -361,7 +361,7 @@ class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest with Predicat
     projections = Seq('b, 'p),
     filter = 'a > 7 && 'a < 9,
     requiredColumns = Seq("b", "a"),
-    pushedFilters = Seq(GreaterThan("a", 7)),
+    pushedFilters = Seq(GreaterThan("a", 7), LessThan("a", 9)),
     inconvertibleFilters = Nil,
     unhandledFilters = Seq('a < 9),
     partitioningFilters = Nil

From 30e743364313d4b81c99de8f9a7170f5bca2771c Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 12 Nov 2015 08:14:08 -0800
Subject: [PATCH 206/518] [SPARK-11673][SQL] Remove the normal Project physical
 operator (and keep TungstenProject)

Also make full outer join being able to produce UnsafeRows.

Author: Reynold Xin <rxin@databricks.com>

Closes #9643 from rxin/SPARK-11673.
---
 .../execution/UnsafeExternalRowSorter.java    |  7 ---
 .../scala/org/apache/spark/sql/Encoder.scala  |  8 +--
 .../spark/sql/catalyst/ScalaReflection.scala  |  6 +--
 .../expressions/EquivalentExpressions.scala   |  2 +-
 .../sql/catalyst/expressions/Projection.scala | 10 ----
 .../apache/spark/sql/execution/Exchange.scala |  5 +-
 .../spark/sql/execution/SparkPlanner.scala    |  6 +--
 .../spark/sql/execution/SparkStrategies.scala | 14 +----
 .../apache/spark/sql/execution/Window.scala   |  6 +--
 .../aggregate/SortBasedAggregate.scala        |  6 +--
 .../aggregate/TungstenAggregate.scala         |  3 +-
 .../spark/sql/execution/basicOperators.scala  | 26 ---------
 .../datasources/DataSourceStrategy.scala      |  3 +-
 .../spark/sql/execution/joins/HashJoin.scala  | 31 +++--------
 .../sql/execution/joins/HashOuterJoin.scala   | 54 +++++++------------
 .../sql/execution/joins/HashSemiJoin.scala    | 25 ++-------
 .../sql/execution/joins/SortMergeJoin.scala   | 41 +++-----------
 .../execution/joins/SortMergeOuterJoin.scala  | 38 +++----------
 .../execution/local/BinaryHashJoinNode.scala  |  6 +--
 .../sql/execution/local/HashJoinNode.scala    | 21 ++------
 .../sql/execution/rowFormatConverters.scala   | 18 ++-----
 .../org/apache/spark/sql/execution/sort.scala |  9 ----
 .../spark/sql/ColumnExpressionSuite.scala     |  3 +-
 .../sql/execution/TungstenSortSuite.scala     |  1 -
 .../execution/local/HashJoinNodeSuite.scala   | 10 +---
 .../execution/HiveTypeCoercionSuite.scala     |  6 ++-
 .../ParquetHadoopFsRelationSuite.scala        |  2 +-
 27 files changed, 80 insertions(+), 287 deletions(-)

diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java b/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
index f7063d1e5c829..3986d6e18f770 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/execution/UnsafeExternalRowSorter.java
@@ -170,13 +170,6 @@ public Iterator<UnsafeRow> sort(Iterator<UnsafeRow> inputIterator) throws IOExce
     return sort();
   }
 
-  /**
-   * Return true if UnsafeExternalRowSorter can sort rows with the given schema, false otherwise.
-   */
-  public static boolean supportsSchema(StructType schema) {
-    return UnsafeProjection.canSupport(schema);
-  }
-
   private static final class RowComparator extends RecordComparator {
     private final Ordering<InternalRow> ordering;
     private final int numFields;
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index 1ff7340557e6c..6134f9e036638 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -17,13 +17,13 @@
 
 package org.apache.spark.sql
 
+import scala.reflect.ClassTag
+
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.types.{ObjectType, StructField, StructType}
 import org.apache.spark.util.Utils
 
-import scala.reflect.ClassTag
-
 /**
  * Used to convert a JVM object of type `T` to and from the internal Spark SQL representation.
  *
@@ -123,9 +123,9 @@ object Encoders {
 
     new ExpressionEncoder[Any](
       schema,
-      false,
+      flat = false,
       extractExpressions,
       constructExpression,
-      ClassTag.apply(cls))
+      ClassTag(cls))
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 0b8a8abd02d67..6d822261b050a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -153,18 +153,18 @@ trait ScalaReflection {
    */
   def constructorFor[T : TypeTag]: Expression = constructorFor(typeOf[T], None)
 
-  protected def constructorFor(
+  private def constructorFor(
       tpe: `Type`,
       path: Option[Expression]): Expression = ScalaReflectionLock.synchronized {
 
     /** Returns the current path with a sub-field extracted. */
-    def addToPath(part: String) =
+    def addToPath(part: String): Expression =
       path
         .map(p => UnresolvedExtractValue(p, expressions.Literal(part)))
         .getOrElse(UnresolvedAttribute(part))
 
     /** Returns the current path with a field at ordinal extracted. */
-    def addToPathOrdinal(ordinal: Int, dataType: DataType) =
+    def addToPathOrdinal(ordinal: Int, dataType: DataType): Expression =
       path
         .map(p => GetStructField(p, StructField(s"_$ordinal", dataType), ordinal))
         .getOrElse(BoundReference(ordinal, dataType, false))
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
index f83df494ba8a6..f7162e420d19a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala
@@ -77,7 +77,7 @@ class EquivalentExpressions {
    * an empty collection if there are none.
    */
   def getEquivalentExprs(e: Expression): Seq[Expression] = {
-    equivalenceMap.get(Expr(e)).getOrElse(mutable.MutableList())
+    equivalenceMap.getOrElse(Expr(e), mutable.MutableList())
   }
 
   /**
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
index 9f0b7821ae74a..053e612f3ecb5 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Projection.scala
@@ -102,16 +102,6 @@ abstract class UnsafeProjection extends Projection {
 
 object UnsafeProjection {
 
-  /*
-   * Returns whether UnsafeProjection can support given StructType, Array[DataType] or
-   * Seq[Expression].
-   */
-  def canSupport(schema: StructType): Boolean = canSupport(schema.fields.map(_.dataType))
-  def canSupport(exprs: Seq[Expression]): Boolean = canSupport(exprs.map(_.dataType).toArray)
-  private def canSupport(types: Array[DataType]): Boolean = {
-    types.forall(GenerateUnsafeProjection.canSupport)
-  }
-
   /**
    * Returns an UnsafeProjection for given StructType.
    */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
index d0e4e068092f9..bc252d98e7144 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
@@ -57,10 +57,7 @@ case class Exchange(
   /**
    * Returns true iff we can support the data type, and we are not doing range partitioning.
    */
-  private lazy val tungstenMode: Boolean = {
-    GenerateUnsafeProjection.canSupport(child.schema) &&
-      !newPartitioning.isInstanceOf[RangePartitioning]
-  }
+  private lazy val tungstenMode: Boolean = !newPartitioning.isInstanceOf[RangePartitioning]
 
   override def outputPartitioning: Partitioning = newPartitioning
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
index cf482ae4a05ee..b7c5476346b2a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
@@ -18,12 +18,10 @@
 package org.apache.spark.sql.execution
 
 import org.apache.spark.SparkContext
-import org.apache.spark.annotation.Experimental
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.execution.datasources.DataSourceStrategy
 
-@Experimental
 class SparkPlanner(val sqlContext: SQLContext) extends SparkStrategies {
   val sparkContext: SparkContext = sqlContext.sparkContext
 
@@ -64,7 +62,7 @@ class SparkPlanner(val sqlContext: SQLContext) extends SparkStrategies {
 
     val projectSet = AttributeSet(projectList.flatMap(_.references))
     val filterSet = AttributeSet(filterPredicates.flatMap(_.references))
-    val filterCondition =
+    val filterCondition: Option[Expression] =
       prunePushedDownFilters(filterPredicates).reduceLeftOption(catalyst.expressions.And)
 
     // Right now we still use a projection even if the only evaluation is applying an alias
@@ -82,7 +80,7 @@ class SparkPlanner(val sqlContext: SQLContext) extends SparkStrategies {
       filterCondition.map(Filter(_, scan)).getOrElse(scan)
     } else {
       val scan = scanBuilder((projectSet ++ filterSet).toSeq)
-      Project(projectList, filterCondition.map(Filter(_, scan)).getOrElse(scan))
+      TungstenProject(projectList, filterCondition.map(Filter(_, scan)).getOrElse(scan))
     }
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index 90989f2cee9a6..a99ae4674bb12 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -309,11 +309,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
      *               if necessary.
      */
     def getSortOperator(sortExprs: Seq[SortOrder], global: Boolean, child: SparkPlan): SparkPlan = {
-      if (TungstenSort.supportsSchema(child.schema)) {
-        execution.TungstenSort(sortExprs, global, child)
-      } else {
-        execution.Sort(sortExprs, global, child)
-      }
+      execution.TungstenSort(sortExprs, global, child)
     }
 
     def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
@@ -347,13 +343,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
       case logical.Sort(sortExprs, global, child) =>
         getSortOperator(sortExprs, global, planLater(child)):: Nil
       case logical.Project(projectList, child) =>
-        // If unsafe mode is enabled and we support these data types in Unsafe, use the
-        // Tungsten project. Otherwise, use the normal project.
-        if (UnsafeProjection.canSupport(projectList) && UnsafeProjection.canSupport(child.schema)) {
-          execution.TungstenProject(projectList, planLater(child)) :: Nil
-        } else {
-          execution.Project(projectList, planLater(child)) :: Nil
-        }
+        execution.TungstenProject(projectList, planLater(child)) :: Nil
       case logical.Filter(condition, child) =>
         execution.Filter(condition, planLater(child)) :: Nil
       case e @ logical.Expand(_, _, child) =>
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Window.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Window.scala
index 53c5ccf8fa37e..b1280c32a6a43 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Window.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Window.scala
@@ -247,11 +247,7 @@ case class Window(
 
         // Get all relevant projections.
         val result = createResultProjection(unboundExpressions)
-        val grouping = if (child.outputsUnsafeRows) {
-          UnsafeProjection.create(partitionSpec, child.output)
-        } else {
-          newProjection(partitionSpec, child.output)
-        }
+        val grouping = UnsafeProjection.create(partitionSpec, child.output)
 
         // Manage the stream and the grouping.
         var nextRow: InternalRow = EmptyRow
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
index fb7f30c2aec99..c8ccbb933df61 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
@@ -78,11 +78,9 @@ case class SortBasedAggregate(
         // so return an empty iterator.
         Iterator[InternalRow]()
       } else {
-        val groupingKeyProjection = if (UnsafeProjection.canSupport(groupingExpressions)) {
+        val groupingKeyProjection =
           UnsafeProjection.create(groupingExpressions, child.output)
-        } else {
-          newMutableProjection(groupingExpressions, child.output)()
-        }
+
         val outputIter = new SortBasedAggregationIterator(
           groupingKeyProjection,
           groupingExpressions.map(_.toAttribute),
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregate.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregate.scala
index 1edde1e5a16d9..920de615e1d86 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregate.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregate.scala
@@ -139,7 +139,6 @@ object TungstenAggregate {
     groupingExpressions: Seq[Expression],
     aggregateBufferAttributes: Seq[Attribute]): Boolean = {
     val aggregationBufferSchema = StructType.fromAttributes(aggregateBufferAttributes)
-    UnsafeFixedWidthAggregationMap.supportsAggregationBufferSchema(aggregationBufferSchema) &&
-      UnsafeProjection.canSupport(groupingExpressions)
+    UnsafeFixedWidthAggregationMap.supportsAggregationBufferSchema(aggregationBufferSchema)
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index 303d636164adb..ae08fb71bf4cb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -30,32 +30,6 @@ import org.apache.spark.util.random.PoissonSampler
 import org.apache.spark.{HashPartitioner, SparkEnv}
 
 
-case class Project(projectList: Seq[NamedExpression], child: SparkPlan) extends UnaryNode {
-  override def output: Seq[Attribute] = projectList.map(_.toAttribute)
-
-  override private[sql] lazy val metrics = Map(
-    "numRows" -> SQLMetrics.createLongMetric(sparkContext, "number of rows"))
-
-  @transient lazy val buildProjection = newMutableProjection(projectList, child.output)
-
-  protected override def doExecute(): RDD[InternalRow] = {
-    val numRows = longMetric("numRows")
-    child.execute().mapPartitions { iter =>
-      val reusableProjection = buildProjection()
-      iter.map { row =>
-        numRows += 1
-        reusableProjection(row)
-      }
-    }
-  }
-
-  override def outputOrdering: Seq[SortOrder] = child.outputOrdering
-}
-
-
-/**
- * A variant of [[Project]] that returns [[UnsafeRow]]s.
- */
 case class TungstenProject(projectList: Seq[NamedExpression], child: SparkPlan) extends UnaryNode {
 
   override private[sql] lazy val metrics = Map(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
index d7c01b6e6f07e..824c89a90eb8a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
@@ -343,7 +343,8 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
         requestedColumns,
         scanBuilder(requestedColumns, candidatePredicates, pushedFilters),
         relation.relation)
-      execution.Project(projects, filterCondition.map(execution.Filter(_, scan)).getOrElse(scan))
+      execution.TungstenProject(
+        projects, filterCondition.map(execution.Filter(_, scan)).getOrElse(scan))
     }
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala
index 997f7f494f4a3..fb961d97c3c3c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala
@@ -44,27 +44,15 @@ trait HashJoin {
 
   override def output: Seq[Attribute] = left.output ++ right.output
 
-  protected[this] def isUnsafeMode: Boolean = {
-    UnsafeProjection.canSupport(buildKeys) && UnsafeProjection.canSupport(self.schema)
-  }
-
-  override def outputsUnsafeRows: Boolean = isUnsafeMode
-  override def canProcessUnsafeRows: Boolean = isUnsafeMode
-  override def canProcessSafeRows: Boolean = !isUnsafeMode
+  override def outputsUnsafeRows: Boolean = true
+  override def canProcessUnsafeRows: Boolean = true
+  override def canProcessSafeRows: Boolean = false
 
   protected def buildSideKeyGenerator: Projection =
-    if (isUnsafeMode) {
-      UnsafeProjection.create(buildKeys, buildPlan.output)
-    } else {
-      newMutableProjection(buildKeys, buildPlan.output)()
-    }
+    UnsafeProjection.create(buildKeys, buildPlan.output)
 
   protected def streamSideKeyGenerator: Projection =
-    if (isUnsafeMode) {
-      UnsafeProjection.create(streamedKeys, streamedPlan.output)
-    } else {
-      newMutableProjection(streamedKeys, streamedPlan.output)()
-    }
+    UnsafeProjection.create(streamedKeys, streamedPlan.output)
 
   protected def hashJoin(
       streamIter: Iterator[InternalRow],
@@ -79,13 +67,8 @@ trait HashJoin {
 
       // Mutable per row objects.
       private[this] val joinRow = new JoinedRow
-      private[this] val resultProjection: (InternalRow) => InternalRow = {
-        if (isUnsafeMode) {
-          UnsafeProjection.create(self.schema)
-        } else {
-          identity[InternalRow]
-        }
-      }
+      private[this] val resultProjection: (InternalRow) => InternalRow =
+        UnsafeProjection.create(self.schema)
 
       private[this] val joinKeys = streamSideKeyGenerator
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashOuterJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashOuterJoin.scala
index 3633f356b014b..ed626fef56af7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashOuterJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashOuterJoin.scala
@@ -64,38 +64,18 @@ trait HashOuterJoin {
         s"HashOuterJoin should not take $x as the JoinType")
   }
 
-  protected[this] def isUnsafeMode: Boolean = {
-    joinType != FullOuter &&
-      UnsafeProjection.canSupport(buildKeys) &&
-      UnsafeProjection.canSupport(self.schema)
-  }
-
-  override def outputsUnsafeRows: Boolean = isUnsafeMode
-  override def canProcessUnsafeRows: Boolean = isUnsafeMode
-  override def canProcessSafeRows: Boolean = !isUnsafeMode
+  override def outputsUnsafeRows: Boolean = true
+  override def canProcessUnsafeRows: Boolean = true
+  override def canProcessSafeRows: Boolean = false
 
   protected def buildKeyGenerator: Projection =
-    if (isUnsafeMode) {
-      UnsafeProjection.create(buildKeys, buildPlan.output)
-    } else {
-      newMutableProjection(buildKeys, buildPlan.output)()
-    }
+    UnsafeProjection.create(buildKeys, buildPlan.output)
 
-  protected[this] def streamedKeyGenerator: Projection = {
-    if (isUnsafeMode) {
-      UnsafeProjection.create(streamedKeys, streamedPlan.output)
-    } else {
-      newProjection(streamedKeys, streamedPlan.output)
-    }
-  }
+  protected[this] def streamedKeyGenerator: Projection =
+    UnsafeProjection.create(streamedKeys, streamedPlan.output)
 
-  protected[this] def resultProjection: InternalRow => InternalRow = {
-    if (isUnsafeMode) {
-      UnsafeProjection.create(self.schema)
-    } else {
-      identity[InternalRow]
-    }
-  }
+  protected[this] def resultProjection: InternalRow => InternalRow =
+    UnsafeProjection.create(self.schema)
 
   @transient private[this] lazy val DUMMY_LIST = CompactBuffer[InternalRow](null)
   @transient protected[this] lazy val EMPTY_LIST = CompactBuffer[InternalRow]()
@@ -173,8 +153,12 @@ trait HashOuterJoin {
   }
 
   protected[this] def fullOuterIterator(
-      key: InternalRow, leftIter: Iterable[InternalRow], rightIter: Iterable[InternalRow],
-      joinedRow: JoinedRow, numOutputRows: LongSQLMetric): Iterator[InternalRow] = {
+      key: InternalRow,
+      leftIter: Iterable[InternalRow],
+      rightIter: Iterable[InternalRow],
+      joinedRow: JoinedRow,
+      resultProjection: InternalRow => InternalRow,
+      numOutputRows: LongSQLMetric): Iterator[InternalRow] = {
     if (!key.anyNull) {
       // Store the positions of records in right, if one of its associated row satisfy
       // the join condition.
@@ -191,7 +175,7 @@ trait HashOuterJoin {
             matched = true
             // if the row satisfy the join condition, add its index into the matched set
             rightMatchedSet.add(idx)
-            joinedRow.copy()
+            resultProjection(joinedRow)
 
         } ++ DUMMY_LIST.filter(_ => !matched).map( _ => {
           // 2. For those unmatched records in left, append additional records with empty right.
@@ -201,7 +185,7 @@ trait HashOuterJoin {
           // of the records in right side.
           // If we didn't get any proper row, then append a single row with empty right.
           numOutputRows += 1
-          joinedRow.withRight(rightNullRow).copy()
+          resultProjection(joinedRow.withRight(rightNullRow))
         })
       } ++ rightIter.zipWithIndex.collect {
         // 3. For those unmatched records in right, append additional records with empty left.
@@ -210,15 +194,15 @@ trait HashOuterJoin {
         // in the matched set.
         case (r, idx) if !rightMatchedSet.contains(idx) =>
           numOutputRows += 1
-          joinedRow(leftNullRow, r).copy()
+          resultProjection(joinedRow(leftNullRow, r))
       }
     } else {
       leftIter.iterator.map[InternalRow] { l =>
         numOutputRows += 1
-        joinedRow(l, rightNullRow).copy()
+        resultProjection(joinedRow(l, rightNullRow))
       } ++ rightIter.iterator.map[InternalRow] { r =>
         numOutputRows += 1
-        joinedRow(leftNullRow, r).copy()
+        resultProjection(joinedRow(leftNullRow, r))
       }
     }
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashSemiJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashSemiJoin.scala
index c7d13e0a72a87..f23a1830e91c1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashSemiJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashSemiJoin.scala
@@ -33,30 +33,15 @@ trait HashSemiJoin {
 
   override def output: Seq[Attribute] = left.output
 
-  protected[this] def supportUnsafe: Boolean = {
-    UnsafeProjection.canSupport(leftKeys) &&
-      UnsafeProjection.canSupport(rightKeys) &&
-      UnsafeProjection.canSupport(left.schema) &&
-      UnsafeProjection.canSupport(right.schema)
-  }
-
-  override def outputsUnsafeRows: Boolean = supportUnsafe
-  override def canProcessUnsafeRows: Boolean = supportUnsafe
-  override def canProcessSafeRows: Boolean = !supportUnsafe
+  override def outputsUnsafeRows: Boolean = true
+  override def canProcessUnsafeRows: Boolean = true
+  override def canProcessSafeRows: Boolean = false
 
   protected def leftKeyGenerator: Projection =
-    if (supportUnsafe) {
-      UnsafeProjection.create(leftKeys, left.output)
-    } else {
-      newMutableProjection(leftKeys, left.output)()
-    }
+    UnsafeProjection.create(leftKeys, left.output)
 
   protected def rightKeyGenerator: Projection =
-    if (supportUnsafe) {
-      UnsafeProjection.create(rightKeys, right.output)
-    } else {
-      newMutableProjection(rightKeys, right.output)()
-    }
+    UnsafeProjection.create(rightKeys, right.output)
 
   @transient private lazy val boundCondition =
     newPredicate(condition.getOrElse(Literal(true)), left.output ++ right.output)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoin.scala
index 7aee8e3dd3fce..4bf7b521c77d3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoin.scala
@@ -53,15 +53,9 @@ case class SortMergeJoin(
   override def requiredChildOrdering: Seq[Seq[SortOrder]] =
     requiredOrders(leftKeys) :: requiredOrders(rightKeys) :: Nil
 
-  protected[this] def isUnsafeMode: Boolean = {
-    UnsafeProjection.canSupport(leftKeys) &&
-      UnsafeProjection.canSupport(rightKeys) &&
-      UnsafeProjection.canSupport(schema)
-  }
-
-  override def outputsUnsafeRows: Boolean = isUnsafeMode
-  override def canProcessUnsafeRows: Boolean = isUnsafeMode
-  override def canProcessSafeRows: Boolean = !isUnsafeMode
+  override def outputsUnsafeRows: Boolean = true
+  override def canProcessUnsafeRows: Boolean = true
+  override def canProcessSafeRows: Boolean = false
 
   private def requiredOrders(keys: Seq[Expression]): Seq[SortOrder] = {
     // This must be ascending in order to agree with the `keyOrdering` defined in `doExecute()`.
@@ -76,26 +70,10 @@ case class SortMergeJoin(
     left.execute().zipPartitions(right.execute()) { (leftIter, rightIter) =>
       new RowIterator {
         // The projection used to extract keys from input rows of the left child.
-        private[this] val leftKeyGenerator = {
-          if (isUnsafeMode) {
-            // It is very important to use UnsafeProjection if input rows are UnsafeRows.
-            // Otherwise, GenerateProjection will cause wrong results.
-            UnsafeProjection.create(leftKeys, left.output)
-          } else {
-            newProjection(leftKeys, left.output)
-          }
-        }
+        private[this] val leftKeyGenerator = UnsafeProjection.create(leftKeys, left.output)
 
         // The projection used to extract keys from input rows of the right child.
-        private[this] val rightKeyGenerator = {
-          if (isUnsafeMode) {
-            // It is very important to use UnsafeProjection if input rows are UnsafeRows.
-            // Otherwise, GenerateProjection will cause wrong results.
-            UnsafeProjection.create(rightKeys, right.output)
-          } else {
-            newProjection(rightKeys, right.output)
-          }
-        }
+        private[this] val rightKeyGenerator = UnsafeProjection.create(rightKeys, right.output)
 
         // An ordering that can be used to compare keys from both sides.
         private[this] val keyOrdering = newNaturalAscendingOrdering(leftKeys.map(_.dataType))
@@ -112,13 +90,8 @@ case class SortMergeJoin(
           numRightRows
         )
         private[this] val joinRow = new JoinedRow
-        private[this] val resultProjection: (InternalRow) => InternalRow = {
-          if (isUnsafeMode) {
-            UnsafeProjection.create(schema)
-          } else {
-            identity[InternalRow]
-          }
-        }
+        private[this] val resultProjection: (InternalRow) => InternalRow =
+          UnsafeProjection.create(schema)
 
         override def advanceNext(): Boolean = {
           if (currentMatchIdx == -1 || currentMatchIdx == currentRightMatches.length) {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
index 5f1590c463836..efaa69c1d3227 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeOuterJoin.scala
@@ -89,31 +89,15 @@ case class SortMergeOuterJoin(
     keys.map(SortOrder(_, Ascending))
   }
 
-  private def isUnsafeMode: Boolean = {
-    UnsafeProjection.canSupport(leftKeys) &&
-      UnsafeProjection.canSupport(rightKeys) &&
-      UnsafeProjection.canSupport(schema)
-  }
-
-  override def outputsUnsafeRows: Boolean = isUnsafeMode
-  override def canProcessUnsafeRows: Boolean = isUnsafeMode
-  override def canProcessSafeRows: Boolean = !isUnsafeMode
+  override def outputsUnsafeRows: Boolean = true
+  override def canProcessUnsafeRows: Boolean = true
+  override def canProcessSafeRows: Boolean = false
 
-  private def createLeftKeyGenerator(): Projection = {
-    if (isUnsafeMode) {
-      UnsafeProjection.create(leftKeys, left.output)
-    } else {
-      newProjection(leftKeys, left.output)
-    }
-  }
+  private def createLeftKeyGenerator(): Projection =
+    UnsafeProjection.create(leftKeys, left.output)
 
-  private def createRightKeyGenerator(): Projection = {
-    if (isUnsafeMode) {
-      UnsafeProjection.create(rightKeys, right.output)
-    } else {
-      newProjection(rightKeys, right.output)
-    }
-  }
+  private def createRightKeyGenerator(): Projection =
+    UnsafeProjection.create(rightKeys, right.output)
 
   override def doExecute(): RDD[InternalRow] = {
     val numLeftRows = longMetric("numLeftRows")
@@ -130,13 +114,7 @@ case class SortMergeOuterJoin(
           (r: InternalRow) => true
         }
       }
-      val resultProj: InternalRow => InternalRow = {
-        if (isUnsafeMode) {
-          UnsafeProjection.create(schema)
-        } else {
-          identity[InternalRow]
-        }
-      }
+      val resultProj: InternalRow => InternalRow = UnsafeProjection.create(schema)
 
       joinType match {
         case LeftOuter =>
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/local/BinaryHashJoinNode.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/local/BinaryHashJoinNode.scala
index 52dcb9e43c4e8..3dcef94095647 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/local/BinaryHashJoinNode.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/local/BinaryHashJoinNode.scala
@@ -50,11 +50,7 @@ case class BinaryHashJoinNode(
   private def buildSideKeyGenerator: Projection = {
     // We are expecting the data types of buildKeys and streamedKeys are the same.
     assert(buildKeys.map(_.dataType) == streamedKeys.map(_.dataType))
-    if (isUnsafeMode) {
-      UnsafeProjection.create(buildKeys, buildNode.output)
-    } else {
-      newMutableProjection(buildKeys, buildNode.output)()
-    }
+    UnsafeProjection.create(buildKeys, buildNode.output)
   }
 
   protected override def doOpen(): Unit = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/local/HashJoinNode.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/local/HashJoinNode.scala
index aef655727fbbb..fd7948ffa9a9b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/local/HashJoinNode.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/local/HashJoinNode.scala
@@ -45,17 +45,8 @@ trait HashJoinNode {
   private[this] var hashed: HashedRelation = _
   private[this] var joinKeys: Projection = _
 
-  protected def isUnsafeMode: Boolean = {
-    UnsafeProjection.canSupport(schema) && UnsafeProjection.canSupport(streamedKeys)
-  }
-
-  private def streamSideKeyGenerator: Projection = {
-    if (isUnsafeMode) {
-      UnsafeProjection.create(streamedKeys, streamedNode.output)
-    } else {
-      newMutableProjection(streamedKeys, streamedNode.output)()
-    }
-  }
+  private def streamSideKeyGenerator: Projection =
+    UnsafeProjection.create(streamedKeys, streamedNode.output)
 
   /**
    * Sets the HashedRelation used by this node. This method needs to be called after
@@ -73,13 +64,7 @@ trait HashJoinNode {
   override def open(): Unit = {
     doOpen()
     joinRow = new JoinedRow
-    resultProjection = {
-      if (isUnsafeMode) {
-        UnsafeProjection.create(schema)
-      } else {
-        identity[InternalRow]
-      }
-    }
+    resultProjection = UnsafeProjection.create(schema)
     joinKeys = streamSideKeyGenerator
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/rowFormatConverters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/rowFormatConverters.scala
index 0e601cd2cab5d..5f8fc2de8b46d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/rowFormatConverters.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/rowFormatConverters.scala
@@ -28,8 +28,6 @@ import org.apache.spark.sql.catalyst.rules.Rule
  */
 case class ConvertToUnsafe(child: SparkPlan) extends UnaryNode {
 
-  require(UnsafeProjection.canSupport(child.schema), s"Cannot convert ${child.schema} to Unsafe")
-
   override def output: Seq[Attribute] = child.output
   override def outputPartitioning: Partitioning = child.outputPartitioning
   override def outputOrdering: Seq[SortOrder] = child.outputOrdering
@@ -97,18 +95,10 @@ private[sql] object EnsureRowFormats extends Rule[SparkPlan] {
     case operator: SparkPlan if handlesBothSafeAndUnsafeRows(operator) =>
       if (operator.children.map(_.outputsUnsafeRows).toSet.size != 1) {
         // If this operator's children produce both unsafe and safe rows,
-        // convert everything unsafe rows if all the schema of them are support by UnsafeRow
-        if (operator.children.forall(c => UnsafeProjection.canSupport(c.schema))) {
-          operator.withNewChildren {
-            operator.children.map {
-              c => if (!c.outputsUnsafeRows) ConvertToUnsafe(c) else c
-            }
-          }
-        } else {
-          operator.withNewChildren {
-            operator.children.map {
-              c => if (c.outputsUnsafeRows) ConvertToSafe(c) else c
-            }
+        // convert everything unsafe rows.
+        operator.withNewChildren {
+          operator.children.map {
+            c => if (!c.outputsUnsafeRows) ConvertToUnsafe(c) else c
           }
         }
       } else {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala
index 1a3832a698b61..47fe70ab154ec 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala
@@ -145,12 +145,3 @@ case class TungstenSort(
   }
 
 }
-
-object TungstenSort {
-  /**
-   * Return true if UnsafeExternalSort can sort rows with the given schema, false otherwise.
-   */
-  def supportsSchema(schema: StructType): Boolean = {
-    UnsafeExternalRowSorter.supportsSchema(schema)
-  }
-}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
index fa559c9c64005..010df2a341589 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
@@ -20,7 +20,7 @@ package org.apache.spark.sql
 import org.apache.spark.sql.catalyst.expressions.NamedExpression
 import org.scalatest.Matchers._
 
-import org.apache.spark.sql.execution.{Project, TungstenProject}
+import org.apache.spark.sql.execution.TungstenProject
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
@@ -615,7 +615,6 @@ class ColumnExpressionSuite extends QueryTest with SharedSQLContext {
 
     def checkNumProjects(df: DataFrame, expectedNumProjects: Int): Unit = {
       val projects = df.queryExecution.executedPlan.collect {
-        case project: Project => project
         case tungstenProject: TungstenProject => tungstenProject
       }
       assert(projects.size === expectedNumProjects)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala
index 85486c08894c9..7c860d1d58d5a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala
@@ -74,7 +74,6 @@ class TungstenSortSuite extends SparkPlanTest with SharedSQLContext {
         sparkContext.parallelize(Random.shuffle(inputData).map(v => Row(v))),
         StructType(StructField("a", dataType, nullable = true) :: Nil)
       )
-      assert(TungstenSort.supportsSchema(inputDf.schema))
       checkThatPlansAgree(
         inputDf,
         plan => ConvertToSafe(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/local/HashJoinNodeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/local/HashJoinNodeSuite.scala
index 44b0d9d4102a1..c30327185e169 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/local/HashJoinNodeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/local/HashJoinNodeSuite.scala
@@ -42,15 +42,7 @@ class HashJoinNodeSuite extends LocalNodeTest {
       buildKeys: Seq[Expression],
       buildNode: LocalNode): HashedRelation = {
 
-    val isUnsafeMode = UnsafeProjection.canSupport(buildKeys)
-
-    val buildSideKeyGenerator =
-      if (isUnsafeMode) {
-        UnsafeProjection.create(buildKeys, buildNode.output)
-      } else {
-        new InterpretedMutableProjection(buildKeys, buildNode.output)
-      }
-
+    val buildSideKeyGenerator = UnsafeProjection.create(buildKeys, buildNode.output)
     buildNode.prepare()
     buildNode.open()
     val hashedRelation = HashedRelation(buildNode, buildSideKeyGenerator)
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveTypeCoercionSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveTypeCoercionSuite.scala
index 197e9bfb02c4e..4cf4e13890294 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveTypeCoercionSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveTypeCoercionSuite.scala
@@ -18,7 +18,7 @@
 package org.apache.spark.sql.hive.execution
 
 import org.apache.spark.sql.catalyst.expressions.{Cast, EqualTo}
-import org.apache.spark.sql.execution.Project
+import org.apache.spark.sql.execution.TungstenProject
 import org.apache.spark.sql.hive.test.TestHive
 
 /**
@@ -43,7 +43,9 @@ class HiveTypeCoercionSuite extends HiveComparisonTest {
 
   test("[SPARK-2210] boolean cast on boolean value should be removed") {
     val q = "select cast(cast(key=0 as boolean) as boolean) from src"
-    val project = TestHive.sql(q).queryExecution.executedPlan.collect { case e: Project => e }.head
+    val project = TestHive.sql(q).queryExecution.executedPlan.collect {
+      case e: TungstenProject => e
+    }.head
 
     // No cast expression introduced
     project.transformAllExpressions { case c: Cast =>
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
index e866493ee6c96..b6db6225805a1 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
@@ -151,7 +151,7 @@ class ParquetHadoopFsRelationSuite extends HadoopFsRelationTest {
       val df = sqlContext.read.parquet(path).filter('a === 0).select('b)
       val physicalPlan = df.queryExecution.executedPlan
 
-      assert(physicalPlan.collect { case p: execution.Project => p }.length === 1)
+      assert(physicalPlan.collect { case p: execution.TungstenProject => p }.length === 1)
       assert(physicalPlan.collect { case p: execution.Filter => p }.length === 1)
     }
   }

From 08660a0bc903a87b72f8ffd8b9b02fd7ee379cf7 Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Thu, 12 Nov 2015 17:23:24 +0100
Subject: [PATCH 207/518] [BUILD][MINOR] Remove non-exist yarnStable module in
 Sbt project

Remove some old yarn related building codes, please review, thanks a lot.

Author: jerryshao <sshao@hortonworks.com>

Closes #9625 from jerryshao/remove-old-module.
---
 project/SparkBuild.scala | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index 570c9e50ed1f4..67724c4e9e411 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -42,9 +42,9 @@ object BuildCommons {
       "streaming-flume", "streaming-kafka", "streaming-mqtt", "streaming-twitter",
       "streaming-zeromq", "launcher", "unsafe", "test-tags").map(ProjectRef(buildLocation, _))
 
-  val optionallyEnabledProjects@Seq(yarn, yarnStable, java8Tests, sparkGangliaLgpl,
+  val optionallyEnabledProjects@Seq(yarn, java8Tests, sparkGangliaLgpl,
     streamingKinesisAsl, dockerIntegrationTests) =
-    Seq("yarn", "yarn-stable", "java8-tests", "ganglia-lgpl", "streaming-kinesis-asl",
+    Seq("yarn", "java8-tests", "ganglia-lgpl", "streaming-kinesis-asl",
       "docker-integration-tests").map(ProjectRef(buildLocation, _))
 
   val assemblyProjects@Seq(assembly, examples, networkYarn, streamingFlumeAssembly, streamingKafkaAssembly, streamingMqttAssembly, streamingKinesisAslAssembly) =
@@ -72,7 +72,6 @@ object SparkBuild extends PomBuild {
   // Provides compatibility for older versions of the Spark build
   def backwardCompatibility = {
     import scala.collection.mutable
-    var isAlphaYarn = false
     var profiles: mutable.Seq[String] = mutable.Seq("sbt")
     // scalastyle:off println
     if (Properties.envOrNone("SPARK_GANGLIA_LGPL").isDefined) {
@@ -85,7 +84,6 @@ object SparkBuild extends PomBuild {
     }
     Properties.envOrNone("SPARK_HADOOP_VERSION") match {
       case Some(v) =>
-        if (v.matches("0.23.*")) isAlphaYarn = true
         println("NOTE: SPARK_HADOOP_VERSION is deprecated, please use -Dhadoop.version=" + v)
         System.setProperty("hadoop.version", v)
       case None =>

From df0e318152165c8e50793aff13aaca5d2d9b8b9d Mon Sep 17 00:00:00 2001
From: Gaurav Kumar <gauravkumar37@gmail.com>
Date: Thu, 12 Nov 2015 12:14:00 -0800
Subject: [PATCH 208/518] Fixed error in scaladoc of convertToCanonicalEdges

The code convertToCanonicalEdges is such that srcIds are smaller than dstIds but the scaladoc suggested otherwise. Have fixed the same.

Author: Gaurav Kumar <gauravkumar37@gmail.com>

Closes #9666 from gauravkumar37/patch-1.
---
 graphx/src/main/scala/org/apache/spark/graphx/GraphOps.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/graphx/src/main/scala/org/apache/spark/graphx/GraphOps.scala b/graphx/src/main/scala/org/apache/spark/graphx/GraphOps.scala
index 9451ff1e5c0e2..9827dfab8684a 100644
--- a/graphx/src/main/scala/org/apache/spark/graphx/GraphOps.scala
+++ b/graphx/src/main/scala/org/apache/spark/graphx/GraphOps.scala
@@ -282,7 +282,7 @@ class GraphOps[VD: ClassTag, ED: ClassTag](graph: Graph[VD, ED]) extends Seriali
    * Convert bi-directional edges into uni-directional ones.
    * Some graph algorithms (e.g., TriangleCount) assume that an input graph
    * has its edges in canonical direction.
-   * This function rewrites the vertex ids of edges so that srcIds are bigger
+   * This function rewrites the vertex ids of edges so that srcIds are smaller
    * than dstIds, and merges the duplicated edges.
    *
    * @param mergeFunc the user defined reduce function which should

From 4fe99c72c60646b1372bb2c089c6fc7c4fa11644 Mon Sep 17 00:00:00 2001
From: Cheng Lian <lian@databricks.com>
Date: Thu, 12 Nov 2015 12:17:51 -0800
Subject: [PATCH 209/518] [SPARK-11191][SQL] Looks up temporary function using
 execution Hive client

When looking up Hive temporary functions, we should always use the `SessionState` within the execution Hive client, since temporary functions are registered there.

Author: Cheng Lian <lian@databricks.com>

Closes #9664 from liancheng/spark-11191.fix-temp-function.
---
 .../HiveThriftServer2Suites.scala             | 45 +++++++++++++++++++
 .../apache/spark/sql/hive/HiveContext.scala   |  2 +-
 .../org/apache/spark/sql/hive/hiveUDFs.scala  | 14 ++++--
 3 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
index 5903b9e71cdd2..eb1895f263d70 100644
--- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
+++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
@@ -463,6 +463,51 @@ class HiveThriftBinaryServerSuite extends HiveThriftJdbcTest {
       assert(conf.get("spark.sql.hive.version") === Some("1.2.1"))
     }
   }
+
+  test("SPARK-11595 ADD JAR with input path having URL scheme") {
+    withJdbcStatement { statement =>
+      val jarPath = "../hive/src/test/resources/TestUDTF.jar"
+      val jarURL = s"file://${System.getProperty("user.dir")}/$jarPath"
+
+      Seq(
+        s"ADD JAR $jarURL",
+        s"""CREATE TEMPORARY FUNCTION udtf_count2
+           |AS 'org.apache.spark.sql.hive.execution.GenericUDTFCount2'
+         """.stripMargin
+      ).foreach(statement.execute)
+
+      val rs1 = statement.executeQuery("DESCRIBE FUNCTION udtf_count2")
+
+      assert(rs1.next())
+      assert(rs1.getString(1) === "Function: udtf_count2")
+
+      assert(rs1.next())
+      assertResult("Class: org.apache.spark.sql.hive.execution.GenericUDTFCount2") {
+        rs1.getString(1)
+      }
+
+      assert(rs1.next())
+      assert(rs1.getString(1) === "Usage: To be added.")
+
+      val dataPath = "../hive/src/test/resources/data/files/kv1.txt"
+
+      Seq(
+        s"CREATE TABLE test_udtf(key INT, value STRING)",
+        s"LOAD DATA LOCAL INPATH '$dataPath' OVERWRITE INTO TABLE test_udtf"
+      ).foreach(statement.execute)
+
+      val rs2 = statement.executeQuery(
+        "SELECT key, cc FROM test_udtf LATERAL VIEW udtf_count2(value) dd AS cc")
+
+      assert(rs2.next())
+      assert(rs2.getInt(1) === 97)
+      assert(rs2.getInt(2) === 500)
+
+      assert(rs2.next())
+      assert(rs2.getInt(1) === 97)
+      assert(rs2.getInt(2) === 500)
+    }
+  }
 }
 
 class HiveThriftHttpServerSuite extends HiveThriftJdbcTest {
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index ba6204633b9ca..0c473799cc991 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -454,7 +454,7 @@ class HiveContext private[hive](
   // Note that HiveUDFs will be overridden by functions registered in this context.
   @transient
   override protected[sql] lazy val functionRegistry: FunctionRegistry =
-    new HiveFunctionRegistry(FunctionRegistry.builtin.copy()) {
+    new HiveFunctionRegistry(FunctionRegistry.builtin.copy(), this) {
       override def lookupFunction(name: String, children: Seq[Expression]): Expression = {
         // Hive Registry need current database to lookup function
         // TODO: the current database of executionHive should be consistent with metadataHive
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
index a9db70119d011..e6fe2ad5f23b6 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
@@ -46,17 +46,23 @@ import org.apache.spark.sql.hive.HiveShim._
 import org.apache.spark.sql.types._
 
 
-private[hive] class HiveFunctionRegistry(underlying: analysis.FunctionRegistry)
+private[hive] class HiveFunctionRegistry(
+    underlying: analysis.FunctionRegistry,
+    hiveContext: HiveContext)
   extends analysis.FunctionRegistry with HiveInspectors {
 
-  def getFunctionInfo(name: String): FunctionInfo = FunctionRegistry.getFunctionInfo(name)
+  def getFunctionInfo(name: String): FunctionInfo = {
+    hiveContext.executionHive.withHiveState {
+      FunctionRegistry.getFunctionInfo(name)
+    }
+  }
 
   override def lookupFunction(name: String, children: Seq[Expression]): Expression = {
     Try(underlying.lookupFunction(name, children)).getOrElse {
       // We only look it up to see if it exists, but do not include it in the HiveUDF since it is
       // not always serializable.
       val functionInfo: FunctionInfo =
-        Option(FunctionRegistry.getFunctionInfo(name.toLowerCase)).getOrElse(
+        Option(getFunctionInfo(name.toLowerCase)).getOrElse(
           throw new AnalysisException(s"undefined function $name"))
 
       val functionClassName = functionInfo.getFunctionClass.getName
@@ -110,7 +116,7 @@ private[hive] class HiveFunctionRegistry(underlying: analysis.FunctionRegistry)
   override def lookupFunction(name: String): Option[ExpressionInfo] = {
     underlying.lookupFunction(name).orElse(
     Try {
-      val info = FunctionRegistry.getFunctionInfo(name)
+      val info = getFunctionInfo(name)
       val annotation = info.getFunctionClass.getAnnotation(classOf[Description])
       if (annotation != null) {
         Some(new ExpressionInfo(

From f5a9526fec284cccd0755d190c91e8d9999f7877 Mon Sep 17 00:00:00 2001
From: hyukjinkwon <gurwls223@gmail.com>
Date: Thu, 12 Nov 2015 12:29:50 -0800
Subject: [PATCH 210/518] [SPARK-10113][SQL] Explicit error message for
 unsigned Parquet logical types

Parquet supports some unsigned datatypes. However, Since Spark does not support unsigned datatypes, it needs to emit an exception with a clear message rather then with the one saying illegal datatype.

Author: hyukjinkwon <gurwls223@gmail.com>

Closes #9646 from HyukjinKwon/SPARK-10113.
---
 .../parquet/CatalystSchemaConverter.scala     |  7 ++++++
 .../datasources/parquet/ParquetIOSuite.scala  | 24 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystSchemaConverter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystSchemaConverter.scala
index 7f3394c20ed3d..f28a18e2756e4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystSchemaConverter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystSchemaConverter.scala
@@ -108,6 +108,9 @@ private[parquet] class CatalystSchemaConverter(
     def typeString =
       if (originalType == null) s"$typeName" else s"$typeName ($originalType)"
 
+    def typeNotSupported() =
+      throw new AnalysisException(s"Parquet type not supported: $typeString")
+
     def typeNotImplemented() =
       throw new AnalysisException(s"Parquet type not yet supported: $typeString")
 
@@ -142,6 +145,9 @@ private[parquet] class CatalystSchemaConverter(
           case INT_32 | null => IntegerType
           case DATE => DateType
           case DECIMAL => makeDecimalType(MAX_PRECISION_FOR_INT32)
+          case UINT_8 => typeNotSupported()
+          case UINT_16 => typeNotSupported()
+          case UINT_32 => typeNotSupported()
           case TIME_MILLIS => typeNotImplemented()
           case _ => illegalType()
         }
@@ -150,6 +156,7 @@ private[parquet] class CatalystSchemaConverter(
         originalType match {
           case INT_64 | null => LongType
           case DECIMAL => makeDecimalType(MAX_PRECISION_FOR_INT64)
+          case UINT_64 => typeNotSupported()
           case TIMESTAMP_MILLIS => typeNotImplemented()
           case _ => illegalType()
         }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index 72744799897be..82a42d68fedc1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -206,6 +206,30 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
     }
   }
 
+  test("SPARK-10113 Support for unsigned Parquet logical types") {
+    val parquetSchema = MessageTypeParser.parseMessageType(
+      """message root {
+        |  required int32 c(UINT_32);
+        |}
+      """.stripMargin)
+
+    withTempPath { location =>
+      val extraMetadata = Map.empty[String, String].asJava
+      val fileMetadata = new FileMetaData(parquetSchema, extraMetadata, "Spark")
+      val path = new Path(location.getCanonicalPath)
+      val footer = List(
+        new Footer(path, new ParquetMetadata(fileMetadata, Collections.emptyList()))
+      ).asJava
+
+      ParquetFileWriter.writeMetadataFile(sparkContext.hadoopConfiguration, path, footer)
+
+      val errorMessage = intercept[Throwable] {
+        sqlContext.read.parquet(path.toString).printSchema()
+      }.toString
+      assert(errorMessage.contains("Parquet type not supported"))
+    }
+  }
+
   test("compression codec") {
     def compressionCodecFor(path: String, codecName: String): String = {
       val codecs = for {

From d292f74831de7e69c852ed26d9c15df85b4fb568 Mon Sep 17 00:00:00 2001
From: JihongMa <linlin200605@gmail.com>
Date: Thu, 12 Nov 2015 13:47:34 -0800
Subject: [PATCH 211/518] [SPARK-11420] Updating Stddev support via Imperative
 Aggregate

switched stddev support from DeclarativeAggregate to ImperativeAggregate.

Author: JihongMa <linlin200605@gmail.com>

Closes #9380 from JihongMA/SPARK-11420.
---
 R/pkg/inst/tests/test_sparkSQL.R              |   4 +-
 python/pyspark/sql/dataframe.py               |   2 +-
 .../catalyst/analysis/HiveTypeCoercion.scala  |   6 +-
 .../expressions/aggregate/Kurtosis.scala      |   4 +-
 .../expressions/aggregate/Skewness.scala      |   4 +-
 .../expressions/aggregate/Stddev.scala        | 128 +++++-------------
 .../org/apache/spark/sql/functions.scala      |   2 +-
 .../spark/sql/DataFrameAggregateSuite.scala   |   4 +-
 .../org/apache/spark/sql/DataFrameSuite.scala |   2 +-
 .../org/apache/spark/sql/SQLQuerySuite.scala  |  11 +-
 10 files changed, 52 insertions(+), 115 deletions(-)

diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 9e453a1e7c2f4..af024e6183a37 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -1007,7 +1007,7 @@ test_that("group by, agg functions", {
   df3 <- agg(gd, age = "stddev")
   expect_is(df3, "DataFrame")
   df3_local <- collect(df3)
-  expect_equal(0, df3_local[df3_local$name == "Andy",][1, 2])
+  expect_true(is.nan(df3_local[df3_local$name == "Andy",][1, 2]))
 
   df4 <- agg(gd, sumAge = sum(df$age))
   expect_is(df4, "DataFrame")
@@ -1038,7 +1038,7 @@ test_that("group by, agg functions", {
   df7 <- agg(gd2, value = "stddev")
   df7_local <- collect(df7)
   expect_true(abs(df7_local[df7_local$name == "ID1",][1, 2] - 6.928203) < 1e-6)
-  expect_equal(0, df7_local[df7_local$name == "ID2",][1, 2])
+  expect_true(is.nan(df7_local[df7_local$name == "ID2",][1, 2]))
 
   mockLines3 <- c("{\"name\":\"Andy\", \"age\":30}",
                   "{\"name\":\"Andy\", \"age\":30}",
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index 0dd75ba7ca820..ad6ad0235a904 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -761,7 +761,7 @@ def describe(self, *cols):
         +-------+------------------+-----+
         |  count|                 2|    2|
         |   mean|               3.5| null|
-        | stddev|2.1213203435596424| null|
+        | stddev|2.1213203435596424|  NaN|
         |    min|                 2|Alice|
         |    max|                 5|  Bob|
         +-------+------------------+-----+
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
index bf2bff0243fa3..92188ee54fd28 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
@@ -297,8 +297,10 @@ object HiveTypeCoercion {
 
       case Sum(e @ StringType()) => Sum(Cast(e, DoubleType))
       case Average(e @ StringType()) => Average(Cast(e, DoubleType))
-      case StddevPop(e @ StringType()) => StddevPop(Cast(e, DoubleType))
-      case StddevSamp(e @ StringType()) => StddevSamp(Cast(e, DoubleType))
+      case StddevPop(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
+        StddevPop(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
+      case StddevSamp(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
+        StddevSamp(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
       case VariancePop(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
         VariancePop(Cast(e, DoubleType), mutableAggBufferOffset, inputAggBufferOffset)
       case VarianceSamp(e @ StringType(), mutableAggBufferOffset, inputAggBufferOffset) =>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
index bae78d98493b8..8fa3aac9f1a51 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
@@ -42,9 +42,11 @@ case class Kurtosis(child: Expression,
       s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
     val m2 = moments(2)
     val m4 = moments(4)
+
     if (n == 0.0 || m2 == 0.0) {
       Double.NaN
-    } else {
+    }
+    else {
       n * m4 / (m2 * m2) - 3.0
     }
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
index c593074fa2479..e1c01a5b82781 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
@@ -41,9 +41,11 @@ case class Skewness(child: Expression,
       s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
     val m2 = moments(2)
     val m3 = moments(3)
+
     if (n == 0.0 || m2 == 0.0) {
       Double.NaN
-    } else {
+    }
+    else {
       math.sqrt(n) * m3 / math.sqrt(m2 * m2 * m2)
     }
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
index 2748009623355..05dd5e3b22543 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
@@ -17,117 +17,55 @@
 
 package org.apache.spark.sql.catalyst.expressions.aggregate
 
-import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
-import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.util.TypeUtils
-import org.apache.spark.sql.types._
 
+case class StddevSamp(child: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0)
+  extends CentralMomentAgg(child) {
 
-// Compute the population standard deviation of a column
-case class StddevPop(child: Expression) extends StddevAgg(child) {
-  override def isSample: Boolean = false
-  override def prettyName: String = "stddev_pop"
-}
-
-
-// Compute the sample standard deviation of a column
-case class StddevSamp(child: Expression) extends StddevAgg(child) {
-  override def isSample: Boolean = true
-  override def prettyName: String = "stddev_samp"
-}
-
-
-// Compute standard deviation based on online algorithm specified here:
-// http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
-abstract class StddevAgg(child: Expression) extends DeclarativeAggregate {
+  def this(child: Expression) = this(child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
 
-  def isSample: Boolean
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
 
-  override def children: Seq[Expression] = child :: Nil
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
 
-  override def nullable: Boolean = true
-
-  override def dataType: DataType = resultType
-
-  override def inputTypes: Seq[AbstractDataType] = Seq(NumericType)
+  override def prettyName: String = "stddev_samp"
 
-  override def checkInputDataTypes(): TypeCheckResult =
-    TypeUtils.checkForNumericExpr(child.dataType, "function stddev")
+  override protected val momentOrder = 2
 
-  private lazy val resultType = DoubleType
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+    require(moments.length == momentOrder + 1,
+      s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
 
-  private lazy val count = AttributeReference("count", resultType)()
-  private lazy val avg = AttributeReference("avg", resultType)()
-  private lazy val mk = AttributeReference("mk", resultType)()
+    if (n == 0.0 || n == 1.0) Double.NaN else math.sqrt(moments(2) / (n - 1.0))
+  }
+}
 
-  override lazy val aggBufferAttributes = count :: avg :: mk :: Nil
+case class StddevPop(
+    child: Expression,
+    mutableAggBufferOffset: Int = 0,
+    inputAggBufferOffset: Int = 0)
+  extends CentralMomentAgg(child) {
 
-  override lazy val initialValues: Seq[Expression] = Seq(
-    /* count = */ Cast(Literal(0), resultType),
-    /* avg = */ Cast(Literal(0), resultType),
-    /* mk = */ Cast(Literal(0), resultType)
-  )
+  def this(child: Expression) = this(child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0)
 
-  override lazy val updateExpressions: Seq[Expression] = {
-    val value = Cast(child, resultType)
-    val newCount = count + Cast(Literal(1), resultType)
+  override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate =
+    copy(mutableAggBufferOffset = newMutableAggBufferOffset)
 
-    // update average
-    // avg = avg + (value - avg)/count
-    val newAvg = avg + (value - avg) / newCount
+  override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate =
+    copy(inputAggBufferOffset = newInputAggBufferOffset)
 
-    // update sum ofference from mean
-    // Mk = Mk + (value - preAvg) * (value - updatedAvg)
-    val newMk = mk + (value - avg) * (value - newAvg)
+  override def prettyName: String = "stddev_pop"
 
-    Seq(
-      /* count = */ If(IsNull(child), count, newCount),
-      /* avg = */ If(IsNull(child), avg, newAvg),
-      /* mk = */ If(IsNull(child), mk, newMk)
-    )
-  }
+  override protected val momentOrder = 2
 
-  override lazy val mergeExpressions: Seq[Expression] = {
-
-    // count merge
-    val newCount = count.left + count.right
-
-    // average merge
-    val newAvg = ((avg.left * count.left) + (avg.right * count.right)) / newCount
-
-    // update sum of square differences
-    val newMk = {
-      val avgDelta = avg.right - avg.left
-      val mkDelta = (avgDelta * avgDelta) * (count.left * count.right) / newCount
-      mk.left + mk.right + mkDelta
-    }
-
-    Seq(
-      /* count = */ If(IsNull(count.left), count.right,
-                       If(IsNull(count.right), count.left, newCount)),
-      /* avg = */ If(IsNull(avg.left), avg.right,
-                     If(IsNull(avg.right), avg.left, newAvg)),
-      /* mk = */ If(IsNull(mk.left), mk.right,
-                    If(IsNull(mk.right), mk.left, newMk))
-    )
-  }
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+    require(moments.length == momentOrder + 1,
+      s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
 
-  override lazy val evaluateExpression: Expression = {
-    // when count == 0, return null
-    // when count == 1, return 0
-    // when count >1
-    // stddev_samp = sqrt (mk/(count -1))
-    // stddev_pop = sqrt (mk/count)
-    val varCol =
-      if (isSample) {
-        mk / Cast(count - Cast(Literal(1), resultType), resultType)
-      } else {
-        mk / count
-      }
-
-    If(EqualTo(count, Cast(Literal(0), resultType)), Cast(Literal(null), resultType),
-      If(EqualTo(count, Cast(Literal(1), resultType)), Cast(Literal(0), resultType),
-        Cast(Sqrt(varCol), resultType)))
+    if (n == 0.0) Double.NaN else math.sqrt(moments(2) / n)
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index b6330e230afef..53cc6e0cda110 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -397,7 +397,7 @@ object functions extends LegacyFunctions {
   def stddev(e: Column): Column = withAggregateFunction { StddevSamp(e.expr) }
 
   /**
-   * Aggregate function: returns the unbiased sample standard deviation of
+   * Aggregate function: returns the sample standard deviation of
    * the expression in a group.
    *
    * @group agg_funcs
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index eb1ee266c5d28..432e8d17623a4 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -195,7 +195,7 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
   }
 
   test("stddev") {
-    val testData2ADev = math.sqrt(4 / 5.0)
+    val testData2ADev = math.sqrt(4.0 / 5.0)
     checkAnswer(
       testData2.agg(stddev('a), stddev_pop('a), stddev_samp('a)),
       Row(testData2ADev, math.sqrt(4 / 6.0), testData2ADev))
@@ -205,7 +205,7 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
     val emptyTableData = Seq.empty[(Int, Int)].toDF("a", "b")
     checkAnswer(
     emptyTableData.agg(stddev('a), stddev_pop('a), stddev_samp('a)),
-    Row(null, null, null))
+    Row(Double.NaN, Double.NaN, Double.NaN))
   }
 
   test("zero sum") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index e4f23fe17b757..35cdab50bdec9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -459,7 +459,7 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
     val emptyDescribeResult = Seq(
       Row("count", "0", "0"),
       Row("mean", null, null),
-      Row("stddev", null, null),
+      Row("stddev", "NaN", "NaN"),
       Row("min", null, null),
       Row("max", null, null))
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index 52a561d2e5452..167aea87de077 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -314,13 +314,6 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
       testCodeGen(
         "SELECT min(key) FROM testData3x",
         Row(1) :: Nil)
-      // STDDEV
-      testCodeGen(
-        "SELECT a, stddev(b), stddev_pop(b) FROM testData2 GROUP BY a",
-        (1 to 3).map(i => Row(i, math.sqrt(0.5), math.sqrt(0.25))))
-      testCodeGen(
-        "SELECT stddev(b), stddev_pop(b), stddev_samp(b) FROM testData2",
-        Row(math.sqrt(1.5 / 5), math.sqrt(1.5 / 6), math.sqrt(1.5 / 5)) :: Nil)
       // Some combinations.
       testCodeGen(
         """
@@ -341,8 +334,8 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
         Row(100, 1, 50.5, 300, 100) :: Nil)
       // Aggregate with Code generation handling all null values
       testCodeGen(
-        "SELECT  sum('a'), avg('a'), stddev('a'), count(null) FROM testData",
-        Row(null, null, null, 0) :: Nil)
+        "SELECT  sum('a'), avg('a'), count(null) FROM testData",
+        Row(null, null, 0) :: Nil)
     } finally {
       sqlContext.dropTempTable("testData3x")
     }

From 767d288b6b33a79d99324b70c2ac079fcf484a50 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Thu, 12 Nov 2015 14:29:16 -0800
Subject: [PATCH 212/518] [SPARK-11655][CORE] Fix deadlock in handling of
 launcher stop().

The stop() callback was trying to close the launcher connection in the
same thread that handles connection data, which ended up causing a
deadlock. So avoid that by dispatching the stop() request in its own
thread.

On top of that, add some exception safety to a few parts of the code,
and use "destroyForcibly" from Java 8 if it's available, to force
kill the child process. The flip side is that "kill()" may not actually
work if running Java 7.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9633 from vanzin/SPARK-11655.
---
 .../spark/launcher/LauncherBackend.scala      | 12 +++++++++--
 .../cluster/SparkDeploySchedulerBackend.scala | 20 ++++++++++---------
 .../spark/launcher/ChildProcAppHandle.java    | 17 ++++++++++++++--
 .../apache/spark/launcher/SparkAppHandle.java |  3 +++
 4 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/launcher/LauncherBackend.scala b/core/src/main/scala/org/apache/spark/launcher/LauncherBackend.scala
index 3ea984c501e02..a5d41a1eeb479 100644
--- a/core/src/main/scala/org/apache/spark/launcher/LauncherBackend.scala
+++ b/core/src/main/scala/org/apache/spark/launcher/LauncherBackend.scala
@@ -21,7 +21,7 @@ import java.net.{InetAddress, Socket}
 
 import org.apache.spark.SPARK_VERSION
 import org.apache.spark.launcher.LauncherProtocol._
-import org.apache.spark.util.ThreadUtils
+import org.apache.spark.util.{ThreadUtils, Utils}
 
 /**
  * A class that can be used to talk to a launcher server. Users should extend this class to
@@ -88,12 +88,20 @@ private[spark] abstract class LauncherBackend {
    */
   protected def onDisconnected() : Unit = { }
 
+  private def fireStopRequest(): Unit = {
+    val thread = LauncherBackend.threadFactory.newThread(new Runnable() {
+      override def run(): Unit = Utils.tryLogNonFatalError {
+        onStopRequest()
+      }
+    })
+    thread.start()
+  }
 
   private class BackendConnection(s: Socket) extends LauncherConnection(s) {
 
     override protected def handle(m: Message): Unit = m match {
       case _: Stop =>
-        onStopRequest()
+        fireStopRequest()
 
       case _ =>
         throw new IllegalArgumentException(s"Unexpected message type: ${m.getClass().getName()}")
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/SparkDeploySchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/SparkDeploySchedulerBackend.scala
index 05d9bc92f228b..5105475c760e2 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/SparkDeploySchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/SparkDeploySchedulerBackend.scala
@@ -191,17 +191,19 @@ private[spark] class SparkDeploySchedulerBackend(
   }
 
   private def stop(finalState: SparkAppHandle.State): Unit = synchronized {
-    stopping = true
+    try {
+      stopping = true
 
-    launcherBackend.setState(finalState)
-    launcherBackend.close()
+      super.stop()
+      client.stop()
 
-    super.stop()
-    client.stop()
-
-    val callback = shutdownCallback
-    if (callback != null) {
-      callback(this)
+      val callback = shutdownCallback
+      if (callback != null) {
+        callback(this)
+      }
+    } finally {
+      launcherBackend.setState(finalState)
+      launcherBackend.close()
     }
   }
 
diff --git a/launcher/src/main/java/org/apache/spark/launcher/ChildProcAppHandle.java b/launcher/src/main/java/org/apache/spark/launcher/ChildProcAppHandle.java
index de50f14fbdc87..1bfda289dec39 100644
--- a/launcher/src/main/java/org/apache/spark/launcher/ChildProcAppHandle.java
+++ b/launcher/src/main/java/org/apache/spark/launcher/ChildProcAppHandle.java
@@ -18,6 +18,7 @@
 package org.apache.spark.launcher;
 
 import java.io.IOException;
+import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.ThreadFactory;
@@ -102,8 +103,20 @@ public synchronized void kill() {
       disconnect();
     }
     if (childProc != null) {
-      childProc.destroy();
-      childProc = null;
+      try {
+        childProc.exitValue();
+      } catch (IllegalThreadStateException e) {
+        // Child is still alive. Try to use Java 8's "destroyForcibly()" if available,
+        // fall back to the old API if it's not there.
+        try {
+          Method destroy = childProc.getClass().getMethod("destroyForcibly");
+          destroy.invoke(childProc);
+        } catch (Exception inner) {
+          childProc.destroy();
+        }
+      } finally {
+        childProc = null;
+      }
     }
   }
 
diff --git a/launcher/src/main/java/org/apache/spark/launcher/SparkAppHandle.java b/launcher/src/main/java/org/apache/spark/launcher/SparkAppHandle.java
index 13dd9f1739fb6..e9caf0b3cb063 100644
--- a/launcher/src/main/java/org/apache/spark/launcher/SparkAppHandle.java
+++ b/launcher/src/main/java/org/apache/spark/launcher/SparkAppHandle.java
@@ -89,6 +89,9 @@ public boolean isFinal() {
    * Tries to kill the underlying application. Implies {@link #disconnect()}. This will not send
    * a {@link #stop()} message to the application, so it's recommended that users first try to
    * stop the application cleanly and only resort to this method if that fails.
+   * <p>
+   * Note that if the application is running as a child process, this method fail to kill the
+   * process when using Java 7. This may happen if, for example, the application is deadlocked.
    */
   void kill();
 

From f0d3b58d91f43697397cdd7a7e7f38cbb7daaa31 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Thu, 12 Nov 2015 14:52:03 -0800
Subject: [PATCH 213/518] [SPARK-11290][STREAMING][TEST-MAVEN] Fix the test for
 maven build

Should not create SparkContext in the constructor of `TrackStateRDDSuite`. This is a follow up PR for #9256 to fix the test for maven build.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9668 from zsxwing/hotfix.
---
 .../spark/streaming/rdd/TrackStateRDDSuite.scala     | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
index fc5f26607ef98..f396b76e8d251 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
@@ -28,11 +28,17 @@ import org.apache.spark.{HashPartitioner, SparkConf, SparkContext, SparkFunSuite
 
 class TrackStateRDDSuite extends SparkFunSuite with BeforeAndAfterAll {
 
-  private var sc = new SparkContext(
-    new SparkConf().setMaster("local").setAppName("TrackStateRDDSuite"))
+  private var sc: SparkContext = null
+
+  override def beforeAll(): Unit = {
+    sc = new SparkContext(
+      new SparkConf().setMaster("local").setAppName("TrackStateRDDSuite"))
+  }
 
   override def afterAll(): Unit = {
-    sc.stop()
+    if (sc != null) {
+      sc.stop()
+    }
   }
 
   test("creation from pair RDD") {

From 380dfcc0dc865d361a97bb045a2ac546dacfdba9 Mon Sep 17 00:00:00 2001
From: Chris Snow <chsnow123@gmail.com>
Date: Thu, 12 Nov 2015 15:42:30 -0800
Subject: [PATCH 214/518] [SPARK-11671] documentation code example typo

Example for sqlContext.createDataDrame from pandas.DataFrame has a typo

Author: Chris Snow <chsnow123@gmail.com>

Closes #9639 from snowch/patch-2.
---
 python/pyspark/sql/context.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/pyspark/sql/context.py b/python/pyspark/sql/context.py
index 924bb6433de0e..5a85ac31025e8 100644
--- a/python/pyspark/sql/context.py
+++ b/python/pyspark/sql/context.py
@@ -415,7 +415,7 @@ def createDataFrame(self, data, schema=None, samplingRatio=None):
 
         >>> sqlContext.createDataFrame(df.toPandas()).collect()  # doctest: +SKIP
         [Row(name=u'Alice', age=1)]
-        >>> sqlContext.createDataFrame(pandas.DataFrame([[1, 2]]).collect())  # doctest: +SKIP
+        >>> sqlContext.createDataFrame(pandas.DataFrame([[1, 2]])).collect()  # doctest: +SKIP
         [Row(0=1, 1=2)]
         """
         if isinstance(data, DataFrame):

From 74c30049a8bf9841eeca48f827572c2044912e21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-Baptiste=20Onofr=C3=A9?= <jbonofre@apache.org>
Date: Thu, 12 Nov 2015 15:46:21 -0800
Subject: [PATCH 215/518] [SPARK-2533] Add locality levels on stage summary
 view
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Author: Jean-Baptiste Onofré <jbonofre@apache.org>

Closes #9487 from jbonofre/SPARK-2533-2.
---
 .../org/apache/spark/ui/jobs/StagePage.scala  | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala b/core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala
index 51425e599e748..1b34ba9f03c44 100644
--- a/core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala
+++ b/core/src/main/scala/org/apache/spark/ui/jobs/StagePage.scala
@@ -28,7 +28,7 @@ import org.apache.commons.lang3.StringEscapeUtils
 
 import org.apache.spark.{InternalAccumulator, SparkConf}
 import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.scheduler.{AccumulableInfo, TaskInfo}
+import org.apache.spark.scheduler.{AccumulableInfo, TaskInfo, TaskLocality}
 import org.apache.spark.ui._
 import org.apache.spark.ui.jobs.UIData._
 import org.apache.spark.util.{Utils, Distribution}
@@ -70,6 +70,21 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
 
   private val displayPeakExecutionMemory = parent.conf.getBoolean("spark.sql.unsafe.enabled", true)
 
+  private def getLocalitySummaryString(stageData: StageUIData): String = {
+    val localities = stageData.taskData.values.map(_.taskInfo.taskLocality)
+    val localityCounts = localities.groupBy(identity).mapValues(_.size)
+    val localityNamesAndCounts = localityCounts.toSeq.map { case (locality, count) =>
+      val localityName = locality match {
+        case TaskLocality.PROCESS_LOCAL => "Process local"
+        case TaskLocality.NODE_LOCAL => "Node local"
+        case TaskLocality.RACK_LOCAL => "Rack local"
+        case TaskLocality.ANY => "Any"
+      }
+      s"$localityName: $count"
+    }
+    localityNamesAndCounts.sorted.mkString("; ")
+  }
+
   def render(request: HttpServletRequest): Seq[Node] = {
     progressListener.synchronized {
       val parameterId = request.getParameter("id")
@@ -129,6 +144,10 @@ private[ui] class StagePage(parent: StagesTab) extends WebUIPage("stage") {
               <strong>Total Time Across All Tasks: </strong>
               {UIUtils.formatDuration(stageData.executorRunTime)}
             </li>
+            <li>
+              <strong>Locality Level Summary: </strong>
+              {getLocalitySummaryString(stageData)}
+            </li>
             {if (stageData.hasInput) {
               <li>
                 <strong>Input Size / Records: </strong>

From cf38fc7551f4743958c2fdc7931affd672755e68 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Thu, 12 Nov 2015 15:47:29 -0800
Subject: [PATCH 216/518] [SPARK-11670] Fix incorrect kryo buffer default value
 in docs

<img width="931" alt="screen shot 2015-11-11 at 1 53 21 pm" src="https://cloud.githubusercontent.com/assets/2133137/11108261/35d183d4-889a-11e5-9572-85e9d6cebd26.png">

Author: Andrew Or <andrew@databricks.com>

Closes #9638 from andrewor14/fix-kryo-docs.
---
 docs/tuning.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/tuning.md b/docs/tuning.md
index 6936912a6be54..879340a01544f 100644
--- a/docs/tuning.md
+++ b/docs/tuning.md
@@ -61,8 +61,8 @@ The [Kryo documentation](https://github.com/EsotericSoftware/kryo) describes mor
 registration options, such as adding custom serialization code.
 
 If your objects are large, you may also need to increase the `spark.kryoserializer.buffer`
-config property. The default is 2, but this value needs to be large enough to hold the *largest*
-object you will serialize.
+[config](configuration.html#compression-and-serialization). This value needs to be large enough
+to hold the *largest* object you will serialize.
 
 Finally, if you don't register your custom classes, Kryo will still work, but it will have to store
 the full class name with each object, which is wasteful.

From 12a0784ac0f314a606f1237e7144eb1355421307 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Thu, 12 Nov 2015 15:48:42 -0800
Subject: [PATCH 217/518] [SPARK-11667] Update dynamic allocation docs to
 reflect supported cluster managers

Author: Andrew Or <andrew@databricks.com>

Closes #9637 from andrewor14/update-da-docs.
---
 docs/job-scheduling.md | 55 +++++++++++++++++++++---------------------
 1 file changed, 27 insertions(+), 28 deletions(-)

diff --git a/docs/job-scheduling.md b/docs/job-scheduling.md
index 8d9c2ba2041b2..a3c34cb6796fa 100644
--- a/docs/job-scheduling.md
+++ b/docs/job-scheduling.md
@@ -56,36 +56,32 @@ provide another approach to share RDDs.
 
 ## Dynamic Resource Allocation
 
-Spark 1.2 introduces the ability to dynamically scale the set of cluster resources allocated to
-your application up and down based on the workload. This means that your application may give
-resources back to the cluster if they are no longer used and request them again later when there
-is demand. This feature is particularly useful if multiple applications share resources in your
-Spark cluster. If a subset of the resources allocated to an application becomes idle, it can be
-returned to the cluster's pool of resources and acquired by other applications. In Spark, dynamic
-resource allocation is performed on the granularity of the executor and can be enabled through
-`spark.dynamicAllocation.enabled`.
-
-This feature is currently disabled by default and available only on [YARN](running-on-yarn.html).
-A future release will extend this to [standalone mode](spark-standalone.html) and
-[Mesos coarse-grained mode](running-on-mesos.html#mesos-run-modes). Note that although Spark on
-Mesos already has a similar notion of dynamic resource sharing in fine-grained mode, enabling
-dynamic allocation allows your Mesos application to take advantage of coarse-grained low-latency
-scheduling while sharing cluster resources efficiently.
+Spark provides a mechanism to dynamically adjust the resources your application occupies based
+on the workload. This means that your application may give resources back to the cluster if they
+are no longer used and request them again later when there is demand. This feature is particularly
+useful if multiple applications share resources in your Spark cluster.
+
+This feature is disabled by default and available on all coarse-grained cluster managers, i.e.
+[standalone mode](spark-standalone.html), [YARN mode](running-on-yarn.html), and
+[Mesos coarse-grained mode](running-on-mesos.html#mesos-run-modes).
 
 ### Configuration and Setup
 
-All configurations used by this feature live under the `spark.dynamicAllocation.*` namespace.
-To enable this feature, your application must set `spark.dynamicAllocation.enabled` to `true`.
-Other relevant configurations are described on the
-[configurations page](configuration.html#dynamic-allocation) and in the subsequent sections in
-detail.
+There are two requirements for using this feature. First, your application must set
+`spark.dynamicAllocation.enabled` to `true`. Second, you must set up an *external shuffle service*
+on each worker node in the same cluster and set `spark.shuffle.service.enabled` to true in your
+application. The purpose of the external shuffle service is to allow executors to be removed
+without deleting shuffle files written by them (more detail described
+[below](job-scheduling.html#graceful-decommission-of-executors)). The way to set up this service
+varies across cluster managers:
+
+In standalone mode, simply start your workers with `spark.shuffle.service.enabled` set to `true`.
 
-Additionally, your application must use an external shuffle service. The purpose of the service is
-to preserve the shuffle files written by executors so the executors can be safely removed (more
-detail described [below](job-scheduling.html#graceful-decommission-of-executors)). To enable
-this service, set `spark.shuffle.service.enabled` to `true`. In YARN, this external shuffle service
-is implemented in `org.apache.spark.yarn.network.YarnShuffleService` that runs in each `NodeManager`
-in your cluster. To start this service, follow these steps:
+In Mesos coarse-grained mode, run `$SPARK_HOME/sbin/start-mesos-shuffle-service.sh` on all
+slave nodes with `spark.shuffle.service.enabled` set to `true`. For instance, you may do so
+through Marathon.
+
+In YARN mode, start the shuffle service on each `NodeManager` as follows:
 
 1. Build Spark with the [YARN profile](building-spark.html). Skip this step if you are using a
 pre-packaged distribution.
@@ -95,10 +91,13 @@ pre-packaged distribution.
 2. Add this jar to the classpath of all `NodeManager`s in your cluster.
 3. In the `yarn-site.xml` on each node, add `spark_shuffle` to `yarn.nodemanager.aux-services`,
 then set `yarn.nodemanager.aux-services.spark_shuffle.class` to
-`org.apache.spark.network.yarn.YarnShuffleService`. Additionally, set all relevant
-`spark.shuffle.service.*` [configurations](configuration.html).
+`org.apache.spark.network.yarn.YarnShuffleService` and `spark.shuffle.service.enabled` to true.
 4. Restart all `NodeManager`s in your cluster.
 
+All other relevant configurations are optional and under the `spark.dynamicAllocation.*` and
+`spark.shuffle.service.*` namespaces. For more detail, see the
+[configurations page](configuration.html#dynamic-allocation).
+
 ### Resource Allocation Policy
 
 At a high level, Spark should relinquish executors when they are no longer used and acquire

From 68ef61bb656bd9c08239726913ca8ab271d52786 Mon Sep 17 00:00:00 2001
From: Chris Snow <chsnow123@gmail.com>
Date: Thu, 12 Nov 2015 15:50:47 -0800
Subject: [PATCH 218/518] [SPARK-11658] simplify documentation for PySpark
 combineByKey

Author: Chris Snow <chsnow123@gmail.com>

Closes #9640 from snowch/patch-3.
---
 python/pyspark/rdd.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/pyspark/rdd.py b/python/pyspark/rdd.py
index 56e892243c79c..4b4d59647b2bc 100644
--- a/python/pyspark/rdd.py
+++ b/python/pyspark/rdd.py
@@ -1760,7 +1760,6 @@ def combineByKey(self, createCombiner, mergeValue, mergeCombiners,
         In addition, users can control the partitioning of the output RDD.
 
         >>> x = sc.parallelize([("a", 1), ("b", 1), ("a", 1)])
-        >>> def f(x): return x
         >>> def add(a, b): return a + str(b)
         >>> sorted(x.combineByKey(str, add, add).collect())
         [('a', '11'), ('b', '1')]

From bc092966f8264c6685b3300461cb79dd6a509ecf Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Thu, 12 Nov 2015 16:43:04 -0800
Subject: [PATCH 219/518] [SPARK-11709] include creation site info in
 SparkContext.assertNotStopped error message

This helps debug issues caused by multiple SparkContext instances. JoshRosen andrewor14

~~~
scala> sc.stop()

scala> sc.parallelize(0 until 10)
java.lang.IllegalStateException: Cannot call methods on a stopped SparkContext.
This stopped SparkContext was created at:

org.apache.spark.SparkContext.<init>(SparkContext.scala:82)
org.apache.spark.repl.SparkILoop.createSparkContext(SparkILoop.scala:1017)
$iwC$$iwC.<init>(<console>:9)
$iwC.<init>(<console>:18)
<init>(<console>:20)
.<init>(<console>:24)
.<clinit>(<console>)
.<init>(<console>:7)
.<clinit>(<console>)
$print(<console>)
sun.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
sun.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:57)
sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
java.lang.reflect.Method.invoke(Method.java:606)
org.apache.spark.repl.SparkIMain$ReadEvalPrint.call(SparkIMain.scala:1065)
org.apache.spark.repl.SparkIMain$Request.loadAndRun(SparkIMain.scala:1340)
org.apache.spark.repl.SparkIMain.loadAndRunReq$1(SparkIMain.scala:840)
org.apache.spark.repl.SparkIMain.interpret(SparkIMain.scala:871)
org.apache.spark.repl.SparkIMain.interpret(SparkIMain.scala:819)
org.apache.spark.repl.SparkILoop.reallyInterpret$1(SparkILoop.scala:857)

The active context was created at:

(No active SparkContext.)
~~~

Author: Xiangrui Meng <meng@databricks.com>

Closes #9675 from mengxr/SPARK-11709.
---
 .../scala/org/apache/spark/SparkContext.scala  | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 43a241686fd7d..4bbd0b038c00f 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -96,7 +96,23 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
 
   private def assertNotStopped(): Unit = {
     if (stopped.get()) {
-      throw new IllegalStateException("Cannot call methods on a stopped SparkContext")
+      val activeContext = SparkContext.activeContext.get()
+      val activeCreationSite =
+        if (activeContext == null) {
+          "(No active SparkContext.)"
+        } else {
+          activeContext.creationSite.longForm
+        }
+      throw new IllegalStateException(
+        s"""Cannot call methods on a stopped SparkContext.
+           |This stopped SparkContext was created at:
+           |
+           |${creationSite.longForm}
+           |
+           |The currently active SparkContext was created at:
+           |
+           |$activeCreationSite
+         """.stripMargin)
     }
   }
 

From dcb896fd8cec83483f700ee985c352be61cdf233 Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Thu, 12 Nov 2015 17:03:19 -0800
Subject: [PATCH 220/518] [SPARK-11712][ML] Make spark.ml LDAModel be abstract

Per discussion in the initial Pipelines LDA PR [https://github.com/apache/spark/pull/9513], we should make LDAModel abstract and create a LocalLDAModel. This code simplification should be done before the 1.6 release to ensure API compatibility in future releases.

CC feynmanliang mengxr

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9678 from jkbradley/lda-pipelines-2.
---
 .../org/apache/spark/ml/clustering/LDA.scala  | 180 +++++++++---------
 .../apache/spark/ml/clustering/LDASuite.scala |   4 +-
 2 files changed, 96 insertions(+), 88 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
index f66233ed3d0f0..92e05815d6a3d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
@@ -314,31 +314,31 @@ private[clustering] trait LDAParams extends Params with HasFeaturesCol with HasM
  * Model fitted by [[LDA]].
  *
  * @param vocabSize  Vocabulary size (number of terms or terms in the vocabulary)
- * @param oldLocalModel  Underlying spark.mllib model.
- *                       If this model was produced by Online LDA, then this is the
- *                       only model representation.
- *                       If this model was produced by EM, then this local
- *                       representation may be built lazily.
  * @param sqlContext  Used to construct local DataFrames for returning query results
  */
 @Since("1.6.0")
 @Experimental
-class LDAModel private[ml] (
+sealed abstract class LDAModel private[ml] (
     @Since("1.6.0") override val uid: String,
     @Since("1.6.0") val vocabSize: Int,
-    @Since("1.6.0") protected var oldLocalModel: Option[OldLocalLDAModel],
     @Since("1.6.0") @transient protected val sqlContext: SQLContext)
   extends Model[LDAModel] with LDAParams with Logging {
 
-  /** Returns underlying spark.mllib model */
+  // NOTE to developers:
+  //  This abstraction should contain all important functionality for basic LDA usage.
+  //  Specializations of this class can contain expert-only functionality.
+
+  /**
+   * Underlying spark.mllib model.
+   * If this model was produced by Online LDA, then this is the only model representation.
+   * If this model was produced by EM, then this local representation may be built lazily.
+   */
   @Since("1.6.0")
-  protected def getModel: OldLDAModel = oldLocalModel match {
-    case Some(m) => m
-    case None =>
-      // Should never happen.
-      throw new RuntimeException("LDAModel required local model format," +
-        " but the underlying model is missing.")
-  }
+  protected def oldLocalModel: OldLocalLDAModel
+
+  /** Returns underlying spark.mllib model, which may be local or distributed */
+  @Since("1.6.0")
+  protected def getModel: OldLDAModel
 
   /**
    * The features for LDA should be a [[Vector]] representing the word counts in a document.
@@ -352,16 +352,17 @@ class LDAModel private[ml] (
   @Since("1.6.0")
   def setSeed(value: Long): this.type = set(seed, value)
 
-  @Since("1.6.0")
-  override def copy(extra: ParamMap): LDAModel = {
-    val copied = new LDAModel(uid, vocabSize, oldLocalModel, sqlContext)
-    copyValues(copied, extra).setParent(parent)
-  }
-
+  /**
+   * Transforms the input dataset.
+   *
+   * WARNING: If this model is an instance of [[DistributedLDAModel]] (produced when [[optimizer]]
+   *          is set to "em"), this involves collecting a large [[topicsMatrix]] to the driver.
+   *          This implementation may be changed in the future.
+   */
   @Since("1.6.0")
   override def transform(dataset: DataFrame): DataFrame = {
     if ($(topicDistributionCol).nonEmpty) {
-      val t = udf(oldLocalModel.get.getTopicDistributionMethod(sqlContext.sparkContext))
+      val t = udf(oldLocalModel.getTopicDistributionMethod(sqlContext.sparkContext))
       dataset.withColumn($(topicDistributionCol), t(col($(featuresCol))))
     } else {
       logWarning("LDAModel.transform was called without any output columns. Set an output column" +
@@ -388,56 +389,50 @@ class LDAModel private[ml] (
    * This is a matrix of size vocabSize x k, where each column is a topic.
    * No guarantees are given about the ordering of the topics.
    *
-   * WARNING: If this model is actually a [[DistributedLDAModel]] instance from EM,
-   *          then this method could involve collecting a large amount of data to the driver
-   *          (on the order of vocabSize x k).
+   * WARNING: If this model is actually a [[DistributedLDAModel]] instance produced by
+   *          the Expectation-Maximization ("em") [[optimizer]], then this method could involve
+   *          collecting a large amount of data to the driver (on the order of vocabSize x k).
    */
   @Since("1.6.0")
-  def topicsMatrix: Matrix = getModel.topicsMatrix
+  def topicsMatrix: Matrix = oldLocalModel.topicsMatrix
 
   /** Indicates whether this instance is of type [[DistributedLDAModel]] */
   @Since("1.6.0")
-  def isDistributed: Boolean = false
+  def isDistributed: Boolean
 
   /**
    * Calculates a lower bound on the log likelihood of the entire corpus.
    *
    * See Equation (16) in the Online LDA paper (Hoffman et al., 2010).
    *
-   * WARNING: If this model was learned via a [[DistributedLDAModel]], this involves collecting
-   *          a large [[topicsMatrix]] to the driver.  This implementation may be changed in the
-   *          future.
+   * WARNING: If this model is an instance of [[DistributedLDAModel]] (produced when [[optimizer]]
+   *          is set to "em"), this involves collecting a large [[topicsMatrix]] to the driver.
+   *          This implementation may be changed in the future.
    *
    * @param dataset  test corpus to use for calculating log likelihood
    * @return variational lower bound on the log likelihood of the entire corpus
    */
   @Since("1.6.0")
-  def logLikelihood(dataset: DataFrame): Double = oldLocalModel match {
-    case Some(m) =>
-      val oldDataset = LDA.getOldDataset(dataset, $(featuresCol))
-      m.logLikelihood(oldDataset)
-    case None =>
-      // Should never happen.
-      throw new RuntimeException("LocalLDAModel.logLikelihood was called," +
-        " but the underlying model is missing.")
+  def logLikelihood(dataset: DataFrame): Double = {
+    val oldDataset = LDA.getOldDataset(dataset, $(featuresCol))
+    oldLocalModel.logLikelihood(oldDataset)
   }
 
   /**
    * Calculate an upper bound bound on perplexity.  (Lower is better.)
    * See Equation (16) in the Online LDA paper (Hoffman et al., 2010).
    *
+   * WARNING: If this model is an instance of [[DistributedLDAModel]] (produced when [[optimizer]]
+   *          is set to "em"), this involves collecting a large [[topicsMatrix]] to the driver.
+   *          This implementation may be changed in the future.
+   *
    * @param dataset test corpus to use for calculating perplexity
    * @return Variational upper bound on log perplexity per token.
    */
   @Since("1.6.0")
-  def logPerplexity(dataset: DataFrame): Double = oldLocalModel match {
-    case Some(m) =>
-      val oldDataset = LDA.getOldDataset(dataset, $(featuresCol))
-      m.logPerplexity(oldDataset)
-    case None =>
-      // Should never happen.
-      throw new RuntimeException("LocalLDAModel.logPerplexity was called," +
-        " but the underlying model is missing.")
+  def logPerplexity(dataset: DataFrame): Double = {
+    val oldDataset = LDA.getOldDataset(dataset, $(featuresCol))
+    oldLocalModel.logPerplexity(oldDataset)
   }
 
   /**
@@ -468,10 +463,43 @@ class LDAModel private[ml] (
 /**
  * :: Experimental ::
  *
- * Distributed model fitted by [[LDA]] using Expectation-Maximization (EM).
+ * Local (non-distributed) model fitted by [[LDA]].
+ *
+ * This model stores the inferred topics only; it does not store info about the training dataset.
+ */
+@Since("1.6.0")
+@Experimental
+class LocalLDAModel private[ml] (
+    uid: String,
+    vocabSize: Int,
+    @Since("1.6.0") override protected val oldLocalModel: OldLocalLDAModel,
+    sqlContext: SQLContext)
+  extends LDAModel(uid, vocabSize, sqlContext) {
+
+  @Since("1.6.0")
+  override def copy(extra: ParamMap): LocalLDAModel = {
+    val copied = new LocalLDAModel(uid, vocabSize, oldLocalModel, sqlContext)
+    copyValues(copied, extra).setParent(parent).asInstanceOf[LocalLDAModel]
+  }
+
+  override protected def getModel: OldLDAModel = oldLocalModel
+
+  @Since("1.6.0")
+  override def isDistributed: Boolean = false
+}
+
+
+/**
+ * :: Experimental ::
+ *
+ * Distributed model fitted by [[LDA]].
+ * This type of model is currently only produced by Expectation-Maximization (EM).
  *
  * This model stores the inferred topics, the full training dataset, and the topic distribution
  * for each training document.
+ *
+ * @param oldLocalModelOption  Used to implement [[oldLocalModel]] as a lazy val, but keeping
+ *                             [[copy()]] cheap.
  */
 @Since("1.6.0")
 @Experimental
@@ -479,59 +507,39 @@ class DistributedLDAModel private[ml] (
     uid: String,
     vocabSize: Int,
     private val oldDistributedModel: OldDistributedLDAModel,
-    sqlContext: SQLContext)
-  extends LDAModel(uid, vocabSize, None, sqlContext) {
+    sqlContext: SQLContext,
+    private var oldLocalModelOption: Option[OldLocalLDAModel])
+  extends LDAModel(uid, vocabSize, sqlContext) {
+
+  override protected def oldLocalModel: OldLocalLDAModel = {
+    if (oldLocalModelOption.isEmpty) {
+      oldLocalModelOption = Some(oldDistributedModel.toLocal)
+    }
+    oldLocalModelOption.get
+  }
+
+  override protected def getModel: OldLDAModel = oldDistributedModel
 
   /**
    * Convert this distributed model to a local representation.  This discards info about the
    * training dataset.
+   *
+   * WARNING: This involves collecting a large [[topicsMatrix]] to the driver.
    */
   @Since("1.6.0")
-  def toLocal: LDAModel = {
-    if (oldLocalModel.isEmpty) {
-      oldLocalModel = Some(oldDistributedModel.toLocal)
-    }
-    new LDAModel(uid, vocabSize, oldLocalModel, sqlContext)
-  }
-
-  @Since("1.6.0")
-  override protected def getModel: OldLDAModel = oldDistributedModel
+  def toLocal: LocalLDAModel = new LocalLDAModel(uid, vocabSize, oldLocalModel, sqlContext)
 
   @Since("1.6.0")
   override def copy(extra: ParamMap): DistributedLDAModel = {
-    val copied = new DistributedLDAModel(uid, vocabSize, oldDistributedModel, sqlContext)
-    if (oldLocalModel.nonEmpty) copied.oldLocalModel = oldLocalModel
+    val copied =
+      new DistributedLDAModel(uid, vocabSize, oldDistributedModel, sqlContext, oldLocalModelOption)
     copyValues(copied, extra).setParent(parent)
     copied
   }
 
-  @Since("1.6.0")
-  override def topicsMatrix: Matrix = {
-    if (oldLocalModel.isEmpty) {
-      oldLocalModel = Some(oldDistributedModel.toLocal)
-    }
-    super.topicsMatrix
-  }
-
   @Since("1.6.0")
   override def isDistributed: Boolean = true
 
-  @Since("1.6.0")
-  override def logLikelihood(dataset: DataFrame): Double = {
-    if (oldLocalModel.isEmpty) {
-      oldLocalModel = Some(oldDistributedModel.toLocal)
-    }
-    super.logLikelihood(dataset)
-  }
-
-  @Since("1.6.0")
-  override def logPerplexity(dataset: DataFrame): Double = {
-    if (oldLocalModel.isEmpty) {
-      oldLocalModel = Some(oldDistributedModel.toLocal)
-    }
-    super.logPerplexity(dataset)
-  }
-
   /**
    * Log likelihood of the observed tokens in the training set,
    * given the current parameter estimates:
@@ -673,9 +681,9 @@ class LDA @Since("1.6.0") (
     val oldModel = oldLDA.run(oldData)
     val newModel = oldModel match {
       case m: OldLocalLDAModel =>
-        new LDAModel(uid, m.vocabSize, Some(m), dataset.sqlContext)
+        new LocalLDAModel(uid, m.vocabSize, m, dataset.sqlContext)
       case m: OldDistributedLDAModel =>
-        new DistributedLDAModel(uid, m.vocabSize, m, dataset.sqlContext)
+        new DistributedLDAModel(uid, m.vocabSize, m, dataset.sqlContext, None)
     }
     copyValues(newModel).setParent(this)
   }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala b/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala
index edb927495e8bf..b634d31cc34f0 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala
@@ -156,7 +156,7 @@ class LDASuite extends SparkFunSuite with MLlibTestSparkContext {
 
     MLTestingUtils.checkCopy(model)
 
-    assert(!model.isInstanceOf[DistributedLDAModel])
+    assert(model.isInstanceOf[LocalLDAModel])
     assert(model.vocabSize === vocabSize)
     assert(model.estimatedDocConcentration.size === k)
     assert(model.topicsMatrix.numRows === vocabSize)
@@ -210,7 +210,7 @@ class LDASuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(model.isDistributed)
 
     val localModel = model.toLocal
-    assert(!localModel.isInstanceOf[DistributedLDAModel])
+    assert(localModel.isInstanceOf[LocalLDAModel])
 
     // training logLikelihood, logPrior
     val ll = model.trainingLogLikelihood

From 41bbd2300472501d69ed46f0407d5ed7cbede4a8 Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Thu, 12 Nov 2015 17:20:30 -0800
Subject: [PATCH 221/518] [SPARK-11654][SQL] add reduce to GroupedDataset

This PR adds a new method, `reduce`, to `GroupedDataset`, which allows similar operations to `reduceByKey` on a traditional `PairRDD`.

```scala
val ds = Seq("abc", "xyz", "hello").toDS()
ds.groupBy(_.length).reduce(_ + _).collect()  // not actually commutative :P

res0: Array(3 -> "abcxyz", 5 -> "hello")
```

While implementing this method and its test cases several more deficiencies were found in our encoder handling.  Specifically, in order to support positional resolution, named resolution and tuple composition, it is important to keep the unresolved encoder around and to use it when constructing new `Datasets` with the same object type but different output attributes.  We now divide the encoder lifecycle into three phases (that mirror the lifecycle of standard expressions) and have checks at various boundaries:

 - Unresoved Encoders: all users facing encoders (those constructed by implicits, static methods, or tuple composition) are unresolved, meaning they have only `UnresolvedAttributes` for named fields and `BoundReferences` for fields accessed by ordinal.
 - Resolved Encoders: internal to a `[Grouped]Dataset` the encoder is resolved, meaning all input has been resolved to a specific `AttributeReference`.  Any encoders that are placed into a logical plan for use in object construction should be resolved.
 - BoundEncoder: Are constructed by physical plans, right before actual conversion from row -> object is performed.

It is left to future work to add explicit checks for resolution and provide good error messages when it fails.  We might also consider enforcing the above constraints in the type system (i.e. `fromRow` only exists on a `ResolvedEncoder`), but we should probably wait before spending too much time on this.

Author: Michael Armbrust <michael@databricks.com>
Author: Wenchen Fan <wenchen@databricks.com>

Closes #9673 from marmbrus/pr/9628.
---
 .../scala/org/apache/spark/sql/Encoder.scala  |  10 +-
 .../catalyst/encoders/ExpressionEncoder.scala | 124 ++++++++++--------
 .../spark/sql/catalyst/encoders/package.scala |  11 +-
 .../expressions/complexTypeExtractors.scala   |   7 +-
 .../plans/logical/basicOperators.scala        |  15 ++-
 .../scala/org/apache/spark/sql/Column.scala   |  43 +++++-
 .../org/apache/spark/sql/DataFrame.scala      |  17 +--
 .../scala/org/apache/spark/sql/Dataset.scala  |  85 ++++++------
 .../org/apache/spark/sql/GroupedDataset.scala |  98 +++++++-------
 .../aggregate/TypedAggregateExpression.scala  |  13 +-
 .../spark/sql/execution/basicOperators.scala  |   7 +-
 .../apache/spark/sql/JavaDatasetSuite.java    |  12 +-
 .../spark/sql/DatasetAggregatorSuite.scala    |  42 ++++++
 .../org/apache/spark/sql/DatasetSuite.scala   |   9 ++
 .../org/apache/spark/sql/QueryTest.scala      |  13 +-
 15 files changed, 309 insertions(+), 197 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index 6134f9e036638..5f619d6c339e3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -84,7 +84,7 @@ object Encoders {
   private def tuple(encoders: Seq[ExpressionEncoder[_]]): ExpressionEncoder[_] = {
     assert(encoders.length > 1)
     // make sure all encoders are resolved, i.e. `Attribute` has been resolved to `BoundReference`.
-    assert(encoders.forall(_.constructExpression.find(_.isInstanceOf[Attribute]).isEmpty))
+    assert(encoders.forall(_.fromRowExpression.find(_.isInstanceOf[Attribute]).isEmpty))
 
     val schema = StructType(encoders.zipWithIndex.map {
       case (e, i) => StructField(s"_${i + 1}", if (e.flat) e.schema.head.dataType else e.schema)
@@ -93,8 +93,8 @@ object Encoders {
     val cls = Utils.getContextOrSparkClassLoader.loadClass(s"scala.Tuple${encoders.size}")
 
     val extractExpressions = encoders.map {
-      case e if e.flat => e.extractExpressions.head
-      case other => CreateStruct(other.extractExpressions)
+      case e if e.flat => e.toRowExpressions.head
+      case other => CreateStruct(other.toRowExpressions)
     }.zipWithIndex.map { case (expr, index) =>
       expr.transformUp {
         case BoundReference(0, t: ObjectType, _) =>
@@ -107,11 +107,11 @@ object Encoders {
 
     val constructExpressions = encoders.zipWithIndex.map { case (enc, index) =>
       if (enc.flat) {
-        enc.constructExpression.transform {
+        enc.fromRowExpression.transform {
           case b: BoundReference => b.copy(ordinal = index)
         }
       } else {
-        enc.constructExpression.transformUp {
+        enc.fromRowExpression.transformUp {
           case BoundReference(ordinal, dt, _) =>
             GetInternalRowField(BoundReference(index, enc.schema, nullable = true), ordinal, dt)
         }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index 294afde5347e2..0d3e4aafb0af4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateSafeProjection, GenerateUnsafeProjection}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.ScalaReflection
-import org.apache.spark.sql.types.{StructField, ObjectType, StructType}
+import org.apache.spark.sql.types.{NullType, StructField, ObjectType, StructType}
 
 /**
  * A factory for constructing encoders that convert objects and primitves to and from the
@@ -61,20 +61,39 @@ object ExpressionEncoder {
 
   /**
    * Given a set of N encoders, constructs a new encoder that produce objects as items in an
-   * N-tuple.  Note that these encoders should first be bound correctly to the combined input
-   * schema.
+   * N-tuple.  Note that these encoders should be unresolved so that information about
+   * name/positional binding is preserved.
    */
   def tuple(encoders: Seq[ExpressionEncoder[_]]): ExpressionEncoder[_] = {
+    encoders.foreach(_.assertUnresolved())
+
     val schema =
       StructType(
-        encoders.zipWithIndex.map { case (e, i) => StructField(s"_${i + 1}", e.schema)})
+        encoders.zipWithIndex.map {
+          case (e, i) => StructField(s"_${i + 1}", if (e.flat) e.schema.head.dataType else e.schema)
+        })
     val cls = Utils.getContextOrSparkClassLoader.loadClass(s"scala.Tuple${encoders.size}")
-    val extractExpressions = encoders.map {
-      case e if e.flat => e.extractExpressions.head
-      case other => CreateStruct(other.extractExpressions)
+
+    // Rebind the encoders to the nested schema.
+    val newConstructExpressions = encoders.zipWithIndex.map {
+      case (e, i) if !e.flat => e.nested(i).fromRowExpression
+      case (e, i) => e.shift(i).fromRowExpression
     }
+
     val constructExpression =
-      NewInstance(cls, encoders.map(_.constructExpression), false, ObjectType(cls))
+      NewInstance(cls, newConstructExpressions, false, ObjectType(cls))
+
+    val input = BoundReference(0, ObjectType(cls), false)
+    val extractExpressions = encoders.zipWithIndex.map {
+      case (e, i) if !e.flat => CreateStruct(e.toRowExpressions.map(_ transformUp {
+        case b: BoundReference =>
+          Invoke(input, s"_${i + 1}", b.dataType, Nil)
+      }))
+      case (e, i) => e.toRowExpressions.head transformUp {
+        case b: BoundReference =>
+          Invoke(input, s"_${i + 1}", b.dataType, Nil)
+      }
+    }
 
     new ExpressionEncoder[Any](
       schema,
@@ -95,35 +114,40 @@ object ExpressionEncoder {
  * A generic encoder for JVM objects.
  *
  * @param schema The schema after converting `T` to a Spark SQL row.
- * @param extractExpressions A set of expressions, one for each top-level field that can be used to
- *                           extract the values from a raw object.
+ * @param toRowExpressions A set of expressions, one for each top-level field that can be used to
+ *                           extract the values from a raw object into an [[InternalRow]].
+ * @param fromRowExpression An expression that will construct an object given an [[InternalRow]].
  * @param clsTag A classtag for `T`.
  */
 case class ExpressionEncoder[T](
     schema: StructType,
     flat: Boolean,
-    extractExpressions: Seq[Expression],
-    constructExpression: Expression,
+    toRowExpressions: Seq[Expression],
+    fromRowExpression: Expression,
     clsTag: ClassTag[T])
   extends Encoder[T] {
 
-  if (flat) require(extractExpressions.size == 1)
+  if (flat) require(toRowExpressions.size == 1)
 
   @transient
-  private lazy val extractProjection = GenerateUnsafeProjection.generate(extractExpressions)
+  private lazy val extractProjection = GenerateUnsafeProjection.generate(toRowExpressions)
   private val inputRow = new GenericMutableRow(1)
 
   @transient
-  private lazy val constructProjection = GenerateSafeProjection.generate(constructExpression :: Nil)
+  private lazy val constructProjection = GenerateSafeProjection.generate(fromRowExpression :: Nil)
 
   /**
    * Returns an encoded version of `t` as a Spark SQL row.  Note that multiple calls to
    * toRow are allowed to return the same actual [[InternalRow]] object.  Thus, the caller should
    * copy the result before making another call if required.
    */
-  def toRow(t: T): InternalRow = {
+  def toRow(t: T): InternalRow = try {
     inputRow(0) = t
     extractProjection(inputRow)
+  } catch {
+    case e: Exception =>
+      throw new RuntimeException(
+        s"Error while encoding: $e\n${toRowExpressions.map(_.treeString).mkString("\n")}", e)
   }
 
   /**
@@ -135,7 +159,20 @@ case class ExpressionEncoder[T](
     constructProjection(row).get(0, ObjectType(clsTag.runtimeClass)).asInstanceOf[T]
   } catch {
     case e: Exception =>
-      throw new RuntimeException(s"Error while decoding: $e\n${constructExpression.treeString}", e)
+      throw new RuntimeException(s"Error while decoding: $e\n${fromRowExpression.treeString}", e)
+  }
+
+  /**
+   * The process of resolution to a given schema throws away information about where a given field
+   * is being bound by ordinal instead of by name.  This method checks to make sure this process
+   * has not been done already in places where we plan to do later composition of encoders.
+   */
+  def assertUnresolved(): Unit = {
+    (fromRowExpression +:  toRowExpressions).foreach(_.foreach {
+      case a: AttributeReference =>
+        sys.error(s"Unresolved encoder expected, but $a was found.")
+      case _ =>
+    })
   }
 
   /**
@@ -143,9 +180,14 @@ case class ExpressionEncoder[T](
    * given schema.
    */
   def resolve(schema: Seq[Attribute]): ExpressionEncoder[T] = {
-    val plan = Project(Alias(constructExpression, "")() :: Nil, LocalRelation(schema))
+    val positionToAttribute = AttributeMap.toIndex(schema)
+    val unbound = fromRowExpression transform {
+      case b: BoundReference => positionToAttribute(b.ordinal)
+    }
+
+    val plan = Project(Alias(unbound, "")() :: Nil, LocalRelation(schema))
     val analyzedPlan = SimpleAnalyzer.execute(plan)
-    copy(constructExpression = analyzedPlan.expressions.head.children.head)
+    copy(fromRowExpression = analyzedPlan.expressions.head.children.head)
   }
 
   /**
@@ -154,39 +196,14 @@ case class ExpressionEncoder[T](
    * resolve before bind.
    */
   def bind(schema: Seq[Attribute]): ExpressionEncoder[T] = {
-    copy(constructExpression = BindReferences.bindReference(constructExpression, schema))
-  }
-
-  /**
-   * Replaces any bound references in the schema with the attributes at the corresponding ordinal
-   * in the provided schema.  This can be used to "relocate" a given encoder to pull values from
-   * a different schema than it was initially bound to.  It can also be used to assign attributes
-   * to ordinal based extraction (i.e. because the input data was a tuple).
-   */
-  def unbind(schema: Seq[Attribute]): ExpressionEncoder[T] = {
-    val positionToAttribute = AttributeMap.toIndex(schema)
-    copy(constructExpression = constructExpression transform {
-      case b: BoundReference => positionToAttribute(b.ordinal)
-    })
+    copy(fromRowExpression = BindReferences.bindReference(fromRowExpression, schema))
   }
 
   /**
-   * Given an encoder that has already been bound to a given schema, returns a new encoder
-   * where the positions are mapped from `oldSchema` to `newSchema`.  This can be used, for example,
-   * when you are trying to use an encoder on grouping keys that were originally part of a larger
-   * row, but now you have projected out only the key expressions.
+   * Returns a new encoder with input columns shifted by `delta` ordinals
    */
-  def rebind(oldSchema: Seq[Attribute], newSchema: Seq[Attribute]): ExpressionEncoder[T] = {
-    val positionToAttribute = AttributeMap.toIndex(oldSchema)
-    val attributeToNewPosition = AttributeMap.byIndex(newSchema)
-    copy(constructExpression = constructExpression transform {
-      case r: BoundReference =>
-        r.copy(ordinal = attributeToNewPosition(positionToAttribute(r.ordinal)))
-    })
-  }
-
   def shift(delta: Int): ExpressionEncoder[T] = {
-    copy(constructExpression = constructExpression transform {
+    copy(fromRowExpression = fromRowExpression transform {
       case r: BoundReference => r.copy(ordinal = r.ordinal + delta)
     })
   }
@@ -196,11 +213,14 @@ case class ExpressionEncoder[T](
    * input row have been modified to pull the object out from a nested struct, instead of the
    * top level fields.
    */
-  def nested(input: Expression = BoundReference(0, schema, true)): ExpressionEncoder[T] = {
-    copy(constructExpression = constructExpression transform {
-      case u: Attribute if u != input =>
+  private def nested(i: Int): ExpressionEncoder[T] = {
+    // We don't always know our input type at this point since it might be unresolved.
+    // We fill in null and it will get unbound to the actual attribute at this position.
+    val input = BoundReference(i, NullType, nullable = true)
+    copy(fromRowExpression = fromRowExpression transformUp {
+      case u: Attribute =>
         UnresolvedExtractValue(input, Literal(u.name))
-      case b: BoundReference if b != input =>
+      case b: BoundReference =>
         GetStructField(
           input,
           StructField(s"i[${b.ordinal}]", b.dataType),
@@ -208,7 +228,7 @@ case class ExpressionEncoder[T](
     })
   }
 
-  protected val attrs = extractExpressions.flatMap(_.collect {
+  protected val attrs = toRowExpressions.flatMap(_.collect {
     case _: UnresolvedAttribute => ""
     case a: Attribute => s"#${a.exprId}"
     case b: BoundReference => s"[${b.ordinal}]"
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/package.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/package.scala
index 2c35adca9c925..9e283f5eb6342 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/package.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/package.scala
@@ -18,10 +18,19 @@
 package org.apache.spark.sql.catalyst
 
 import org.apache.spark.sql.Encoder
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
 
 package object encoders {
+  /**
+   * Returns an internal encoder object that can be used to serialize / deserialize JVM objects
+   * into Spark SQL rows.  The implicit encoder should always be unresolved (i.e. have no attribute
+   * references from a specific schema.)  This requirement allows us to preserve whether a given
+   * object type is being bound by name or by ordinal when doing resolution.
+   */
   private[sql] def encoderFor[A : Encoder]: ExpressionEncoder[A] = implicitly[Encoder[A]] match {
-    case e: ExpressionEncoder[A] => e
+    case e: ExpressionEncoder[A] =>
+      e.assertUnresolved()
+      e
     case _ => sys.error(s"Only expression encoders are supported today")
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
index 41cd0a104a1f5..f871b737fff3a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
@@ -97,11 +97,16 @@ object ExtractValue {
  * Returns the value of fields in the Struct `child`.
  *
  * No need to do type checking since it is handled by [[ExtractValue]].
+ * TODO: Unify with [[GetInternalRowField]], remove the need to specify a [[StructField]].
  */
 case class GetStructField(child: Expression, field: StructField, ordinal: Int)
   extends UnaryExpression {
 
-  override def dataType: DataType = field.dataType
+  override def dataType: DataType = child.dataType match {
+    case s: StructType => s(ordinal).dataType
+    // This is a hack to avoid breaking existing code until we remove the need for the struct field
+    case _ => field.dataType
+  }
   override def nullable: Boolean = child.nullable || field.nullable
   override def toString: String = s"$child.${field.name}"
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index 32b09b59af436..d9f046efce0bf 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -483,9 +483,12 @@ case class MapPartitions[T, U](
 
 /** Factory for constructing new `AppendColumn` nodes. */
 object AppendColumn {
-  def apply[T : Encoder, U : Encoder](func: T => U, child: LogicalPlan): AppendColumn[T, U] = {
+  def apply[T, U : Encoder](
+      func: T => U,
+      tEncoder: ExpressionEncoder[T],
+      child: LogicalPlan): AppendColumn[T, U] = {
     val attrs = encoderFor[U].schema.toAttributes
-    new AppendColumn[T, U](func, encoderFor[T], encoderFor[U], attrs, child)
+    new AppendColumn[T, U](func, tEncoder, encoderFor[U], attrs, child)
   }
 }
 
@@ -506,14 +509,16 @@ case class AppendColumn[T, U](
 
 /** Factory for constructing new `MapGroups` nodes. */
 object MapGroups {
-  def apply[K : Encoder, T : Encoder, U : Encoder](
+  def apply[K, T, U : Encoder](
       func: (K, Iterator[T]) => TraversableOnce[U],
+      kEncoder: ExpressionEncoder[K],
+      tEncoder: ExpressionEncoder[T],
       groupingAttributes: Seq[Attribute],
       child: LogicalPlan): MapGroups[K, T, U] = {
     new MapGroups(
       func,
-      encoderFor[K],
-      encoderFor[T],
+      kEncoder,
+      tEncoder,
       encoderFor[U],
       groupingAttributes,
       encoderFor[U].schema.toAttributes,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index f0f275e91f1a3..929224460dc09 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -17,13 +17,15 @@
 
 package org.apache.spark.sql
 
+import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
+
 import scala.language.implicitConversions
 
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.Logging
 import org.apache.spark.sql.functions.lit
 import org.apache.spark.sql.catalyst.analysis._
-import org.apache.spark.sql.catalyst.encoders.encoderFor
+import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor}
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.util.DataTypeParser
 import org.apache.spark.sql.types._
@@ -45,7 +47,25 @@ private[sql] object Column {
  *           checked by the analyzer instead of the compiler (i.e. `expr("sum(...)")`).
  * @tparam U The output type of this column.
  */
-class TypedColumn[-T, U](expr: Expression, val encoder: Encoder[U]) extends Column(expr)
+class TypedColumn[-T, U](
+    expr: Expression,
+    private[sql] val encoder: ExpressionEncoder[U]) extends Column(expr) {
+
+  /**
+   * Inserts the specific input type and schema into any expressions that are expected to operate
+   * on a decoded object.
+   */
+  private[sql] def withInputType(
+      inputEncoder: ExpressionEncoder[_],
+      schema: Seq[Attribute]): TypedColumn[T, U] = {
+    new TypedColumn[T, U] (expr transform {
+      case ta: TypedAggregateExpression if ta.aEncoder.isEmpty =>
+        ta.copy(
+          aEncoder = Some(inputEncoder.asInstanceOf[ExpressionEncoder[Any]]),
+          children = schema)
+    }, encoder)
+  }
+}
 
 /**
  * :: Experimental ::
@@ -73,6 +93,25 @@ class Column(protected[sql] val expr: Expression) extends Logging {
   /** Creates a column based on the given expression. */
   private def withExpr(newExpr: Expression): Column = new Column(newExpr)
 
+  /**
+   * Returns the expression for this column either with an existing or auto assigned name.
+   */
+  private[sql] def named: NamedExpression = expr match {
+    // Wrap UnresolvedAttribute with UnresolvedAlias, as when we resolve UnresolvedAttribute, we
+    // will remove intermediate Alias for ExtractValue chain, and we need to alias it again to
+    // make it a NamedExpression.
+    case u: UnresolvedAttribute => UnresolvedAlias(u)
+
+    case expr: NamedExpression => expr
+
+    // Leave an unaliased generator with an empty list of names since the analyzer will generate
+    // the correct defaults after the nested expression's type has been resolved.
+    case explode: Explode => MultiAlias(explode, Nil)
+    case jt: JsonTuple => MultiAlias(jt, Nil)
+
+    case expr: Expression => Alias(expr, expr.prettyString)()
+  }
+
   override def toString: String = expr.prettyString
 
   override def equals(that: Any): Boolean = that match {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index a492099b9392b..3ba4ba18d2122 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -735,22 +735,7 @@ class DataFrame private[sql](
    */
   @scala.annotation.varargs
   def select(cols: Column*): DataFrame = withPlan {
-    val namedExpressions = cols.map {
-      // Wrap UnresolvedAttribute with UnresolvedAlias, as when we resolve UnresolvedAttribute, we
-      // will remove intermediate Alias for ExtractValue chain, and we need to alias it again to
-      // make it a NamedExpression.
-      case Column(u: UnresolvedAttribute) => UnresolvedAlias(u)
-
-      case Column(expr: NamedExpression) => expr
-
-      // Leave an unaliased generator with an empty list of names since the analyzer will generate
-      // the correct defaults after the nested expression's type has been resolved.
-      case Column(explode: Explode) => MultiAlias(explode, Nil)
-      case Column(jt: JsonTuple) => MultiAlias(jt, Nil)
-
-      case Column(expr: Expression) => Alias(expr, expr.prettyString)()
-    }
-    Project(namedExpressions.toSeq, logicalPlan)
+    Project(cols.map(_.named), logicalPlan)
   }
 
   /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 87dae6b331593..b930e4661c1a6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -29,7 +29,6 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedAlias
 import org.apache.spark.sql.catalyst.plans.Inner
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.{Queryable, QueryExecution}
-import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
 import org.apache.spark.sql.types.StructType
 
 /**
@@ -63,15 +62,20 @@ import org.apache.spark.sql.types.StructType
 class Dataset[T] private[sql](
     @transient val sqlContext: SQLContext,
     @transient val queryExecution: QueryExecution,
-    unresolvedEncoder: Encoder[T]) extends Queryable with Serializable {
+    tEncoder: Encoder[T]) extends Queryable with Serializable {
+
+  /**
+   * An unresolved version of the internal encoder for the type of this dataset.  This one is marked
+   * implicit so that we can use it when constructing new [[Dataset]] objects that have the same
+   * object type (that will be possibly resolved to a different schema).
+   */
+  private implicit val unresolvedTEncoder: ExpressionEncoder[T] = encoderFor(tEncoder)
 
   /** The encoder for this [[Dataset]] that has been resolved to its output schema. */
-  private[sql] implicit val encoder: ExpressionEncoder[T] = unresolvedEncoder match {
-    case e: ExpressionEncoder[T] => e.resolve(queryExecution.analyzed.output)
-    case _ => throw new IllegalArgumentException("Only expression encoders are currently supported")
-  }
+  private[sql] val resolvedTEncoder: ExpressionEncoder[T] =
+    unresolvedTEncoder.resolve(queryExecution.analyzed.output)
 
-  private implicit def classTag = encoder.clsTag
+  private implicit def classTag = resolvedTEncoder.clsTag
 
   private[sql] def this(sqlContext: SQLContext, plan: LogicalPlan)(implicit encoder: Encoder[T]) =
     this(sqlContext, new QueryExecution(sqlContext, plan), encoder)
@@ -81,7 +85,7 @@ class Dataset[T] private[sql](
    *
    * @since 1.6.0
    */
-  def schema: StructType = encoder.schema
+  def schema: StructType = resolvedTEncoder.schema
 
   /* ************* *
    *  Conversions  *
@@ -134,7 +138,7 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def rdd: RDD[T] = {
-    val tEnc = encoderFor[T]
+    val tEnc = resolvedTEncoder
     val input = queryExecution.analyzed.output
     queryExecution.toRdd.mapPartitions { iter =>
       val bound = tEnc.bind(input)
@@ -195,7 +199,7 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def mapPartitions[U : Encoder](func: Iterator[T] => Iterator[U]): Dataset[U] = {
-    new Dataset(
+    new Dataset[U](
       sqlContext,
       MapPartitions[T, U](
         func,
@@ -295,12 +299,12 @@ class Dataset[T] private[sql](
    */
   def groupBy[K : Encoder](func: T => K): GroupedDataset[K, T] = {
     val inputPlan = queryExecution.analyzed
-    val withGroupingKey = AppendColumn(func, inputPlan)
+    val withGroupingKey = AppendColumn(func, resolvedTEncoder, inputPlan)
     val executed = sqlContext.executePlan(withGroupingKey)
 
     new GroupedDataset(
-      encoderFor[K].resolve(withGroupingKey.newColumns),
-      encoderFor[T].bind(inputPlan.output),
+      encoderFor[K],
+      encoderFor[T],
       executed,
       inputPlan.output,
       withGroupingKey.newColumns)
@@ -360,7 +364,15 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def select[U1: Encoder](c1: TypedColumn[T, U1]): Dataset[U1] = {
-    new Dataset[U1](sqlContext, Project(Alias(withEncoder(c1).expr, "_1")() :: Nil, logicalPlan))
+    // We use an unbound encoder since the expression will make up its own schema.
+    // TODO: This probably doesn't work if we are relying on reordering of the input class fields.
+    new Dataset[U1](
+      sqlContext,
+      Project(
+        c1.withInputType(
+          resolvedTEncoder.bind(queryExecution.analyzed.output),
+          queryExecution.analyzed.output).named :: Nil,
+        logicalPlan))
   }
 
   /**
@@ -369,28 +381,14 @@ class Dataset[T] private[sql](
    * that cast appropriately for the user facing interface.
    */
   protected def selectUntyped(columns: TypedColumn[_, _]*): Dataset[_] = {
-    val withEncoders = columns.map(withEncoder)
-    val aliases = withEncoders.zipWithIndex.map { case (c, i) => Alias(c.expr, s"_${i + 1}")() }
-    val unresolvedPlan = Project(aliases, logicalPlan)
-    val execution = new QueryExecution(sqlContext, unresolvedPlan)
-    // Rebind the encoders to the nested schema that will be produced by the select.
-    val encoders = withEncoders.map(_.encoder.asInstanceOf[ExpressionEncoder[_]]).zip(aliases).map {
-      case (e: ExpressionEncoder[_], a) if !e.flat =>
-        e.nested(a.toAttribute).resolve(execution.analyzed.output)
-      case (e, a) =>
-        e.unbind(a.toAttribute :: Nil).resolve(execution.analyzed.output)
-    }
-    new Dataset(sqlContext, execution, ExpressionEncoder.tuple(encoders))
-  }
+    val encoders = columns.map(_.encoder)
+    // We use an unbound encoder since the expression will make up its own schema.
+    // TODO: This probably doesn't work if we are relying on reordering of the input class fields.
+    val namedColumns =
+      columns.map(_.withInputType(unresolvedTEncoder, queryExecution.analyzed.output).named)
+    val execution = new QueryExecution(sqlContext, Project(namedColumns, logicalPlan))
 
-  private def withEncoder(c: TypedColumn[_, _]): TypedColumn[_, _] = {
-    val e = c.expr transform {
-      case ta: TypedAggregateExpression if ta.aEncoder.isEmpty =>
-        ta.copy(
-          aEncoder = Some(encoder.asInstanceOf[ExpressionEncoder[Any]]),
-          children = queryExecution.analyzed.output)
-    }
-    new TypedColumn(e, c.encoder)
+    new Dataset(sqlContext, execution, ExpressionEncoder.tuple(encoders))
   }
 
   /**
@@ -497,23 +495,18 @@ class Dataset[T] private[sql](
     val left = this.logicalPlan
     val right = other.logicalPlan
 
-    val leftData = this.encoder match {
+    val leftData = this.unresolvedTEncoder match {
       case e if e.flat => Alias(left.output.head, "_1")()
       case _ => Alias(CreateStruct(left.output), "_1")()
     }
-    val rightData = other.encoder match {
+    val rightData = other.unresolvedTEncoder match {
       case e if e.flat => Alias(right.output.head, "_2")()
       case _ => Alias(CreateStruct(right.output), "_2")()
     }
-    val leftEncoder =
-      if (encoder.flat) encoder else encoder.nested(leftData.toAttribute)
-    val rightEncoder =
-      if (other.encoder.flat) other.encoder else other.encoder.nested(rightData.toAttribute)
-    implicit val tuple2Encoder: Encoder[(T, U)] =
-      ExpressionEncoder.tuple(
-        leftEncoder,
-        rightEncoder.rebind(right.output, left.output ++ right.output))
 
+
+    implicit val tuple2Encoder: Encoder[(T, U)] =
+      ExpressionEncoder.tuple(this.unresolvedTEncoder, other.unresolvedTEncoder)
     withPlan[(T, U)](other) { (left, right) =>
       Project(
         leftData :: rightData :: Nil,
@@ -580,7 +573,7 @@ class Dataset[T] private[sql](
   private[sql] def logicalPlan = queryExecution.analyzed
 
   private[sql] def withPlan(f: LogicalPlan => LogicalPlan): Dataset[T] =
-    new Dataset[T](sqlContext, sqlContext.executePlan(f(logicalPlan)), encoder)
+    new Dataset[T](sqlContext, sqlContext.executePlan(f(logicalPlan)), tEncoder)
 
   private[sql] def withPlan[R : Encoder](
       other: Dataset[_])(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 61e2a9545069b..ae1272ae531fb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -17,20 +17,16 @@
 
 package org.apache.spark.sql
 
-import java.util.{Iterator => JIterator}
 
 import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
-import org.apache.spark.api.java.function.{Function2 => JFunction2, Function3 => JFunction3, _}
-import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttribute}
+import org.apache.spark.api.java.function._
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor}
-import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression, Alias, Attribute}
+import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.catalyst.plans.logical._
-import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
 import org.apache.spark.sql.execution.QueryExecution
 
-
 /**
  * :: Experimental ::
  * A [[Dataset]] has been logically grouped by a user specified grouping key.  Users should not
@@ -44,23 +40,21 @@ import org.apache.spark.sql.execution.QueryExecution
  */
 @Experimental
 class GroupedDataset[K, T] private[sql](
-    private val kEncoder: Encoder[K],
-    private val tEncoder: Encoder[T],
-    queryExecution: QueryExecution,
+    kEncoder: Encoder[K],
+    tEncoder: Encoder[T],
+    val queryExecution: QueryExecution,
     private val dataAttributes: Seq[Attribute],
     private val groupingAttributes: Seq[Attribute]) extends Serializable {
 
-  private implicit val kEnc = kEncoder match {
-    case e: ExpressionEncoder[K] => e.unbind(groupingAttributes).resolve(groupingAttributes)
-    case other =>
-      throw new UnsupportedOperationException("Only expression encoders are currently supported")
-  }
+  // Similar to [[Dataset]], we use unresolved encoders for later composition and resolved encoders
+  // when constructing new logical plans that will operate on the output of the current
+  // queryexecution.
 
-  private implicit val tEnc = tEncoder match {
-    case e: ExpressionEncoder[T] => e.resolve(dataAttributes)
-    case other =>
-      throw new UnsupportedOperationException("Only expression encoders are currently supported")
-  }
+  private implicit val unresolvedKEncoder = encoderFor(kEncoder)
+  private implicit val unresolvedTEncoder = encoderFor(tEncoder)
+
+  private val resolvedKEncoder = unresolvedKEncoder.resolve(groupingAttributes)
+  private val resolvedTEncoder = unresolvedTEncoder.resolve(dataAttributes)
 
   /** Encoders for built in aggregations. */
   private implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder[Long](flat = true)
@@ -79,7 +73,7 @@ class GroupedDataset[K, T] private[sql](
   def asKey[L : Encoder]: GroupedDataset[L, T] =
     new GroupedDataset(
       encoderFor[L],
-      tEncoder,
+      unresolvedTEncoder,
       queryExecution,
       dataAttributes,
       groupingAttributes)
@@ -95,7 +89,7 @@ class GroupedDataset[K, T] private[sql](
   }
 
   /**
-   * Applies the given function to each group of data.  For each unique group, the function will
+   * Applies the given function to each group of data.  For each unique group, the function  will
    * be passed the group key and an iterator that contains all of the elements in the group. The
    * function can return an iterator containing elements of an arbitrary type which will be returned
    * as a new [[Dataset]].
@@ -108,7 +102,12 @@ class GroupedDataset[K, T] private[sql](
   def flatMap[U : Encoder](f: (K, Iterator[T]) => TraversableOnce[U]): Dataset[U] = {
     new Dataset[U](
       sqlContext,
-      MapGroups(f, groupingAttributes, logicalPlan))
+      MapGroups(
+        f,
+        resolvedKEncoder,
+        resolvedTEncoder,
+        groupingAttributes,
+        logicalPlan))
   }
 
   def flatMap[U](f: FlatMapGroupFunction[K, T, U], encoder: Encoder[U]): Dataset[U] = {
@@ -127,15 +126,28 @@ class GroupedDataset[K, T] private[sql](
    */
   def map[U : Encoder](f: (K, Iterator[T]) => U): Dataset[U] = {
     val func = (key: K, it: Iterator[T]) => Iterator(f(key, it))
-    new Dataset[U](
-      sqlContext,
-      MapGroups(func, groupingAttributes, logicalPlan))
+    flatMap(func)
   }
 
   def map[U](f: MapGroupFunction[K, T, U], encoder: Encoder[U]): Dataset[U] = {
     map((key, data) => f.call(key, data.asJava))(encoder)
   }
 
+  /**
+   * Reduces the elements of each group of data using the specified binary function.
+   * The given function must be commutative and associative or the result may be non-deterministic.
+   */
+  def reduce(f: (T, T) => T): Dataset[(K, T)] = {
+    val func = (key: K, it: Iterator[T]) => Iterator(key -> it.reduce(f))
+
+    implicit val resultEncoder = ExpressionEncoder.tuple(unresolvedKEncoder, unresolvedTEncoder)
+    flatMap(func)
+  }
+
+  def reduce(f: ReduceFunction[T]): Dataset[(K, T)] = {
+    reduce(f.call _)
+  }
+
   // To ensure valid overloading.
   protected def agg(expr: Column, exprs: Column*): DataFrame =
     groupedData.agg(expr, exprs: _*)
@@ -147,37 +159,17 @@ class GroupedDataset[K, T] private[sql](
    * TODO: does not handle aggrecations that return nonflat results,
    */
   protected def aggUntyped(columns: TypedColumn[_, _]*): Dataset[_] = {
-    val aliases = (groupingAttributes ++ columns.map(_.expr)).map {
-      case u: UnresolvedAttribute => UnresolvedAlias(u)
-      case expr: NamedExpression => expr
-      case expr: Expression => Alias(expr, expr.prettyString)()
-    }
-
-    val unresolvedPlan = Aggregate(groupingAttributes, aliases, logicalPlan)
-
-    // Fill in the input encoders for any aggregators in the plan.
-    val withEncoders = unresolvedPlan transformAllExpressions {
-      case ta: TypedAggregateExpression if ta.aEncoder.isEmpty =>
-        ta.copy(
-          aEncoder = Some(tEnc.asInstanceOf[ExpressionEncoder[Any]]),
-          children = dataAttributes)
-    }
-    val execution = new QueryExecution(sqlContext, withEncoders)
-
-    val columnEncoders = columns.map(_.encoder.asInstanceOf[ExpressionEncoder[_]])
-
-    // Rebind the encoders to the nested schema that will be produced by the aggregation.
-    val encoders = (kEnc +: columnEncoders).zip(execution.analyzed.output).map {
-      case (e: ExpressionEncoder[_], a) if !e.flat =>
-        e.nested(a).resolve(execution.analyzed.output)
-      case (e, a) =>
-        e.unbind(a :: Nil).resolve(execution.analyzed.output)
-    }
+    val encoders = columns.map(_.encoder)
+    val namedColumns =
+      columns.map(
+        _.withInputType(resolvedTEncoder.bind(dataAttributes), dataAttributes).named)
+    val aggregate = Aggregate(groupingAttributes, groupingAttributes ++ namedColumns, logicalPlan)
+    val execution = new QueryExecution(sqlContext, aggregate)
 
     new Dataset(
       sqlContext,
       execution,
-      ExpressionEncoder.tuple(encoders))
+      ExpressionEncoder.tuple(unresolvedKEncoder +: encoders))
   }
 
   /**
@@ -230,7 +222,7 @@ class GroupedDataset[K, T] private[sql](
   def cogroup[U, R : Encoder](
       other: GroupedDataset[K, U])(
       f: (K, Iterator[T], Iterator[U]) => TraversableOnce[R]): Dataset[R] = {
-    implicit def uEnc: Encoder[U] = other.tEncoder
+    implicit def uEnc: Encoder[U] = other.unresolvedTEncoder
     new Dataset[R](
       sqlContext,
       CoGroup(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
index dfcbac8687b3e..3f2775896bb8c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
@@ -55,7 +55,7 @@ case class TypedAggregateExpression(
     aEncoder: Option[ExpressionEncoder[Any]],
     bEncoder: ExpressionEncoder[Any],
     cEncoder: ExpressionEncoder[Any],
-    children: Seq[Expression],
+    children: Seq[Attribute],
     mutableAggBufferOffset: Int,
     inputAggBufferOffset: Int)
   extends ImperativeAggregate with Logging {
@@ -78,8 +78,7 @@ case class TypedAggregateExpression(
 
   override lazy val resolved: Boolean = aEncoder.isDefined
 
-  override lazy val inputTypes: Seq[DataType] =
-    aEncoder.map(_.schema.map(_.dataType)).getOrElse(Nil)
+  override lazy val inputTypes: Seq[DataType] = Nil
 
   override val aggBufferSchema: StructType = bEncoder.schema
 
@@ -90,12 +89,8 @@ case class TypedAggregateExpression(
   override val inputAggBufferAttributes: Seq[AttributeReference] =
     aggBufferAttributes.map(_.newInstance())
 
-  lazy val inputAttributes = aEncoder.get.schema.toAttributes
-  lazy val inputMapping = AttributeMap(inputAttributes.zip(children))
-  lazy val boundA =
-    aEncoder.get.copy(constructExpression = aEncoder.get.constructExpression transform {
-      case a: AttributeReference => inputMapping(a)
-    })
+  // We let the dataset do the binding for us.
+  lazy val boundA = aEncoder.get
 
   val bAttributes = bEncoder.schema.toAttributes
   lazy val boundB = bEncoder.resolve(bAttributes).bind(bAttributes)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index ae08fb71bf4cb..ed82c9a6a3770 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -311,6 +311,10 @@ case class AppendColumns[T, U](
     newColumns: Seq[Attribute],
     child: SparkPlan) extends UnaryNode {
 
+  // We are using an unsafe combiner.
+  override def canProcessSafeRows: Boolean = false
+  override def canProcessUnsafeRows: Boolean = true
+
   override def output: Seq[Attribute] = child.output ++ newColumns
 
   override protected def doExecute(): RDD[InternalRow] = {
@@ -349,11 +353,12 @@ case class MapGroups[K, T, U](
     child.execute().mapPartitions { iter =>
       val grouped = GroupedIterator(iter, groupingAttributes, child.output)
       val groupKeyEncoder = kEncoder.bind(groupingAttributes)
+      val groupDataEncoder = tEncoder.bind(child.output)
 
       grouped.flatMap { case (key, rowIter) =>
         val result = func(
           groupKeyEncoder.fromRow(key),
-          rowIter.map(tEncoder.fromRow))
+          rowIter.map(groupDataEncoder.fromRow))
         result.map(uEncoder.toRow)
       }
     }
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index 33d8388f615ae..46169ca07d715 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -157,7 +157,6 @@ public Integer call(Integer v1, Integer v2) throws Exception {
     Assert.assertEquals(6, reduced);
   }
 
-  @Test
   public void testGroupBy() {
     List<String> data = Arrays.asList("a", "foo", "bar");
     Dataset<String> ds = context.createDataset(data, Encoders.STRING());
@@ -196,6 +195,17 @@ public Iterable<String> call(Integer key, Iterator<String> values) throws Except
 
     Assert.assertEquals(Arrays.asList("1a", "3foobar"), flatMapped.collectAsList());
 
+    Dataset<Tuple2<Integer, String>> reduced = grouped.reduce(new ReduceFunction<String>() {
+      @Override
+      public String call(String v1, String v2) throws Exception {
+        return v1 + v2;
+      }
+    });
+
+    Assert.assertEquals(
+      Arrays.asList(tuple2(1, "a"), tuple2(3, "foobar")),
+      reduced.collectAsList());
+
     List<Integer> data2 = Arrays.asList(2, 6, 10);
     Dataset<Integer> ds2 = context.createDataset(data2, Encoders.INT());
     GroupedDataset<Integer, Integer> grouped2 = ds2.groupBy(new MapFunction<Integer, Integer>() {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
index 378cd365276b3..20896efdfec16 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -67,6 +67,28 @@ object ComplexResultAgg extends Aggregator[(String, Int), (Long, Long), (Long, L
   override def finish(reduction: (Long, Long)): (Long, Long) = reduction
 }
 
+case class AggData(a: Int, b: String)
+object ClassInputAgg extends Aggregator[AggData, Int, Int] with Serializable {
+  /** A zero value for this aggregation. Should satisfy the property that any b + zero = b */
+  override def zero: Int = 0
+
+  /**
+   * Combine two values to produce a new value.  For performance, the function may modify `b` and
+   * return it instead of constructing new object for b.
+   */
+  override def reduce(b: Int, a: AggData): Int = b + a.a
+
+  /**
+   * Transform the output of the reduction.
+   */
+  override def finish(reduction: Int): Int = reduction
+
+  /**
+   * Merge two intermediate values
+   */
+  override def merge(b1: Int, b2: Int): Int = b1 + b2
+}
+
 class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
 
   import testImplicits._
@@ -123,4 +145,24 @@ class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
       ds.select(sum((i: Int) => i), sum((i: Int) => i * 2)),
       11 -> 22)
   }
+
+  test("typed aggregation: class input") {
+    val ds = Seq(AggData(1, "one"), AggData(2, "two")).toDS()
+
+    checkAnswer(
+      ds.select(ClassInputAgg.toColumn),
+      3)
+  }
+
+  test("typed aggregation: class input with reordering") {
+    val ds = sql("SELECT 'one' AS b, 1 as a").as[AggData]
+
+    checkAnswer(
+      ds.select(ClassInputAgg.toColumn),
+      1)
+
+    checkAnswer(
+      ds.groupBy(_.b).agg(ClassInputAgg.toColumn),
+      ("one", 1))
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 621148528714f..c23dd46d3767b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -218,6 +218,15 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       "a", "30", "b", "3", "c", "1")
   }
 
+  test("groupBy function, reduce") {
+    val ds = Seq("abc", "xyz", "hello").toDS()
+    val agged = ds.groupBy(_.length).reduce(_ + _)
+
+    checkAnswer(
+      agged,
+      3 -> "abcxyz", 5 -> "hello")
+  }
+
   test("groupBy columns, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
index 7a8b7ae5bf265..b5417b195f396 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
@@ -82,18 +82,21 @@ abstract class QueryTest extends PlanTest {
         fail(
           s"""
              |Exception collecting dataset as objects
-             |${ds.encoder}
-             |${ds.encoder.constructExpression.treeString}
+             |${ds.resolvedTEncoder}
+             |${ds.resolvedTEncoder.fromRowExpression.treeString}
              |${ds.queryExecution}
            """.stripMargin, e)
     }
 
     if (decoded != expectedAnswer.toSet) {
+      val expected = expectedAnswer.toSet.toSeq.map((a: Any) => a.toString).sorted
+      val actual = decoded.toSet.toSeq.map((a: Any) => a.toString).sorted
+
+      val comparision = sideBySide("expected" +: expected, "spark" +: actual).mkString("\n")
       fail(
         s"""Decoded objects do not match expected objects:
-           |Expected: ${expectedAnswer.toSet.toSeq.map((a: Any) => a.toString).sorted}
-            |Actual ${decoded.toSet.toSeq.map((a: Any) => a.toString).sorted}
-            |${ds.encoder.constructExpression.treeString}
+            |$comparision
+            |${ds.resolvedTEncoder.fromRowExpression.treeString}
          """.stripMargin)
     }
   }

From 0f1d00a905614bb5eebf260566dbcb831158d445 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Thu, 12 Nov 2015 17:48:43 -0800
Subject: [PATCH 222/518] [SPARK-11663][STREAMING] Add Java API for
 trackStateByKey

TODO
- [x] Add Java API
- [x] Add API tests
- [x] Add a function test

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9636 from zsxwing/java-track.
---
 .../spark/api/java/function/Function4.java    |  27 +++
 .../JavaStatefulNetworkWordCount.java         |  45 ++--
 .../streaming/StatefulNetworkWordCount.scala  |   2 +-
 .../apache/spark/streaming/Java8APISuite.java |  43 ++++
 .../org/apache/spark/streaming/State.scala    |  25 ++-
 .../apache/spark/streaming/StateSpec.scala    |  84 +++++--
 .../streaming/api/java/JavaPairDStream.scala  |  46 +++-
 .../api/java/JavaTrackStateDStream.scala      |  44 ++++
 .../streaming/dstream/TrackStateDStream.scala |   1 +
 .../spark/streaming/rdd/TrackStateRDD.scala   |   4 +-
 .../spark/streaming/util/StateMap.scala       |   6 +-
 .../streaming/JavaTrackStateByKeySuite.java   | 210 ++++++++++++++++++
 12 files changed, 485 insertions(+), 52 deletions(-)
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/Function4.java
 create mode 100644 streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaTrackStateDStream.scala
 create mode 100644 streaming/src/test/java/org/apache/spark/streaming/JavaTrackStateByKeySuite.java

diff --git a/core/src/main/java/org/apache/spark/api/java/function/Function4.java b/core/src/main/java/org/apache/spark/api/java/function/Function4.java
new file mode 100644
index 0000000000000..fd727d64863d7
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/Function4.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+
+/**
+ * A four-argument function that takes arguments of type T1, T2, T3 and T4 and returns an R.
+ */
+public interface Function4<T1, T2, T3, T4, R> extends Serializable {
+  public R call(T1 v1, T2 v2, T3 v3, T4 v4) throws Exception;
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/streaming/JavaStatefulNetworkWordCount.java b/examples/src/main/java/org/apache/spark/examples/streaming/JavaStatefulNetworkWordCount.java
index 99b63a2590ae2..c400e4237abe3 100644
--- a/examples/src/main/java/org/apache/spark/examples/streaming/JavaStatefulNetworkWordCount.java
+++ b/examples/src/main/java/org/apache/spark/examples/streaming/JavaStatefulNetworkWordCount.java
@@ -26,18 +26,15 @@
 import com.google.common.base.Optional;
 import com.google.common.collect.Lists;
 
-import org.apache.spark.HashPartitioner;
 import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.function.*;
 import org.apache.spark.api.java.JavaPairRDD;
 import org.apache.spark.api.java.StorageLevels;
-import org.apache.spark.api.java.function.FlatMapFunction;
-import org.apache.spark.api.java.function.Function2;
-import org.apache.spark.api.java.function.PairFunction;
 import org.apache.spark.streaming.Durations;
-import org.apache.spark.streaming.api.java.JavaDStream;
-import org.apache.spark.streaming.api.java.JavaPairDStream;
-import org.apache.spark.streaming.api.java.JavaReceiverInputDStream;
-import org.apache.spark.streaming.api.java.JavaStreamingContext;
+import org.apache.spark.streaming.State;
+import org.apache.spark.streaming.StateSpec;
+import org.apache.spark.streaming.Time;
+import org.apache.spark.streaming.api.java.*;
 
 /**
  * Counts words cumulatively in UTF8 encoded, '\n' delimited text received from the network every
@@ -63,25 +60,12 @@ public static void main(String[] args) {
 
     StreamingExamples.setStreamingLogLevels();
 
-    // Update the cumulative count function
-    final Function2<List<Integer>, Optional<Integer>, Optional<Integer>> updateFunction =
-        new Function2<List<Integer>, Optional<Integer>, Optional<Integer>>() {
-          @Override
-          public Optional<Integer> call(List<Integer> values, Optional<Integer> state) {
-            Integer newSum = state.or(0);
-            for (Integer value : values) {
-              newSum += value;
-            }
-            return Optional.of(newSum);
-          }
-        };
-
     // Create the context with a 1 second batch size
     SparkConf sparkConf = new SparkConf().setAppName("JavaStatefulNetworkWordCount");
     JavaStreamingContext ssc = new JavaStreamingContext(sparkConf, Durations.seconds(1));
     ssc.checkpoint(".");
 
-    // Initial RDD input to updateStateByKey
+    // Initial RDD input to trackStateByKey
     @SuppressWarnings("unchecked")
     List<Tuple2<String, Integer>> tuples = Arrays.asList(new Tuple2<String, Integer>("hello", 1),
             new Tuple2<String, Integer>("world", 1));
@@ -105,9 +89,22 @@ public Tuple2<String, Integer> call(String s) {
           }
         });
 
+    // Update the cumulative count function
+    final Function4<Time, String, Optional<Integer>, State<Integer>, Optional<Tuple2<String, Integer>>> trackStateFunc =
+        new Function4<Time, String, Optional<Integer>, State<Integer>, Optional<Tuple2<String, Integer>>>() {
+
+          @Override
+          public Optional<Tuple2<String, Integer>> call(Time time, String word, Optional<Integer> one, State<Integer> state) {
+            int sum = one.or(0) + (state.exists() ? state.get() : 0);
+            Tuple2<String, Integer> output = new Tuple2<String, Integer>(word, sum);
+            state.update(sum);
+            return Optional.of(output);
+          }
+        };
+
     // This will give a Dstream made of state (which is the cumulative count of the words)
-    JavaPairDStream<String, Integer> stateDstream = wordsDstream.updateStateByKey(updateFunction,
-            new HashPartitioner(ssc.sparkContext().defaultParallelism()), initialRDD);
+    JavaTrackStateDStream<String, Integer, Integer, Tuple2<String, Integer>> stateDstream =
+        wordsDstream.trackStateByKey(StateSpec.function(trackStateFunc).initialState(initialRDD));
 
     stateDstream.print();
     ssc.start();
diff --git a/examples/src/main/scala/org/apache/spark/examples/streaming/StatefulNetworkWordCount.scala b/examples/src/main/scala/org/apache/spark/examples/streaming/StatefulNetworkWordCount.scala
index be2ae0b473363..a4f847f118b2c 100644
--- a/examples/src/main/scala/org/apache/spark/examples/streaming/StatefulNetworkWordCount.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/streaming/StatefulNetworkWordCount.scala
@@ -49,7 +49,7 @@ object StatefulNetworkWordCount {
     val ssc = new StreamingContext(sparkConf, Seconds(1))
     ssc.checkpoint(".")
 
-    // Initial RDD input to updateStateByKey
+    // Initial RDD input to trackStateByKey
     val initialRDD = ssc.sparkContext.parallelize(List(("hello", 1), ("world", 1)))
 
     // Create a ReceiverInputDStream on target ip:port and count the
diff --git a/extras/java8-tests/src/test/java/org/apache/spark/streaming/Java8APISuite.java b/extras/java8-tests/src/test/java/org/apache/spark/streaming/Java8APISuite.java
index 73091cfe2c09e..163ae92c12c6d 100644
--- a/extras/java8-tests/src/test/java/org/apache/spark/streaming/Java8APISuite.java
+++ b/extras/java8-tests/src/test/java/org/apache/spark/streaming/Java8APISuite.java
@@ -31,9 +31,12 @@
 import org.apache.spark.HashPartitioner;
 import org.apache.spark.api.java.JavaPairRDD;
 import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.function.Function2;
+import org.apache.spark.api.java.function.Function4;
 import org.apache.spark.api.java.function.PairFunction;
 import org.apache.spark.streaming.api.java.JavaDStream;
 import org.apache.spark.streaming.api.java.JavaPairDStream;
+import org.apache.spark.streaming.api.java.JavaTrackStateDStream;
 
 /**
  * Most of these tests replicate org.apache.spark.streaming.JavaAPISuite using java 8
@@ -831,4 +834,44 @@ public void testFlatMapValues() {
     Assert.assertEquals(expected, result);
   }
 
+  /**
+   * This test is only for testing the APIs. It's not necessary to run it.
+   */
+  public void testTrackStateByAPI() {
+    JavaPairRDD<String, Boolean> initialRDD = null;
+    JavaPairDStream<String, Integer> wordsDstream = null;
+
+    JavaTrackStateDStream<String, Integer, Boolean, Double> stateDstream =
+        wordsDstream.trackStateByKey(
+            StateSpec.<String, Integer, Boolean, Double> function((time, key, value, state) -> {
+              // Use all State's methods here
+              state.exists();
+              state.get();
+              state.isTimingOut();
+              state.remove();
+              state.update(true);
+              return Optional.of(2.0);
+            }).initialState(initialRDD)
+                .numPartitions(10)
+                .partitioner(new HashPartitioner(10))
+                .timeout(Durations.seconds(10)));
+
+    JavaPairDStream<String, Boolean> emittedRecords = stateDstream.stateSnapshots();
+
+    JavaTrackStateDStream<String, Integer, Boolean, Double> stateDstream2 =
+        wordsDstream.trackStateByKey(
+            StateSpec.<String, Integer, Boolean, Double>function((value, state) -> {
+              state.exists();
+              state.get();
+              state.isTimingOut();
+              state.remove();
+              state.update(true);
+              return 2.0;
+            }).initialState(initialRDD)
+                .numPartitions(10)
+                .partitioner(new HashPartitioner(10))
+                .timeout(Durations.seconds(10)));
+
+    JavaPairDStream<String, Boolean> emittedRecords2 = stateDstream2.stateSnapshots();
+  }
 }
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/State.scala b/streaming/src/main/scala/org/apache/spark/streaming/State.scala
index 7dd1b72f80499..604e64fc61630 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/State.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/State.scala
@@ -50,9 +50,30 @@ import org.apache.spark.annotation.Experimental
  *
  * }}}
  *
- * Java example:
+ * Java example of using `State`:
  * {{{
- *      TODO(@zsxwing)
+ *    // A tracking function that maintains an integer state and return a String
+ *   Function2<Optional<Integer>, State<Integer>, Optional<String>> trackStateFunc =
+ *       new Function2<Optional<Integer>, State<Integer>, Optional<String>>() {
+ *
+ *         @Override
+ *         public Optional<String> call(Optional<Integer> one, State<Integer> state) {
+ *           if (state.exists()) {
+ *             int existingState = state.get(); // Get the existing state
+ *             boolean shouldRemove = ...; // Decide whether to remove the state
+ *             if (shouldRemove) {
+ *               state.remove(); // Remove the state
+ *             } else {
+ *               int newState = ...;
+ *               state.update(newState); // Set the new state
+ *             }
+ *           } else {
+ *             int initialState = ...; // Set the initial state
+ *             state.update(initialState);
+ *           }
+ *           // return something
+ *         }
+ *       };
  * }}}
  */
 @Experimental
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala b/streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala
index c9fe35e74c1c7..bea5b9df20b53 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala
@@ -17,15 +17,14 @@
 
 package org.apache.spark.streaming
 
-import scala.reflect.ClassTag
-
+import com.google.common.base.Optional
 import org.apache.spark.annotation.Experimental
-import org.apache.spark.api.java.JavaPairRDD
+import org.apache.spark.api.java.{JavaPairRDD, JavaUtils}
+import org.apache.spark.api.java.function.{Function2 => JFunction2, Function4 => JFunction4}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.util.ClosureCleaner
 import org.apache.spark.{HashPartitioner, Partitioner}
 
-
 /**
  * :: Experimental ::
  * Abstract class representing all the specifications of the DStream transformation
@@ -49,12 +48,12 @@ import org.apache.spark.{HashPartitioner, Partitioner}
  *
  * Example in Java:
  * {{{
- *    StateStateSpec[KeyType, ValueType, StateType, EmittedDataType] spec =
- *      StateStateSpec.function[KeyType, ValueType, StateType, EmittedDataType](trackingFunction)
+ *    StateSpec<KeyType, ValueType, StateType, EmittedDataType> spec =
+ *      StateSpec.<KeyType, ValueType, StateType, EmittedDataType>function(trackingFunction)
  *                    .numPartition(10);
  *
- *    JavaDStream[EmittedDataType] emittedRecordDStream =
- *      javaPairDStream.trackStateByKey[StateType, EmittedDataType](spec);
+ *    JavaTrackStateDStream<KeyType, ValueType, StateType, EmittedType> emittedRecordDStream =
+ *      javaPairDStream.<StateType, EmittedDataType>trackStateByKey(spec);
  * }}}
  */
 @Experimental
@@ -92,6 +91,7 @@ sealed abstract class StateSpec[KeyType, ValueType, StateType, EmittedType] exte
 /**
  * :: Experimental ::
  * Builder object for creating instances of [[org.apache.spark.streaming.StateSpec StateSpec]]
+ * that is used for specifying the parameters of the DStream transformation `trackStateByKey`
  * that is used for specifying the parameters of the DStream transformation
  * `trackStateByKey` operation of a
  * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
@@ -103,28 +103,27 @@ sealed abstract class StateSpec[KeyType, ValueType, StateType, EmittedType] exte
  *      ...
  *    }
  *
- *    val spec = StateSpec.function(trackingFunction).numPartitions(10)
- *
- *    val emittedRecordDStream = keyValueDStream.trackStateByKey[StateType, EmittedDataType](spec)
+ *    val emittedRecordDStream = keyValueDStream.trackStateByKey[StateType, EmittedDataType](
+ *        StateSpec.function(trackingFunction).numPartitions(10))
  * }}}
  *
  * Example in Java:
  * {{{
- *    StateStateSpec[KeyType, ValueType, StateType, EmittedDataType] spec =
- *      StateStateSpec.function[KeyType, ValueType, StateType, EmittedDataType](trackingFunction)
+ *    StateSpec<KeyType, ValueType, StateType, EmittedDataType> spec =
+ *      StateSpec.<KeyType, ValueType, StateType, EmittedDataType>function(trackingFunction)
  *                    .numPartition(10);
  *
- *    JavaDStream[EmittedDataType] emittedRecordDStream =
- *      javaPairDStream.trackStateByKey[StateType, EmittedDataType](spec);
+ *    JavaTrackStateDStream<KeyType, ValueType, StateType, EmittedType> emittedRecordDStream =
+ *      javaPairDStream.<StateType, EmittedDataType>trackStateByKey(spec);
  * }}}
  */
 @Experimental
 object StateSpec {
   /**
    * Create a [[org.apache.spark.streaming.StateSpec StateSpec]] for setting all the specifications
-   * `trackStateByKey` operation on a
-   * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
-   * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
+   * of the `trackStateByKey` operation on a
+   * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]].
+   *
    * @param trackingFunction The function applied on every data item to manage the associated state
    *                         and generate the emitted data
    * @tparam KeyType      Class of the keys
@@ -141,9 +140,9 @@ object StateSpec {
 
   /**
    * Create a [[org.apache.spark.streaming.StateSpec StateSpec]] for setting all the specifications
-   * `trackStateByKey` operation on a
-   * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
-   * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
+   * of the `trackStateByKey` operation on a
+   * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]].
+   *
    * @param trackingFunction The function applied on every data item to manage the associated state
    *                         and generate the emitted data
    * @tparam ValueType    Class of the values
@@ -160,6 +159,48 @@ object StateSpec {
       }
     new StateSpecImpl(wrappedFunction)
   }
+
+  /**
+   * Create a [[org.apache.spark.streaming.StateSpec StateSpec]] for setting all
+   * the specifications of the `trackStateByKey` operation on a
+   * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]].
+   *
+   * @param javaTrackingFunction The function applied on every data item to manage the associated
+   *                             state and generate the emitted data
+   * @tparam KeyType      Class of the keys
+   * @tparam ValueType    Class of the values
+   * @tparam StateType    Class of the states data
+   * @tparam EmittedType  Class of the emitted data
+   */
+  def function[KeyType, ValueType, StateType, EmittedType](javaTrackingFunction:
+      JFunction4[Time, KeyType, Optional[ValueType], State[StateType], Optional[EmittedType]]):
+    StateSpec[KeyType, ValueType, StateType, EmittedType] = {
+    val trackingFunc = (time: Time, k: KeyType, v: Option[ValueType], s: State[StateType]) => {
+      val t = javaTrackingFunction.call(time, k, JavaUtils.optionToOptional(v), s)
+      Option(t.orNull)
+    }
+    StateSpec.function(trackingFunc)
+  }
+
+  /**
+   * Create a [[org.apache.spark.streaming.StateSpec StateSpec]] for setting all the specifications
+   * of the `trackStateByKey` operation on a
+   * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]].
+   *
+   * @param javaTrackingFunction The function applied on every data item to manage the associated
+   *                             state and generate the emitted data
+   * @tparam ValueType    Class of the values
+   * @tparam StateType    Class of the states data
+   * @tparam EmittedType  Class of the emitted data
+   */
+  def function[KeyType, ValueType, StateType, EmittedType](
+      javaTrackingFunction: JFunction2[Optional[ValueType], State[StateType], EmittedType]):
+    StateSpec[KeyType, ValueType, StateType, EmittedType] = {
+    val trackingFunc = (v: Option[ValueType], s: State[StateType]) => {
+      javaTrackingFunction.call(Optional.fromNullable(v.get), s)
+    }
+    StateSpec.function(trackingFunc)
+  }
 }
 
 
@@ -184,7 +225,6 @@ case class StateSpecImpl[K, V, S, T](
     this
   }
 
-
   override def numPartitions(numPartitions: Int): this.type = {
     this.partitioner(new HashPartitioner(numPartitions))
     this
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaPairDStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaPairDStream.scala
index e2aec6c2f63e7..70e32b383e458 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaPairDStream.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaPairDStream.scala
@@ -28,8 +28,10 @@ import com.google.common.base.Optional
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.mapred.{JobConf, OutputFormat}
 import org.apache.hadoop.mapreduce.{OutputFormat => NewOutputFormat}
+
 import org.apache.spark.Partitioner
-import org.apache.spark.api.java.{JavaPairRDD, JavaUtils}
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.api.java.{JavaPairRDD, JavaSparkContext, JavaUtils}
 import org.apache.spark.api.java.JavaPairRDD._
 import org.apache.spark.api.java.JavaSparkContext.fakeClassTag
 import org.apache.spark.api.java.function.{Function => JFunction, Function2 => JFunction2}
@@ -426,6 +428,48 @@ class JavaPairDStream[K, V](val dstream: DStream[(K, V)])(
     )
   }
 
+  /**
+   * :: Experimental ::
+   * Return a new [[JavaDStream]] of data generated by combining the key-value data in `this` stream
+   * with a continuously updated per-key state. The user-provided state tracking function is
+   * applied on each keyed data item along with its corresponding state. The function can choose to
+   * update/remove the state and return a transformed data, which forms the
+   * [[JavaTrackStateDStream]].
+   *
+   * The specifications of this transformation is made through the
+   * [[org.apache.spark.streaming.StateSpec StateSpec]] class. Besides the tracking function, there
+   * are a number of optional parameters - initial state data, number of partitions, timeouts, etc.
+   * See the [[org.apache.spark.streaming.StateSpec StateSpec]] for more details.
+   *
+   * Example of using `trackStateByKey`:
+   * {{{
+   *   // A tracking function that maintains an integer state and return a String
+   *   Function2<Optional<Integer>, State<Integer>, Optional<String>> trackStateFunc =
+   *       new Function2<Optional<Integer>, State<Integer>, Optional<String>>() {
+   *
+   *         @Override
+   *         public Optional<String> call(Optional<Integer> one, State<Integer> state) {
+   *           // Check if state exists, accordingly update/remove state and return transformed data
+   *         }
+   *       };
+   *
+   *    JavaTrackStateDStream<Integer, Integer, Integer, String> trackStateDStream =
+   *        keyValueDStream.<Integer, String>trackStateByKey(
+   *                 StateSpec.function(trackStateFunc).numPartitions(10));
+   * }}}
+   *
+   * @param spec          Specification of this transformation
+   * @tparam StateType    Class type of the state
+   * @tparam EmittedType  Class type of the tranformed data return by the tracking function
+   */
+  @Experimental
+  def trackStateByKey[StateType, EmittedType](spec: StateSpec[K, V, StateType, EmittedType]):
+    JavaTrackStateDStream[K, V, StateType, EmittedType] = {
+    new JavaTrackStateDStream(dstream.trackStateByKey(spec)(
+      JavaSparkContext.fakeClassTag,
+      JavaSparkContext.fakeClassTag))
+  }
+
   private def convertUpdateStateFunction[S](in: JFunction2[JList[V], Optional[S], Optional[S]]):
   (Seq[V], Option[S]) => Option[S] = {
     val scalaFunc: (Seq[V], Option[S]) => Option[S] = (values, state) => {
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaTrackStateDStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaTrackStateDStream.scala
new file mode 100644
index 0000000000000..f459930d0660b
--- /dev/null
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaTrackStateDStream.scala
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming.api.java
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.api.java.JavaSparkContext
+import org.apache.spark.streaming.dstream.TrackStateDStream
+
+/**
+ * :: Experimental ::
+ * [[JavaDStream]] representing the stream of records emitted by the tracking function in the
+ * `trackStateByKey` operation on a [[JavaPairDStream]]. Additionally, it also gives access to the
+ * stream of state snapshots, that is, the state data of all keys after a batch has updated them.
+ *
+ * @tparam KeyType Class of the state key
+ * @tparam ValueType Class of the state value
+ * @tparam StateType Class of the state
+ * @tparam EmittedType Class of the emitted records
+ */
+@Experimental
+class JavaTrackStateDStream[KeyType, ValueType, StateType, EmittedType](
+    dstream: TrackStateDStream[KeyType, ValueType, StateType, EmittedType])
+  extends JavaDStream[EmittedType](dstream)(JavaSparkContext.fakeClassTag) {
+
+  def stateSnapshots(): JavaPairDStream[KeyType, StateType] =
+    new JavaPairDStream(dstream.stateSnapshots())(
+      JavaSparkContext.fakeClassTag,
+      JavaSparkContext.fakeClassTag)
+}
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
index 58d89c93bcbef..98e881e6ae115 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
@@ -35,6 +35,7 @@ import org.apache.spark.streaming.rdd.{TrackStateRDD, TrackStateRDDRecord}
  * all keys after a batch has updated them.
  *
  * @tparam KeyType Class of the state key
+ * @tparam ValueType Class of the state value
  * @tparam StateType Class of the state data
  * @tparam EmittedType Class of the emitted records
  */
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala b/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
index ed7cea26d0608..fc51496be47bf 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
@@ -70,12 +70,14 @@ private[streaming] class TrackStateRDDPartition(
  *                           in the `prevStateRDD` to create `this` RDD
  * @param trackingFunction The function that will be used to update state and return new data
  * @param batchTime        The time of the batch to which this RDD belongs to. Use to update
+ * @param timeoutThresholdTime The time to indicate which keys are timeout
  */
 private[streaming] class TrackStateRDD[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
     private var prevStateRDD: RDD[TrackStateRDDRecord[K, S, T]],
     private var partitionedDataRDD: RDD[(K, V)],
     trackingFunction: (Time, K, Option[V], State[S]) => Option[T],
-    batchTime: Time, timeoutThresholdTime: Option[Long]
+    batchTime: Time,
+    timeoutThresholdTime: Option[Long]
   ) extends RDD[TrackStateRDDRecord[K, S, T]](
     partitionedDataRDD.sparkContext,
     List(
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
index ed622ef7bf700..34287c3e00908 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
@@ -267,7 +267,11 @@ private[streaming] class OpenHashMapBasedStateMap[K: ClassTag, S: ClassTag](
 
     // Read the data of the delta
     val deltaMapSize = inputStream.readInt()
-    deltaMap = new OpenHashMap[K, StateInfo[S]]()
+    deltaMap = if (deltaMapSize != 0) {
+        new OpenHashMap[K, StateInfo[S]](deltaMapSize)
+      } else {
+        new OpenHashMap[K, StateInfo[S]](initialCapacity)
+      }
     var deltaMapCount = 0
     while (deltaMapCount < deltaMapSize) {
       val key = inputStream.readObject().asInstanceOf[K]
diff --git a/streaming/src/test/java/org/apache/spark/streaming/JavaTrackStateByKeySuite.java b/streaming/src/test/java/org/apache/spark/streaming/JavaTrackStateByKeySuite.java
new file mode 100644
index 0000000000000..eac4cdd14a683
--- /dev/null
+++ b/streaming/src/test/java/org/apache/spark/streaming/JavaTrackStateByKeySuite.java
@@ -0,0 +1,210 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.streaming;
+
+import java.io.Serializable;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+
+import scala.Tuple2;
+
+import com.google.common.base.Optional;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.streaming.api.java.JavaDStream;
+import org.apache.spark.util.ManualClock;
+import org.junit.Assert;
+import org.junit.Test;
+
+import org.apache.spark.HashPartitioner;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.function.Function2;
+import org.apache.spark.api.java.function.Function4;
+import org.apache.spark.streaming.api.java.JavaPairDStream;
+import org.apache.spark.streaming.api.java.JavaTrackStateDStream;
+
+public class JavaTrackStateByKeySuite extends LocalJavaStreamingContext implements Serializable {
+
+  /**
+   * This test is only for testing the APIs. It's not necessary to run it.
+   */
+  public void testAPI() {
+    JavaPairRDD<String, Boolean> initialRDD = null;
+    JavaPairDStream<String, Integer> wordsDstream = null;
+
+    final Function4<Time, String, Optional<Integer>, State<Boolean>, Optional<Double>>
+        trackStateFunc =
+        new Function4<Time, String, Optional<Integer>, State<Boolean>, Optional<Double>>() {
+
+          @Override
+          public Optional<Double> call(
+              Time time, String word, Optional<Integer> one, State<Boolean> state) {
+            // Use all State's methods here
+            state.exists();
+            state.get();
+            state.isTimingOut();
+            state.remove();
+            state.update(true);
+            return Optional.of(2.0);
+          }
+        };
+
+    JavaTrackStateDStream<String, Integer, Boolean, Double> stateDstream =
+        wordsDstream.trackStateByKey(
+            StateSpec.function(trackStateFunc)
+                .initialState(initialRDD)
+                .numPartitions(10)
+                .partitioner(new HashPartitioner(10))
+                .timeout(Durations.seconds(10)));
+
+    JavaPairDStream<String, Boolean> emittedRecords = stateDstream.stateSnapshots();
+
+    final Function2<Optional<Integer>, State<Boolean>, Double> trackStateFunc2 =
+        new Function2<Optional<Integer>, State<Boolean>, Double>() {
+
+          @Override
+          public Double call(Optional<Integer> one, State<Boolean> state) {
+            // Use all State's methods here
+            state.exists();
+            state.get();
+            state.isTimingOut();
+            state.remove();
+            state.update(true);
+            return 2.0;
+          }
+        };
+
+    JavaTrackStateDStream<String, Integer, Boolean, Double> stateDstream2 =
+        wordsDstream.trackStateByKey(
+            StateSpec.<String, Integer, Boolean, Double> function(trackStateFunc2)
+                .initialState(initialRDD)
+                .numPartitions(10)
+                .partitioner(new HashPartitioner(10))
+                .timeout(Durations.seconds(10)));
+
+    JavaPairDStream<String, Boolean> emittedRecords2 = stateDstream2.stateSnapshots();
+  }
+
+  @Test
+  public void testBasicFunction() {
+    List<List<String>> inputData = Arrays.asList(
+        Collections.<String>emptyList(),
+        Arrays.asList("a"),
+        Arrays.asList("a", "b"),
+        Arrays.asList("a", "b", "c"),
+        Arrays.asList("a", "b"),
+        Arrays.asList("a"),
+        Collections.<String>emptyList()
+    );
+
+    List<Set<Integer>> outputData = Arrays.asList(
+        Collections.<Integer>emptySet(),
+        Sets.newHashSet(1),
+        Sets.newHashSet(2, 1),
+        Sets.newHashSet(3, 2, 1),
+        Sets.newHashSet(4, 3),
+        Sets.newHashSet(5),
+        Collections.<Integer>emptySet()
+    );
+
+    List<Set<Tuple2<String, Integer>>> stateData = Arrays.asList(
+        Collections.<Tuple2<String, Integer>>emptySet(),
+        Sets.newHashSet(new Tuple2<String, Integer>("a", 1)),
+        Sets.newHashSet(new Tuple2<String, Integer>("a", 2), new Tuple2<String, Integer>("b", 1)),
+        Sets.newHashSet(
+            new Tuple2<String, Integer>("a", 3),
+            new Tuple2<String, Integer>("b", 2),
+            new Tuple2<String, Integer>("c", 1)),
+        Sets.newHashSet(
+            new Tuple2<String, Integer>("a", 4),
+            new Tuple2<String, Integer>("b", 3),
+            new Tuple2<String, Integer>("c", 1)),
+        Sets.newHashSet(
+            new Tuple2<String, Integer>("a", 5),
+            new Tuple2<String, Integer>("b", 3),
+            new Tuple2<String, Integer>("c", 1)),
+        Sets.newHashSet(
+            new Tuple2<String, Integer>("a", 5),
+            new Tuple2<String, Integer>("b", 3),
+            new Tuple2<String, Integer>("c", 1))
+    );
+
+    Function2<Optional<Integer>, State<Integer>, Integer> trackStateFunc =
+        new Function2<Optional<Integer>, State<Integer>, Integer>() {
+
+          @Override
+          public Integer call(Optional<Integer> value, State<Integer> state) throws Exception {
+            int sum = value.or(0) + (state.exists() ? state.get() : 0);
+            state.update(sum);
+            return sum;
+          }
+        };
+    testOperation(
+        inputData,
+        StateSpec.<String, Integer, Integer, Integer>function(trackStateFunc),
+        outputData,
+        stateData);
+  }
+
+  private <K, S, T> void testOperation(
+      List<List<K>> input,
+      StateSpec<K, Integer, S, T> trackStateSpec,
+      List<Set<T>> expectedOutputs,
+      List<Set<Tuple2<K, S>>> expectedStateSnapshots) {
+    int numBatches = expectedOutputs.size();
+    JavaDStream<K> inputStream = JavaTestUtils.attachTestInputStream(ssc, input, 2);
+    JavaTrackStateDStream<K, Integer, S, T> trackeStateStream =
+        JavaPairDStream.fromJavaDStream(inputStream.map(new Function<K, Tuple2<K, Integer>>() {
+          @Override
+          public Tuple2<K, Integer> call(K x) throws Exception {
+            return new Tuple2<K, Integer>(x, 1);
+          }
+        })).trackStateByKey(trackStateSpec);
+
+    final List<Set<T>> collectedOutputs =
+        Collections.synchronizedList(Lists.<Set<T>>newArrayList());
+    trackeStateStream.foreachRDD(new Function<JavaRDD<T>, Void>() {
+      @Override
+      public Void call(JavaRDD<T> rdd) throws Exception {
+        collectedOutputs.add(Sets.newHashSet(rdd.collect()));
+        return null;
+      }
+    });
+    final List<Set<Tuple2<K, S>>> collectedStateSnapshots =
+        Collections.synchronizedList(Lists.<Set<Tuple2<K, S>>>newArrayList());
+    trackeStateStream.stateSnapshots().foreachRDD(new Function<JavaPairRDD<K, S>, Void>() {
+      @Override
+      public Void call(JavaPairRDD<K, S> rdd) throws Exception {
+        collectedStateSnapshots.add(Sets.newHashSet(rdd.collect()));
+        return null;
+      }
+    });
+    BatchCounter batchCounter = new BatchCounter(ssc.ssc());
+    ssc.start();
+    ((ManualClock) ssc.ssc().scheduler().clock())
+        .advance(ssc.ssc().progressListener().batchDuration() * numBatches + 1);
+    batchCounter.waitUntilBatchesCompleted(numBatches, 10000);
+
+    Assert.assertEquals(expectedOutputs, collectedOutputs);
+    Assert.assertEquals(expectedStateSnapshots, collectedStateSnapshots);
+  }
+}

From 7786f9cc0790d27854a1e184f66a9b4df4d040a2 Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Thu, 12 Nov 2015 18:03:23 -0800
Subject: [PATCH 223/518] [SPARK-11419][STREAMING] Parallel recovery for
 FileBasedWriteAheadLog + minor recovery tweaks

The support for closing WriteAheadLog files after writes was just merged in. Closing every file after a write is a very expensive operation as it creates many small files on S3. It's not necessary to enable it on HDFS anyway.

However, when you have many small files on S3, recovery takes very long. In addition, files start stacking up pretty quickly, and deletes may not be able to keep up, therefore deletes can also be parallelized.

This PR adds support for the two parallelization steps mentioned above, in addition to a couple more failures I encountered during recovery.

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9373 from brkyvz/par-recovery.
---
 .../streaming/scheduler/JobScheduler.scala    |  6 +-
 .../util/FileBasedWriteAheadLog.scala         | 78 +++++++++++-----
 .../FileBasedWriteAheadLogRandomReader.scala  |  2 +-
 .../util/FileBasedWriteAheadLogReader.scala   | 17 +++-
 .../spark/streaming/util/HdfsUtils.scala      | 24 ++++-
 .../streaming/ReceivedBlockTrackerSuite.scala | 91 ++++++++++++++++++-
 .../streaming/util/WriteAheadLogSuite.scala   | 87 +++++++++++++++++-
 7 files changed, 268 insertions(+), 37 deletions(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/JobScheduler.scala b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/JobScheduler.scala
index 2480b4ec093e2..1ed6fb0aa9d52 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/scheduler/JobScheduler.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/scheduler/JobScheduler.scala
@@ -88,8 +88,10 @@ class JobScheduler(val ssc: StreamingContext) extends Logging {
     if (eventLoop == null) return // scheduler has already been stopped
     logDebug("Stopping JobScheduler")
 
-    // First, stop receiving
-    receiverTracker.stop(processAllReceivedData)
+    if (receiverTracker != null) {
+      // First, stop receiving
+      receiverTracker.stop(processAllReceivedData)
+    }
 
     // Second, stop generating jobs. If it has to process all received data,
     // then this will wait for all the processing through JobScheduler to be over.
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala
index bc3f2486c21fd..72705f1a9c010 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala
@@ -17,10 +17,12 @@
 package org.apache.spark.streaming.util
 
 import java.nio.ByteBuffer
+import java.util.concurrent.ThreadPoolExecutor
 import java.util.{Iterator => JIterator}
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
+import scala.collection.parallel.ThreadPoolTaskSupport
 import scala.concurrent.{Await, ExecutionContext, Future}
 import scala.language.postfixOps
 
@@ -57,8 +59,8 @@ private[streaming] class FileBasedWriteAheadLog(
   private val callerNameTag = getCallerName.map(c => s" for $c").getOrElse("")
 
   private val threadpoolName = s"WriteAheadLogManager $callerNameTag"
-  implicit private val executionContext = ExecutionContext.fromExecutorService(
-    ThreadUtils.newDaemonSingleThreadExecutor(threadpoolName))
+  private val threadpool = ThreadUtils.newDaemonCachedThreadPool(threadpoolName, 20)
+  private val executionContext = ExecutionContext.fromExecutorService(threadpool)
   override protected val logName = s"WriteAheadLogManager $callerNameTag"
 
   private var currentLogPath: Option[String] = None
@@ -124,13 +126,19 @@ private[streaming] class FileBasedWriteAheadLog(
    */
   def readAll(): JIterator[ByteBuffer] = synchronized {
     val logFilesToRead = pastLogs.map{ _.path} ++ currentLogPath
-    logInfo("Reading from the logs: " + logFilesToRead.mkString("\n"))
-
-    logFilesToRead.iterator.map { file =>
+    logInfo("Reading from the logs:\n" + logFilesToRead.mkString("\n"))
+    def readFile(file: String): Iterator[ByteBuffer] = {
       logDebug(s"Creating log reader with $file")
       val reader = new FileBasedWriteAheadLogReader(file, hadoopConf)
       CompletionIterator[ByteBuffer, Iterator[ByteBuffer]](reader, reader.close _)
-    }.flatten.asJava
+    }
+    if (!closeFileAfterWrite) {
+      logFilesToRead.iterator.map(readFile).flatten.asJava
+    } else {
+      // For performance gains, it makes sense to parallelize the recovery if
+      // closeFileAfterWrite = true
+      seqToParIterator(threadpool, logFilesToRead, readFile).asJava
+    }
   }
 
   /**
@@ -146,30 +154,33 @@ private[streaming] class FileBasedWriteAheadLog(
    * asynchronously.
    */
   def clean(threshTime: Long, waitForCompletion: Boolean): Unit = {
-    val oldLogFiles = synchronized { pastLogs.filter { _.endTime < threshTime } }
+    val oldLogFiles = synchronized {
+      val expiredLogs = pastLogs.filter { _.endTime < threshTime }
+      pastLogs --= expiredLogs
+      expiredLogs
+    }
     logInfo(s"Attempting to clear ${oldLogFiles.size} old log files in $logDirectory " +
       s"older than $threshTime: ${oldLogFiles.map { _.path }.mkString("\n")}")
 
-    def deleteFiles() {
-      oldLogFiles.foreach { logInfo =>
-        try {
-          val path = new Path(logInfo.path)
-          val fs = HdfsUtils.getFileSystemForPath(path, hadoopConf)
-          fs.delete(path, true)
-          synchronized { pastLogs -= logInfo }
-          logDebug(s"Cleared log file $logInfo")
-        } catch {
-          case ex: Exception =>
-            logWarning(s"Error clearing write ahead log file $logInfo", ex)
-        }
+    def deleteFile(walInfo: LogInfo): Unit = {
+      try {
+        val path = new Path(walInfo.path)
+        val fs = HdfsUtils.getFileSystemForPath(path, hadoopConf)
+        fs.delete(path, true)
+        logDebug(s"Cleared log file $walInfo")
+      } catch {
+        case ex: Exception =>
+          logWarning(s"Error clearing write ahead log file $walInfo", ex)
       }
       logInfo(s"Cleared log files in $logDirectory older than $threshTime")
     }
-    if (!executionContext.isShutdown) {
-      val f = Future { deleteFiles() }
-      if (waitForCompletion) {
-        import scala.concurrent.duration._
-        Await.ready(f, 1 second)
+    oldLogFiles.foreach { logInfo =>
+      if (!executionContext.isShutdown) {
+        val f = Future { deleteFile(logInfo) }(executionContext)
+        if (waitForCompletion) {
+          import scala.concurrent.duration._
+          Await.ready(f, 1 second)
+        }
       }
     }
   }
@@ -251,4 +262,23 @@ private[streaming] object FileBasedWriteAheadLog {
       }
     }.sortBy { _.startTime }
   }
+
+  /**
+   * This creates an iterator from a parallel collection, by keeping at most `n` objects in memory
+   * at any given time, where `n` is the size of the thread pool. This is crucial for use cases
+   * where we create `FileBasedWriteAheadLogReader`s during parallel recovery. We don't want to
+   * open up `k` streams altogether where `k` is the size of the Seq that we want to parallelize.
+   */
+  def seqToParIterator[I, O](
+      tpool: ThreadPoolExecutor,
+      source: Seq[I],
+      handler: I => Iterator[O]): Iterator[O] = {
+    val taskSupport = new ThreadPoolTaskSupport(tpool)
+    val groupSize = tpool.getMaximumPoolSize.max(8)
+    source.grouped(groupSize).flatMap { group =>
+      val parallelCollection = group.par
+      parallelCollection.tasksupport = taskSupport
+      parallelCollection.map(handler)
+    }.flatten
+  }
 }
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogRandomReader.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogRandomReader.scala
index f7168229ec15a..56d4977da0b51 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogRandomReader.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogRandomReader.scala
@@ -30,7 +30,7 @@ private[streaming] class FileBasedWriteAheadLogRandomReader(path: String, conf:
   extends Closeable {
 
   private val instream = HdfsUtils.getInputStream(path, conf)
-  private var closed = false
+  private var closed = (instream == null) // the file may be deleted as we're opening the stream
 
   def read(segment: FileBasedWriteAheadLogSegment): ByteBuffer = synchronized {
     assertOpen()
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogReader.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogReader.scala
index c3bb59f3fef94..a375c0729534b 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogReader.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLogReader.scala
@@ -16,7 +16,7 @@
  */
 package org.apache.spark.streaming.util
 
-import java.io.{Closeable, EOFException}
+import java.io.{IOException, Closeable, EOFException}
 import java.nio.ByteBuffer
 
 import org.apache.hadoop.conf.Configuration
@@ -32,7 +32,7 @@ private[streaming] class FileBasedWriteAheadLogReader(path: String, conf: Config
   extends Iterator[ByteBuffer] with Closeable with Logging {
 
   private val instream = HdfsUtils.getInputStream(path, conf)
-  private var closed = false
+  private var closed = (instream == null) // the file may be deleted as we're opening the stream
   private var nextItem: Option[ByteBuffer] = None
 
   override def hasNext: Boolean = synchronized {
@@ -55,6 +55,19 @@ private[streaming] class FileBasedWriteAheadLogReader(path: String, conf: Config
           logDebug("Error reading next item, EOF reached", e)
           close()
           false
+        case e: IOException =>
+          logWarning("Error while trying to read data. If the file was deleted, " +
+            "this should be okay.", e)
+          close()
+          if (HdfsUtils.checkFileExists(path, conf)) {
+            // If file exists, this could be a legitimate error
+            throw e
+          } else {
+            // File was deleted. This can occur when the daemon cleanup thread takes time to
+            // delete the file during recovery.
+            false
+          }
+
         case e: Exception =>
           logWarning("Error while trying to read data from HDFS.", e)
           close()
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/HdfsUtils.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/HdfsUtils.scala
index f60688f173c44..13a765d035ee8 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/HdfsUtils.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/HdfsUtils.scala
@@ -16,6 +16,8 @@
  */
 package org.apache.spark.streaming.util
 
+import java.io.IOException
+
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs._
 
@@ -42,8 +44,19 @@ private[streaming] object HdfsUtils {
   def getInputStream(path: String, conf: Configuration): FSDataInputStream = {
     val dfsPath = new Path(path)
     val dfs = getFileSystemForPath(dfsPath, conf)
-    val instream = dfs.open(dfsPath)
-    instream
+    if (dfs.isFile(dfsPath)) {
+      try {
+        dfs.open(dfsPath)
+      } catch {
+        case e: IOException =>
+          // If we are really unlucky, the file may be deleted as we're opening the stream.
+          // This can happen as clean up is performed by daemon threads that may be left over from
+          // previous runs.
+          if (!dfs.isFile(dfsPath)) null else throw e
+      }
+    } else {
+      null
+    }
   }
 
   def checkState(state: Boolean, errorMsg: => String) {
@@ -71,4 +84,11 @@ private[streaming] object HdfsUtils {
       case _ => fs
     }
   }
+
+  /** Check if the file exists at the given path. */
+  def checkFileExists(path: String, conf: Configuration): Boolean = {
+    val hdpPath = new Path(path)
+    val fs = getFileSystemForPath(hdpPath, conf)
+    fs.isFile(hdpPath)
+  }
 }
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala
index f793a12843b2f..7db17abb7947c 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.streaming
 
 import java.io.File
+import java.nio.ByteBuffer
 
 import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.duration._
@@ -32,7 +33,7 @@ import org.apache.spark.{Logging, SparkConf, SparkException, SparkFunSuite}
 import org.apache.spark.storage.StreamBlockId
 import org.apache.spark.streaming.receiver.BlockManagerBasedStoreResult
 import org.apache.spark.streaming.scheduler._
-import org.apache.spark.streaming.util.{WriteAheadLogUtils, FileBasedWriteAheadLogReader}
+import org.apache.spark.streaming.util._
 import org.apache.spark.streaming.util.WriteAheadLogSuite._
 import org.apache.spark.util.{Clock, ManualClock, SystemClock, Utils}
 
@@ -207,6 +208,75 @@ class ReceivedBlockTrackerSuite
     tracker1.isWriteAheadLogEnabled should be (false)
   }
 
+  test("parallel file deletion in FileBasedWriteAheadLog is robust to deletion error") {
+    conf.set("spark.streaming.driver.writeAheadLog.rollingIntervalSecs", "1")
+    require(WriteAheadLogUtils.getRollingIntervalSecs(conf, isDriver = true) === 1)
+
+    val addBlocks = generateBlockInfos()
+    val batch1 = addBlocks.slice(0, 1)
+    val batch2 = addBlocks.slice(1, 3)
+    val batch3 = addBlocks.slice(3, addBlocks.length)
+
+    assert(getWriteAheadLogFiles().length === 0)
+
+    // list of timestamps for files
+    val t = Seq.tabulate(5)(i => i * 1000)
+
+    writeEventsManually(getLogFileName(t(0)), Seq(createBatchCleanup(t(0))))
+    assert(getWriteAheadLogFiles().length === 1)
+
+    // The goal is to create several log files which should have been cleaned up.
+    // If we face any issue during recovery, because these old files exist, then we need to make
+    // deletion more robust rather than a parallelized operation where we fire and forget
+    val batch1Allocation = createBatchAllocation(t(1), batch1)
+    writeEventsManually(getLogFileName(t(1)), batch1.map(BlockAdditionEvent) :+ batch1Allocation)
+
+    writeEventsManually(getLogFileName(t(2)), Seq(createBatchCleanup(t(1))))
+
+    val batch2Allocation = createBatchAllocation(t(3), batch2)
+    writeEventsManually(getLogFileName(t(3)), batch2.map(BlockAdditionEvent) :+ batch2Allocation)
+
+    writeEventsManually(getLogFileName(t(4)), batch3.map(BlockAdditionEvent))
+
+    // We should have 5 different log files as we called `writeEventsManually` with 5 different
+    // timestamps
+    assert(getWriteAheadLogFiles().length === 5)
+
+    // Create the tracker to recover from the log files. We're going to ask the tracker to clean
+    // things up, and then we're going to rewrite that data, and recover using a different tracker.
+    // They should have identical data no matter what
+    val tracker = createTracker(recoverFromWriteAheadLog = true, clock = new ManualClock(t(4)))
+
+    def compareTrackers(base: ReceivedBlockTracker, subject: ReceivedBlockTracker): Unit = {
+      subject.getBlocksOfBatchAndStream(t(3), streamId) should be(
+        base.getBlocksOfBatchAndStream(t(3), streamId))
+      subject.getBlocksOfBatchAndStream(t(1), streamId) should be(
+        base.getBlocksOfBatchAndStream(t(1), streamId))
+      subject.getBlocksOfBatchAndStream(t(0), streamId) should be(Nil)
+    }
+
+    // ask the tracker to clean up some old files
+    tracker.cleanupOldBatches(t(3), waitForCompletion = true)
+    assert(getWriteAheadLogFiles().length === 3)
+
+    val tracker2 = createTracker(recoverFromWriteAheadLog = true, clock = new ManualClock(t(4)))
+    compareTrackers(tracker, tracker2)
+
+    // rewrite first file
+    writeEventsManually(getLogFileName(t(0)), Seq(createBatchCleanup(t(0))))
+    assert(getWriteAheadLogFiles().length === 4)
+    // make sure trackers are consistent
+    val tracker3 = createTracker(recoverFromWriteAheadLog = true, clock = new ManualClock(t(4)))
+    compareTrackers(tracker, tracker3)
+
+    // rewrite second file
+    writeEventsManually(getLogFileName(t(1)), batch1.map(BlockAdditionEvent) :+ batch1Allocation)
+    assert(getWriteAheadLogFiles().length === 5)
+    // make sure trackers are consistent
+    val tracker4 = createTracker(recoverFromWriteAheadLog = true, clock = new ManualClock(t(4)))
+    compareTrackers(tracker, tracker4)
+  }
+
   /**
    * Create tracker object with the optional provided clock. Use fake clock if you
    * want to control time by manually incrementing it to test log clean.
@@ -228,11 +298,30 @@ class ReceivedBlockTrackerSuite
       BlockManagerBasedStoreResult(StreamBlockId(streamId, math.abs(Random.nextInt)), Some(0L))))
   }
 
+  /**
+   * Write received block tracker events to a file manually.
+   */
+  def writeEventsManually(filePath: String, events: Seq[ReceivedBlockTrackerLogEvent]): Unit = {
+    val writer = HdfsUtils.getOutputStream(filePath, hadoopConf)
+    events.foreach { event =>
+      val bytes = Utils.serialize(event)
+      writer.writeInt(bytes.size)
+      writer.write(bytes)
+    }
+    writer.close()
+  }
+
   /** Get all the data written in the given write ahead log file. */
   def getWrittenLogData(logFile: String): Seq[ReceivedBlockTrackerLogEvent] = {
     getWrittenLogData(Seq(logFile))
   }
 
+  /** Get the log file name for the given log start time. */
+  def getLogFileName(time: Long, rollingIntervalSecs: Int = 1): String = {
+    checkpointDirectory.toString + File.separator + "receivedBlockMetadata" +
+      File.separator + s"log-$time-${time + rollingIntervalSecs * 1000}"
+  }
+
   /**
    * Get all the data written in the given write ahead log files. By default, it will read all
    * files in the test log directory.
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
index 9e13f25c2efea..4273fd7dda8be 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
@@ -19,7 +19,8 @@ package org.apache.spark.streaming.util
 import java.io._
 import java.nio.ByteBuffer
 import java.util.{Iterator => JIterator}
-import java.util.concurrent.ThreadPoolExecutor
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.{TimeUnit, CountDownLatch, ThreadPoolExecutor}
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
@@ -32,15 +33,13 @@ import org.apache.hadoop.fs.Path
 import org.mockito.Matchers.{eq => meq}
 import org.mockito.Matchers._
 import org.mockito.Mockito._
-import org.mockito.invocation.InvocationOnMock
-import org.mockito.stubbing.Answer
 import org.scalatest.concurrent.Eventually
 import org.scalatest.concurrent.Eventually._
 import org.scalatest.{PrivateMethodTester, BeforeAndAfterEach, BeforeAndAfter}
 import org.scalatest.mock.MockitoSugar
 
 import org.apache.spark.streaming.scheduler._
-import org.apache.spark.util.{ThreadUtils, ManualClock, Utils}
+import org.apache.spark.util.{CompletionIterator, ThreadUtils, ManualClock, Utils}
 import org.apache.spark.{SparkConf, SparkFunSuite}
 
 /** Common tests for WriteAheadLogs that we would like to test with different configurations. */
@@ -198,6 +197,64 @@ class FileBasedWriteAheadLogSuite
 
   import WriteAheadLogSuite._
 
+  test("FileBasedWriteAheadLog - seqToParIterator") {
+    /*
+     If the setting `closeFileAfterWrite` is enabled, we start generating a very large number of
+     files. This causes recovery to take a very long time. In order to make it quicker, we
+     parallelized the reading of these files. This test makes sure that we limit the number of
+     open files to the size of the number of threads in our thread pool rather than the size of
+     the list of files.
+     */
+    val numThreads = 8
+    val tpool = ThreadUtils.newDaemonFixedThreadPool(numThreads, "wal-test-thread-pool")
+    class GetMaxCounter {
+      private val value = new AtomicInteger()
+      @volatile private var max: Int = 0
+      def increment(): Unit = synchronized {
+        val atInstant = value.incrementAndGet()
+        if (atInstant > max) max = atInstant
+      }
+      def decrement(): Unit = synchronized { value.decrementAndGet() }
+      def get(): Int = synchronized { value.get() }
+      def getMax(): Int = synchronized { max }
+    }
+    try {
+      // If Jenkins is slow, we may not have a chance to run many threads simultaneously. Having
+      // a latch will make sure that all the threads can be launched altogether.
+      val latch = new CountDownLatch(1)
+      val testSeq = 1 to 1000
+      val counter = new GetMaxCounter()
+      def handle(value: Int): Iterator[Int] = {
+        new CompletionIterator[Int, Iterator[Int]](Iterator(value)) {
+          counter.increment()
+          // block so that other threads also launch
+          latch.await(10, TimeUnit.SECONDS)
+          override def completion() { counter.decrement() }
+        }
+      }
+      @volatile var collected: Seq[Int] = Nil
+      val t = new Thread() {
+        override def run() {
+          // run the calculation on a separate thread so that we can release the latch
+          val iterator = FileBasedWriteAheadLog.seqToParIterator[Int, Int](tpool, testSeq, handle)
+          collected = iterator.toSeq
+        }
+      }
+      t.start()
+      eventually(Eventually.timeout(10.seconds)) {
+        // make sure we are doing a parallel computation!
+        assert(counter.getMax() > 1)
+      }
+      latch.countDown()
+      t.join(10000)
+      assert(collected === testSeq)
+      // make sure we didn't open too many Iterators
+      assert(counter.getMax() <= numThreads)
+    } finally {
+      tpool.shutdownNow()
+    }
+  }
+
   test("FileBasedWriteAheadLogWriter - writing data") {
     val dataToWrite = generateRandomData()
     val segments = writeDataUsingWriter(testFile, dataToWrite)
@@ -259,6 +316,26 @@ class FileBasedWriteAheadLogSuite
     assert(readDataUsingReader(testFile) === (dataToWrite.dropRight(1)))
   }
 
+  test("FileBasedWriteAheadLogReader - handles errors when file doesn't exist") {
+    // Write data manually for testing the sequential reader
+    val dataToWrite = generateRandomData()
+    writeDataUsingWriter(testFile, dataToWrite)
+    val tFile = new File(testFile)
+    assert(tFile.exists())
+    // Verify the data can be read and is same as the one correctly written
+    assert(readDataUsingReader(testFile) === dataToWrite)
+
+    tFile.delete()
+    assert(!tFile.exists())
+
+    val reader = new FileBasedWriteAheadLogReader(testFile, hadoopConf)
+    assert(!reader.hasNext)
+    reader.close()
+
+    // Verify that no exception is thrown if file doesn't exist
+    assert(readDataUsingReader(testFile) === Nil)
+  }
+
   test("FileBasedWriteAheadLogRandomReader - reading data using random reader") {
     // Write data manually for testing the random reader
     val writtenData = generateRandomData()
@@ -581,7 +658,7 @@ object WriteAheadLogSuite {
       closeFileAfterWrite: Boolean,
       allowBatching: Boolean): Seq[String] = {
     val wal = createWriteAheadLog(logDirectory, closeFileAfterWrite, allowBatching)
-    val data = wal.readAll().asScala.map(byteBufferToString).toSeq
+    val data = wal.readAll().asScala.map(byteBufferToString).toArray
     wal.close()
     data
   }

From e4e46b20f6475f8e148d5326f7c88c57850d46a1 Mon Sep 17 00:00:00 2001
From: Tathagata Das <tathagata.das1565@gmail.com>
Date: Thu, 12 Nov 2015 19:02:49 -0800
Subject: [PATCH 224/518] [SPARK-11681][STREAMING] Correctly update state
 timestamp even when state is not updated

Bug: Timestamp is not updated if there is data but the corresponding state is not updated. This is wrong, and timeout is defined as "no data for a while", not "not state update for a while".

Fix: Update timestamp when timestamp when timeout is specified, otherwise no need.
Also refactored the code for better testability and added unit tests.

Author: Tathagata Das <tathagata.das1565@gmail.com>

Closes #9648 from tdas/SPARK-11681.
---
 .../spark/streaming/rdd/TrackStateRDD.scala   | 105 ++++++++------
 .../streaming/rdd/TrackStateRDDSuite.scala    | 136 +++++++++++++++++-
 2 files changed, 192 insertions(+), 49 deletions(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala b/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
index fc51496be47bf..7050378d0feb0 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/rdd/TrackStateRDD.scala
@@ -32,8 +32,51 @@ import org.apache.spark._
  * Record storing the keyed-state [[TrackStateRDD]]. Each record contains a [[StateMap]] and a
  * sequence of records returned by the tracking function of `trackStateByKey`.
  */
-private[streaming] case class TrackStateRDDRecord[K, S, T](
-    var stateMap: StateMap[K, S], var emittedRecords: Seq[T])
+private[streaming] case class TrackStateRDDRecord[K, S, E](
+    var stateMap: StateMap[K, S], var emittedRecords: Seq[E])
+
+private[streaming] object TrackStateRDDRecord {
+  def updateRecordWithData[K: ClassTag, V: ClassTag, S: ClassTag, E: ClassTag](
+    prevRecord: Option[TrackStateRDDRecord[K, S, E]],
+    dataIterator: Iterator[(K, V)],
+    updateFunction: (Time, K, Option[V], State[S]) => Option[E],
+    batchTime: Time,
+    timeoutThresholdTime: Option[Long],
+    removeTimedoutData: Boolean
+  ): TrackStateRDDRecord[K, S, E] = {
+    // Create a new state map by cloning the previous one (if it exists) or by creating an empty one
+    val newStateMap = prevRecord.map { _.stateMap.copy() }. getOrElse { new EmptyStateMap[K, S]() }
+
+    val emittedRecords = new ArrayBuffer[E]
+    val wrappedState = new StateImpl[S]()
+
+    // Call the tracking function on each record in the data iterator, and accordingly
+    // update the states touched, and collect the data returned by the tracking function
+    dataIterator.foreach { case (key, value) =>
+      wrappedState.wrap(newStateMap.get(key))
+      val emittedRecord = updateFunction(batchTime, key, Some(value), wrappedState)
+      if (wrappedState.isRemoved) {
+        newStateMap.remove(key)
+      } else if (wrappedState.isUpdated || timeoutThresholdTime.isDefined) {
+        newStateMap.put(key, wrappedState.get(), batchTime.milliseconds)
+      }
+      emittedRecords ++= emittedRecord
+    }
+
+    // Get the timed out state records, call the tracking function on each and collect the
+    // data returned
+    if (removeTimedoutData && timeoutThresholdTime.isDefined) {
+      newStateMap.getByTime(timeoutThresholdTime.get).foreach { case (key, state, _) =>
+        wrappedState.wrapTiminoutState(state)
+        val emittedRecord = updateFunction(batchTime, key, None, wrappedState)
+        emittedRecords ++= emittedRecord
+        newStateMap.remove(key)
+      }
+    }
+
+    TrackStateRDDRecord(newStateMap, emittedRecords)
+  }
+}
 
 /**
  * Partition of the [[TrackStateRDD]], which depends on corresponding partitions of prev state
@@ -72,16 +115,16 @@ private[streaming] class TrackStateRDDPartition(
  * @param batchTime        The time of the batch to which this RDD belongs to. Use to update
  * @param timeoutThresholdTime The time to indicate which keys are timeout
  */
-private[streaming] class TrackStateRDD[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
-    private var prevStateRDD: RDD[TrackStateRDDRecord[K, S, T]],
+private[streaming] class TrackStateRDD[K: ClassTag, V: ClassTag, S: ClassTag, E: ClassTag](
+    private var prevStateRDD: RDD[TrackStateRDDRecord[K, S, E]],
     private var partitionedDataRDD: RDD[(K, V)],
-    trackingFunction: (Time, K, Option[V], State[S]) => Option[T],
+    trackingFunction: (Time, K, Option[V], State[S]) => Option[E],
     batchTime: Time,
     timeoutThresholdTime: Option[Long]
-  ) extends RDD[TrackStateRDDRecord[K, S, T]](
+  ) extends RDD[TrackStateRDDRecord[K, S, E]](
     partitionedDataRDD.sparkContext,
     List(
-      new OneToOneDependency[TrackStateRDDRecord[K, S, T]](prevStateRDD),
+      new OneToOneDependency[TrackStateRDDRecord[K, S, E]](prevStateRDD),
       new OneToOneDependency(partitionedDataRDD))
   ) {
 
@@ -98,7 +141,7 @@ private[streaming] class TrackStateRDD[K: ClassTag, V: ClassTag, S: ClassTag, T:
   }
 
   override def compute(
-      partition: Partition, context: TaskContext): Iterator[TrackStateRDDRecord[K, S, T]] = {
+      partition: Partition, context: TaskContext): Iterator[TrackStateRDDRecord[K, S, E]] = {
 
     val stateRDDPartition = partition.asInstanceOf[TrackStateRDDPartition]
     val prevStateRDDIterator = prevStateRDD.iterator(
@@ -106,42 +149,16 @@ private[streaming] class TrackStateRDD[K: ClassTag, V: ClassTag, S: ClassTag, T:
     val dataIterator = partitionedDataRDD.iterator(
       stateRDDPartition.partitionedDataRDDPartition, context)
 
-    // Create a new state map by cloning the previous one (if it exists) or by creating an empty one
-    val newStateMap = if (prevStateRDDIterator.hasNext) {
-      prevStateRDDIterator.next().stateMap.copy()
-    } else {
-      new EmptyStateMap[K, S]()
-    }
-
-    val emittedRecords = new ArrayBuffer[T]
-    val wrappedState = new StateImpl[S]()
-
-    // Call the tracking function on each record in the data RDD partition, and accordingly
-    // update the states touched, and the data returned by the tracking function.
-    dataIterator.foreach { case (key, value) =>
-      wrappedState.wrap(newStateMap.get(key))
-      val emittedRecord = trackingFunction(batchTime, key, Some(value), wrappedState)
-      if (wrappedState.isRemoved) {
-        newStateMap.remove(key)
-      } else if (wrappedState.isUpdated) {
-        newStateMap.put(key, wrappedState.get(), batchTime.milliseconds)
-      }
-      emittedRecords ++= emittedRecord
-    }
-
-    // If the RDD is expected to be doing a full scan of all the data in the StateMap,
-    // then use this opportunity to filter out those keys that have timed out.
-    // For each of them call the tracking function.
-    if (doFullScan && timeoutThresholdTime.isDefined) {
-      newStateMap.getByTime(timeoutThresholdTime.get).foreach { case (key, state, _) =>
-        wrappedState.wrapTiminoutState(state)
-        val emittedRecord = trackingFunction(batchTime, key, None, wrappedState)
-        emittedRecords ++= emittedRecord
-        newStateMap.remove(key)
-      }
-    }
-
-    Iterator(TrackStateRDDRecord(newStateMap, emittedRecords))
+    val prevRecord = if (prevStateRDDIterator.hasNext) Some(prevStateRDDIterator.next()) else None
+    val newRecord = TrackStateRDDRecord.updateRecordWithData(
+      prevRecord,
+      dataIterator,
+      trackingFunction,
+      batchTime,
+      timeoutThresholdTime,
+      removeTimedoutData = doFullScan // remove timedout data only when full scan is enabled
+    )
+    Iterator(newRecord)
   }
 
   override protected def getPartitions: Array[Partition] = {
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
index f396b76e8d251..19ef5a14f8ab4 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
@@ -23,6 +23,7 @@ import scala.reflect.ClassTag
 import org.scalatest.BeforeAndAfterAll
 
 import org.apache.spark.rdd.RDD
+import org.apache.spark.streaming.util.OpenHashMapBasedStateMap
 import org.apache.spark.streaming.{Time, State}
 import org.apache.spark.{HashPartitioner, SparkConf, SparkContext, SparkFunSuite}
 
@@ -52,6 +53,131 @@ class TrackStateRDDSuite extends SparkFunSuite with BeforeAndAfterAll {
     assert(rdd.partitioner === Some(partitioner))
   }
 
+  test("updating state and generating emitted data in TrackStateRecord") {
+
+    val initialTime = 1000L
+    val updatedTime = 2000L
+    val thresholdTime = 1500L
+    @volatile var functionCalled = false
+
+    /**
+     * Assert that applying given data on a prior record generates correct updated record, with
+     * correct state map and emitted data
+     */
+    def assertRecordUpdate(
+        initStates: Iterable[Int],
+        data: Iterable[String],
+        expectedStates: Iterable[(Int, Long)],
+        timeoutThreshold: Option[Long] = None,
+        removeTimedoutData: Boolean = false,
+        expectedOutput: Iterable[Int] = None,
+        expectedTimingOutStates: Iterable[Int] = None,
+        expectedRemovedStates: Iterable[Int] = None
+      ): Unit = {
+      val initialStateMap = new OpenHashMapBasedStateMap[String, Int]()
+      initStates.foreach { s => initialStateMap.put("key", s, initialTime) }
+      functionCalled = false
+      val record = TrackStateRDDRecord[String, Int, Int](initialStateMap, Seq.empty)
+      val dataIterator = data.map { v => ("key", v) }.iterator
+      val removedStates = new ArrayBuffer[Int]
+      val timingOutStates = new ArrayBuffer[Int]
+      /**
+       * Tracking function that updates/removes state based on instructions in the data, and
+       * return state (when instructed or when state is timing out).
+       */
+      def testFunc(t: Time, key: String, data: Option[String], state: State[Int]): Option[Int] = {
+        functionCalled = true
+
+        assert(t.milliseconds === updatedTime, "tracking func called with wrong time")
+
+        data match {
+          case Some("noop") =>
+            None
+          case Some("get-state") =>
+            Some(state.getOption().getOrElse(-1))
+          case Some("update-state") =>
+            if (state.exists) state.update(state.get + 1) else state.update(0)
+            None
+          case Some("remove-state") =>
+            removedStates += state.get()
+            state.remove()
+            None
+          case None =>
+            assert(state.isTimingOut() === true, "State is not timing out when data = None")
+            timingOutStates += state.get()
+            None
+          case _ =>
+            fail("Unexpected test data")
+        }
+      }
+
+      val updatedRecord = TrackStateRDDRecord.updateRecordWithData[String, String, Int, Int](
+        Some(record), dataIterator, testFunc,
+        Time(updatedTime), timeoutThreshold, removeTimedoutData)
+
+      val updatedStateData = updatedRecord.stateMap.getAll().map { x => (x._2, x._3) }
+      assert(updatedStateData.toSet === expectedStates.toSet,
+        "states do not match after updating the TrackStateRecord")
+
+      assert(updatedRecord.emittedRecords.toSet === expectedOutput.toSet,
+        "emitted data do not match after updating the TrackStateRecord")
+
+      assert(timingOutStates.toSet === expectedTimingOutStates.toSet, "timing out states do not " +
+        "match those that were expected to do so while updating the TrackStateRecord")
+
+      assert(removedStates.toSet === expectedRemovedStates.toSet, "removed states do not " +
+        "match those that were expected to do so while updating the TrackStateRecord")
+
+    }
+
+    // No data, no state should be changed, function should not be called,
+    assertRecordUpdate(initStates = Nil, data = None, expectedStates = Nil)
+    assert(functionCalled === false)
+    assertRecordUpdate(initStates = Seq(0), data = None, expectedStates = Seq((0, initialTime)))
+    assert(functionCalled === false)
+
+    // Data present, function should be called irrespective of whether state exists
+    assertRecordUpdate(initStates = Seq(0), data = Seq("noop"),
+      expectedStates = Seq((0, initialTime)))
+    assert(functionCalled === true)
+    assertRecordUpdate(initStates = None, data = Some("noop"), expectedStates = None)
+    assert(functionCalled === true)
+
+    // Function called with right state data
+    assertRecordUpdate(initStates = None, data = Seq("get-state"),
+      expectedStates = None, expectedOutput = Seq(-1))
+    assertRecordUpdate(initStates = Seq(123), data = Seq("get-state"),
+      expectedStates = Seq((123, initialTime)), expectedOutput = Seq(123))
+
+    // Update state and timestamp, when timeout not present
+    assertRecordUpdate(initStates = Nil, data = Seq("update-state"),
+      expectedStates = Seq((0, updatedTime)))
+    assertRecordUpdate(initStates = Seq(0), data = Seq("update-state"),
+      expectedStates = Seq((1, updatedTime)))
+
+    // Remove state
+    assertRecordUpdate(initStates = Seq(345), data = Seq("remove-state"),
+      expectedStates = Nil, expectedRemovedStates = Seq(345))
+
+    // State strictly older than timeout threshold should be timed out
+    assertRecordUpdate(initStates = Seq(123), data = Nil,
+      timeoutThreshold = Some(initialTime), removeTimedoutData = true,
+      expectedStates = Seq((123, initialTime)), expectedTimingOutStates = Nil)
+
+    assertRecordUpdate(initStates = Seq(123), data = Nil,
+      timeoutThreshold = Some(initialTime + 1), removeTimedoutData = true,
+      expectedStates = Nil, expectedTimingOutStates = Seq(123))
+
+    // State should not be timed out after it has received data
+    assertRecordUpdate(initStates = Seq(123), data = Seq("noop"),
+      timeoutThreshold = Some(initialTime + 1), removeTimedoutData = true,
+      expectedStates = Seq((123, updatedTime)), expectedTimingOutStates = Nil)
+    assertRecordUpdate(initStates = Seq(123), data = Seq("remove-state"),
+      timeoutThreshold = Some(initialTime + 1), removeTimedoutData = true,
+      expectedStates = Nil, expectedTimingOutStates = Nil, expectedRemovedStates = Seq(123))
+
+  }
+
   test("states generated by TrackStateRDD") {
     val initStates = Seq(("k1", 0), ("k2", 0))
     val initTime = 123
@@ -148,9 +274,8 @@ class TrackStateRDDSuite extends SparkFunSuite with BeforeAndAfterAll {
     val rdd7 = testStateUpdates(                      // should remove k2's state
       rdd6, Seq(("k2", 2), ("k0", 2), ("k3", 1)), Set(("k3", 0, updateTime)))
 
-    val rdd8 = testStateUpdates(
-      rdd7, Seq(("k3", 2)), Set()                     //
-    )
+    val rdd8 = testStateUpdates(                      // should remove k3's state
+      rdd7, Seq(("k3", 2)), Set())
   }
 
   /** Assert whether the `trackStateByKey` operation generates expected results */
@@ -176,7 +301,7 @@ class TrackStateRDDSuite extends SparkFunSuite with BeforeAndAfterAll {
 
     // Persist to make sure that it gets computed only once and we can track precisely how many
     // state keys the computing touched
-    newStateRDD.persist()
+    newStateRDD.persist().count()
     assertRDD(newStateRDD, expectedStates, expectedEmittedRecords)
     newStateRDD
   }
@@ -188,7 +313,8 @@ class TrackStateRDDSuite extends SparkFunSuite with BeforeAndAfterAll {
       expectedEmittedRecords: Set[T]): Unit = {
     val states = trackStateRDD.flatMap { _.stateMap.getAll() }.collect().toSet
     val emittedRecords = trackStateRDD.flatMap { _.emittedRecords }.collect().toSet
-    assert(states === expectedStates, "states after track state operation were not as expected")
+    assert(states === expectedStates,
+      "states after track state operation were not as expected")
     assert(emittedRecords === expectedEmittedRecords,
       "emitted records after track state operation were not as expected")
   }

From e71c07557c39e2f74bd20d2ab3a2fca88aa5dfbb Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Thu, 12 Nov 2015 20:01:13 -0800
Subject: [PATCH 225/518] [SPARK-11672][ML] flaky spark.ml read/write tests

We set `sqlContext = null` in `afterAll`. However, this doesn't change `SQLContext.activeContext`  and then `SQLContext.getOrCreate` might use the `SparkContext` from previous test suite and hence causes the error. This PR calls `clearActive` in `beforeAll` and `afterAll` to avoid using an old context from other test suites.

cc: yhuai

Author: Xiangrui Meng <meng@databricks.com>

Closes #9677 from mengxr/SPARK-11672.2.
---
 .../org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java   | 4 ++--
 .../spark/ml/classification/LogisticRegressionSuite.scala     | 2 +-
 .../scala/org/apache/spark/ml/feature/BinarizerSuite.scala    | 2 +-
 .../scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala | 2 +-
 .../org/apache/spark/mllib/util/MLlibTestSparkContext.scala   | 2 ++
 5 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java b/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
index 4f7aeac1ec54c..c39538014be81 100644
--- a/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
+++ b/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
@@ -23,7 +23,7 @@
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
-import org.junit.Ignore;
+import org.junit.Test;
 
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.SQLContext;
@@ -50,7 +50,7 @@ public void tearDown() {
     Utils.deleteRecursively(tempDir);
   }
 
-  @Ignore // SPARK-11672
+  @Test
   public void testDefaultReadWrite() throws IOException {
     String uid = "my_params";
     MyParams instance = new MyParams(uid);
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
index e4c2f1baa4fa1..51b06b7eb6d53 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
@@ -872,7 +872,7 @@ class LogisticRegressionSuite
     assert(model1a0.intercept ~== model1b.intercept absTol 1E-3)
   }
 
-  ignore("read/write") { // SPARK-11672
+  test("read/write") {
     // Set some Params to make sure set Params are serialized.
     val lr = new LogisticRegression()
       .setElasticNetParam(0.1)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
index a66fe03281935..9dfa1439cc303 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
@@ -68,7 +68,7 @@ class BinarizerSuite extends SparkFunSuite with MLlibTestSparkContext with Defau
     }
   }
 
-  ignore("read/write") { // SPARK-11672
+  test("read/write") {
     val binarizer = new Binarizer()
       .setInputCol("feature")
       .setOutputCol("binarized_feature")
diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
index 44e09c38f9375..cac4bd9aa3ab8 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
@@ -105,7 +105,7 @@ object MyParams extends Readable[MyParams] {
 class DefaultReadWriteSuite extends SparkFunSuite with MLlibTestSparkContext
   with DefaultReadWriteTest {
 
-  ignore("default read/write") { // SPARK-11672
+  test("default read/write") {
     val myParams = new MyParams("my_params")
     testDefaultReadWrite(myParams)
   }
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/util/MLlibTestSparkContext.scala b/mllib/src/test/scala/org/apache/spark/mllib/util/MLlibTestSparkContext.scala
index 5d1796ef65722..998ee48186558 100644
--- a/mllib/src/test/scala/org/apache/spark/mllib/util/MLlibTestSparkContext.scala
+++ b/mllib/src/test/scala/org/apache/spark/mllib/util/MLlibTestSparkContext.scala
@@ -32,11 +32,13 @@ trait MLlibTestSparkContext extends BeforeAndAfterAll { self: Suite =>
       .setMaster("local[2]")
       .setAppName("MLlibUnitTest")
     sc = new SparkContext(conf)
+    SQLContext.clearActive()
     sqlContext = new SQLContext(sc)
   }
 
   override def afterAll() {
     sqlContext = null
+    SQLContext.clearActive()
     if (sc != null) {
       sc.stop()
     }

From ed04846e144db5bdab247c0e1fe2a47b99155c82 Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Thu, 12 Nov 2015 20:02:49 -0800
Subject: [PATCH 226/518] [SPARK-11263][SPARKR] lintr Throws Warnings on
 Commented Code in Documentation

Clean out hundreds of `style: Commented code should be removed.` from lintr

Like these:
```
/opt/spark-1.6.0-bin-hadoop2.6/R/pkg/R/DataFrame.R:513:3: style: Commented code should be removed.
# sc <- sparkR.init()
  ^~~~~~~~~~~~~~~~~~~
/opt/spark-1.6.0-bin-hadoop2.6/R/pkg/R/DataFrame.R:514:3: style: Commented code should be removed.
# sqlContext <- sparkRSQL.init(sc)
  ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/opt/spark-1.6.0-bin-hadoop2.6/R/pkg/R/DataFrame.R:515:3: style: Commented code should be removed.
# path <- "path/to/file.json"
  ^~~~~~~~~~~~~~~~~~~~~~~~~~~
```

tried without export or rdname, neither work
instead, added this `#' noRd` to suppress .Rd file generation

also updated `family` for DataFrame functions for longer descriptive text instead of `dataframe_funcs`
![image](https://cloud.githubusercontent.com/assets/8969467/10933937/17bf5b1e-8291-11e5-9777-40fc632105dc.png)

this covers *most* of 'Commented code' but I left out a few that looks legitimate.

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9463 from felixcheung/rlintr.
---
 R/pkg/R/DataFrame.R        |  232 +++---
 R/pkg/R/RDD.R              | 1585 ++++++++++++++++++------------------
 R/pkg/R/SQLContext.R       |   66 +-
 R/pkg/R/context.R          |  235 +++---
 R/pkg/R/generics.R         |   18 -
 R/pkg/R/pairRDD.R          |  910 +++++++++++----------
 R/pkg/R/sparkR.R           |    3 +-
 R/pkg/inst/profile/shell.R |    2 +-
 8 files changed, 1539 insertions(+), 1512 deletions(-)

diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index cc868069d1e5a..fd105ba5bc9bb 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -25,7 +25,7 @@ setOldClass("jobj")
 #' @title S4 class that represents a DataFrame
 #' @description DataFrames can be created using functions like \link{createDataFrame},
 #'              \link{jsonFile}, \link{table} etc.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname DataFrame
 #' @docType class
 #'
@@ -68,7 +68,7 @@ dataFrame <- function(sdf, isCached = FALSE) {
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname printSchema
 #' @name printSchema
 #' @export
@@ -93,7 +93,7 @@ setMethod("printSchema",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname schema
 #' @name schema
 #' @export
@@ -117,7 +117,7 @@ setMethod("schema",
 #'
 #' @param x A SparkSQL DataFrame
 #' @param extended Logical. If extended is False, explain() only prints the physical plan.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname explain
 #' @name explain
 #' @export
@@ -148,7 +148,7 @@ setMethod("explain",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname isLocal
 #' @name isLocal
 #' @export
@@ -173,7 +173,7 @@ setMethod("isLocal",
 #' @param x A SparkSQL DataFrame
 #' @param numRows The number of rows to print. Defaults to 20.
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname showDF
 #' @name showDF
 #' @export
@@ -198,7 +198,7 @@ setMethod("showDF",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname show
 #' @name show
 #' @export
@@ -225,7 +225,7 @@ setMethod("show", "DataFrame",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname dtypes
 #' @name dtypes
 #' @export
@@ -251,7 +251,7 @@ setMethod("dtypes",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname columns
 #' @name columns
 #' @aliases names
@@ -272,7 +272,7 @@ setMethod("columns",
             })
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname columns
 #' @name names
 setMethod("names",
@@ -281,7 +281,7 @@ setMethod("names",
             columns(x)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname columns
 #' @name names<-
 setMethod("names<-",
@@ -300,7 +300,7 @@ setMethod("names<-",
 #' @param x A SparkSQL DataFrame
 #' @param tableName A character vector containing the name of the table
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname registerTempTable
 #' @name registerTempTable
 #' @export
@@ -328,7 +328,7 @@ setMethod("registerTempTable",
 #' @param overwrite A logical argument indicating whether or not to overwrite
 #' the existing rows in the table.
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname insertInto
 #' @name insertInto
 #' @export
@@ -353,7 +353,7 @@ setMethod("insertInto",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname cache
 #' @name cache
 #' @export
@@ -381,7 +381,7 @@ setMethod("cache",
 #'
 #' @param x The DataFrame to persist
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname persist
 #' @name persist
 #' @export
@@ -409,7 +409,7 @@ setMethod("persist",
 #' @param x The DataFrame to unpersist
 #' @param blocking Whether to block until all blocks are deleted
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname unpersist-methods
 #' @name unpersist
 #' @export
@@ -437,7 +437,7 @@ setMethod("unpersist",
 #' @param x A SparkSQL DataFrame
 #' @param numPartitions The number of partitions to use.
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname repartition
 #' @name repartition
 #' @export
@@ -456,25 +456,24 @@ setMethod("repartition",
             dataFrame(sdf)
           })
 
-# toJSON
-#
-# Convert the rows of a DataFrame into JSON objects and return an RDD where
-# each element contains a JSON string.
-#
-# @param x A SparkSQL DataFrame
-# @return A StringRRDD of JSON objects
-#
-# @family dataframe_funcs
-# @rdname tojson
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# sqlContext <- sparkRSQL.init(sc)
-# path <- "path/to/file.json"
-# df <- jsonFile(sqlContext, path)
-# newRDD <- toJSON(df)
-#}
+#' toJSON
+#'
+#' Convert the rows of a DataFrame into JSON objects and return an RDD where
+#' each element contains a JSON string.
+#'
+#' @param x A SparkSQL DataFrame
+#' @return A StringRRDD of JSON objects
+#' @family DataFrame functions
+#' @rdname tojson
+#' @noRd
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' sqlContext <- sparkRSQL.init(sc)
+#' path <- "path/to/file.json"
+#' df <- jsonFile(sqlContext, path)
+#' newRDD <- toJSON(df)
+#'}
 setMethod("toJSON",
           signature(x = "DataFrame"),
           function(x) {
@@ -491,7 +490,7 @@ setMethod("toJSON",
 #' @param x A SparkSQL DataFrame
 #' @param path The directory where the file is saved
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname saveAsParquetFile
 #' @name saveAsParquetFile
 #' @export
@@ -515,7 +514,7 @@ setMethod("saveAsParquetFile",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname distinct
 #' @name distinct
 #' @export
@@ -538,7 +537,7 @@ setMethod("distinct",
 #
 #' @description Returns a new DataFrame containing distinct rows in this DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname unique
 #' @name unique
 #' @aliases distinct
@@ -556,7 +555,7 @@ setMethod("unique",
 #' @param withReplacement Sampling with replacement or not
 #' @param fraction The (rough) sample target fraction
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname sample
 #' @aliases sample_frac
 #' @export
@@ -580,7 +579,7 @@ setMethod("sample",
             dataFrame(sdf)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname sample
 #' @name sample_frac
 setMethod("sample_frac",
@@ -596,7 +595,7 @@ setMethod("sample_frac",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname count
 #' @name count
 #' @aliases nrow
@@ -620,7 +619,7 @@ setMethod("count",
 #'
 #' @name nrow
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname nrow
 #' @aliases count
 setMethod("nrow",
@@ -633,7 +632,7 @@ setMethod("nrow",
 #'
 #' @param x a SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname ncol
 #' @name ncol
 #' @export
@@ -654,7 +653,7 @@ setMethod("ncol",
 #' Returns the dimentions (number of rows and columns) of a DataFrame
 #' @param x a SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname dim
 #' @name dim
 #' @export
@@ -678,7 +677,7 @@ setMethod("dim",
 #' @param stringsAsFactors (Optional) A logical indicating whether or not string columns
 #' should be converted to factors. FALSE by default.
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname collect
 #' @name collect
 #' @export
@@ -746,7 +745,7 @@ setMethod("collect",
 #' @param num The number of rows to return
 #' @return A new DataFrame containing the number of rows specified.
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname limit
 #' @name limit
 #' @export
@@ -767,7 +766,7 @@ setMethod("limit",
 
 #' Take the first NUM rows of a DataFrame and return a the results as a data.frame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname take
 #' @name take
 #' @export
@@ -796,7 +795,7 @@ setMethod("take",
 #' @param num The number of rows to return. Default is 6.
 #' @return A data.frame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname head
 #' @name head
 #' @export
@@ -819,7 +818,7 @@ setMethod("head",
 #'
 #' @param x A SparkSQL DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname first
 #' @name first
 #' @export
@@ -837,23 +836,21 @@ setMethod("first",
             take(x, 1)
           })
 
-# toRDD
-#
-# Converts a Spark DataFrame to an RDD while preserving column names.
-#
-# @param x A Spark DataFrame
-#
-# @family dataframe_funcs
-# @rdname DataFrame
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# sqlContext <- sparkRSQL.init(sc)
-# path <- "path/to/file.json"
-# df <- jsonFile(sqlContext, path)
-# rdd <- toRDD(df)
-# }
+#' toRDD
+#'
+#' Converts a Spark DataFrame to an RDD while preserving column names.
+#'
+#' @param x A Spark DataFrame
+#'
+#' @noRd
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' sqlContext <- sparkRSQL.init(sc)
+#' path <- "path/to/file.json"
+#' df <- jsonFile(sqlContext, path)
+#' rdd <- toRDD(df)
+#'}
 setMethod("toRDD",
           signature(x = "DataFrame"),
           function(x) {
@@ -874,7 +871,7 @@ setMethod("toRDD",
 #' @return a GroupedData
 #' @seealso GroupedData
 #' @aliases group_by
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname groupBy
 #' @name groupBy
 #' @export
@@ -899,7 +896,7 @@ setMethod("groupBy",
              groupedData(sgd)
            })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname groupBy
 #' @name group_by
 setMethod("group_by",
@@ -913,7 +910,7 @@ setMethod("group_by",
 #' Compute aggregates by specifying a list of columns
 #'
 #' @param x a DataFrame
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname agg
 #' @name agg
 #' @aliases summarize
@@ -924,7 +921,7 @@ setMethod("agg",
             agg(groupBy(x), ...)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname agg
 #' @name summarize
 setMethod("summarize",
@@ -940,8 +937,8 @@ setMethod("summarize",
 # the requested map function.                                                     #
 ###################################################################################
 
-# @family dataframe_funcs
-# @rdname lapply
+#' @rdname lapply
+#' @noRd
 setMethod("lapply",
           signature(X = "DataFrame", FUN = "function"),
           function(X, FUN) {
@@ -949,24 +946,25 @@ setMethod("lapply",
             lapply(rdd, FUN)
           })
 
-# @family dataframe_funcs
-# @rdname lapply
+#' @rdname lapply
+#' @noRd
 setMethod("map",
           signature(X = "DataFrame", FUN = "function"),
           function(X, FUN) {
             lapply(X, FUN)
           })
 
-# @family dataframe_funcs
-# @rdname flatMap
+#' @rdname flatMap
+#' @noRd
 setMethod("flatMap",
           signature(X = "DataFrame", FUN = "function"),
           function(X, FUN) {
             rdd <- toRDD(X)
             flatMap(rdd, FUN)
           })
-# @family dataframe_funcs
-# @rdname lapplyPartition
+
+#' @rdname lapplyPartition
+#' @noRd
 setMethod("lapplyPartition",
           signature(X = "DataFrame", FUN = "function"),
           function(X, FUN) {
@@ -974,16 +972,16 @@ setMethod("lapplyPartition",
             lapplyPartition(rdd, FUN)
           })
 
-# @family dataframe_funcs
-# @rdname lapplyPartition
+#' @rdname lapplyPartition
+#' @noRd
 setMethod("mapPartitions",
           signature(X = "DataFrame", FUN = "function"),
           function(X, FUN) {
             lapplyPartition(X, FUN)
           })
 
-# @family dataframe_funcs
-# @rdname foreach
+#' @rdname foreach
+#' @noRd
 setMethod("foreach",
           signature(x = "DataFrame", func = "function"),
           function(x, func) {
@@ -991,8 +989,8 @@ setMethod("foreach",
             foreach(rdd, func)
           })
 
-# @family dataframe_funcs
-# @rdname foreach
+#' @rdname foreach
+#' @noRd
 setMethod("foreachPartition",
           signature(x = "DataFrame", func = "function"),
           function(x, func) {
@@ -1091,7 +1089,7 @@ setMethod("[", signature(x = "DataFrame", i = "Column"),
 #' @param select expression for the single Column or a list of columns to select from the DataFrame
 #' @return A new DataFrame containing only the rows that meet the condition with selected columns
 #' @export
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname subset
 #' @name subset
 #' @aliases [
@@ -1122,7 +1120,7 @@ setMethod("subset", signature(x = "DataFrame"),
 #' @param col A list of columns or single Column or name
 #' @return A new DataFrame with selected columns
 #' @export
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname select
 #' @name select
 #' @family subsetting functions
@@ -1150,7 +1148,7 @@ setMethod("select", signature(x = "DataFrame", col = "character"),
             }
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname select
 #' @export
 setMethod("select", signature(x = "DataFrame", col = "Column"),
@@ -1162,7 +1160,7 @@ setMethod("select", signature(x = "DataFrame", col = "Column"),
             dataFrame(sdf)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname select
 #' @export
 setMethod("select",
@@ -1187,7 +1185,7 @@ setMethod("select",
 #' @param expr A string containing a SQL expression
 #' @param ... Additional expressions
 #' @return A DataFrame
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname selectExpr
 #' @name selectExpr
 #' @export
@@ -1215,7 +1213,7 @@ setMethod("selectExpr",
 #' @param colName A string containing the name of the new column.
 #' @param col A Column expression.
 #' @return A DataFrame with the new column added.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname withColumn
 #' @name withColumn
 #' @aliases mutate transform
@@ -1241,7 +1239,7 @@ setMethod("withColumn",
 #' @param .data A DataFrame
 #' @param col a named argument of the form name = col
 #' @return A new DataFrame with the new columns added.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname withColumn
 #' @name mutate
 #' @aliases withColumn transform
@@ -1275,7 +1273,7 @@ setMethod("mutate",
           })
 
 #' @export
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname withColumn
 #' @name transform
 #' @aliases withColumn mutate
@@ -1293,7 +1291,7 @@ setMethod("transform",
 #' @param existingCol The name of the column you want to change.
 #' @param newCol The new column name.
 #' @return A DataFrame with the column name changed.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname withColumnRenamed
 #' @name withColumnRenamed
 #' @export
@@ -1325,7 +1323,7 @@ setMethod("withColumnRenamed",
 #' @param x A DataFrame
 #' @param newCol A named pair of the form new_column_name = existing_column
 #' @return A DataFrame with the column name changed.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname withColumnRenamed
 #' @name rename
 #' @aliases withColumnRenamed
@@ -1370,7 +1368,7 @@ setClassUnion("characterOrColumn", c("character", "Column"))
 #' @param decreasing A logical argument indicating sorting order for columns when
 #'                   a character vector is specified for col
 #' @return A DataFrame where all elements are sorted.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname arrange
 #' @name arrange
 #' @aliases orderby
@@ -1397,7 +1395,7 @@ setMethod("arrange",
             dataFrame(sdf)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname arrange
 #' @export
 setMethod("arrange",
@@ -1429,7 +1427,7 @@ setMethod("arrange",
             do.call("arrange", c(x, jcols))
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname arrange
 #' @name orderby
 setMethod("orderBy",
@@ -1446,7 +1444,7 @@ setMethod("orderBy",
 #' @param condition The condition to filter on. This may either be a Column expression
 #' or a string containing a SQL statement
 #' @return A DataFrame containing only the rows that meet the condition.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname filter
 #' @name filter
 #' @family subsetting functions
@@ -1470,7 +1468,7 @@ setMethod("filter",
             dataFrame(sdf)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname filter
 #' @name where
 setMethod("where",
@@ -1491,7 +1489,7 @@ setMethod("where",
 #' 'inner', 'outer', 'full', 'fullouter', leftouter', 'left_outer', 'left',
 #' 'right_outer', 'rightouter', 'right', and 'leftsemi'. The default joinType is "inner".
 #' @return A DataFrame containing the result of the join operation.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname join
 #' @name join
 #' @export
@@ -1550,7 +1548,7 @@ setMethod("join",
 #'   be returned. If all.x is set to FALSE and all.y is set to TRUE, a right
 #'   outer join will be returned. If all.x and all.y are set to TRUE, a full
 #'   outer join will be returned.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname merge
 #' @export
 #' @examples
@@ -1682,7 +1680,7 @@ generateAliasesForIntersectedCols <- function (x, intersectedColNames, suffix) {
 #' @param x A Spark DataFrame
 #' @param y A Spark DataFrame
 #' @return A DataFrame containing the result of the union.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname unionAll
 #' @name unionAll
 #' @export
@@ -1705,7 +1703,7 @@ setMethod("unionAll",
 #'
 #' @description Returns a new DataFrame containing rows of all parameters.
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname rbind
 #' @name rbind
 #' @aliases unionAll
@@ -1727,7 +1725,7 @@ setMethod("rbind",
 #' @param x A Spark DataFrame
 #' @param y A Spark DataFrame
 #' @return A DataFrame containing the result of the intersect.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname intersect
 #' @name intersect
 #' @export
@@ -1754,7 +1752,7 @@ setMethod("intersect",
 #' @param x A Spark DataFrame
 #' @param y A Spark DataFrame
 #' @return A DataFrame containing the result of the except operation.
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname except
 #' @name except
 #' @export
@@ -1794,7 +1792,7 @@ setMethod("except",
 #' @param source A name for external data source
 #' @param mode One of 'append', 'overwrite', 'error', 'ignore' save mode
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname write.df
 #' @name write.df
 #' @aliases saveDF
@@ -1830,7 +1828,7 @@ setMethod("write.df",
             callJMethod(df@sdf, "save", source, jmode, options)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname write.df
 #' @name saveDF
 #' @export
@@ -1861,7 +1859,7 @@ setMethod("saveDF",
 #' @param source A name for external data source
 #' @param mode One of 'append', 'overwrite', 'error', 'ignore' save mode
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname saveAsTable
 #' @name saveAsTable
 #' @export
@@ -1902,7 +1900,7 @@ setMethod("saveAsTable",
 #' @param col A string of name
 #' @param ... Additional expressions
 #' @return A DataFrame
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname describe
 #' @name describe
 #' @aliases summary
@@ -1925,7 +1923,7 @@ setMethod("describe",
             dataFrame(sdf)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname describe
 #' @name describe
 setMethod("describe",
@@ -1940,7 +1938,7 @@ setMethod("describe",
 #'
 #' @description Computes statistics for numeric columns of the DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname summary
 #' @name summary
 setMethod("summary",
@@ -1965,7 +1963,7 @@ setMethod("summary",
 #' @param cols Optional list of column names to consider.
 #' @return A DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname nafunctions
 #' @name dropna
 #' @aliases na.omit
@@ -1995,7 +1993,7 @@ setMethod("dropna",
             dataFrame(sdf)
           })
 
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname nafunctions
 #' @name na.omit
 #' @export
@@ -2023,7 +2021,7 @@ setMethod("na.omit",
 #'             column is simply ignored.
 #' @return A DataFrame
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname nafunctions
 #' @name fillna
 #' @export
@@ -2087,7 +2085,7 @@ setMethod("fillna",
 #' @title Download data from a DataFrame into a data.frame
 #' @param x a DataFrame
 #' @return a data.frame
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname as.data.frame
 #' @examples \dontrun{
 #'
@@ -2108,7 +2106,7 @@ setMethod("as.data.frame",
 #' the DataFrame is searched by R when evaluating a variable, so columns in
 #' the DataFrame can be accessed by simply giving their names.
 #'
-#' @family dataframe_funcs
+#' @family DataFrame functions
 #' @rdname attach
 #' @title Attach DataFrame to R search path
 #' @param what (DataFrame) The DataFrame to attach
diff --git a/R/pkg/R/RDD.R b/R/pkg/R/RDD.R
index 051e441d4e063..47945c2825da9 100644
--- a/R/pkg/R/RDD.R
+++ b/R/pkg/R/RDD.R
@@ -19,16 +19,15 @@
 
 setOldClass("jobj")
 
-# @title S4 class that represents an RDD
-# @description RDD can be created using functions like
-#              \code{parallelize}, \code{textFile} etc.
-# @rdname RDD
-# @seealso parallelize, textFile
-#
-# @slot env An R environment that stores bookkeeping states of the RDD
-# @slot jrdd Java object reference to the backing JavaRDD
-# to an RDD
-# @export
+#' @title S4 class that represents an RDD
+#' @description RDD can be created using functions like
+#'              \code{parallelize}, \code{textFile} etc.
+#' @rdname RDD
+#' @seealso parallelize, textFile
+#' @slot env An R environment that stores bookkeeping states of the RDD
+#' @slot jrdd Java object reference to the backing JavaRDD
+#' to an RDD
+#' @noRd
 setClass("RDD",
          slots = list(env = "environment",
                       jrdd = "jobj"))
@@ -111,14 +110,13 @@ setMethod("initialize", "PipelinedRDD", function(.Object, prev, func, jrdd_val)
   .Object
 })
 
-# @rdname RDD
-# @export
-#
-# @param jrdd Java object reference to the backing JavaRDD
-# @param serializedMode Use "byte" if the RDD stores data serialized in R, "string" if the RDD
-# stores strings, and "row" if the RDD stores the rows of a DataFrame
-# @param isCached TRUE if the RDD is cached
-# @param isCheckpointed TRUE if the RDD has been checkpointed
+#' @rdname RDD
+#' @noRd
+#' @param jrdd Java object reference to the backing JavaRDD
+#' @param serializedMode Use "byte" if the RDD stores data serialized in R, "string" if the RDD
+#' stores strings, and "row" if the RDD stores the rows of a DataFrame
+#' @param isCached TRUE if the RDD is cached
+#' @param isCheckpointed TRUE if the RDD has been checkpointed
 RDD <- function(jrdd, serializedMode = "byte", isCached = FALSE,
                 isCheckpointed = FALSE) {
   new("RDD", jrdd, serializedMode, isCached, isCheckpointed)
@@ -201,19 +199,20 @@ setValidity("RDD",
 
 ############ Actions and Transformations ############
 
-# Persist an RDD
-#
-# Persist this RDD with the default storage level (MEMORY_ONLY).
-#
-# @param x The RDD to cache
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10, 2L)
-# cache(rdd)
-#}
-# @rdname cache-methods
-# @aliases cache,RDD-method
+#' Persist an RDD
+#'
+#' Persist this RDD with the default storage level (MEMORY_ONLY).
+#'
+#' @param x The RDD to cache
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10, 2L)
+#' cache(rdd)
+#'}
+#' @rdname cache-methods
+#' @aliases cache,RDD-method
+#' @noRd
 setMethod("cache",
           signature(x = "RDD"),
           function(x) {
@@ -222,22 +221,23 @@ setMethod("cache",
             x
           })
 
-# Persist an RDD
-#
-# Persist this RDD with the specified storage level. For details of the
-# supported storage levels, refer to
-# http://spark.apache.org/docs/latest/programming-guide.html#rdd-persistence.
-#
-# @param x The RDD to persist
-# @param newLevel The new storage level to be assigned
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10, 2L)
-# persist(rdd, "MEMORY_AND_DISK")
-#}
-# @rdname persist
-# @aliases persist,RDD-method
+#' Persist an RDD
+#'
+#' Persist this RDD with the specified storage level. For details of the
+#' supported storage levels, refer to
+#' http://spark.apache.org/docs/latest/programming-guide.html#rdd-persistence.
+#'
+#' @param x The RDD to persist
+#' @param newLevel The new storage level to be assigned
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10, 2L)
+#' persist(rdd, "MEMORY_AND_DISK")
+#'}
+#' @rdname persist
+#' @aliases persist,RDD-method
+#' @noRd
 setMethod("persist",
           signature(x = "RDD", newLevel = "character"),
           function(x, newLevel = "MEMORY_ONLY") {
@@ -246,21 +246,22 @@ setMethod("persist",
             x
           })
 
-# Unpersist an RDD
-#
-# Mark the RDD as non-persistent, and remove all blocks for it from memory and
-# disk.
-#
-# @param x The RDD to unpersist
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10, 2L)
-# cache(rdd) # rdd@@env$isCached == TRUE
-# unpersist(rdd) # rdd@@env$isCached == FALSE
-#}
-# @rdname unpersist-methods
-# @aliases unpersist,RDD-method
+#' Unpersist an RDD
+#'
+#' Mark the RDD as non-persistent, and remove all blocks for it from memory and
+#' disk.
+#'
+#' @param x The RDD to unpersist
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10, 2L)
+#' cache(rdd) # rdd@@env$isCached == TRUE
+#' unpersist(rdd) # rdd@@env$isCached == FALSE
+#'}
+#' @rdname unpersist-methods
+#' @aliases unpersist,RDD-method
+#' @noRd
 setMethod("unpersist",
           signature(x = "RDD"),
           function(x) {
@@ -269,24 +270,25 @@ setMethod("unpersist",
             x
           })
 
-# Checkpoint an RDD
-#
-# Mark this RDD for checkpointing. It will be saved to a file inside the
-# checkpoint directory set with setCheckpointDir() and all references to its
-# parent RDDs will be removed. This function must be called before any job has
-# been executed on this RDD. It is strongly recommended that this RDD is
-# persisted in memory, otherwise saving it on a file will require recomputation.
-#
-# @param x The RDD to checkpoint
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# setCheckpointDir(sc, "checkpoint")
-# rdd <- parallelize(sc, 1:10, 2L)
-# checkpoint(rdd)
-#}
-# @rdname checkpoint-methods
-# @aliases checkpoint,RDD-method
+#' Checkpoint an RDD
+#'
+#' Mark this RDD for checkpointing. It will be saved to a file inside the
+#' checkpoint directory set with setCheckpointDir() and all references to its
+#' parent RDDs will be removed. This function must be called before any job has
+#' been executed on this RDD. It is strongly recommended that this RDD is
+#' persisted in memory, otherwise saving it on a file will require recomputation.
+#'
+#' @param x The RDD to checkpoint
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' setCheckpointDir(sc, "checkpoint")
+#' rdd <- parallelize(sc, 1:10, 2L)
+#' checkpoint(rdd)
+#'}
+#' @rdname checkpoint-methods
+#' @aliases checkpoint,RDD-method
+#' @noRd
 setMethod("checkpoint",
           signature(x = "RDD"),
           function(x) {
@@ -296,18 +298,19 @@ setMethod("checkpoint",
             x
           })
 
-# Gets the number of partitions of an RDD
-#
-# @param x A RDD.
-# @return the number of partitions of rdd as an integer.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10, 2L)
-# numPartitions(rdd)  # 2L
-#}
-# @rdname numPartitions
-# @aliases numPartitions,RDD-method
+#' Gets the number of partitions of an RDD
+#'
+#' @param x A RDD.
+#' @return the number of partitions of rdd as an integer.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10, 2L)
+#' numPartitions(rdd)  # 2L
+#'}
+#' @rdname numPartitions
+#' @aliases numPartitions,RDD-method
+#' @noRd
 setMethod("numPartitions",
           signature(x = "RDD"),
           function(x) {
@@ -316,24 +319,25 @@ setMethod("numPartitions",
             callJMethod(partitions, "size")
           })
 
-# Collect elements of an RDD
-#
-# @description
-# \code{collect} returns a list that contains all of the elements in this RDD.
-#
-# @param x The RDD to collect
-# @param ... Other optional arguments to collect
-# @param flatten FALSE if the list should not flattened
-# @return a list containing elements in the RDD
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10, 2L)
-# collect(rdd) # list from 1 to 10
-# collectPartition(rdd, 0L) # list from 1 to 5
-#}
-# @rdname collect-methods
-# @aliases collect,RDD-method
+#' Collect elements of an RDD
+#'
+#' @description
+#' \code{collect} returns a list that contains all of the elements in this RDD.
+#'
+#' @param x The RDD to collect
+#' @param ... Other optional arguments to collect
+#' @param flatten FALSE if the list should not flattened
+#' @return a list containing elements in the RDD
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10, 2L)
+#' collect(rdd) # list from 1 to 10
+#' collectPartition(rdd, 0L) # list from 1 to 5
+#'}
+#' @rdname collect-methods
+#' @aliases collect,RDD-method
+#' @noRd
 setMethod("collect",
           signature(x = "RDD"),
           function(x, flatten = TRUE) {
@@ -344,12 +348,13 @@ setMethod("collect",
           })
 
 
-# @description
-# \code{collectPartition} returns a list that contains all of the elements
-# in the specified partition of the RDD.
-# @param partitionId the partition to collect (starts from 0)
-# @rdname collect-methods
-# @aliases collectPartition,integer,RDD-method
+#' @description
+#' \code{collectPartition} returns a list that contains all of the elements
+#' in the specified partition of the RDD.
+#' @param partitionId the partition to collect (starts from 0)
+#' @rdname collect-methods
+#' @aliases collectPartition,integer,RDD-method
+#' @noRd
 setMethod("collectPartition",
           signature(x = "RDD", partitionId = "integer"),
           function(x, partitionId) {
@@ -362,17 +367,18 @@ setMethod("collectPartition",
               serializedMode = getSerializedMode(x))
           })
 
-# @description
-# \code{collectAsMap} returns a named list as a map that contains all of the elements
-# in a key-value pair RDD.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(list(1, 2), list(3, 4)), 2L)
-# collectAsMap(rdd) # list(`1` = 2, `3` = 4)
-#}
-# @rdname collect-methods
-# @aliases collectAsMap,RDD-method
+#' @description
+#' \code{collectAsMap} returns a named list as a map that contains all of the elements
+#' in a key-value pair RDD.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(list(1, 2), list(3, 4)), 2L)
+#' collectAsMap(rdd) # list(`1` = 2, `3` = 4)
+#'}
+#' @rdname collect-methods
+#' @aliases collectAsMap,RDD-method
+#' @noRd
 setMethod("collectAsMap",
           signature(x = "RDD"),
           function(x) {
@@ -382,19 +388,20 @@ setMethod("collectAsMap",
             as.list(map)
           })
 
-# Return the number of elements in the RDD.
-#
-# @param x The RDD to count
-# @return number of elements in the RDD.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# count(rdd) # 10
-# length(rdd) # Same as count
-#}
-# @rdname count
-# @aliases count,RDD-method
+#' Return the number of elements in the RDD.
+#'
+#' @param x The RDD to count
+#' @return number of elements in the RDD.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' count(rdd) # 10
+#' length(rdd) # Same as count
+#'}
+#' @rdname count
+#' @aliases count,RDD-method
+#' @noRd
 setMethod("count",
           signature(x = "RDD"),
           function(x) {
@@ -406,31 +413,32 @@ setMethod("count",
             sum(as.integer(vals))
           })
 
-# Return the number of elements in the RDD
-# @export
-# @rdname count
+#' Return the number of elements in the RDD
+#' @rdname count
+#' @noRd
 setMethod("length",
           signature(x = "RDD"),
           function(x) {
             count(x)
           })
 
-# Return the count of each unique value in this RDD as a list of
-# (value, count) pairs.
-#
-# Same as countByValue in Spark.
-#
-# @param x The RDD to count
-# @return list of (value, count) pairs, where count is number of each unique
-# value in rdd.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, c(1,2,3,2,1))
-# countByValue(rdd) # (1,2L), (2,2L), (3,1L)
-#}
-# @rdname countByValue
-# @aliases countByValue,RDD-method
+#' Return the count of each unique value in this RDD as a list of
+#' (value, count) pairs.
+#'
+#' Same as countByValue in Spark.
+#'
+#' @param x The RDD to count
+#' @return list of (value, count) pairs, where count is number of each unique
+#' value in rdd.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, c(1,2,3,2,1))
+#' countByValue(rdd) # (1,2L), (2,2L), (3,1L)
+#'}
+#' @rdname countByValue
+#' @aliases countByValue,RDD-method
+#' @noRd
 setMethod("countByValue",
           signature(x = "RDD"),
           function(x) {
@@ -438,23 +446,24 @@ setMethod("countByValue",
             collect(reduceByKey(ones, `+`, numPartitions(x)))
           })
 
-# Apply a function to all elements
-#
-# This function creates a new RDD by applying the given transformation to all
-# elements of the given RDD
-#
-# @param X The RDD to apply the transformation.
-# @param FUN the transformation to apply on each element
-# @return a new RDD created by the transformation.
-# @rdname lapply
-# @aliases lapply
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# multiplyByTwo <- lapply(rdd, function(x) { x * 2 })
-# collect(multiplyByTwo) # 2,4,6...
-#}
+#' Apply a function to all elements
+#'
+#' This function creates a new RDD by applying the given transformation to all
+#' elements of the given RDD
+#'
+#' @param X The RDD to apply the transformation.
+#' @param FUN the transformation to apply on each element
+#' @return a new RDD created by the transformation.
+#' @rdname lapply
+#' @noRd
+#' @aliases lapply
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' multiplyByTwo <- lapply(rdd, function(x) { x * 2 })
+#' collect(multiplyByTwo) # 2,4,6...
+#'}
 setMethod("lapply",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
@@ -464,31 +473,33 @@ setMethod("lapply",
             lapplyPartitionsWithIndex(X, func)
           })
 
-# @rdname lapply
-# @aliases map,RDD,function-method
+#' @rdname lapply
+#' @aliases map,RDD,function-method
+#' @noRd
 setMethod("map",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
             lapply(X, FUN)
           })
 
-# Flatten results after apply a function to all elements
-#
-# This function return a new RDD by first applying a function to all
-# elements of this RDD, and then flattening the results.
-#
-# @param X The RDD to apply the transformation.
-# @param FUN the transformation to apply on each element
-# @return a new RDD created by the transformation.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# multiplyByTwo <- flatMap(rdd, function(x) { list(x*2, x*10) })
-# collect(multiplyByTwo) # 2,20,4,40,6,60...
-#}
-# @rdname flatMap
-# @aliases flatMap,RDD,function-method
+#' Flatten results after apply a function to all elements
+#'
+#' This function return a new RDD by first applying a function to all
+#' elements of this RDD, and then flattening the results.
+#'
+#' @param X The RDD to apply the transformation.
+#' @param FUN the transformation to apply on each element
+#' @return a new RDD created by the transformation.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' multiplyByTwo <- flatMap(rdd, function(x) { list(x*2, x*10) })
+#' collect(multiplyByTwo) # 2,20,4,40,6,60...
+#'}
+#' @rdname flatMap
+#' @aliases flatMap,RDD,function-method
+#' @noRd
 setMethod("flatMap",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
@@ -501,83 +512,88 @@ setMethod("flatMap",
             lapplyPartition(X, partitionFunc)
           })
 
-# Apply a function to each partition of an RDD
-#
-# Return a new RDD by applying a function to each partition of this RDD.
-#
-# @param X The RDD to apply the transformation.
-# @param FUN the transformation to apply on each partition.
-# @return a new RDD created by the transformation.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# partitionSum <- lapplyPartition(rdd, function(part) { Reduce("+", part) })
-# collect(partitionSum) # 15, 40
-#}
-# @rdname lapplyPartition
-# @aliases lapplyPartition,RDD,function-method
+#' Apply a function to each partition of an RDD
+#'
+#' Return a new RDD by applying a function to each partition of this RDD.
+#'
+#' @param X The RDD to apply the transformation.
+#' @param FUN the transformation to apply on each partition.
+#' @return a new RDD created by the transformation.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' partitionSum <- lapplyPartition(rdd, function(part) { Reduce("+", part) })
+#' collect(partitionSum) # 15, 40
+#'}
+#' @rdname lapplyPartition
+#' @aliases lapplyPartition,RDD,function-method
+#' @noRd
 setMethod("lapplyPartition",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
             lapplyPartitionsWithIndex(X, function(s, part) { FUN(part) })
           })
 
-# mapPartitions is the same as lapplyPartition.
-#
-# @rdname lapplyPartition
-# @aliases mapPartitions,RDD,function-method
+#' mapPartitions is the same as lapplyPartition.
+#'
+#' @rdname lapplyPartition
+#' @aliases mapPartitions,RDD,function-method
+#' @noRd
 setMethod("mapPartitions",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
             lapplyPartition(X, FUN)
           })
 
-# Return a new RDD by applying a function to each partition of this RDD, while
-# tracking the index of the original partition.
-#
-# @param X The RDD to apply the transformation.
-# @param FUN the transformation to apply on each partition; takes the partition
-#        index and a list of elements in the particular partition.
-# @return a new RDD created by the transformation.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10, 5L)
-# prod <- lapplyPartitionsWithIndex(rdd, function(partIndex, part) {
-#                                          partIndex * Reduce("+", part) })
-# collect(prod, flatten = FALSE) # 0, 7, 22, 45, 76
-#}
-# @rdname lapplyPartitionsWithIndex
-# @aliases lapplyPartitionsWithIndex,RDD,function-method
+#' Return a new RDD by applying a function to each partition of this RDD, while
+#' tracking the index of the original partition.
+#'
+#' @param X The RDD to apply the transformation.
+#' @param FUN the transformation to apply on each partition; takes the partition
+#'        index and a list of elements in the particular partition.
+#' @return a new RDD created by the transformation.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10, 5L)
+#' prod <- lapplyPartitionsWithIndex(rdd, function(partIndex, part) {
+#'                                          partIndex * Reduce("+", part) })
+#' collect(prod, flatten = FALSE) # 0, 7, 22, 45, 76
+#'}
+#' @rdname lapplyPartitionsWithIndex
+#' @aliases lapplyPartitionsWithIndex,RDD,function-method
+#' @noRd
 setMethod("lapplyPartitionsWithIndex",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
             PipelinedRDD(X, FUN)
           })
 
-# @rdname lapplyPartitionsWithIndex
-# @aliases mapPartitionsWithIndex,RDD,function-method
+#' @rdname lapplyPartitionsWithIndex
+#' @aliases mapPartitionsWithIndex,RDD,function-method
+#' @noRd
 setMethod("mapPartitionsWithIndex",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
             lapplyPartitionsWithIndex(X, FUN)
           })
 
-# This function returns a new RDD containing only the elements that satisfy
-# a predicate (i.e. returning TRUE in a given logical function).
-# The same as `filter()' in Spark.
-#
-# @param x The RDD to be filtered.
-# @param f A unary predicate function.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# unlist(collect(filterRDD(rdd, function (x) { x < 3 }))) # c(1, 2)
-#}
-# @rdname filterRDD
-# @aliases filterRDD,RDD,function-method
+#' This function returns a new RDD containing only the elements that satisfy
+#' a predicate (i.e. returning TRUE in a given logical function).
+#' The same as `filter()' in Spark.
+#'
+#' @param x The RDD to be filtered.
+#' @param f A unary predicate function.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' unlist(collect(filterRDD(rdd, function (x) { x < 3 }))) # c(1, 2)
+#'}
+#' @rdname filterRDD
+#' @aliases filterRDD,RDD,function-method
+#' @noRd
 setMethod("filterRDD",
           signature(x = "RDD", f = "function"),
           function(x, f) {
@@ -587,30 +603,32 @@ setMethod("filterRDD",
             lapplyPartition(x, filter.func)
           })
 
-# @rdname filterRDD
-# @aliases Filter
+#' @rdname filterRDD
+#' @aliases Filter
+#' @noRd
 setMethod("Filter",
           signature(f = "function", x = "RDD"),
           function(f, x) {
             filterRDD(x, f)
           })
 
-# Reduce across elements of an RDD.
-#
-# This function reduces the elements of this RDD using the
-# specified commutative and associative binary operator.
-#
-# @param x The RDD to reduce
-# @param func Commutative and associative function to apply on elements
-#             of the RDD.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# reduce(rdd, "+") # 55
-#}
-# @rdname reduce
-# @aliases reduce,RDD,ANY-method
+#' Reduce across elements of an RDD.
+#'
+#' This function reduces the elements of this RDD using the
+#' specified commutative and associative binary operator.
+#'
+#' @param x The RDD to reduce
+#' @param func Commutative and associative function to apply on elements
+#'             of the RDD.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' reduce(rdd, "+") # 55
+#'}
+#' @rdname reduce
+#' @aliases reduce,RDD,ANY-method
+#' @noRd
 setMethod("reduce",
           signature(x = "RDD", func = "ANY"),
           function(x, func) {
@@ -624,70 +642,74 @@ setMethod("reduce",
             Reduce(func, partitionList)
           })
 
-# Get the maximum element of an RDD.
-#
-# @param x The RDD to get the maximum element from
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# maximum(rdd) # 10
-#}
-# @rdname maximum
-# @aliases maximum,RDD
+#' Get the maximum element of an RDD.
+#'
+#' @param x The RDD to get the maximum element from
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' maximum(rdd) # 10
+#'}
+#' @rdname maximum
+#' @aliases maximum,RDD
+#' @noRd
 setMethod("maximum",
           signature(x = "RDD"),
           function(x) {
             reduce(x, max)
           })
 
-# Get the minimum element of an RDD.
-#
-# @param x The RDD to get the minimum element from
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# minimum(rdd) # 1
-#}
-# @rdname minimum
-# @aliases minimum,RDD
+#' Get the minimum element of an RDD.
+#'
+#' @param x The RDD to get the minimum element from
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' minimum(rdd) # 1
+#'}
+#' @rdname minimum
+#' @aliases minimum,RDD
+#' @noRd
 setMethod("minimum",
           signature(x = "RDD"),
           function(x) {
             reduce(x, min)
           })
 
-# Add up the elements in an RDD.
-#
-# @param x The RDD to add up the elements in
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# sumRDD(rdd) # 55
-#}
-# @rdname sumRDD
-# @aliases sumRDD,RDD
+#' Add up the elements in an RDD.
+#'
+#' @param x The RDD to add up the elements in
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' sumRDD(rdd) # 55
+#'}
+#' @rdname sumRDD
+#' @aliases sumRDD,RDD
+#' @noRd
 setMethod("sumRDD",
           signature(x = "RDD"),
           function(x) {
             reduce(x, "+")
           })
 
-# Applies a function to all elements in an RDD, and force evaluation.
-#
-# @param x The RDD to apply the function
-# @param func The function to be applied.
-# @return invisible NULL.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# foreach(rdd, function(x) { save(x, file=...) })
-#}
-# @rdname foreach
-# @aliases foreach,RDD,function-method
+#' Applies a function to all elements in an RDD, and force evaluation.
+#'
+#' @param x The RDD to apply the function
+#' @param func The function to be applied.
+#' @return invisible NULL.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' foreach(rdd, function(x) { save(x, file=...) })
+#'}
+#' @rdname foreach
+#' @aliases foreach,RDD,function-method
+#' @noRd
 setMethod("foreach",
           signature(x = "RDD", func = "function"),
           function(x, func) {
@@ -698,37 +720,39 @@ setMethod("foreach",
             invisible(collect(mapPartitions(x, partition.func)))
           })
 
-# Applies a function to each partition in an RDD, and force evaluation.
-#
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# foreachPartition(rdd, function(part) { save(part, file=...); NULL })
-#}
-# @rdname foreach
-# @aliases foreachPartition,RDD,function-method
+#' Applies a function to each partition in an RDD, and force evaluation.
+#'
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' foreachPartition(rdd, function(part) { save(part, file=...); NULL })
+#'}
+#' @rdname foreach
+#' @aliases foreachPartition,RDD,function-method
+#' @noRd
 setMethod("foreachPartition",
           signature(x = "RDD", func = "function"),
           function(x, func) {
             invisible(collect(mapPartitions(x, func)))
           })
 
-# Take elements from an RDD.
-#
-# This function takes the first NUM elements in the RDD and
-# returns them in a list.
-#
-# @param x The RDD to take elements from
-# @param num Number of elements to take
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# take(rdd, 2L) # list(1, 2)
-#}
-# @rdname take
-# @aliases take,RDD,numeric-method
+#' Take elements from an RDD.
+#'
+#' This function takes the first NUM elements in the RDD and
+#' returns them in a list.
+#'
+#' @param x The RDD to take elements from
+#' @param num Number of elements to take
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' take(rdd, 2L) # list(1, 2)
+#'}
+#' @rdname take
+#' @aliases take,RDD,numeric-method
+#' @noRd
 setMethod("take",
           signature(x = "RDD", num = "numeric"),
           function(x, num) {
@@ -763,39 +787,40 @@ setMethod("take",
           })
 
 
-# First
-#
-# Return the first element of an RDD
-#
-# @rdname first
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# first(rdd)
-# }
+#' First
+#'
+#' Return the first element of an RDD
+#'
+#' @rdname first
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' first(rdd)
+#' }
+#' @noRd
 setMethod("first",
           signature(x = "RDD"),
           function(x) {
             take(x, 1)[[1]]
           })
 
-# Removes the duplicates from RDD.
-#
-# This function returns a new RDD containing the distinct elements in the
-# given RDD. The same as `distinct()' in Spark.
-#
-# @param x The RDD to remove duplicates from.
-# @param numPartitions Number of partitions to create.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, c(1,2,2,3,3,3))
-# sort(unlist(collect(distinct(rdd)))) # c(1, 2, 3)
-#}
-# @rdname distinct
-# @aliases distinct,RDD-method
+#' Removes the duplicates from RDD.
+#'
+#' This function returns a new RDD containing the distinct elements in the
+#' given RDD. The same as `distinct()' in Spark.
+#'
+#' @param x The RDD to remove duplicates from.
+#' @param numPartitions Number of partitions to create.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, c(1,2,2,3,3,3))
+#' sort(unlist(collect(distinct(rdd)))) # c(1, 2, 3)
+#'}
+#' @rdname distinct
+#' @aliases distinct,RDD-method
+#' @noRd
 setMethod("distinct",
           signature(x = "RDD"),
           function(x, numPartitions = SparkR:::numPartitions(x)) {
@@ -807,24 +832,25 @@ setMethod("distinct",
             resRDD
           })
 
-# Return an RDD that is a sampled subset of the given RDD.
-#
-# The same as `sample()' in Spark. (We rename it due to signature
-# inconsistencies with the `sample()' function in R's base package.)
-#
-# @param x The RDD to sample elements from
-# @param withReplacement Sampling with replacement or not
-# @param fraction The (rough) sample target fraction
-# @param seed Randomness seed value
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# collect(sampleRDD(rdd, FALSE, 0.5, 1618L)) # ~5 distinct elements
-# collect(sampleRDD(rdd, TRUE, 0.5, 9L)) # ~5 elements possibly with duplicates
-#}
-# @rdname sampleRDD
-# @aliases sampleRDD,RDD
+#' Return an RDD that is a sampled subset of the given RDD.
+#'
+#' The same as `sample()' in Spark. (We rename it due to signature
+#' inconsistencies with the `sample()' function in R's base package.)
+#'
+#' @param x The RDD to sample elements from
+#' @param withReplacement Sampling with replacement or not
+#' @param fraction The (rough) sample target fraction
+#' @param seed Randomness seed value
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' collect(sampleRDD(rdd, FALSE, 0.5, 1618L)) # ~5 distinct elements
+#' collect(sampleRDD(rdd, TRUE, 0.5, 9L)) # ~5 elements possibly with duplicates
+#'}
+#' @rdname sampleRDD
+#' @aliases sampleRDD,RDD
+#' @noRd
 setMethod("sampleRDD",
           signature(x = "RDD", withReplacement = "logical",
                     fraction = "numeric", seed = "integer"),
@@ -868,23 +894,24 @@ setMethod("sampleRDD",
             lapplyPartitionsWithIndex(x, samplingFunc)
           })
 
-# Return a list of the elements that are a sampled subset of the given RDD.
-#
-# @param x The RDD to sample elements from
-# @param withReplacement Sampling with replacement or not
-# @param num Number of elements to return
-# @param seed Randomness seed value
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:100)
-# # exactly 5 elements sampled, which may not be distinct
-# takeSample(rdd, TRUE, 5L, 1618L)
-# # exactly 5 distinct elements sampled
-# takeSample(rdd, FALSE, 5L, 16181618L)
-#}
-# @rdname takeSample
-# @aliases takeSample,RDD
+#' Return a list of the elements that are a sampled subset of the given RDD.
+#'
+#' @param x The RDD to sample elements from
+#' @param withReplacement Sampling with replacement or not
+#' @param num Number of elements to return
+#' @param seed Randomness seed value
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:100)
+#' # exactly 5 elements sampled, which may not be distinct
+#' takeSample(rdd, TRUE, 5L, 1618L)
+#' # exactly 5 distinct elements sampled
+#' takeSample(rdd, FALSE, 5L, 16181618L)
+#'}
+#' @rdname takeSample
+#' @aliases takeSample,RDD
+#' @noRd
 setMethod("takeSample", signature(x = "RDD", withReplacement = "logical",
                                   num = "integer", seed = "integer"),
           function(x, withReplacement, num, seed) {
@@ -931,18 +958,19 @@ setMethod("takeSample", signature(x = "RDD", withReplacement = "logical",
             base::sample(samples)[1:total]
           })
 
-# Creates tuples of the elements in this RDD by applying a function.
-#
-# @param x The RDD.
-# @param func The function to be applied.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(1, 2, 3))
-# collect(keyBy(rdd, function(x) { x*x })) # list(list(1, 1), list(4, 2), list(9, 3))
-#}
-# @rdname keyBy
-# @aliases keyBy,RDD
+#' Creates tuples of the elements in this RDD by applying a function.
+#'
+#' @param x The RDD.
+#' @param func The function to be applied.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(1, 2, 3))
+#' collect(keyBy(rdd, function(x) { x*x })) # list(list(1, 1), list(4, 2), list(9, 3))
+#'}
+#' @rdname keyBy
+#' @aliases keyBy,RDD
+#' @noRd
 setMethod("keyBy",
           signature(x = "RDD", func = "function"),
           function(x, func) {
@@ -952,44 +980,46 @@ setMethod("keyBy",
             lapply(x, apply.func)
           })
 
-# Return a new RDD that has exactly numPartitions partitions.
-# Can increase or decrease the level of parallelism in this RDD. Internally,
-# this uses a shuffle to redistribute data.
-# If you are decreasing the number of partitions in this RDD, consider using
-# coalesce, which can avoid performing a shuffle.
-#
-# @param x The RDD.
-# @param numPartitions Number of partitions to create.
-# @seealso coalesce
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(1, 2, 3, 4, 5, 6, 7), 4L)
-# numPartitions(rdd)                   # 4
-# numPartitions(repartition(rdd, 2L))  # 2
-#}
-# @rdname repartition
-# @aliases repartition,RDD
+#' Return a new RDD that has exactly numPartitions partitions.
+#' Can increase or decrease the level of parallelism in this RDD. Internally,
+#' this uses a shuffle to redistribute data.
+#' If you are decreasing the number of partitions in this RDD, consider using
+#' coalesce, which can avoid performing a shuffle.
+#'
+#' @param x The RDD.
+#' @param numPartitions Number of partitions to create.
+#' @seealso coalesce
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(1, 2, 3, 4, 5, 6, 7), 4L)
+#' numPartitions(rdd)                   # 4
+#' numPartitions(repartition(rdd, 2L))  # 2
+#'}
+#' @rdname repartition
+#' @aliases repartition,RDD
+#' @noRd
 setMethod("repartition",
           signature(x = "RDD", numPartitions = "numeric"),
           function(x, numPartitions) {
             coalesce(x, numPartitions, TRUE)
           })
 
-# Return a new RDD that is reduced into numPartitions partitions.
-#
-# @param x The RDD.
-# @param numPartitions Number of partitions to create.
-# @seealso repartition
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(1, 2, 3, 4, 5), 3L)
-# numPartitions(rdd)               # 3
-# numPartitions(coalesce(rdd, 1L)) # 1
-#}
-# @rdname coalesce
-# @aliases coalesce,RDD
+#' Return a new RDD that is reduced into numPartitions partitions.
+#'
+#' @param x The RDD.
+#' @param numPartitions Number of partitions to create.
+#' @seealso repartition
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(1, 2, 3, 4, 5), 3L)
+#' numPartitions(rdd)               # 3
+#' numPartitions(coalesce(rdd, 1L)) # 1
+#'}
+#' @rdname coalesce
+#' @aliases coalesce,RDD
+#' @noRd
 setMethod("coalesce",
            signature(x = "RDD", numPartitions = "numeric"),
            function(x, numPartitions, shuffle = FALSE) {
@@ -1013,19 +1043,20 @@ setMethod("coalesce",
              }
            })
 
-# Save this RDD as a SequenceFile of serialized objects.
-#
-# @param x The RDD to save
-# @param path The directory where the file is saved
-# @seealso objectFile
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:3)
-# saveAsObjectFile(rdd, "/tmp/sparkR-tmp")
-#}
-# @rdname saveAsObjectFile
-# @aliases saveAsObjectFile,RDD
+#' Save this RDD as a SequenceFile of serialized objects.
+#'
+#' @param x The RDD to save
+#' @param path The directory where the file is saved
+#' @seealso objectFile
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:3)
+#' saveAsObjectFile(rdd, "/tmp/sparkR-tmp")
+#'}
+#' @rdname saveAsObjectFile
+#' @aliases saveAsObjectFile,RDD
+#' @noRd
 setMethod("saveAsObjectFile",
           signature(x = "RDD", path = "character"),
           function(x, path) {
@@ -1038,18 +1069,19 @@ setMethod("saveAsObjectFile",
             invisible(callJMethod(getJRDD(x), "saveAsObjectFile", path))
           })
 
-# Save this RDD as a text file, using string representations of elements.
-#
-# @param x The RDD to save
-# @param path The directory where the partitions of the text file are saved
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:3)
-# saveAsTextFile(rdd, "/tmp/sparkR-tmp")
-#}
-# @rdname saveAsTextFile
-# @aliases saveAsTextFile,RDD
+#' Save this RDD as a text file, using string representations of elements.
+#'
+#' @param x The RDD to save
+#' @param path The directory where the partitions of the text file are saved
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:3)
+#' saveAsTextFile(rdd, "/tmp/sparkR-tmp")
+#'}
+#' @rdname saveAsTextFile
+#' @aliases saveAsTextFile,RDD
+#' @noRd
 setMethod("saveAsTextFile",
           signature(x = "RDD", path = "character"),
           function(x, path) {
@@ -1062,21 +1094,22 @@ setMethod("saveAsTextFile",
               callJMethod(getJRDD(stringRdd, serializedMode = "string"), "saveAsTextFile", path))
           })
 
-# Sort an RDD by the given key function.
-#
-# @param x An RDD to be sorted.
-# @param func A function used to compute the sort key for each element.
-# @param ascending A flag to indicate whether the sorting is ascending or descending.
-# @param numPartitions Number of partitions to create.
-# @return An RDD where all elements are sorted.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(3, 2, 1))
-# collect(sortBy(rdd, function(x) { x })) # list (1, 2, 3)
-#}
-# @rdname sortBy
-# @aliases sortBy,RDD,RDD-method
+#' Sort an RDD by the given key function.
+#'
+#' @param x An RDD to be sorted.
+#' @param func A function used to compute the sort key for each element.
+#' @param ascending A flag to indicate whether the sorting is ascending or descending.
+#' @param numPartitions Number of partitions to create.
+#' @return An RDD where all elements are sorted.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(3, 2, 1))
+#' collect(sortBy(rdd, function(x) { x })) # list (1, 2, 3)
+#'}
+#' @rdname sortBy
+#' @aliases sortBy,RDD,RDD-method
+#' @noRd
 setMethod("sortBy",
           signature(x = "RDD", func = "function"),
           function(x, func, ascending = TRUE, numPartitions = SparkR:::numPartitions(x)) {
@@ -1138,97 +1171,95 @@ takeOrderedElem <- function(x, num, ascending = TRUE) {
   resList
 }
 
-# Returns the first N elements from an RDD in ascending order.
-#
-# @param x An RDD.
-# @param num Number of elements to return.
-# @return The first N elements from the RDD in ascending order.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(10, 1, 2, 9, 3, 4, 5, 6, 7))
-# takeOrdered(rdd, 6L) # list(1, 2, 3, 4, 5, 6)
-#}
-# @rdname takeOrdered
-# @aliases takeOrdered,RDD,RDD-method
+#' Returns the first N elements from an RDD in ascending order.
+#'
+#' @param x An RDD.
+#' @param num Number of elements to return.
+#' @return The first N elements from the RDD in ascending order.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(10, 1, 2, 9, 3, 4, 5, 6, 7))
+#' takeOrdered(rdd, 6L) # list(1, 2, 3, 4, 5, 6)
+#'}
+#' @rdname takeOrdered
+#' @aliases takeOrdered,RDD,RDD-method
+#' @noRd
 setMethod("takeOrdered",
           signature(x = "RDD", num = "integer"),
           function(x, num) {
             takeOrderedElem(x, num)
           })
 
-# Returns the top N elements from an RDD.
-#
-# @param x An RDD.
-# @param num Number of elements to return.
-# @return The top N elements from the RDD.
-# @rdname top
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(10, 1, 2, 9, 3, 4, 5, 6, 7))
-# top(rdd, 6L) # list(10, 9, 7, 6, 5, 4)
-#}
-# @rdname top
-# @aliases top,RDD,RDD-method
+#' Returns the top N elements from an RDD.
+#'
+#' @param x An RDD.
+#' @param num Number of elements to return.
+#' @return The top N elements from the RDD.
+#' @rdname top
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(10, 1, 2, 9, 3, 4, 5, 6, 7))
+#' top(rdd, 6L) # list(10, 9, 7, 6, 5, 4)
+#'}
+#' @aliases top,RDD,RDD-method
+#' @noRd
 setMethod("top",
           signature(x = "RDD", num = "integer"),
           function(x, num) {
             takeOrderedElem(x, num, FALSE)
           })
 
-# Fold an RDD using a given associative function and a neutral "zero value".
-#
-# Aggregate the elements of each partition, and then the results for all the
-# partitions, using a given associative function and a neutral "zero value".
-#
-# @param x An RDD.
-# @param zeroValue A neutral "zero value".
-# @param op An associative function for the folding operation.
-# @return The folding result.
-# @rdname fold
-# @seealso reduce
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(1, 2, 3, 4, 5))
-# fold(rdd, 0, "+") # 15
-#}
-# @rdname fold
-# @aliases fold,RDD,RDD-method
+#' Fold an RDD using a given associative function and a neutral "zero value".
+#'
+#' Aggregate the elements of each partition, and then the results for all the
+#' partitions, using a given associative function and a neutral "zero value".
+#'
+#' @param x An RDD.
+#' @param zeroValue A neutral "zero value".
+#' @param op An associative function for the folding operation.
+#' @return The folding result.
+#' @rdname fold
+#' @seealso reduce
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(1, 2, 3, 4, 5))
+#' fold(rdd, 0, "+") # 15
+#'}
+#' @aliases fold,RDD,RDD-method
+#' @noRd
 setMethod("fold",
           signature(x = "RDD", zeroValue = "ANY", op = "ANY"),
           function(x, zeroValue, op) {
             aggregateRDD(x, zeroValue, op, op)
           })
 
-# Aggregate an RDD using the given combine functions and a neutral "zero value".
-#
-# Aggregate the elements of each partition, and then the results for all the
-# partitions, using given combine functions and a neutral "zero value".
-#
-# @param x An RDD.
-# @param zeroValue A neutral "zero value".
-# @param seqOp A function to aggregate the RDD elements. It may return a different
-#              result type from the type of the RDD elements.
-# @param combOp A function to aggregate results of seqOp.
-# @return The aggregation result.
-# @rdname aggregateRDD
-# @seealso reduce
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(1, 2, 3, 4))
-# zeroValue <- list(0, 0)
-# seqOp <- function(x, y) { list(x[[1]] + y, x[[2]] + 1) }
-# combOp <- function(x, y) { list(x[[1]] + y[[1]], x[[2]] + y[[2]]) }
-# aggregateRDD(rdd, zeroValue, seqOp, combOp) # list(10, 4)
-#}
-# @rdname aggregateRDD
-# @aliases aggregateRDD,RDD,RDD-method
+#' Aggregate an RDD using the given combine functions and a neutral "zero value".
+#'
+#' Aggregate the elements of each partition, and then the results for all the
+#' partitions, using given combine functions and a neutral "zero value".
+#'
+#' @param x An RDD.
+#' @param zeroValue A neutral "zero value".
+#' @param seqOp A function to aggregate the RDD elements. It may return a different
+#'              result type from the type of the RDD elements.
+#' @param combOp A function to aggregate results of seqOp.
+#' @return The aggregation result.
+#' @rdname aggregateRDD
+#' @seealso reduce
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(1, 2, 3, 4))
+#' zeroValue <- list(0, 0)
+#' seqOp <- function(x, y) { list(x[[1]] + y, x[[2]] + 1) }
+#' combOp <- function(x, y) { list(x[[1]] + y[[1]], x[[2]] + y[[2]]) }
+#' aggregateRDD(rdd, zeroValue, seqOp, combOp) # list(10, 4)
+#'}
+#' @aliases aggregateRDD,RDD,RDD-method
+#' @noRd
 setMethod("aggregateRDD",
           signature(x = "RDD", zeroValue = "ANY", seqOp = "ANY", combOp = "ANY"),
           function(x, zeroValue, seqOp, combOp) {
@@ -1241,25 +1272,24 @@ setMethod("aggregateRDD",
             Reduce(combOp, partitionList, zeroValue)
           })
 
-# Pipes elements to a forked external process.
-#
-# The same as 'pipe()' in Spark.
-#
-# @param x The RDD whose elements are piped to the forked external process.
-# @param command The command to fork an external process.
-# @param env A named list to set environment variables of the external process.
-# @return A new RDD created by piping all elements to a forked external process.
-# @rdname pipeRDD
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# collect(pipeRDD(rdd, "more")
-# Output: c("1", "2", ..., "10")
-#}
-# @rdname pipeRDD
-# @aliases pipeRDD,RDD,character-method
+#' Pipes elements to a forked external process.
+#'
+#' The same as 'pipe()' in Spark.
+#'
+#' @param x The RDD whose elements are piped to the forked external process.
+#' @param command The command to fork an external process.
+#' @param env A named list to set environment variables of the external process.
+#' @return A new RDD created by piping all elements to a forked external process.
+#' @rdname pipeRDD
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' collect(pipeRDD(rdd, "more")
+#' Output: c("1", "2", ..., "10")
+#'}
+#' @aliases pipeRDD,RDD,character-method
+#' @noRd
 setMethod("pipeRDD",
           signature(x = "RDD", command = "character"),
           function(x, command, env = list()) {
@@ -1274,42 +1304,40 @@ setMethod("pipeRDD",
             lapplyPartition(x, func)
           })
 
-# TODO: Consider caching the name in the RDD's environment
-# Return an RDD's name.
-#
-# @param x The RDD whose name is returned.
-# @rdname name
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(1,2,3))
-# name(rdd) # NULL (if not set before)
-#}
-# @rdname name
-# @aliases name,RDD
+#' TODO: Consider caching the name in the RDD's environment
+#' Return an RDD's name.
+#'
+#' @param x The RDD whose name is returned.
+#' @rdname name
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(1,2,3))
+#' name(rdd) # NULL (if not set before)
+#'}
+#' @aliases name,RDD
+#' @noRd
 setMethod("name",
           signature(x = "RDD"),
           function(x) {
             callJMethod(getJRDD(x), "name")
           })
 
-# Set an RDD's name.
-#
-# @param x The RDD whose name is to be set.
-# @param name The RDD name to be set.
-# @return a new RDD renamed.
-# @rdname setName
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(1,2,3))
-# setName(rdd, "myRDD")
-# name(rdd) # "myRDD"
-#}
-# @rdname setName
-# @aliases setName,RDD
+#' Set an RDD's name.
+#'
+#' @param x The RDD whose name is to be set.
+#' @param name The RDD name to be set.
+#' @return a new RDD renamed.
+#' @rdname setName
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(1,2,3))
+#' setName(rdd, "myRDD")
+#' name(rdd) # "myRDD"
+#'}
+#' @aliases setName,RDD
+#' @noRd
 setMethod("setName",
           signature(x = "RDD", name = "character"),
           function(x, name) {
@@ -1317,25 +1345,26 @@ setMethod("setName",
             x
           })
 
-# Zip an RDD with generated unique Long IDs.
-#
-# Items in the kth partition will get ids k, n+k, 2*n+k, ..., where
-# n is the number of partitions. So there may exist gaps, but this
-# method won't trigger a spark job, which is different from
-# zipWithIndex.
-#
-# @param x An RDD to be zipped.
-# @return An RDD with zipped items.
-# @seealso zipWithIndex
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list("a", "b", "c", "d", "e"), 3L)
-# collect(zipWithUniqueId(rdd))
-# # list(list("a", 0), list("b", 3), list("c", 1), list("d", 4), list("e", 2))
-#}
-# @rdname zipWithUniqueId
-# @aliases zipWithUniqueId,RDD
+#' Zip an RDD with generated unique Long IDs.
+#'
+#' Items in the kth partition will get ids k, n+k, 2*n+k, ..., where
+#' n is the number of partitions. So there may exist gaps, but this
+#' method won't trigger a spark job, which is different from
+#' zipWithIndex.
+#'
+#' @param x An RDD to be zipped.
+#' @return An RDD with zipped items.
+#' @seealso zipWithIndex
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list("a", "b", "c", "d", "e"), 3L)
+#' collect(zipWithUniqueId(rdd))
+#' # list(list("a", 0), list("b", 3), list("c", 1), list("d", 4), list("e", 2))
+#'}
+#' @rdname zipWithUniqueId
+#' @aliases zipWithUniqueId,RDD
+#' @noRd
 setMethod("zipWithUniqueId",
           signature(x = "RDD"),
           function(x) {
@@ -1354,28 +1383,29 @@ setMethod("zipWithUniqueId",
             lapplyPartitionsWithIndex(x, partitionFunc)
           })
 
-# Zip an RDD with its element indices.
-#
-# The ordering is first based on the partition index and then the
-# ordering of items within each partition. So the first item in
-# the first partition gets index 0, and the last item in the last
-# partition receives the largest index.
-#
-# This method needs to trigger a Spark job when this RDD contains
-# more than one partition.
-#
-# @param x An RDD to be zipped.
-# @return An RDD with zipped items.
-# @seealso zipWithUniqueId
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list("a", "b", "c", "d", "e"), 3L)
-# collect(zipWithIndex(rdd))
-# # list(list("a", 0), list("b", 1), list("c", 2), list("d", 3), list("e", 4))
-#}
-# @rdname zipWithIndex
-# @aliases zipWithIndex,RDD
+#' Zip an RDD with its element indices.
+#'
+#' The ordering is first based on the partition index and then the
+#' ordering of items within each partition. So the first item in
+#' the first partition gets index 0, and the last item in the last
+#' partition receives the largest index.
+#'
+#' This method needs to trigger a Spark job when this RDD contains
+#' more than one partition.
+#'
+#' @param x An RDD to be zipped.
+#' @return An RDD with zipped items.
+#' @seealso zipWithUniqueId
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list("a", "b", "c", "d", "e"), 3L)
+#' collect(zipWithIndex(rdd))
+#' # list(list("a", 0), list("b", 1), list("c", 2), list("d", 3), list("e", 4))
+#'}
+#' @rdname zipWithIndex
+#' @aliases zipWithIndex,RDD
+#' @noRd
 setMethod("zipWithIndex",
           signature(x = "RDD"),
           function(x) {
@@ -1407,20 +1437,21 @@ setMethod("zipWithIndex",
            lapplyPartitionsWithIndex(x, partitionFunc)
          })
 
-# Coalesce all elements within each partition of an RDD into a list.
-#
-# @param x An RDD.
-# @return An RDD created by coalescing all elements within
-#         each partition into a list.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, as.list(1:4), 2L)
-# collect(glom(rdd))
-# # list(list(1, 2), list(3, 4))
-#}
-# @rdname glom
-# @aliases glom,RDD
+#' Coalesce all elements within each partition of an RDD into a list.
+#'
+#' @param x An RDD.
+#' @return An RDD created by coalescing all elements within
+#'         each partition into a list.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, as.list(1:4), 2L)
+#' collect(glom(rdd))
+#' # list(list(1, 2), list(3, 4))
+#'}
+#' @rdname glom
+#' @aliases glom,RDD
+#' @noRd
 setMethod("glom",
           signature(x = "RDD"),
           function(x) {
@@ -1433,21 +1464,22 @@ setMethod("glom",
 
 ############ Binary Functions #############
 
-# Return the union RDD of two RDDs.
-# The same as union() in Spark.
-#
-# @param x An RDD.
-# @param y An RDD.
-# @return a new RDD created by performing the simple union (witout removing
-# duplicates) of two input RDDs.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:3)
-# unionRDD(rdd, rdd) # 1, 2, 3, 1, 2, 3
-#}
-# @rdname unionRDD
-# @aliases unionRDD,RDD,RDD-method
+#' Return the union RDD of two RDDs.
+#' The same as union() in Spark.
+#'
+#' @param x An RDD.
+#' @param y An RDD.
+#' @return a new RDD created by performing the simple union (witout removing
+#' duplicates) of two input RDDs.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:3)
+#' unionRDD(rdd, rdd) # 1, 2, 3, 1, 2, 3
+#'}
+#' @rdname unionRDD
+#' @aliases unionRDD,RDD,RDD-method
+#' @noRd
 setMethod("unionRDD",
           signature(x = "RDD", y = "RDD"),
           function(x, y) {
@@ -1464,27 +1496,28 @@ setMethod("unionRDD",
             union.rdd
           })
 
-# Zip an RDD with another RDD.
-#
-# Zips this RDD with another one, returning key-value pairs with the
-# first element in each RDD second element in each RDD, etc. Assumes
-# that the two RDDs have the same number of partitions and the same
-# number of elements in each partition (e.g. one was made through
-# a map on the other).
-#
-# @param x An RDD to be zipped.
-# @param other Another RDD to be zipped.
-# @return An RDD zipped from the two RDDs.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, 0:4)
-# rdd2 <- parallelize(sc, 1000:1004)
-# collect(zipRDD(rdd1, rdd2))
-# # list(list(0, 1000), list(1, 1001), list(2, 1002), list(3, 1003), list(4, 1004))
-#}
-# @rdname zipRDD
-# @aliases zipRDD,RDD
+#' Zip an RDD with another RDD.
+#'
+#' Zips this RDD with another one, returning key-value pairs with the
+#' first element in each RDD second element in each RDD, etc. Assumes
+#' that the two RDDs have the same number of partitions and the same
+#' number of elements in each partition (e.g. one was made through
+#' a map on the other).
+#'
+#' @param x An RDD to be zipped.
+#' @param other Another RDD to be zipped.
+#' @return An RDD zipped from the two RDDs.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, 0:4)
+#' rdd2 <- parallelize(sc, 1000:1004)
+#' collect(zipRDD(rdd1, rdd2))
+#' # list(list(0, 1000), list(1, 1001), list(2, 1002), list(3, 1003), list(4, 1004))
+#'}
+#' @rdname zipRDD
+#' @aliases zipRDD,RDD
+#' @noRd
 setMethod("zipRDD",
           signature(x = "RDD", other = "RDD"),
           function(x, other) {
@@ -1503,24 +1536,25 @@ setMethod("zipRDD",
             mergePartitions(rdd, TRUE)
           })
 
-# Cartesian product of this RDD and another one.
-#
-# Return the Cartesian product of this RDD and another one,
-# that is, the RDD of all pairs of elements (a, b) where a
-# is in this and b is in other.
-#
-# @param x An RDD.
-# @param other An RDD.
-# @return A new RDD which is the Cartesian product of these two RDDs.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:2)
-# sortByKey(cartesian(rdd, rdd))
-# # list(list(1, 1), list(1, 2), list(2, 1), list(2, 2))
-#}
-# @rdname cartesian
-# @aliases cartesian,RDD,RDD-method
+#' Cartesian product of this RDD and another one.
+#'
+#' Return the Cartesian product of this RDD and another one,
+#' that is, the RDD of all pairs of elements (a, b) where a
+#' is in this and b is in other.
+#'
+#' @param x An RDD.
+#' @param other An RDD.
+#' @return A new RDD which is the Cartesian product of these two RDDs.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:2)
+#' sortByKey(cartesian(rdd, rdd))
+#' # list(list(1, 1), list(1, 2), list(2, 1), list(2, 2))
+#'}
+#' @rdname cartesian
+#' @aliases cartesian,RDD,RDD-method
+#' @noRd
 setMethod("cartesian",
           signature(x = "RDD", other = "RDD"),
           function(x, other) {
@@ -1533,24 +1567,25 @@ setMethod("cartesian",
             mergePartitions(rdd, FALSE)
           })
 
-# Subtract an RDD with another RDD.
-#
-# Return an RDD with the elements from this that are not in other.
-#
-# @param x An RDD.
-# @param other An RDD.
-# @param numPartitions Number of the partitions in the result RDD.
-# @return An RDD with the elements from this that are not in other.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, list(1, 1, 2, 2, 3, 4))
-# rdd2 <- parallelize(sc, list(2, 4))
-# collect(subtract(rdd1, rdd2))
-# # list(1, 1, 3)
-#}
-# @rdname subtract
-# @aliases subtract,RDD
+#' Subtract an RDD with another RDD.
+#'
+#' Return an RDD with the elements from this that are not in other.
+#'
+#' @param x An RDD.
+#' @param other An RDD.
+#' @param numPartitions Number of the partitions in the result RDD.
+#' @return An RDD with the elements from this that are not in other.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, list(1, 1, 2, 2, 3, 4))
+#' rdd2 <- parallelize(sc, list(2, 4))
+#' collect(subtract(rdd1, rdd2))
+#' # list(1, 1, 3)
+#'}
+#' @rdname subtract
+#' @aliases subtract,RDD
+#' @noRd
 setMethod("subtract",
           signature(x = "RDD", other = "RDD"),
           function(x, other, numPartitions = SparkR:::numPartitions(x)) {
@@ -1560,28 +1595,29 @@ setMethod("subtract",
             keys(subtractByKey(rdd1, rdd2, numPartitions))
           })
 
-# Intersection of this RDD and another one.
-#
-# Return the intersection of this RDD and another one.
-# The output will not contain any duplicate elements,
-# even if the input RDDs did. Performs a hash partition
-# across the cluster.
-# Note that this method performs a shuffle internally.
-#
-# @param x An RDD.
-# @param other An RDD.
-# @param numPartitions The number of partitions in the result RDD.
-# @return An RDD which is the intersection of these two RDDs.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, list(1, 10, 2, 3, 4, 5))
-# rdd2 <- parallelize(sc, list(1, 6, 2, 3, 7, 8))
-# collect(sortBy(intersection(rdd1, rdd2), function(x) { x }))
-# # list(1, 2, 3)
-#}
-# @rdname intersection
-# @aliases intersection,RDD
+#' Intersection of this RDD and another one.
+#'
+#' Return the intersection of this RDD and another one.
+#' The output will not contain any duplicate elements,
+#' even if the input RDDs did. Performs a hash partition
+#' across the cluster.
+#' Note that this method performs a shuffle internally.
+#'
+#' @param x An RDD.
+#' @param other An RDD.
+#' @param numPartitions The number of partitions in the result RDD.
+#' @return An RDD which is the intersection of these two RDDs.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, list(1, 10, 2, 3, 4, 5))
+#' rdd2 <- parallelize(sc, list(1, 6, 2, 3, 7, 8))
+#' collect(sortBy(intersection(rdd1, rdd2), function(x) { x }))
+#' # list(1, 2, 3)
+#'}
+#' @rdname intersection
+#' @aliases intersection,RDD
+#' @noRd
 setMethod("intersection",
           signature(x = "RDD", other = "RDD"),
           function(x, other, numPartitions = SparkR:::numPartitions(x)) {
@@ -1597,26 +1633,27 @@ setMethod("intersection",
             keys(filterRDD(cogroup(rdd1, rdd2, numPartitions = numPartitions), filterFunction))
           })
 
-# Zips an RDD's partitions with one (or more) RDD(s).
-# Same as zipPartitions in Spark.
-#
-# @param ... RDDs to be zipped.
-# @param func A function to transform zipped partitions.
-# @return A new RDD by applying a function to the zipped partitions.
-#         Assumes that all the RDDs have the *same number of partitions*, but
-#         does *not* require them to have the same number of elements in each partition.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, 1:2, 2L)  # 1, 2
-# rdd2 <- parallelize(sc, 1:4, 2L)  # 1:2, 3:4
-# rdd3 <- parallelize(sc, 1:6, 2L)  # 1:3, 4:6
-# collect(zipPartitions(rdd1, rdd2, rdd3,
-#                       func = function(x, y, z) { list(list(x, y, z))} ))
-# # list(list(1, c(1,2), c(1,2,3)), list(2, c(3,4), c(4,5,6)))
-#}
-# @rdname zipRDD
-# @aliases zipPartitions,RDD
+#' Zips an RDD's partitions with one (or more) RDD(s).
+#' Same as zipPartitions in Spark.
+#'
+#' @param ... RDDs to be zipped.
+#' @param func A function to transform zipped partitions.
+#' @return A new RDD by applying a function to the zipped partitions.
+#'         Assumes that all the RDDs have the *same number of partitions*, but
+#'         does *not* require them to have the same number of elements in each partition.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, 1:2, 2L)  # 1, 2
+#' rdd2 <- parallelize(sc, 1:4, 2L)  # 1:2, 3:4
+#' rdd3 <- parallelize(sc, 1:6, 2L)  # 1:3, 4:6
+#' collect(zipPartitions(rdd1, rdd2, rdd3,
+#'                       func = function(x, y, z) { list(list(x, y, z))} ))
+#' # list(list(1, c(1,2), c(1,2,3)), list(2, c(3,4), c(4,5,6)))
+#'}
+#' @rdname zipRDD
+#' @aliases zipPartitions,RDD
+#' @noRd
 setMethod("zipPartitions",
           "RDD",
           function(..., func) {
diff --git a/R/pkg/R/SQLContext.R b/R/pkg/R/SQLContext.R
index 1bf025cce4376..fd013fdb304df 100644
--- a/R/pkg/R/SQLContext.R
+++ b/R/pkg/R/SQLContext.R
@@ -144,7 +144,6 @@ createDataFrame <- function(sqlContext, data, schema = NULL, samplingRatio = 1.0
   }
 
   stopifnot(class(schema) == "structType")
-  # schemaString <- tojson(schema)
 
   jrdd <- getJRDD(lapply(rdd, function(x) x), "row")
   srdd <- callJMethod(jrdd, "rdd")
@@ -160,22 +159,21 @@ as.DataFrame <- function(sqlContext, data, schema = NULL, samplingRatio = 1.0) {
   createDataFrame(sqlContext, data, schema, samplingRatio)
 }
 
-# toDF
-#
-# Converts an RDD to a DataFrame by infer the types.
-#
-# @param x An RDD
-#
-# @rdname DataFrame
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# sqlContext <- sparkRSQL.init(sc)
-# rdd <- lapply(parallelize(sc, 1:10), function(x) list(a=x, b=as.character(x)))
-# df <- toDF(rdd)
-# }
-
+#' toDF
+#'
+#' Converts an RDD to a DataFrame by infer the types.
+#'
+#' @param x An RDD
+#'
+#' @rdname DataFrame
+#' @noRd
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' sqlContext <- sparkRSQL.init(sc)
+#' rdd <- lapply(parallelize(sc, 1:10), function(x) list(a=x, b=as.character(x)))
+#' df <- toDF(rdd)
+#'}
 setGeneric("toDF", function(x, ...) { standardGeneric("toDF") })
 
 setMethod("toDF", signature(x = "RDD"),
@@ -217,23 +215,23 @@ jsonFile <- function(sqlContext, path) {
 }
 
 
-# JSON RDD
-#
-# Loads an RDD storing one JSON object per string as a DataFrame.
-#
-# @param sqlContext SQLContext to use
-# @param rdd An RDD of JSON string
-# @param schema A StructType object to use as schema
-# @param samplingRatio The ratio of simpling used to infer the schema
-# @return A DataFrame
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# sqlContext <- sparkRSQL.init(sc)
-# rdd <- texFile(sc, "path/to/json")
-# df <- jsonRDD(sqlContext, rdd)
-# }
+#' JSON RDD
+#'
+#' Loads an RDD storing one JSON object per string as a DataFrame.
+#'
+#' @param sqlContext SQLContext to use
+#' @param rdd An RDD of JSON string
+#' @param schema A StructType object to use as schema
+#' @param samplingRatio The ratio of simpling used to infer the schema
+#' @return A DataFrame
+#' @noRd
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' sqlContext <- sparkRSQL.init(sc)
+#' rdd <- texFile(sc, "path/to/json")
+#' df <- jsonRDD(sqlContext, rdd)
+#'}
 
 # TODO: support schema
 jsonRDD <- function(sqlContext, rdd, schema = NULL, samplingRatio = 1.0) {
diff --git a/R/pkg/R/context.R b/R/pkg/R/context.R
index 720990e1c6087..471bec1eacf03 100644
--- a/R/pkg/R/context.R
+++ b/R/pkg/R/context.R
@@ -25,23 +25,23 @@ getMinPartitions <- function(sc, minPartitions) {
   as.integer(minPartitions)
 }
 
-# Create an RDD from a text file.
-#
-# This function reads a text file from HDFS, a local file system (available on all
-# nodes), or any Hadoop-supported file system URI, and creates an
-# RDD of strings from it.
-#
-# @param sc SparkContext to use
-# @param path Path of file to read. A vector of multiple paths is allowed.
-# @param minPartitions Minimum number of partitions to be created. If NULL, the default
-#  value is chosen based on available parallelism.
-# @return RDD where each item is of type \code{character}
-# @export
-# @examples
-#\dontrun{
-#  sc <- sparkR.init()
-#  lines <- textFile(sc, "myfile.txt")
-#}
+#' Create an RDD from a text file.
+#'
+#' This function reads a text file from HDFS, a local file system (available on all
+#' nodes), or any Hadoop-supported file system URI, and creates an
+#' RDD of strings from it.
+#'
+#' @param sc SparkContext to use
+#' @param path Path of file to read. A vector of multiple paths is allowed.
+#' @param minPartitions Minimum number of partitions to be created. If NULL, the default
+#'  value is chosen based on available parallelism.
+#' @return RDD where each item is of type \code{character}
+#' @noRd
+#' @examples
+#'\dontrun{
+#'  sc <- sparkR.init()
+#'  lines <- textFile(sc, "myfile.txt")
+#'}
 textFile <- function(sc, path, minPartitions = NULL) {
   # Allow the user to have a more flexible definiton of the text file path
   path <- suppressWarnings(normalizePath(path))
@@ -53,23 +53,23 @@ textFile <- function(sc, path, minPartitions = NULL) {
   RDD(jrdd, "string")
 }
 
-# Load an RDD saved as a SequenceFile containing serialized objects.
-#
-# The file to be loaded should be one that was previously generated by calling
-# saveAsObjectFile() of the RDD class.
-#
-# @param sc SparkContext to use
-# @param path Path of file to read. A vector of multiple paths is allowed.
-# @param minPartitions Minimum number of partitions to be created. If NULL, the default
-#  value is chosen based on available parallelism.
-# @return RDD containing serialized R objects.
-# @seealso saveAsObjectFile
-# @export
-# @examples
-#\dontrun{
-#  sc <- sparkR.init()
-#  rdd <- objectFile(sc, "myfile")
-#}
+#' Load an RDD saved as a SequenceFile containing serialized objects.
+#'
+#' The file to be loaded should be one that was previously generated by calling
+#' saveAsObjectFile() of the RDD class.
+#'
+#' @param sc SparkContext to use
+#' @param path Path of file to read. A vector of multiple paths is allowed.
+#' @param minPartitions Minimum number of partitions to be created. If NULL, the default
+#'  value is chosen based on available parallelism.
+#' @return RDD containing serialized R objects.
+#' @seealso saveAsObjectFile
+#' @noRd
+#' @examples
+#'\dontrun{
+#'  sc <- sparkR.init()
+#'  rdd <- objectFile(sc, "myfile")
+#'}
 objectFile <- function(sc, path, minPartitions = NULL) {
   # Allow the user to have a more flexible definiton of the text file path
   path <- suppressWarnings(normalizePath(path))
@@ -81,24 +81,24 @@ objectFile <- function(sc, path, minPartitions = NULL) {
   RDD(jrdd, "byte")
 }
 
-# Create an RDD from a homogeneous list or vector.
-#
-# This function creates an RDD from a local homogeneous list in R. The elements
-# in the list are split into \code{numSlices} slices and distributed to nodes
-# in the cluster.
-#
-# @param sc SparkContext to use
-# @param coll collection to parallelize
-# @param numSlices number of partitions to create in the RDD
-# @return an RDD created from this collection
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10, 2)
-# # The RDD should contain 10 elements
-# length(rdd)
-#}
+#' Create an RDD from a homogeneous list or vector.
+#'
+#' This function creates an RDD from a local homogeneous list in R. The elements
+#' in the list are split into \code{numSlices} slices and distributed to nodes
+#' in the cluster.
+#'
+#' @param sc SparkContext to use
+#' @param coll collection to parallelize
+#' @param numSlices number of partitions to create in the RDD
+#' @return an RDD created from this collection
+#' @noRd
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10, 2)
+#' # The RDD should contain 10 elements
+#' length(rdd)
+#'}
 parallelize <- function(sc, coll, numSlices = 1) {
   # TODO: bound/safeguard numSlices
   # TODO: unit tests for if the split works for all primitives
@@ -133,33 +133,32 @@ parallelize <- function(sc, coll, numSlices = 1) {
   RDD(jrdd, "byte")
 }
 
-# Include this specified package on all workers
-#
-# This function can be used to include a package on all workers before the
-# user's code is executed. This is useful in scenarios where other R package
-# functions are used in a function passed to functions like \code{lapply}.
-# NOTE: The package is assumed to be installed on every node in the Spark
-# cluster.
-#
-# @param sc SparkContext to use
-# @param pkg Package name
-#
-# @export
-# @examples
-#\dontrun{
-#  library(Matrix)
-#
-#  sc <- sparkR.init()
-#  # Include the matrix library we will be using
-#  includePackage(sc, Matrix)
-#
-#  generateSparse <- function(x) {
-#    sparseMatrix(i=c(1, 2, 3), j=c(1, 2, 3), x=c(1, 2, 3))
-#  }
-#
-#  rdd <- lapplyPartition(parallelize(sc, 1:2, 2L), generateSparse)
-#  collect(rdd)
-#}
+#' Include this specified package on all workers
+#'
+#' This function can be used to include a package on all workers before the
+#' user's code is executed. This is useful in scenarios where other R package
+#' functions are used in a function passed to functions like \code{lapply}.
+#' NOTE: The package is assumed to be installed on every node in the Spark
+#' cluster.
+#'
+#' @param sc SparkContext to use
+#' @param pkg Package name
+#' @noRd
+#' @examples
+#'\dontrun{
+#'  library(Matrix)
+#'
+#'  sc <- sparkR.init()
+#'  # Include the matrix library we will be using
+#'  includePackage(sc, Matrix)
+#'
+#'  generateSparse <- function(x) {
+#'    sparseMatrix(i=c(1, 2, 3), j=c(1, 2, 3), x=c(1, 2, 3))
+#'  }
+#'
+#'  rdd <- lapplyPartition(parallelize(sc, 1:2, 2L), generateSparse)
+#'  collect(rdd)
+#'}
 includePackage <- function(sc, pkg) {
   pkg <- as.character(substitute(pkg))
   if (exists(".packages", .sparkREnv)) {
@@ -171,30 +170,30 @@ includePackage <- function(sc, pkg) {
   .sparkREnv$.packages <- packages
 }
 
-# @title Broadcast a variable to all workers
-#
-# @description
-# Broadcast a read-only variable to the cluster, returning a \code{Broadcast}
-# object for reading it in distributed functions.
-#
-# @param sc Spark Context to use
-# @param object Object to be broadcast
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:2, 2L)
-#
-# # Large Matrix object that we want to broadcast
-# randomMat <- matrix(nrow=100, ncol=10, data=rnorm(1000))
-# randomMatBr <- broadcast(sc, randomMat)
-#
-# # Use the broadcast variable inside the function
-# useBroadcast <- function(x) {
-#   sum(value(randomMatBr) * x)
-# }
-# sumRDD <- lapply(rdd, useBroadcast)
-#}
+#' @title Broadcast a variable to all workers
+#'
+#' @description
+#' Broadcast a read-only variable to the cluster, returning a \code{Broadcast}
+#' object for reading it in distributed functions.
+#'
+#' @param sc Spark Context to use
+#' @param object Object to be broadcast
+#' @noRd
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:2, 2L)
+#'
+#' # Large Matrix object that we want to broadcast
+#' randomMat <- matrix(nrow=100, ncol=10, data=rnorm(1000))
+#' randomMatBr <- broadcast(sc, randomMat)
+#'
+#' # Use the broadcast variable inside the function
+#' useBroadcast <- function(x) {
+#'   sum(value(randomMatBr) * x)
+#' }
+#' sumRDD <- lapply(rdd, useBroadcast)
+#'}
 broadcast <- function(sc, object) {
   objName <- as.character(substitute(object))
   serializedObj <- serialize(object, connection = NULL)
@@ -205,21 +204,21 @@ broadcast <- function(sc, object) {
   Broadcast(id, object, jBroadcast, objName)
 }
 
-# @title Set the checkpoint directory
-#
-# Set the directory under which RDDs are going to be checkpointed. The
-# directory must be a HDFS path if running on a cluster.
-#
-# @param sc Spark Context to use
-# @param dirName Directory path
-# @export
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# setCheckpointDir(sc, "~/checkpoint")
-# rdd <- parallelize(sc, 1:2, 2L)
-# checkpoint(rdd)
-#}
+#' @title Set the checkpoint directory
+#'
+#' Set the directory under which RDDs are going to be checkpointed. The
+#' directory must be a HDFS path if running on a cluster.
+#'
+#' @param sc Spark Context to use
+#' @param dirName Directory path
+#' @noRd
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' setCheckpointDir(sc, "~/checkpoint")
+#' rdd <- parallelize(sc, 1:2, 2L)
+#' checkpoint(rdd)
+#'}
 setCheckpointDir <- function(sc, dirName) {
   invisible(callJMethod(sc, "setCheckpointDir", suppressWarnings(normalizePath(dirName))))
 }
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 92ad4ee8685ee..612e639f8ad99 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -88,12 +88,8 @@ setGeneric("flatMap", function(X, FUN) { standardGeneric("flatMap") })
 # @export
 setGeneric("fold", function(x, zeroValue, op) { standardGeneric("fold") })
 
-# @rdname foreach
-# @export
 setGeneric("foreach", function(x, func) { standardGeneric("foreach") })
 
-# @rdname foreach
-# @export
 setGeneric("foreachPartition", function(x, func) { standardGeneric("foreachPartition") })
 
 # The jrdd accessor function.
@@ -107,27 +103,17 @@ setGeneric("glom", function(x) { standardGeneric("glom") })
 # @export
 setGeneric("keyBy", function(x, func) { standardGeneric("keyBy") })
 
-# @rdname lapplyPartition
-# @export
 setGeneric("lapplyPartition", function(X, FUN) { standardGeneric("lapplyPartition") })
 
-# @rdname lapplyPartitionsWithIndex
-# @export
 setGeneric("lapplyPartitionsWithIndex",
            function(X, FUN) {
              standardGeneric("lapplyPartitionsWithIndex")
            })
 
-# @rdname lapply
-# @export
 setGeneric("map", function(X, FUN) { standardGeneric("map") })
 
-# @rdname lapplyPartition
-# @export
 setGeneric("mapPartitions", function(X, FUN) { standardGeneric("mapPartitions") })
 
-# @rdname lapplyPartitionsWithIndex
-# @export
 setGeneric("mapPartitionsWithIndex",
            function(X, FUN) { standardGeneric("mapPartitionsWithIndex") })
 
@@ -563,12 +549,8 @@ setGeneric("summarize", function(x,...) { standardGeneric("summarize") })
 #' @export
 setGeneric("summary", function(object, ...) { standardGeneric("summary") })
 
-# @rdname tojson
-# @export
 setGeneric("toJSON", function(x) { standardGeneric("toJSON") })
 
-#' @rdname DataFrame
-#' @export
 setGeneric("toRDD", function(x) { standardGeneric("toRDD") })
 
 #' @rdname unionAll
diff --git a/R/pkg/R/pairRDD.R b/R/pkg/R/pairRDD.R
index 199c3fd6ab1b2..991bea4d2022d 100644
--- a/R/pkg/R/pairRDD.R
+++ b/R/pkg/R/pairRDD.R
@@ -21,23 +21,24 @@ NULL
 
 ############ Actions and Transformations ############
 
-# Look up elements of a key in an RDD
-#
-# @description
-# \code{lookup} returns a list of values in this RDD for key key.
-#
-# @param x The RDD to collect
-# @param key The key to look up for
-# @return a list of values in this RDD for key key
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# pairs <- list(c(1, 1), c(2, 2), c(1, 3))
-# rdd <- parallelize(sc, pairs)
-# lookup(rdd, 1) # list(1, 3)
-#}
-# @rdname lookup
-# @aliases lookup,RDD-method
+#' Look up elements of a key in an RDD
+#'
+#' @description
+#' \code{lookup} returns a list of values in this RDD for key key.
+#'
+#' @param x The RDD to collect
+#' @param key The key to look up for
+#' @return a list of values in this RDD for key key
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' pairs <- list(c(1, 1), c(2, 2), c(1, 3))
+#' rdd <- parallelize(sc, pairs)
+#' lookup(rdd, 1) # list(1, 3)
+#'}
+#' @rdname lookup
+#' @aliases lookup,RDD-method
+#' @noRd
 setMethod("lookup",
           signature(x = "RDD", key = "ANY"),
           function(x, key) {
@@ -49,21 +50,22 @@ setMethod("lookup",
             collect(valsRDD)
           })
 
-# Count the number of elements for each key, and return the result to the
-# master as lists of (key, count) pairs.
-#
-# Same as countByKey in Spark.
-#
-# @param x The RDD to count keys.
-# @return list of (key, count) pairs, where count is number of each key in rdd.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(c("a", 1), c("b", 1), c("a", 1)))
-# countByKey(rdd) # ("a", 2L), ("b", 1L)
-#}
-# @rdname countByKey
-# @aliases countByKey,RDD-method
+#' Count the number of elements for each key, and return the result to the
+#' master as lists of (key, count) pairs.
+#'
+#' Same as countByKey in Spark.
+#'
+#' @param x The RDD to count keys.
+#' @return list of (key, count) pairs, where count is number of each key in rdd.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(c("a", 1), c("b", 1), c("a", 1)))
+#' countByKey(rdd) # ("a", 2L), ("b", 1L)
+#'}
+#' @rdname countByKey
+#' @aliases countByKey,RDD-method
+#' @noRd
 setMethod("countByKey",
           signature(x = "RDD"),
           function(x) {
@@ -71,17 +73,18 @@ setMethod("countByKey",
             countByValue(keys)
           })
 
-# Return an RDD with the keys of each tuple.
-#
-# @param x The RDD from which the keys of each tuple is returned.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(list(1, 2), list(3, 4)))
-# collect(keys(rdd)) # list(1, 3)
-#}
-# @rdname keys
-# @aliases keys,RDD
+#' Return an RDD with the keys of each tuple.
+#'
+#' @param x The RDD from which the keys of each tuple is returned.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(list(1, 2), list(3, 4)))
+#' collect(keys(rdd)) # list(1, 3)
+#'}
+#' @rdname keys
+#' @aliases keys,RDD
+#' @noRd
 setMethod("keys",
           signature(x = "RDD"),
           function(x) {
@@ -91,17 +94,18 @@ setMethod("keys",
             lapply(x, func)
           })
 
-# Return an RDD with the values of each tuple.
-#
-# @param x The RDD from which the values of each tuple is returned.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(list(1, 2), list(3, 4)))
-# collect(values(rdd)) # list(2, 4)
-#}
-# @rdname values
-# @aliases values,RDD
+#' Return an RDD with the values of each tuple.
+#'
+#' @param x The RDD from which the values of each tuple is returned.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(list(1, 2), list(3, 4)))
+#' collect(values(rdd)) # list(2, 4)
+#'}
+#' @rdname values
+#' @aliases values,RDD
+#' @noRd
 setMethod("values",
           signature(x = "RDD"),
           function(x) {
@@ -111,23 +115,24 @@ setMethod("values",
             lapply(x, func)
           })
 
-# Applies a function to all values of the elements, without modifying the keys.
-#
-# The same as `mapValues()' in Spark.
-#
-# @param X The RDD to apply the transformation.
-# @param FUN the transformation to apply on the value of each element.
-# @return a new RDD created by the transformation.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:10)
-# makePairs <- lapply(rdd, function(x) { list(x, x) })
-# collect(mapValues(makePairs, function(x) { x * 2) })
-# Output: list(list(1,2), list(2,4), list(3,6), ...)
-#}
-# @rdname mapValues
-# @aliases mapValues,RDD,function-method
+#' Applies a function to all values of the elements, without modifying the keys.
+#'
+#' The same as `mapValues()' in Spark.
+#'
+#' @param X The RDD to apply the transformation.
+#' @param FUN the transformation to apply on the value of each element.
+#' @return a new RDD created by the transformation.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:10)
+#' makePairs <- lapply(rdd, function(x) { list(x, x) })
+#' collect(mapValues(makePairs, function(x) { x * 2) })
+#' Output: list(list(1,2), list(2,4), list(3,6), ...)
+#'}
+#' @rdname mapValues
+#' @aliases mapValues,RDD,function-method
+#' @noRd
 setMethod("mapValues",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
@@ -137,23 +142,24 @@ setMethod("mapValues",
             lapply(X, func)
           })
 
-# Pass each value in the key-value pair RDD through a flatMap function without
-# changing the keys; this also retains the original RDD's partitioning.
-#
-# The same as 'flatMapValues()' in Spark.
-#
-# @param X The RDD to apply the transformation.
-# @param FUN the transformation to apply on the value of each element.
-# @return a new RDD created by the transformation.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(list(1, c(1,2)), list(2, c(3,4))))
-# collect(flatMapValues(rdd, function(x) { x }))
-# Output: list(list(1,1), list(1,2), list(2,3), list(2,4))
-#}
-# @rdname flatMapValues
-# @aliases flatMapValues,RDD,function-method
+#' Pass each value in the key-value pair RDD through a flatMap function without
+#' changing the keys; this also retains the original RDD's partitioning.
+#'
+#' The same as 'flatMapValues()' in Spark.
+#'
+#' @param X The RDD to apply the transformation.
+#' @param FUN the transformation to apply on the value of each element.
+#' @return a new RDD created by the transformation.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(list(1, c(1,2)), list(2, c(3,4))))
+#' collect(flatMapValues(rdd, function(x) { x }))
+#' Output: list(list(1,1), list(1,2), list(2,3), list(2,4))
+#'}
+#' @rdname flatMapValues
+#' @aliases flatMapValues,RDD,function-method
+#' @noRd
 setMethod("flatMapValues",
           signature(X = "RDD", FUN = "function"),
           function(X, FUN) {
@@ -165,38 +171,34 @@ setMethod("flatMapValues",
 
 ############ Shuffle Functions ############
 
-# Partition an RDD by key
-#
-# This function operates on RDDs where every element is of the form list(K, V) or c(K, V).
-# For each element of this RDD, the partitioner is used to compute a hash
-# function and the RDD is partitioned using this hash value.
-#
-# @param x The RDD to partition. Should be an RDD where each element is
-#             list(K, V) or c(K, V).
-# @param numPartitions Number of partitions to create.
-# @param ... Other optional arguments to partitionBy.
-#
-# @param partitionFunc The partition function to use. Uses a default hashCode
-#                      function if not provided
-# @return An RDD partitioned using the specified partitioner.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
-# rdd <- parallelize(sc, pairs)
-# parts <- partitionBy(rdd, 2L)
-# collectPartition(parts, 0L) # First partition should contain list(1, 2) and list(1, 4)
-#}
-# @rdname partitionBy
-# @aliases partitionBy,RDD,integer-method
+#' Partition an RDD by key
+#'
+#' This function operates on RDDs where every element is of the form list(K, V) or c(K, V).
+#' For each element of this RDD, the partitioner is used to compute a hash
+#' function and the RDD is partitioned using this hash value.
+#'
+#' @param x The RDD to partition. Should be an RDD where each element is
+#'             list(K, V) or c(K, V).
+#' @param numPartitions Number of partitions to create.
+#' @param ... Other optional arguments to partitionBy.
+#'
+#' @param partitionFunc The partition function to use. Uses a default hashCode
+#'                      function if not provided
+#' @return An RDD partitioned using the specified partitioner.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
+#' rdd <- parallelize(sc, pairs)
+#' parts <- partitionBy(rdd, 2L)
+#' collectPartition(parts, 0L) # First partition should contain list(1, 2) and list(1, 4)
+#'}
+#' @rdname partitionBy
+#' @aliases partitionBy,RDD,integer-method
+#' @noRd
 setMethod("partitionBy",
           signature(x = "RDD", numPartitions = "numeric"),
           function(x, numPartitions, partitionFunc = hashCode) {
-
-            #if (missing(partitionFunc)) {
-            #  partitionFunc <- hashCode
-            #}
-
             partitionFunc <- cleanClosure(partitionFunc)
             serializedHashFuncBytes <- serialize(partitionFunc, connection = NULL)
 
@@ -233,27 +235,28 @@ setMethod("partitionBy",
             RDD(r, serializedMode = "byte")
           })
 
-# Group values by key
-#
-# This function operates on RDDs where every element is of the form list(K, V) or c(K, V).
-# and group values for each key in the RDD into a single sequence.
-#
-# @param x The RDD to group. Should be an RDD where each element is
-#             list(K, V) or c(K, V).
-# @param numPartitions Number of partitions to create.
-# @return An RDD where each element is list(K, list(V))
-# @seealso reduceByKey
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
-# rdd <- parallelize(sc, pairs)
-# parts <- groupByKey(rdd, 2L)
-# grouped <- collect(parts)
-# grouped[[1]] # Should be a list(1, list(2, 4))
-#}
-# @rdname groupByKey
-# @aliases groupByKey,RDD,integer-method
+#' Group values by key
+#'
+#' This function operates on RDDs where every element is of the form list(K, V) or c(K, V).
+#' and group values for each key in the RDD into a single sequence.
+#'
+#' @param x The RDD to group. Should be an RDD where each element is
+#'             list(K, V) or c(K, V).
+#' @param numPartitions Number of partitions to create.
+#' @return An RDD where each element is list(K, list(V))
+#' @seealso reduceByKey
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
+#' rdd <- parallelize(sc, pairs)
+#' parts <- groupByKey(rdd, 2L)
+#' grouped <- collect(parts)
+#' grouped[[1]] # Should be a list(1, list(2, 4))
+#'}
+#' @rdname groupByKey
+#' @aliases groupByKey,RDD,integer-method
+#' @noRd
 setMethod("groupByKey",
           signature(x = "RDD", numPartitions = "numeric"),
           function(x, numPartitions) {
@@ -291,28 +294,29 @@ setMethod("groupByKey",
             lapplyPartition(shuffled, groupVals)
           })
 
-# Merge values by key
-#
-# This function operates on RDDs where every element is of the form list(K, V) or c(K, V).
-# and merges the values for each key using an associative reduce function.
-#
-# @param x The RDD to reduce by key. Should be an RDD where each element is
-#             list(K, V) or c(K, V).
-# @param combineFunc The associative reduce function to use.
-# @param numPartitions Number of partitions to create.
-# @return An RDD where each element is list(K, V') where V' is the merged
-#         value
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
-# rdd <- parallelize(sc, pairs)
-# parts <- reduceByKey(rdd, "+", 2L)
-# reduced <- collect(parts)
-# reduced[[1]] # Should be a list(1, 6)
-#}
-# @rdname reduceByKey
-# @aliases reduceByKey,RDD,integer-method
+#'  Merge values by key
+#'
+#' This function operates on RDDs where every element is of the form list(K, V) or c(K, V).
+#' and merges the values for each key using an associative reduce function.
+#'
+#' @param x The RDD to reduce by key. Should be an RDD where each element is
+#'             list(K, V) or c(K, V).
+#' @param combineFunc The associative reduce function to use.
+#' @param numPartitions Number of partitions to create.
+#' @return An RDD where each element is list(K, V') where V' is the merged
+#'         value
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
+#' rdd <- parallelize(sc, pairs)
+#' parts <- reduceByKey(rdd, "+", 2L)
+#' reduced <- collect(parts)
+#' reduced[[1]] # Should be a list(1, 6)
+#'}
+#' @rdname reduceByKey
+#' @aliases reduceByKey,RDD,integer-method
+#' @noRd
 setMethod("reduceByKey",
           signature(x = "RDD", combineFunc = "ANY", numPartitions = "numeric"),
           function(x, combineFunc, numPartitions) {
@@ -332,27 +336,28 @@ setMethod("reduceByKey",
             lapplyPartition(shuffled, reduceVals)
           })
 
-# Merge values by key locally
-#
-# This function operates on RDDs where every element is of the form list(K, V) or c(K, V).
-# and merges the values for each key using an associative reduce function, but return the
-# results immediately to the driver as an R list.
-#
-# @param x The RDD to reduce by key. Should be an RDD where each element is
-#             list(K, V) or c(K, V).
-# @param combineFunc The associative reduce function to use.
-# @return A list of elements of type list(K, V') where V' is the merged value for each key
-# @seealso reduceByKey
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
-# rdd <- parallelize(sc, pairs)
-# reduced <- reduceByKeyLocally(rdd, "+")
-# reduced # list(list(1, 6), list(1.1, 3))
-#}
-# @rdname reduceByKeyLocally
-# @aliases reduceByKeyLocally,RDD,integer-method
+#' Merge values by key locally
+#'
+#' This function operates on RDDs where every element is of the form list(K, V) or c(K, V).
+#' and merges the values for each key using an associative reduce function, but return the
+#' results immediately to the driver as an R list.
+#'
+#' @param x The RDD to reduce by key. Should be an RDD where each element is
+#'             list(K, V) or c(K, V).
+#' @param combineFunc The associative reduce function to use.
+#' @return A list of elements of type list(K, V') where V' is the merged value for each key
+#' @seealso reduceByKey
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
+#' rdd <- parallelize(sc, pairs)
+#' reduced <- reduceByKeyLocally(rdd, "+")
+#' reduced # list(list(1, 6), list(1.1, 3))
+#'}
+#' @rdname reduceByKeyLocally
+#' @aliases reduceByKeyLocally,RDD,integer-method
+#' @noRd
 setMethod("reduceByKeyLocally",
           signature(x = "RDD", combineFunc = "ANY"),
           function(x, combineFunc) {
@@ -384,41 +389,40 @@ setMethod("reduceByKeyLocally",
             convertEnvsToList(merged[[1]], merged[[2]])
           })
 
-# Combine values by key
-#
-# Generic function to combine the elements for each key using a custom set of
-# aggregation functions. Turns an RDD[(K, V)] into a result of type RDD[(K, C)],
-# for a "combined type" C. Note that V and C can be different -- for example, one
-# might group an RDD of type (Int, Int) into an RDD of type (Int, Seq[Int]).
-
-# Users provide three functions:
-# \itemize{
-#   \item createCombiner, which turns a V into a C (e.g., creates a one-element list)
-#   \item mergeValue, to merge a V into a C (e.g., adds it to the end of a list) -
-#   \item mergeCombiners, to combine two C's into a single one (e.g., concatentates
-#    two lists).
-# }
-#
-# @param x The RDD to combine. Should be an RDD where each element is
-#             list(K, V) or c(K, V).
-# @param createCombiner Create a combiner (C) given a value (V)
-# @param mergeValue Merge the given value (V) with an existing combiner (C)
-# @param mergeCombiners Merge two combiners and return a new combiner
-# @param numPartitions Number of partitions to create.
-# @return An RDD where each element is list(K, C) where C is the combined type
-#
-# @seealso groupByKey, reduceByKey
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
-# rdd <- parallelize(sc, pairs)
-# parts <- combineByKey(rdd, function(x) { x }, "+", "+", 2L)
-# combined <- collect(parts)
-# combined[[1]] # Should be a list(1, 6)
-#}
-# @rdname combineByKey
-# @aliases combineByKey,RDD,ANY,ANY,ANY,integer-method
+#' Combine values by key
+#'
+#' Generic function to combine the elements for each key using a custom set of
+#' aggregation functions. Turns an RDD[(K, V)] into a result of type RDD[(K, C)],
+#' for a "combined type" C. Note that V and C can be different -- for example, one
+#' might group an RDD of type (Int, Int) into an RDD of type (Int, Seq[Int]).
+#' Users provide three functions:
+#' \itemize{
+#'   \item createCombiner, which turns a V into a C (e.g., creates a one-element list)
+#'   \item mergeValue, to merge a V into a C (e.g., adds it to the end of a list) -
+#'   \item mergeCombiners, to combine two C's into a single one (e.g., concatentates
+#'    two lists).
+#' }
+#'
+#' @param x The RDD to combine. Should be an RDD where each element is
+#'             list(K, V) or c(K, V).
+#' @param createCombiner Create a combiner (C) given a value (V)
+#' @param mergeValue Merge the given value (V) with an existing combiner (C)
+#' @param mergeCombiners Merge two combiners and return a new combiner
+#' @param numPartitions Number of partitions to create.
+#' @return An RDD where each element is list(K, C) where C is the combined type
+#' @seealso groupByKey, reduceByKey
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' pairs <- list(list(1, 2), list(1.1, 3), list(1, 4))
+#' rdd <- parallelize(sc, pairs)
+#' parts <- combineByKey(rdd, function(x) { x }, "+", "+", 2L)
+#' combined <- collect(parts)
+#' combined[[1]] # Should be a list(1, 6)
+#'}
+#' @rdname combineByKey
+#' @aliases combineByKey,RDD,ANY,ANY,ANY,integer-method
+#' @noRd
 setMethod("combineByKey",
           signature(x = "RDD", createCombiner = "ANY", mergeValue = "ANY",
                     mergeCombiners = "ANY", numPartitions = "numeric"),
@@ -450,36 +454,37 @@ setMethod("combineByKey",
             lapplyPartition(shuffled, mergeAfterShuffle)
           })
 
-# Aggregate a pair RDD by each key.
-#
-# Aggregate the values of each key in an RDD, using given combine functions
-# and a neutral "zero value". This function can return a different result type,
-# U, than the type of the values in this RDD, V. Thus, we need one operation
-# for merging a V into a U and one operation for merging two U's, The former
-# operation is used for merging values within a partition, and the latter is
-# used for merging values between partitions. To avoid memory allocation, both
-# of these functions are allowed to modify and return their first argument
-# instead of creating a new U.
-#
-# @param x An RDD.
-# @param zeroValue A neutral "zero value".
-# @param seqOp A function to aggregate the values of each key. It may return
-#              a different result type from the type of the values.
-# @param combOp A function to aggregate results of seqOp.
-# @return An RDD containing the aggregation result.
-# @seealso foldByKey, combineByKey
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(list(1, 1), list(1, 2), list(2, 3), list(2, 4)))
-# zeroValue <- list(0, 0)
-# seqOp <- function(x, y) { list(x[[1]] + y, x[[2]] + 1) }
-# combOp <- function(x, y) { list(x[[1]] + y[[1]], x[[2]] + y[[2]]) }
-# aggregateByKey(rdd, zeroValue, seqOp, combOp, 2L)
-#   # list(list(1, list(3, 2)), list(2, list(7, 2)))
-#}
-# @rdname aggregateByKey
-# @aliases aggregateByKey,RDD,ANY,ANY,ANY,integer-method
+#' Aggregate a pair RDD by each key.
+#'
+#' Aggregate the values of each key in an RDD, using given combine functions
+#' and a neutral "zero value". This function can return a different result type,
+#' U, than the type of the values in this RDD, V. Thus, we need one operation
+#' for merging a V into a U and one operation for merging two U's, The former
+#' operation is used for merging values within a partition, and the latter is
+#' used for merging values between partitions. To avoid memory allocation, both
+#' of these functions are allowed to modify and return their first argument
+#' instead of creating a new U.
+#'
+#' @param x An RDD.
+#' @param zeroValue A neutral "zero value".
+#' @param seqOp A function to aggregate the values of each key. It may return
+#'              a different result type from the type of the values.
+#' @param combOp A function to aggregate results of seqOp.
+#' @return An RDD containing the aggregation result.
+#' @seealso foldByKey, combineByKey
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(list(1, 1), list(1, 2), list(2, 3), list(2, 4)))
+#' zeroValue <- list(0, 0)
+#' seqOp <- function(x, y) { list(x[[1]] + y, x[[2]] + 1) }
+#' combOp <- function(x, y) { list(x[[1]] + y[[1]], x[[2]] + y[[2]]) }
+#' aggregateByKey(rdd, zeroValue, seqOp, combOp, 2L)
+#'   # list(list(1, list(3, 2)), list(2, list(7, 2)))
+#'}
+#' @rdname aggregateByKey
+#' @aliases aggregateByKey,RDD,ANY,ANY,ANY,integer-method
+#' @noRd
 setMethod("aggregateByKey",
           signature(x = "RDD", zeroValue = "ANY", seqOp = "ANY",
                     combOp = "ANY", numPartitions = "numeric"),
@@ -491,26 +496,27 @@ setMethod("aggregateByKey",
             combineByKey(x, createCombiner, seqOp, combOp, numPartitions)
           })
 
-# Fold a pair RDD by each key.
-#
-# Aggregate the values of each key in an RDD, using an associative function "func"
-# and a neutral "zero value" which may be added to the result an arbitrary
-# number of times, and must not change the result (e.g., 0 for addition, or
-# 1 for multiplication.).
-#
-# @param x An RDD.
-# @param zeroValue A neutral "zero value".
-# @param func An associative function for folding values of each key.
-# @return An RDD containing the aggregation result.
-# @seealso aggregateByKey, combineByKey
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(list(1, 1), list(1, 2), list(2, 3), list(2, 4)))
-# foldByKey(rdd, 0, "+", 2L) # list(list(1, 3), list(2, 7))
-#}
-# @rdname foldByKey
-# @aliases foldByKey,RDD,ANY,ANY,integer-method
+#' Fold a pair RDD by each key.
+#'
+#' Aggregate the values of each key in an RDD, using an associative function "func"
+#' and a neutral "zero value" which may be added to the result an arbitrary
+#' number of times, and must not change the result (e.g., 0 for addition, or
+#' 1 for multiplication.).
+#'
+#' @param x An RDD.
+#' @param zeroValue A neutral "zero value".
+#' @param func An associative function for folding values of each key.
+#' @return An RDD containing the aggregation result.
+#' @seealso aggregateByKey, combineByKey
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(list(1, 1), list(1, 2), list(2, 3), list(2, 4)))
+#' foldByKey(rdd, 0, "+", 2L) # list(list(1, 3), list(2, 7))
+#'}
+#' @rdname foldByKey
+#' @aliases foldByKey,RDD,ANY,ANY,integer-method
+#' @noRd
 setMethod("foldByKey",
           signature(x = "RDD", zeroValue = "ANY",
                     func = "ANY", numPartitions = "numeric"),
@@ -520,28 +526,29 @@ setMethod("foldByKey",
 
 ############ Binary Functions #############
 
-# Join two RDDs
-#
-# @description
-# \code{join} This function joins two RDDs where every element is of the form list(K, V).
-# The key types of the two RDDs should be the same.
-#
-# @param x An RDD to be joined. Should be an RDD where each element is
-#             list(K, V).
-# @param y An RDD to be joined. Should be an RDD where each element is
-#             list(K, V).
-# @param numPartitions Number of partitions to create.
-# @return a new RDD containing all pairs of elements with matching keys in
-#         two input RDDs.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, list(list(1, 1), list(2, 4)))
-# rdd2 <- parallelize(sc, list(list(1, 2), list(1, 3)))
-# join(rdd1, rdd2, 2L) # list(list(1, list(1, 2)), list(1, list(1, 3))
-#}
-# @rdname join-methods
-# @aliases join,RDD,RDD-method
+#' Join two RDDs
+#'
+#' @description
+#' \code{join} This function joins two RDDs where every element is of the form list(K, V).
+#' The key types of the two RDDs should be the same.
+#'
+#' @param x An RDD to be joined. Should be an RDD where each element is
+#'             list(K, V).
+#' @param y An RDD to be joined. Should be an RDD where each element is
+#'             list(K, V).
+#' @param numPartitions Number of partitions to create.
+#' @return a new RDD containing all pairs of elements with matching keys in
+#'         two input RDDs.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, list(list(1, 1), list(2, 4)))
+#' rdd2 <- parallelize(sc, list(list(1, 2), list(1, 3)))
+#' join(rdd1, rdd2, 2L) # list(list(1, list(1, 2)), list(1, list(1, 3))
+#'}
+#' @rdname join-methods
+#' @aliases join,RDD,RDD-method
+#' @noRd
 setMethod("join",
           signature(x = "RDD", y = "RDD"),
           function(x, y, numPartitions) {
@@ -556,30 +563,31 @@ setMethod("join",
                                     doJoin)
           })
 
-# Left outer join two RDDs
-#
-# @description
-# \code{leftouterjoin} This function left-outer-joins two RDDs where every element is of
-# the form list(K, V). The key types of the two RDDs should be the same.
-#
-# @param x An RDD to be joined. Should be an RDD where each element is
-#             list(K, V).
-# @param y An RDD to be joined. Should be an RDD where each element is
-#             list(K, V).
-# @param numPartitions Number of partitions to create.
-# @return For each element (k, v) in x, the resulting RDD will either contain
-#         all pairs (k, (v, w)) for (k, w) in rdd2, or the pair (k, (v, NULL))
-#         if no elements in rdd2 have key k.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, list(list(1, 1), list(2, 4)))
-# rdd2 <- parallelize(sc, list(list(1, 2), list(1, 3)))
-# leftOuterJoin(rdd1, rdd2, 2L)
-# # list(list(1, list(1, 2)), list(1, list(1, 3)), list(2, list(4, NULL)))
-#}
-# @rdname join-methods
-# @aliases leftOuterJoin,RDD,RDD-method
+#' Left outer join two RDDs
+#'
+#' @description
+#' \code{leftouterjoin} This function left-outer-joins two RDDs where every element is of
+#' the form list(K, V). The key types of the two RDDs should be the same.
+#'
+#' @param x An RDD to be joined. Should be an RDD where each element is
+#'             list(K, V).
+#' @param y An RDD to be joined. Should be an RDD where each element is
+#'             list(K, V).
+#' @param numPartitions Number of partitions to create.
+#' @return For each element (k, v) in x, the resulting RDD will either contain
+#'         all pairs (k, (v, w)) for (k, w) in rdd2, or the pair (k, (v, NULL))
+#'         if no elements in rdd2 have key k.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, list(list(1, 1), list(2, 4)))
+#' rdd2 <- parallelize(sc, list(list(1, 2), list(1, 3)))
+#' leftOuterJoin(rdd1, rdd2, 2L)
+#' # list(list(1, list(1, 2)), list(1, list(1, 3)), list(2, list(4, NULL)))
+#'}
+#' @rdname join-methods
+#' @aliases leftOuterJoin,RDD,RDD-method
+#' @noRd
 setMethod("leftOuterJoin",
           signature(x = "RDD", y = "RDD", numPartitions = "numeric"),
           function(x, y, numPartitions) {
@@ -593,30 +601,31 @@ setMethod("leftOuterJoin",
             joined <- flatMapValues(groupByKey(unionRDD(xTagged, yTagged), numPartitions), doJoin)
           })
 
-# Right outer join two RDDs
-#
-# @description
-# \code{rightouterjoin} This function right-outer-joins two RDDs where every element is of
-# the form list(K, V). The key types of the two RDDs should be the same.
-#
-# @param x An RDD to be joined. Should be an RDD where each element is
-#             list(K, V).
-# @param y An RDD to be joined. Should be an RDD where each element is
-#             list(K, V).
-# @param numPartitions Number of partitions to create.
-# @return For each element (k, w) in y, the resulting RDD will either contain
-#         all pairs (k, (v, w)) for (k, v) in x, or the pair (k, (NULL, w))
-#         if no elements in x have key k.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, list(list(1, 2), list(1, 3)))
-# rdd2 <- parallelize(sc, list(list(1, 1), list(2, 4)))
-# rightOuterJoin(rdd1, rdd2, 2L)
-# # list(list(1, list(2, 1)), list(1, list(3, 1)), list(2, list(NULL, 4)))
-#}
-# @rdname join-methods
-# @aliases rightOuterJoin,RDD,RDD-method
+#' Right outer join two RDDs
+#'
+#' @description
+#' \code{rightouterjoin} This function right-outer-joins two RDDs where every element is of
+#' the form list(K, V). The key types of the two RDDs should be the same.
+#'
+#' @param x An RDD to be joined. Should be an RDD where each element is
+#'             list(K, V).
+#' @param y An RDD to be joined. Should be an RDD where each element is
+#'             list(K, V).
+#' @param numPartitions Number of partitions to create.
+#' @return For each element (k, w) in y, the resulting RDD will either contain
+#'         all pairs (k, (v, w)) for (k, v) in x, or the pair (k, (NULL, w))
+#'         if no elements in x have key k.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, list(list(1, 2), list(1, 3)))
+#' rdd2 <- parallelize(sc, list(list(1, 1), list(2, 4)))
+#' rightOuterJoin(rdd1, rdd2, 2L)
+#' # list(list(1, list(2, 1)), list(1, list(3, 1)), list(2, list(NULL, 4)))
+#'}
+#' @rdname join-methods
+#' @aliases rightOuterJoin,RDD,RDD-method
+#' @noRd
 setMethod("rightOuterJoin",
           signature(x = "RDD", y = "RDD", numPartitions = "numeric"),
           function(x, y, numPartitions) {
@@ -630,33 +639,34 @@ setMethod("rightOuterJoin",
             joined <- flatMapValues(groupByKey(unionRDD(xTagged, yTagged), numPartitions), doJoin)
           })
 
-# Full outer join two RDDs
-#
-# @description
-# \code{fullouterjoin} This function full-outer-joins two RDDs where every element is of
-# the form list(K, V). The key types of the two RDDs should be the same.
-#
-# @param x An RDD to be joined. Should be an RDD where each element is
-#             list(K, V).
-# @param y An RDD to be joined. Should be an RDD where each element is
-#             list(K, V).
-# @param numPartitions Number of partitions to create.
-# @return For each element (k, v) in x and (k, w) in y, the resulting RDD
-#         will contain all pairs (k, (v, w)) for both (k, v) in x and
-#         (k, w) in y, or the pair (k, (NULL, w))/(k, (v, NULL)) if no elements
-#         in x/y have key k.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, list(list(1, 2), list(1, 3), list(3, 3)))
-# rdd2 <- parallelize(sc, list(list(1, 1), list(2, 4)))
-# fullOuterJoin(rdd1, rdd2, 2L) # list(list(1, list(2, 1)),
-#                               #      list(1, list(3, 1)),
-#                               #      list(2, list(NULL, 4)))
-#                               #      list(3, list(3, NULL)),
-#}
-# @rdname join-methods
-# @aliases fullOuterJoin,RDD,RDD-method
+#' Full outer join two RDDs
+#'
+#' @description
+#' \code{fullouterjoin} This function full-outer-joins two RDDs where every element is of
+#' the form list(K, V). The key types of the two RDDs should be the same.
+#'
+#' @param x An RDD to be joined. Should be an RDD where each element is
+#'             list(K, V).
+#' @param y An RDD to be joined. Should be an RDD where each element is
+#'             list(K, V).
+#' @param numPartitions Number of partitions to create.
+#' @return For each element (k, v) in x and (k, w) in y, the resulting RDD
+#'         will contain all pairs (k, (v, w)) for both (k, v) in x and
+#'         (k, w) in y, or the pair (k, (NULL, w))/(k, (v, NULL)) if no elements
+#'         in x/y have key k.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, list(list(1, 2), list(1, 3), list(3, 3)))
+#' rdd2 <- parallelize(sc, list(list(1, 1), list(2, 4)))
+#' fullOuterJoin(rdd1, rdd2, 2L) # list(list(1, list(2, 1)),
+#'                               #      list(1, list(3, 1)),
+#'                               #      list(2, list(NULL, 4)))
+#'                               #      list(3, list(3, NULL)),
+#'}
+#' @rdname join-methods
+#' @aliases fullOuterJoin,RDD,RDD-method
+#' @noRd
 setMethod("fullOuterJoin",
           signature(x = "RDD", y = "RDD", numPartitions = "numeric"),
           function(x, y, numPartitions) {
@@ -670,23 +680,24 @@ setMethod("fullOuterJoin",
             joined <- flatMapValues(groupByKey(unionRDD(xTagged, yTagged), numPartitions), doJoin)
           })
 
-# For each key k in several RDDs, return a resulting RDD that
-# whose values are a list of values for the key in all RDDs.
-#
-# @param ... Several RDDs.
-# @param numPartitions Number of partitions to create.
-# @return a new RDD containing all pairs of elements with values in a list
-# in all RDDs.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, list(list(1, 1), list(2, 4)))
-# rdd2 <- parallelize(sc, list(list(1, 2), list(1, 3)))
-# cogroup(rdd1, rdd2, numPartitions = 2L)
-# # list(list(1, list(1, list(2, 3))), list(2, list(list(4), list()))
-#}
-# @rdname cogroup
-# @aliases cogroup,RDD-method
+#' For each key k in several RDDs, return a resulting RDD that
+#' whose values are a list of values for the key in all RDDs.
+#'
+#' @param ... Several RDDs.
+#' @param numPartitions Number of partitions to create.
+#' @return a new RDD containing all pairs of elements with values in a list
+#' in all RDDs.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, list(list(1, 1), list(2, 4)))
+#' rdd2 <- parallelize(sc, list(list(1, 2), list(1, 3)))
+#' cogroup(rdd1, rdd2, numPartitions = 2L)
+#' # list(list(1, list(1, list(2, 3))), list(2, list(list(4), list()))
+#'}
+#' @rdname cogroup
+#' @aliases cogroup,RDD-method
+#' @noRd
 setMethod("cogroup",
           "RDD",
           function(..., numPartitions) {
@@ -722,20 +733,21 @@ setMethod("cogroup",
                                      group.func)
           })
 
-# Sort a (k, v) pair RDD by k.
-#
-# @param x A (k, v) pair RDD to be sorted.
-# @param ascending A flag to indicate whether the sorting is ascending or descending.
-# @param numPartitions Number of partitions to create.
-# @return An RDD where all (k, v) pair elements are sorted.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, list(list(3, 1), list(2, 2), list(1, 3)))
-# collect(sortByKey(rdd)) # list (list(1, 3), list(2, 2), list(3, 1))
-#}
-# @rdname sortByKey
-# @aliases sortByKey,RDD,RDD-method
+#' Sort a (k, v) pair RDD by k.
+#'
+#' @param x A (k, v) pair RDD to be sorted.
+#' @param ascending A flag to indicate whether the sorting is ascending or descending.
+#' @param numPartitions Number of partitions to create.
+#' @return An RDD where all (k, v) pair elements are sorted.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, list(list(3, 1), list(2, 2), list(1, 3)))
+#' collect(sortByKey(rdd)) # list (list(1, 3), list(2, 2), list(3, 1))
+#'}
+#' @rdname sortByKey
+#' @aliases sortByKey,RDD,RDD-method
+#' @noRd
 setMethod("sortByKey",
           signature(x = "RDD"),
           function(x, ascending = TRUE, numPartitions = SparkR:::numPartitions(x)) {
@@ -784,25 +796,26 @@ setMethod("sortByKey",
             lapplyPartition(newRDD, partitionFunc)
           })
 
-# Subtract a pair RDD with another pair RDD.
-#
-# Return an RDD with the pairs from x whose keys are not in other.
-#
-# @param x An RDD.
-# @param other An RDD.
-# @param numPartitions Number of the partitions in the result RDD.
-# @return An RDD with the pairs from x whose keys are not in other.
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd1 <- parallelize(sc, list(list("a", 1), list("b", 4),
-#                              list("b", 5), list("a", 2)))
-# rdd2 <- parallelize(sc, list(list("a", 3), list("c", 1)))
-# collect(subtractByKey(rdd1, rdd2))
-# # list(list("b", 4), list("b", 5))
-#}
-# @rdname subtractByKey
-# @aliases subtractByKey,RDD
+#' Subtract a pair RDD with another pair RDD.
+#'
+#' Return an RDD with the pairs from x whose keys are not in other.
+#'
+#' @param x An RDD.
+#' @param other An RDD.
+#' @param numPartitions Number of the partitions in the result RDD.
+#' @return An RDD with the pairs from x whose keys are not in other.
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd1 <- parallelize(sc, list(list("a", 1), list("b", 4),
+#'                              list("b", 5), list("a", 2)))
+#' rdd2 <- parallelize(sc, list(list("a", 3), list("c", 1)))
+#' collect(subtractByKey(rdd1, rdd2))
+#' # list(list("b", 4), list("b", 5))
+#'}
+#' @rdname subtractByKey
+#' @aliases subtractByKey,RDD
+#' @noRd
 setMethod("subtractByKey",
           signature(x = "RDD", other = "RDD"),
           function(x, other, numPartitions = SparkR:::numPartitions(x)) {
@@ -818,41 +831,42 @@ setMethod("subtractByKey",
                           function (v) { v[[1]] })
           })
 
-# Return a subset of this RDD sampled by key.
-#
-# @description
-# \code{sampleByKey} Create a sample of this RDD using variable sampling rates
-# for different keys as specified by fractions, a key to sampling rate map.
-#
-# @param x The RDD to sample elements by key, where each element is
-#             list(K, V) or c(K, V).
-# @param withReplacement Sampling with replacement or not
-# @param fraction The (rough) sample target fraction
-# @param seed Randomness seed value
-# @examples
-#\dontrun{
-# sc <- sparkR.init()
-# rdd <- parallelize(sc, 1:3000)
-# pairs <- lapply(rdd, function(x) { if (x %% 3 == 0) list("a", x)
-#                                    else { if (x %% 3 == 1) list("b", x) else list("c", x) }})
-# fractions <- list(a = 0.2, b = 0.1, c = 0.3)
-# sample <- sampleByKey(pairs, FALSE, fractions, 1618L)
-# 100 < length(lookup(sample, "a")) && 300 > length(lookup(sample, "a")) # TRUE
-# 50 < length(lookup(sample, "b")) && 150 > length(lookup(sample, "b")) # TRUE
-# 200 < length(lookup(sample, "c")) && 400 > length(lookup(sample, "c")) # TRUE
-# lookup(sample, "a")[which.min(lookup(sample, "a"))] >= 0 # TRUE
-# lookup(sample, "a")[which.max(lookup(sample, "a"))] <= 2000 # TRUE
-# lookup(sample, "b")[which.min(lookup(sample, "b"))] >= 0 # TRUE
-# lookup(sample, "b")[which.max(lookup(sample, "b"))] <= 2000 # TRUE
-# lookup(sample, "c")[which.min(lookup(sample, "c"))] >= 0 # TRUE
-# lookup(sample, "c")[which.max(lookup(sample, "c"))] <= 2000 # TRUE
-# fractions <- list(a = 0.2, b = 0.1, c = 0.3, d = 0.4)
-# sample <- sampleByKey(pairs, FALSE, fractions, 1618L) # Key "d" will be ignored
-# fractions <- list(a = 0.2, b = 0.1)
-# sample <- sampleByKey(pairs, FALSE, fractions, 1618L) # KeyError: "c"
-#}
-# @rdname sampleByKey
-# @aliases sampleByKey,RDD-method
+#' Return a subset of this RDD sampled by key.
+#'
+#' @description
+#' \code{sampleByKey} Create a sample of this RDD using variable sampling rates
+#' for different keys as specified by fractions, a key to sampling rate map.
+#'
+#' @param x The RDD to sample elements by key, where each element is
+#'             list(K, V) or c(K, V).
+#' @param withReplacement Sampling with replacement or not
+#' @param fraction The (rough) sample target fraction
+#' @param seed Randomness seed value
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' rdd <- parallelize(sc, 1:3000)
+#' pairs <- lapply(rdd, function(x) { if (x %% 3 == 0) list("a", x)
+#'                                    else { if (x %% 3 == 1) list("b", x) else list("c", x) }})
+#' fractions <- list(a = 0.2, b = 0.1, c = 0.3)
+#' sample <- sampleByKey(pairs, FALSE, fractions, 1618L)
+#' 100 < length(lookup(sample, "a")) && 300 > length(lookup(sample, "a")) # TRUE
+#' 50 < length(lookup(sample, "b")) && 150 > length(lookup(sample, "b")) # TRUE
+#' 200 < length(lookup(sample, "c")) && 400 > length(lookup(sample, "c")) # TRUE
+#' lookup(sample, "a")[which.min(lookup(sample, "a"))] >= 0 # TRUE
+#' lookup(sample, "a")[which.max(lookup(sample, "a"))] <= 2000 # TRUE
+#' lookup(sample, "b")[which.min(lookup(sample, "b"))] >= 0 # TRUE
+#' lookup(sample, "b")[which.max(lookup(sample, "b"))] <= 2000 # TRUE
+#' lookup(sample, "c")[which.min(lookup(sample, "c"))] >= 0 # TRUE
+#' lookup(sample, "c")[which.max(lookup(sample, "c"))] <= 2000 # TRUE
+#' fractions <- list(a = 0.2, b = 0.1, c = 0.3, d = 0.4)
+#' sample <- sampleByKey(pairs, FALSE, fractions, 1618L) # Key "d" will be ignored
+#' fractions <- list(a = 0.2, b = 0.1)
+#' sample <- sampleByKey(pairs, FALSE, fractions, 1618L) # KeyError: "c"
+#'}
+#' @rdname sampleByKey
+#' @aliases sampleByKey,RDD-method
+#' @noRd
 setMethod("sampleByKey",
           signature(x = "RDD", withReplacement = "logical",
                     fractions = "vector", seed = "integer"),
diff --git a/R/pkg/R/sparkR.R b/R/pkg/R/sparkR.R
index 004d08e74e1cd..ebe2b2b8dc1d0 100644
--- a/R/pkg/R/sparkR.R
+++ b/R/pkg/R/sparkR.R
@@ -34,7 +34,6 @@ connExists <- function(env) {
 sparkR.stop <- function() {
   env <- .sparkREnv
   if (exists(".sparkRCon", envir = env)) {
-    # cat("Stopping SparkR\n")
     if (exists(".sparkRjsc", envir = env)) {
       sc <- get(".sparkRjsc", envir = env)
       callJMethod(sc, "stop")
@@ -78,7 +77,7 @@ sparkR.stop <- function() {
 #' Initialize a new Spark Context.
 #'
 #' This function initializes a new SparkContext. For details on how to initialize
-#' and use SparkR, refer to SparkR programming guide at 
+#' and use SparkR, refer to SparkR programming guide at
 #' \url{http://spark.apache.org/docs/latest/sparkr.html#starting-up-sparkcontext-sqlcontext}.
 #'
 #' @param master The Spark master URL.
diff --git a/R/pkg/inst/profile/shell.R b/R/pkg/inst/profile/shell.R
index 7189f1a260934..90a3761e41f82 100644
--- a/R/pkg/inst/profile/shell.R
+++ b/R/pkg/inst/profile/shell.R
@@ -38,7 +38,7 @@
   if (nchar(sparkVer) == 0) {
     cat("\n")
   } else {
-    cat("   version ", sparkVer, "\n") 
+    cat("   version ", sparkVer, "\n")
   }
   cat("    /_/", "\n")
   cat("\n")

From 2035ed392e0a9c18ff9c176a7b0f0097ed1276df Mon Sep 17 00:00:00 2001
From: Lewuathe <lewuathe@me.com>
Date: Thu, 12 Nov 2015 20:09:42 -0800
Subject: [PATCH 227/518] [SPARK-11717] Ignore R session and history files from
 git

see: https://issues.apache.org/jira/browse/SPARK-11717

SparkR generates R session data and history files under current directory.
It might be useful to ignore these files even running SparkR on spark directory for test or development.

Author: Lewuathe <lewuathe@me.com>

Closes #9681 from Lewuathe/SPARK-11717.
---
 .gitignore | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.gitignore b/.gitignore
index debad77ec2ad3..08f2d8f7543f0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -74,3 +74,7 @@ metastore/
 warehouse/
 TempStatsStore/
 sql/hive-thriftserver/test_warehouses
+
+# For R session data
+.RHistory
+.RData

From ea5ae2705afa4eaadd4192c37d74c97364378cf9 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 12 Nov 2015 21:29:43 -0800
Subject: [PATCH 228/518] [SPARK-11629][ML][PYSPARK][DOC] Python example code
 for Multilayer Perceptron Classification

Add Python example code for Multilayer Perceptron Classification, and make example code in user guide document testable. mengxr

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9594 from yanboliang/spark-11629.
---
 docs/ml-ann.md                                | 71 ++----------------
 ...MultilayerPerceptronClassifierExample.java | 74 +++++++++++++++++++
 .../multilayer_perceptron_classification.py   | 56 ++++++++++++++
 ...ultilayerPerceptronClassifierExample.scala | 71 ++++++++++++++++++
 4 files changed, 206 insertions(+), 66 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java
 create mode 100644 examples/src/main/python/ml/multilayer_perceptron_classification.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala

diff --git a/docs/ml-ann.md b/docs/ml-ann.md
index d5ddd92af1e96..6e763e8f41568 100644
--- a/docs/ml-ann.md
+++ b/docs/ml-ann.md
@@ -48,76 +48,15 @@ MLPC employes backpropagation for learning the model. We use logistic loss funct
 <div class="codetabs">
 
 <div data-lang="scala" markdown="1">
-
-{% highlight scala %}
-import org.apache.spark.ml.classification.MultilayerPerceptronClassifier
-import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
-import org.apache.spark.mllib.util.MLUtils
-import org.apache.spark.sql.Row
-
-// Load training data
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt").toDF()
-// Split the data into train and test
-val splits = data.randomSplit(Array(0.6, 0.4), seed = 1234L)
-val train = splits(0)
-val test = splits(1)
-// specify layers for the neural network: 
-// input layer of size 4 (features), two intermediate of size 5 and 4 and output of size 3 (classes)
-val layers = Array[Int](4, 5, 4, 3)
-// create the trainer and set its parameters
-val trainer = new MultilayerPerceptronClassifier()
-  .setLayers(layers)
-  .setBlockSize(128)
-  .setSeed(1234L)
-  .setMaxIter(100)
-// train the model
-val model = trainer.fit(train)
-// compute precision on the test set
-val result = model.transform(test)
-val predictionAndLabels = result.select("prediction", "label")
-val evaluator = new MulticlassClassificationEvaluator()
-  .setMetricName("precision")
-println("Precision:" + evaluator.evaluate(predictionAndLabels))
-{% endhighlight %}
-
+{% include_example scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
+{% include_example java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java %}
+</div>
 
-{% highlight java %}
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.ml.classification.MultilayerPerceptronClassificationModel;
-import org.apache.spark.ml.classification.MultilayerPerceptronClassifier;
-import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-
-// Load training data
-String path = "data/mllib/sample_multiclass_classification_data.txt";
-JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD();
-DataFrame dataFrame = sqlContext.createDataFrame(data, LabeledPoint.class);
-// Split the data into train and test
-DataFrame[] splits = dataFrame.randomSplit(new double[]{0.6, 0.4}, 1234L);
-DataFrame train = splits[0];
-DataFrame test = splits[1];
-// specify layers for the neural network:
-// input layer of size 4 (features), two intermediate of size 5 and 4 and output of size 3 (classes)
-int[] layers = new int[] {4, 5, 4, 3};
-// create the trainer and set its parameters
-MultilayerPerceptronClassifier trainer = new MultilayerPerceptronClassifier()
-  .setLayers(layers)
-  .setBlockSize(128)
-  .setSeed(1234L)
-  .setMaxIter(100);
-// train the model
-MultilayerPerceptronClassificationModel model = trainer.fit(train);
-// compute precision on the test set
-DataFrame result = model.transform(test);
-DataFrame predictionAndLabels = result.select("prediction", "label");
-MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
-  .setMetricName("precision");
-System.out.println("Precision = " + evaluator.evaluate(predictionAndLabels));
-{% endhighlight %}
+<div data-lang="python" markdown="1">
+{% include_example python/ml/multilayer_perceptron_classification.py %}
 </div>
 
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java
new file mode 100644
index 0000000000000..f48e1339c5007
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java
@@ -0,0 +1,74 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+// $example on$
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.ml.classification.MultilayerPerceptronClassificationModel;
+import org.apache.spark.ml.classification.MultilayerPerceptronClassifier;
+import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.util.MLUtils;
+import org.apache.spark.sql.DataFrame;
+// $example off$
+
+/**
+ * An example for Multilayer Perceptron Classification.
+ */
+public class JavaMultilayerPerceptronClassifierExample {
+
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaMultilayerPerceptronClassifierExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext jsql = new SQLContext(jsc);
+
+    // $example on$
+    // Load training data
+    String path = "data/mllib/sample_multiclass_classification_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), path).toJavaRDD();
+    DataFrame dataFrame = jsql.createDataFrame(data, LabeledPoint.class);
+    // Split the data into train and test
+    DataFrame[] splits = dataFrame.randomSplit(new double[]{0.6, 0.4}, 1234L);
+    DataFrame train = splits[0];
+    DataFrame test = splits[1];
+    // specify layers for the neural network:
+    // input layer of size 4 (features), two intermediate of size 5 and 4
+    // and output of size 3 (classes)
+    int[] layers = new int[] {4, 5, 4, 3};
+    // create the trainer and set its parameters
+    MultilayerPerceptronClassifier trainer = new MultilayerPerceptronClassifier()
+      .setLayers(layers)
+      .setBlockSize(128)
+      .setSeed(1234L)
+      .setMaxIter(100);
+    // train the model
+    MultilayerPerceptronClassificationModel model = trainer.fit(train);
+    // compute precision on the test set
+    DataFrame result = model.transform(test);
+    DataFrame predictionAndLabels = result.select("prediction", "label");
+    MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
+      .setMetricName("precision");
+    System.out.println("Precision = " + evaluator.evaluate(predictionAndLabels));
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/python/ml/multilayer_perceptron_classification.py b/examples/src/main/python/ml/multilayer_perceptron_classification.py
new file mode 100644
index 0000000000000..d8ef9f39e3fa5
--- /dev/null
+++ b/examples/src/main/python/ml/multilayer_perceptron_classification.py
@@ -0,0 +1,56 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import print_function
+
+from pyspark import SparkContext
+from pyspark.sql import SQLContext
+# $example on$
+from pyspark.ml.classification import MultilayerPerceptronClassifier
+from pyspark.ml.evaluation import MulticlassClassificationEvaluator
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+
+    sc = SparkContext(appName="multilayer_perceptron_classification_example")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load training data
+    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt")\
+        .toDF()
+    # Split the data into train and test
+    splits = data.randomSplit([0.6, 0.4], 1234)
+    train = splits[0]
+    test = splits[1]
+    # specify layers for the neural network:
+    # input layer of size 4 (features), two intermediate of size 5 and 4
+    # and output of size 3 (classes)
+    layers = [4, 5, 4, 3]
+    # create the trainer and set its parameters
+    trainer = MultilayerPerceptronClassifier(maxIter=100, layers=layers, blockSize=128, seed=1234)
+    # train the model
+    model = trainer.fit(train)
+    # compute precision on the test set
+    result = model.transform(test)
+    predictionAndLabels = result.select("prediction", "label")
+    evaluator = MulticlassClassificationEvaluator(metricName="precision")
+    print("Precision:" + str(evaluator.evaluate(predictionAndLabels)))
+    # $example off$
+
+    sc.stop()
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala
new file mode 100644
index 0000000000000..99d5f35b5a565
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+import org.apache.spark.{SparkContext, SparkConf}
+import org.apache.spark.sql.SQLContext
+// $example on$
+import org.apache.spark.ml.classification.MultilayerPerceptronClassifier
+import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+
+/**
+ * An example for Multilayer Perceptron Classification.
+ */
+object MultilayerPerceptronClassifierExample {
+
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("MultilayerPerceptronClassifierExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+    import sqlContext.implicits._
+
+    // $example on$
+    // Load training data
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt")
+      .toDF()
+    // Split the data into train and test
+    val splits = data.randomSplit(Array(0.6, 0.4), seed = 1234L)
+    val train = splits(0)
+    val test = splits(1)
+    // specify layers for the neural network:
+    // input layer of size 4 (features), two intermediate of size 5 and 4
+    // and output of size 3 (classes)
+    val layers = Array[Int](4, 5, 4, 3)
+    // create the trainer and set its parameters
+    val trainer = new MultilayerPerceptronClassifier()
+      .setLayers(layers)
+      .setBlockSize(128)
+      .setSeed(1234L)
+      .setMaxIter(100)
+    // train the model
+    val model = trainer.fit(train)
+    // compute precision on the test set
+    val result = model.transform(test)
+    val predictionAndLabels = result.select("prediction", "label")
+    val evaluator = new MulticlassClassificationEvaluator()
+      .setMetricName("precision")
+    println("Precision:" + evaluator.evaluate(predictionAndLabels))
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:off println

From ad960885bfee7850c18eb5338546cecf2b2e9876 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Thu, 12 Nov 2015 22:44:57 -0800
Subject: [PATCH 229/518] [SPARK-8029] Robust shuffle writer

Currently, all the shuffle writer will write to target path directly, the file could be corrupted by other attempt of the same partition on the same executor. They should write to temporary file then rename to target path, as what we do in output committer. In order to make the rename atomic, the temporary file should be created in the same local directory (FileSystem).

This PR is based on #9214 , thanks to squito . Closes #9214

Author: Davies Liu <davies@databricks.com>

Closes #9610 from davies/safe_shuffle.
---
 .../sort/BypassMergeSortShuffleWriter.java    |   9 +-
 .../shuffle/sort/UnsafeShuffleWriter.java     |  13 +-
 .../shuffle/FileShuffleBlockResolver.scala    |  17 +--
 .../shuffle/IndexShuffleBlockResolver.scala   | 102 ++++++++++++++--
 .../shuffle/hash/HashShuffleWriter.scala      |  25 ++++
 .../shuffle/sort/SortShuffleWriter.scala      |  11 +-
 .../apache/spark/storage/BlockManager.scala   |   9 +-
 .../spark/storage/DiskBlockObjectWriter.scala |   5 +-
 .../scala/org/apache/spark/util/Utils.scala   |  12 +-
 .../util/collection/ExternalSorter.scala      |   1 -
 .../sort/IndexShuffleBlockResolverSuite.scala | 114 ++++++++++++++++++
 .../sort/UnsafeShuffleWriterSuite.java        |   9 +-
 .../map/AbstractBytesToBytesMapSuite.java     |   3 +-
 .../sort/UnsafeExternalSorterSuite.java       |   3 +-
 .../scala/org/apache/spark/ShuffleSuite.scala | 107 +++++++++++++++-
 .../BypassMergeSortShuffleWriterSuite.scala   |  14 ++-
 16 files changed, 402 insertions(+), 52 deletions(-)
 create mode 100644 core/src/test/java/org/apache/spark/shuffle/sort/IndexShuffleBlockResolverSuite.scala

diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java
index ee82d679935c0..a1a1fb01426a0 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java
@@ -125,7 +125,7 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
     assert (partitionWriters == null);
     if (!records.hasNext()) {
       partitionLengths = new long[numPartitions];
-      shuffleBlockResolver.writeIndexFile(shuffleId, mapId, partitionLengths);
+      shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, null);
       mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
       return;
     }
@@ -155,9 +155,10 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
       writer.commitAndClose();
     }
 
-    partitionLengths =
-      writePartitionedFile(shuffleBlockResolver.getDataFile(shuffleId, mapId));
-    shuffleBlockResolver.writeIndexFile(shuffleId, mapId, partitionLengths);
+    File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);
+    File tmp = Utils.tempFileWith(output);
+    partitionLengths = writePartitionedFile(tmp);
+    shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);
     mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
   }
 
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
index 6a0a89e81c321..744c3008ca50e 100644
--- a/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
@@ -41,7 +41,7 @@
 import org.apache.spark.executor.ShuffleWriteMetrics;
 import org.apache.spark.io.CompressionCodec;
 import org.apache.spark.io.CompressionCodec$;
-import org.apache.spark.io.LZFCompressionCodec;
+import org.apache.spark.memory.TaskMemoryManager;
 import org.apache.spark.network.util.LimitedInputStream;
 import org.apache.spark.scheduler.MapStatus;
 import org.apache.spark.scheduler.MapStatus$;
@@ -53,7 +53,7 @@
 import org.apache.spark.storage.BlockManager;
 import org.apache.spark.storage.TimeTrackingOutputStream;
 import org.apache.spark.unsafe.Platform;
-import org.apache.spark.memory.TaskMemoryManager;
+import org.apache.spark.util.Utils;
 
 @Private
 public class UnsafeShuffleWriter<K, V> extends ShuffleWriter<K, V> {
@@ -206,8 +206,10 @@ void closeAndWriteOutput() throws IOException {
     final SpillInfo[] spills = sorter.closeAndGetSpills();
     sorter = null;
     final long[] partitionLengths;
+    final File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);
+    final File tmp = Utils.tempFileWith(output);
     try {
-      partitionLengths = mergeSpills(spills);
+      partitionLengths = mergeSpills(spills, tmp);
     } finally {
       for (SpillInfo spill : spills) {
         if (spill.file.exists() && ! spill.file.delete()) {
@@ -215,7 +217,7 @@ void closeAndWriteOutput() throws IOException {
         }
       }
     }
-    shuffleBlockResolver.writeIndexFile(shuffleId, mapId, partitionLengths);
+    shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);
     mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
   }
 
@@ -248,8 +250,7 @@ void forceSorterToSpill() throws IOException {
    *
    * @return the partition lengths in the merged file.
    */
-  private long[] mergeSpills(SpillInfo[] spills) throws IOException {
-    final File outputFile = shuffleBlockResolver.getDataFile(shuffleId, mapId);
+  private long[] mergeSpills(SpillInfo[] spills, File outputFile) throws IOException {
     final boolean compressionEnabled = sparkConf.getBoolean("spark.shuffle.compress", true);
     final CompressionCodec compressionCodec = CompressionCodec$.MODULE$.createCodec(sparkConf);
     final boolean fastMergeEnabled =
diff --git a/core/src/main/scala/org/apache/spark/shuffle/FileShuffleBlockResolver.scala b/core/src/main/scala/org/apache/spark/shuffle/FileShuffleBlockResolver.scala
index cd253a78c2b19..39fadd8783518 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/FileShuffleBlockResolver.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/FileShuffleBlockResolver.scala
@@ -21,13 +21,13 @@ import java.util.concurrent.ConcurrentLinkedQueue
 
 import scala.collection.JavaConverters._
 
-import org.apache.spark.{Logging, SparkConf, SparkEnv}
 import org.apache.spark.executor.ShuffleWriteMetrics
 import org.apache.spark.network.buffer.{FileSegmentManagedBuffer, ManagedBuffer}
 import org.apache.spark.network.netty.SparkTransportConf
 import org.apache.spark.serializer.Serializer
 import org.apache.spark.storage._
-import org.apache.spark.util.{MetadataCleaner, MetadataCleanerType, TimeStampedHashMap}
+import org.apache.spark.util.{MetadataCleaner, MetadataCleanerType, TimeStampedHashMap, Utils}
+import org.apache.spark.{Logging, SparkConf, SparkEnv}
 
 /** A group of writers for a ShuffleMapTask, one writer per reducer. */
 private[spark] trait ShuffleWriterGroup {
@@ -84,17 +84,8 @@ private[spark] class FileShuffleBlockResolver(conf: SparkConf)
         Array.tabulate[DiskBlockObjectWriter](numReducers) { bucketId =>
           val blockId = ShuffleBlockId(shuffleId, mapId, bucketId)
           val blockFile = blockManager.diskBlockManager.getFile(blockId)
-          // Because of previous failures, the shuffle file may already exist on this machine.
-          // If so, remove it.
-          if (blockFile.exists) {
-            if (blockFile.delete()) {
-              logInfo(s"Removed existing shuffle file $blockFile")
-            } else {
-              logWarning(s"Failed to remove existing shuffle file $blockFile")
-            }
-          }
-          blockManager.getDiskWriter(blockId, blockFile, serializerInstance, bufferSize,
-            writeMetrics)
+          val tmp = Utils.tempFileWith(blockFile)
+          blockManager.getDiskWriter(blockId, tmp, serializerInstance, bufferSize, writeMetrics)
         }
       }
       // Creating the file to write to and creating a disk writer both involve interacting with
diff --git a/core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala b/core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala
index 5e4c2b5d0a5c4..05b1eed7f3bef 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala
@@ -21,13 +21,12 @@ import java.io._
 
 import com.google.common.io.ByteStreams
 
-import org.apache.spark.{SparkConf, SparkEnv, Logging}
 import org.apache.spark.network.buffer.{FileSegmentManagedBuffer, ManagedBuffer}
 import org.apache.spark.network.netty.SparkTransportConf
+import org.apache.spark.shuffle.IndexShuffleBlockResolver.NOOP_REDUCE_ID
 import org.apache.spark.storage._
 import org.apache.spark.util.Utils
-
-import IndexShuffleBlockResolver.NOOP_REDUCE_ID
+import org.apache.spark.{SparkEnv, Logging, SparkConf}
 
 /**
  * Create and maintain the shuffle blocks' mapping between logic block and physical file location.
@@ -40,10 +39,13 @@ import IndexShuffleBlockResolver.NOOP_REDUCE_ID
  */
 // Note: Changes to the format in this file should be kept in sync with
 // org.apache.spark.network.shuffle.ExternalShuffleBlockResolver#getSortBasedShuffleBlockData().
-private[spark] class IndexShuffleBlockResolver(conf: SparkConf) extends ShuffleBlockResolver
+private[spark] class IndexShuffleBlockResolver(
+    conf: SparkConf,
+    _blockManager: BlockManager = null)
+  extends ShuffleBlockResolver
   with Logging {
 
-  private lazy val blockManager = SparkEnv.get.blockManager
+  private lazy val blockManager = Option(_blockManager).getOrElse(SparkEnv.get.blockManager)
 
   private val transportConf = SparkTransportConf.fromSparkConf(conf)
 
@@ -74,14 +76,69 @@ private[spark] class IndexShuffleBlockResolver(conf: SparkConf) extends ShuffleB
     }
   }
 
+  /**
+   * Check whether the given index and data files match each other.
+   * If so, return the partition lengths in the data file. Otherwise return null.
+   */
+  private def checkIndexAndDataFile(index: File, data: File, blocks: Int): Array[Long] = {
+    // the index file should have `block + 1` longs as offset.
+    if (index.length() != (blocks + 1) * 8) {
+      return null
+    }
+    val lengths = new Array[Long](blocks)
+    // Read the lengths of blocks
+    val in = try {
+      new DataInputStream(new BufferedInputStream(new FileInputStream(index)))
+    } catch {
+      case e: IOException =>
+        return null
+    }
+    try {
+      // Convert the offsets into lengths of each block
+      var offset = in.readLong()
+      if (offset != 0L) {
+        return null
+      }
+      var i = 0
+      while (i < blocks) {
+        val off = in.readLong()
+        lengths(i) = off - offset
+        offset = off
+        i += 1
+      }
+    } catch {
+      case e: IOException =>
+        return null
+    } finally {
+      in.close()
+    }
+
+    // the size of data file should match with index file
+    if (data.length() == lengths.sum) {
+      lengths
+    } else {
+      null
+    }
+  }
+
   /**
    * Write an index file with the offsets of each block, plus a final offset at the end for the
    * end of the output file. This will be used by getBlockData to figure out where each block
    * begins and ends.
+   *
+   * It will commit the data and index file as an atomic operation, use the existing ones, or
+   * replace them with new ones.
+   *
+   * Note: the `lengths` will be updated to match the existing index file if use the existing ones.
    * */
-  def writeIndexFile(shuffleId: Int, mapId: Int, lengths: Array[Long]): Unit = {
+  def writeIndexFileAndCommit(
+      shuffleId: Int,
+      mapId: Int,
+      lengths: Array[Long],
+      dataTmp: File): Unit = {
     val indexFile = getIndexFile(shuffleId, mapId)
-    val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexFile)))
+    val indexTmp = Utils.tempFileWith(indexFile)
+    val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexTmp)))
     Utils.tryWithSafeFinally {
       // We take in lengths of each block, need to convert it to offsets.
       var offset = 0L
@@ -93,6 +150,37 @@ private[spark] class IndexShuffleBlockResolver(conf: SparkConf) extends ShuffleB
     } {
       out.close()
     }
+
+    val dataFile = getDataFile(shuffleId, mapId)
+    // There is only one IndexShuffleBlockResolver per executor, this synchronization make sure
+    // the following check and rename are atomic.
+    synchronized {
+      val existingLengths = checkIndexAndDataFile(indexFile, dataFile, lengths.length)
+      if (existingLengths != null) {
+        // Another attempt for the same task has already written our map outputs successfully,
+        // so just use the existing partition lengths and delete our temporary map outputs.
+        System.arraycopy(existingLengths, 0, lengths, 0, lengths.length)
+        if (dataTmp != null && dataTmp.exists()) {
+          dataTmp.delete()
+        }
+        indexTmp.delete()
+      } else {
+        // This is the first successful attempt in writing the map outputs for this task,
+        // so override any existing index and data files with the ones we wrote.
+        if (indexFile.exists()) {
+          indexFile.delete()
+        }
+        if (dataFile.exists()) {
+          dataFile.delete()
+        }
+        if (!indexTmp.renameTo(indexFile)) {
+          throw new IOException("fail to rename file " + indexTmp + " to " + indexFile)
+        }
+        if (dataTmp != null && dataTmp.exists() && !dataTmp.renameTo(dataFile)) {
+          throw new IOException("fail to rename file " + dataTmp + " to " + dataFile)
+        }
+      }
+    }
   }
 
   override def getBlockData(blockId: ShuffleBlockId): ManagedBuffer = {
diff --git a/core/src/main/scala/org/apache/spark/shuffle/hash/HashShuffleWriter.scala b/core/src/main/scala/org/apache/spark/shuffle/hash/HashShuffleWriter.scala
index 41df70c602c30..412bf70000da7 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/hash/HashShuffleWriter.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/hash/HashShuffleWriter.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.shuffle.hash
 
+import java.io.IOException
+
 import org.apache.spark._
 import org.apache.spark.executor.ShuffleWriteMetrics
 import org.apache.spark.scheduler.MapStatus
@@ -106,6 +108,29 @@ private[spark] class HashShuffleWriter[K, V](
       writer.commitAndClose()
       writer.fileSegment().length
     }
+    // rename all shuffle files to final paths
+    // Note: there is only one ShuffleBlockResolver in executor
+    shuffleBlockResolver.synchronized {
+      shuffle.writers.zipWithIndex.foreach { case (writer, i) =>
+        val output = blockManager.diskBlockManager.getFile(writer.blockId)
+        if (sizes(i) > 0) {
+          if (output.exists()) {
+            // Use length of existing file and delete our own temporary one
+            sizes(i) = output.length()
+            writer.file.delete()
+          } else {
+            // Commit by renaming our temporary file to something the fetcher expects
+            if (!writer.file.renameTo(output)) {
+              throw new IOException(s"fail to rename ${writer.file} to $output")
+            }
+          }
+        } else {
+          if (output.exists()) {
+            output.delete()
+          }
+        }
+      }
+    }
     MapStatus(blockManager.shuffleServerId, sizes)
   }
 
diff --git a/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala b/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala
index 808317b017a0f..f83cf8859e581 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/sort/SortShuffleWriter.scala
@@ -20,8 +20,9 @@ package org.apache.spark.shuffle.sort
 import org.apache.spark._
 import org.apache.spark.executor.ShuffleWriteMetrics
 import org.apache.spark.scheduler.MapStatus
-import org.apache.spark.shuffle.{IndexShuffleBlockResolver, ShuffleWriter, BaseShuffleHandle}
+import org.apache.spark.shuffle.{BaseShuffleHandle, IndexShuffleBlockResolver, ShuffleWriter}
 import org.apache.spark.storage.ShuffleBlockId
+import org.apache.spark.util.Utils
 import org.apache.spark.util.collection.ExternalSorter
 
 private[spark] class SortShuffleWriter[K, V, C](
@@ -65,11 +66,11 @@ private[spark] class SortShuffleWriter[K, V, C](
     // Don't bother including the time to open the merged output file in the shuffle write time,
     // because it just opens a single file, so is typically too fast to measure accurately
     // (see SPARK-3570).
-    val outputFile = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
+    val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
+    val tmp = Utils.tempFileWith(output)
     val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
-    val partitionLengths = sorter.writePartitionedFile(blockId, outputFile)
-    shuffleBlockResolver.writeIndexFile(dep.shuffleId, mapId, partitionLengths)
-
+    val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
+    shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
     mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
   }
 
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
index c374b93766225..661c706af32b1 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
@@ -21,10 +21,10 @@ import java.io._
 import java.nio.{ByteBuffer, MappedByteBuffer}
 
 import scala.collection.mutable.{ArrayBuffer, HashMap}
-import scala.concurrent.{ExecutionContext, Await, Future}
 import scala.concurrent.duration._
-import scala.util.control.NonFatal
+import scala.concurrent.{Await, ExecutionContext, Future}
 import scala.util.Random
+import scala.util.control.NonFatal
 
 import sun.nio.ch.DirectBuffer
 
@@ -38,9 +38,8 @@ import org.apache.spark.network.netty.SparkTransportConf
 import org.apache.spark.network.shuffle.ExternalShuffleClient
 import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo
 import org.apache.spark.rpc.RpcEnv
-import org.apache.spark.serializer.{SerializerInstance, Serializer}
+import org.apache.spark.serializer.{Serializer, SerializerInstance}
 import org.apache.spark.shuffle.ShuffleManager
-import org.apache.spark.shuffle.hash.HashShuffleManager
 import org.apache.spark.util._
 
 private[spark] sealed trait BlockValues
@@ -660,7 +659,7 @@ private[spark] class BlockManager(
     val compressStream: OutputStream => OutputStream = wrapForCompression(blockId, _)
     val syncWrites = conf.getBoolean("spark.shuffle.sync", false)
     new DiskBlockObjectWriter(file, serializerInstance, bufferSize, compressStream,
-      syncWrites, writeMetrics)
+      syncWrites, writeMetrics, blockId)
   }
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala b/core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala
index 80d426fadc65e..e2dd80f243930 100644
--- a/core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala
+++ b/core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala
@@ -34,14 +34,15 @@ import org.apache.spark.util.Utils
  * reopened again.
  */
 private[spark] class DiskBlockObjectWriter(
-    file: File,
+    val file: File,
     serializerInstance: SerializerInstance,
     bufferSize: Int,
     compressStream: OutputStream => OutputStream,
     syncWrites: Boolean,
     // These write metrics concurrently shared with other active DiskBlockObjectWriters who
     // are themselves performing writes. All updates must be relative.
-    writeMetrics: ShuffleWriteMetrics)
+    writeMetrics: ShuffleWriteMetrics,
+    val blockId: BlockId = null)
   extends OutputStream
   with Logging {
 
diff --git a/core/src/main/scala/org/apache/spark/util/Utils.scala b/core/src/main/scala/org/apache/spark/util/Utils.scala
index 316c194ff3454..1b3acb8ef7f51 100644
--- a/core/src/main/scala/org/apache/spark/util/Utils.scala
+++ b/core/src/main/scala/org/apache/spark/util/Utils.scala
@@ -21,8 +21,8 @@ import java.io._
 import java.lang.management.ManagementFactory
 import java.net._
 import java.nio.ByteBuffer
-import java.util.{Properties, Locale, Random, UUID}
 import java.util.concurrent._
+import java.util.{Locale, Properties, Random, UUID}
 import javax.net.ssl.HttpsURLConnection
 
 import scala.collection.JavaConverters._
@@ -30,7 +30,7 @@ import scala.collection.Map
 import scala.collection.mutable.ArrayBuffer
 import scala.io.Source
 import scala.reflect.ClassTag
-import scala.util.{Failure, Success, Try}
+import scala.util.Try
 import scala.util.control.{ControlThrowable, NonFatal}
 
 import com.google.common.io.{ByteStreams, Files}
@@ -42,7 +42,6 @@ import org.apache.hadoop.security.UserGroupInformation
 import org.apache.log4j.PropertyConfigurator
 import org.eclipse.jetty.util.MultiException
 import org.json4s._
-
 import tachyon.TachyonURI
 import tachyon.client.{TachyonFS, TachyonFile}
 
@@ -2169,6 +2168,13 @@ private[spark] object Utils extends Logging {
     val resource = createResource
     try f.apply(resource) finally resource.close()
   }
+
+  /**
+   * Returns a path of temporary file which is in the same directory with `path`.
+   */
+  def tempFileWith(path: File): File = {
+    new File(path.getAbsolutePath + "." + UUID.randomUUID())
+  }
 }
 
 /**
diff --git a/core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala b/core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala
index bd6844d045cad..2440139ac95e9 100644
--- a/core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala
+++ b/core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala
@@ -638,7 +638,6 @@ private[spark] class ExternalSorter[K, V, C](
    * called by the SortShuffleWriter.
    *
    * @param blockId block ID to write to. The index file will be blockId.name + ".index".
-   * @param context a TaskContext for a running Spark task, for us to update shuffle metrics.
    * @return array of lengths, in bytes, of each partition of the file (used by map output tracker)
    */
   def writePartitionedFile(
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/IndexShuffleBlockResolverSuite.scala b/core/src/test/java/org/apache/spark/shuffle/sort/IndexShuffleBlockResolverSuite.scala
new file mode 100644
index 0000000000000..0b19861fc41ee
--- /dev/null
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/IndexShuffleBlockResolverSuite.scala
@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.sort
+
+import java.io.{File, FileInputStream, FileOutputStream}
+
+import org.mockito.Answers.RETURNS_SMART_NULLS
+import org.mockito.Matchers._
+import org.mockito.Mockito._
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
+import org.mockito.{Mock, MockitoAnnotations}
+import org.scalatest.BeforeAndAfterEach
+
+import org.apache.spark.shuffle.IndexShuffleBlockResolver
+import org.apache.spark.storage._
+import org.apache.spark.util.Utils
+import org.apache.spark.{SparkConf, SparkFunSuite}
+
+
+class IndexShuffleBlockResolverSuite extends SparkFunSuite with BeforeAndAfterEach {
+
+  @Mock(answer = RETURNS_SMART_NULLS) private var blockManager: BlockManager = _
+  @Mock(answer = RETURNS_SMART_NULLS) private var diskBlockManager: DiskBlockManager = _
+
+  private var tempDir: File = _
+  private val conf: SparkConf = new SparkConf(loadDefaults = false)
+
+  override def beforeEach(): Unit = {
+    tempDir = Utils.createTempDir()
+    MockitoAnnotations.initMocks(this)
+
+    when(blockManager.diskBlockManager).thenReturn(diskBlockManager)
+    when(diskBlockManager.getFile(any[BlockId])).thenAnswer(
+      new Answer[File] {
+        override def answer(invocation: InvocationOnMock): File = {
+          new File(tempDir, invocation.getArguments.head.toString)
+        }
+      })
+  }
+
+  override def afterEach(): Unit = {
+    Utils.deleteRecursively(tempDir)
+  }
+
+  test("commit shuffle files multiple times") {
+    val lengths = Array[Long](10, 0, 20)
+    val resolver = new IndexShuffleBlockResolver(conf, blockManager)
+    val dataTmp = File.createTempFile("shuffle", null, tempDir)
+    val out = new FileOutputStream(dataTmp)
+    out.write(new Array[Byte](30))
+    out.close()
+    resolver.writeIndexFileAndCommit(1, 2, lengths, dataTmp)
+
+    val dataFile = resolver.getDataFile(1, 2)
+    assert(dataFile.exists())
+    assert(dataFile.length() === 30)
+    assert(!dataTmp.exists())
+
+    val dataTmp2 = File.createTempFile("shuffle", null, tempDir)
+    val out2 = new FileOutputStream(dataTmp2)
+    val lengths2 = new Array[Long](3)
+    out2.write(Array[Byte](1))
+    out2.write(new Array[Byte](29))
+    out2.close()
+    resolver.writeIndexFileAndCommit(1, 2, lengths2, dataTmp2)
+    assert(lengths2.toSeq === lengths.toSeq)
+    assert(dataFile.exists())
+    assert(dataFile.length() === 30)
+    assert(!dataTmp2.exists())
+
+    // The dataFile should be the previous one
+    val in = new FileInputStream(dataFile)
+    val firstByte = new Array[Byte](1)
+    in.read(firstByte)
+    assert(firstByte(0) === 0)
+
+    // remove data file
+    dataFile.delete()
+
+    val dataTmp3 = File.createTempFile("shuffle", null, tempDir)
+    val out3 = new FileOutputStream(dataTmp3)
+    val lengths3 = Array[Long](10, 10, 15)
+    out3.write(Array[Byte](2))
+    out3.write(new Array[Byte](34))
+    out3.close()
+    resolver.writeIndexFileAndCommit(1, 2, lengths3, dataTmp3)
+    assert(lengths3.toSeq != lengths.toSeq)
+    assert(dataFile.exists())
+    assert(dataFile.length() === 35)
+    assert(!dataTmp2.exists())
+
+    // The dataFile should be the previous one
+    val in2 = new FileInputStream(dataFile)
+    val firstByte2 = new Array[Byte](1)
+    in2.read(firstByte2)
+    assert(firstByte2(0) === 2)
+  }
+}
diff --git a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
index 0e0eca515afc1..bc85918c59aab 100644
--- a/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
+++ b/core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java
@@ -130,7 +130,8 @@ public DiskBlockObjectWriter answer(InvocationOnMock invocationOnMock) throws Th
           (Integer) args[3],
           new CompressStream(),
           false,
-          (ShuffleWriteMetrics) args[4]
+          (ShuffleWriteMetrics) args[4],
+          (BlockId) args[0]
         );
       }
     });
@@ -169,9 +170,13 @@ public OutputStream answer(InvocationOnMock invocation) throws Throwable {
       @Override
       public Void answer(InvocationOnMock invocationOnMock) throws Throwable {
         partitionSizesInMergedFile = (long[]) invocationOnMock.getArguments()[2];
+        File tmp = (File) invocationOnMock.getArguments()[3];
+        mergedOutputFile.delete();
+        tmp.renameTo(mergedOutputFile);
         return null;
       }
-    }).when(shuffleBlockResolver).writeIndexFile(anyInt(), anyInt(), any(long[].class));
+    }).when(shuffleBlockResolver)
+      .writeIndexFileAndCommit(anyInt(), anyInt(), any(long[].class), any(File.class));
 
     when(diskBlockManager.createTempShuffleBlock()).thenAnswer(
       new Answer<Tuple2<TempShuffleBlockId, File>>() {
diff --git a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
index 3bca790f30870..d87a1d2a56d99 100644
--- a/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
+++ b/core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java
@@ -117,7 +117,8 @@ public DiskBlockObjectWriter answer(InvocationOnMock invocationOnMock) throws Th
           (Integer) args[3],
           new CompressStream(),
           false,
-          (ShuffleWriteMetrics) args[4]
+          (ShuffleWriteMetrics) args[4],
+          (BlockId) args[0]
         );
       }
     });
diff --git a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
index 11c3a7be38875..a1c9f6fab8e65 100644
--- a/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
+++ b/core/src/test/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorterSuite.java
@@ -130,7 +130,8 @@ public DiskBlockObjectWriter answer(InvocationOnMock invocationOnMock) throws Th
           (Integer) args[3],
           new CompressStream(),
           false,
-          (ShuffleWriteMetrics) args[4]
+          (ShuffleWriteMetrics) args[4],
+          (BlockId) args[0]
         );
       }
     });
diff --git a/core/src/test/scala/org/apache/spark/ShuffleSuite.scala b/core/src/test/scala/org/apache/spark/ShuffleSuite.scala
index 4a0877d86f2c6..0de10ae485378 100644
--- a/core/src/test/scala/org/apache/spark/ShuffleSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ShuffleSuite.scala
@@ -17,12 +17,16 @@
 
 package org.apache.spark
 
+import java.util.concurrent.{Callable, Executors, ExecutorService, CyclicBarrier}
+
 import org.scalatest.Matchers
 
 import org.apache.spark.ShuffleSuite.NonJavaSerializableClass
+import org.apache.spark.memory.TaskMemoryManager
 import org.apache.spark.rdd.{CoGroupedRDD, OrderedRDDFunctions, RDD, ShuffledRDD, SubtractedRDD}
-import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
+import org.apache.spark.scheduler.{MyRDD, MapStatus, SparkListener, SparkListenerTaskEnd}
 import org.apache.spark.serializer.KryoSerializer
+import org.apache.spark.shuffle.ShuffleWriter
 import org.apache.spark.storage.{ShuffleDataBlockId, ShuffleBlockId}
 import org.apache.spark.util.MutablePair
 
@@ -317,6 +321,107 @@ abstract class ShuffleSuite extends SparkFunSuite with Matchers with LocalSparkC
     assert(metrics.bytesWritten === metrics.byresRead)
     assert(metrics.bytesWritten > 0)
   }
+
+  test("multiple simultaneous attempts for one task (SPARK-8029)") {
+    sc = new SparkContext("local", "test", conf)
+    val mapTrackerMaster = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster]
+    val manager = sc.env.shuffleManager
+
+    val taskMemoryManager = new TaskMemoryManager(sc.env.memoryManager, 0L)
+    val metricsSystem = sc.env.metricsSystem
+    val shuffleMapRdd = new MyRDD(sc, 1, Nil)
+    val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(1))
+    val shuffleHandle = manager.registerShuffle(0, 1, shuffleDep)
+
+    // first attempt -- its successful
+    val writer1 = manager.getWriter[Int, Int](shuffleHandle, 0,
+      new TaskContextImpl(0, 0, 0L, 0, taskMemoryManager, metricsSystem,
+        InternalAccumulator.create(sc)))
+    val data1 = (1 to 10).map { x => x -> x}
+
+    // second attempt -- also successful.  We'll write out different data,
+    // just to simulate the fact that the records may get written differently
+    // depending on what gets spilled, what gets combined, etc.
+    val writer2 = manager.getWriter[Int, Int](shuffleHandle, 0,
+      new TaskContextImpl(0, 0, 1L, 0, taskMemoryManager, metricsSystem,
+        InternalAccumulator.create(sc)))
+    val data2 = (11 to 20).map { x => x -> x}
+
+    // interleave writes of both attempts -- we want to test that both attempts can occur
+    // simultaneously, and everything is still OK
+
+    def writeAndClose(
+      writer: ShuffleWriter[Int, Int])(
+      iter: Iterator[(Int, Int)]): Option[MapStatus] = {
+      val files = writer.write(iter)
+      writer.stop(true)
+    }
+    val interleaver = new InterleaveIterators(
+      data1, writeAndClose(writer1), data2, writeAndClose(writer2))
+    val (mapOutput1, mapOutput2) = interleaver.run()
+
+    // check that we can read the map output and it has the right data
+    assert(mapOutput1.isDefined)
+    assert(mapOutput2.isDefined)
+    assert(mapOutput1.get.location === mapOutput2.get.location)
+    assert(mapOutput1.get.getSizeForBlock(0) === mapOutput1.get.getSizeForBlock(0))
+
+    // register one of the map outputs -- doesn't matter which one
+    mapOutput1.foreach { case mapStatus =>
+      mapTrackerMaster.registerMapOutputs(0, Array(mapStatus))
+    }
+
+    val reader = manager.getReader[Int, Int](shuffleHandle, 0, 1,
+      new TaskContextImpl(1, 0, 2L, 0, taskMemoryManager, metricsSystem,
+        InternalAccumulator.create(sc)))
+    val readData = reader.read().toIndexedSeq
+    assert(readData === data1.toIndexedSeq || readData === data2.toIndexedSeq)
+
+    manager.unregisterShuffle(0)
+  }
+}
+
+/**
+ * Utility to help tests make sure that we can process two different iterators simultaneously
+ * in different threads.  This makes sure that in your test, you don't completely process data1 with
+ * f1 before processing data2 with f2 (or vice versa).  It adds a barrier so that the functions only
+ * process one element, before pausing to wait for the other function to "catch up".
+ */
+class InterleaveIterators[T, R](
+  data1: Seq[T],
+  f1: Iterator[T] => R,
+  data2: Seq[T],
+  f2: Iterator[T] => R) {
+
+  require(data1.size == data2.size)
+
+  val barrier = new CyclicBarrier(2)
+  class BarrierIterator[E](id: Int, sub: Iterator[E]) extends Iterator[E] {
+    def hasNext: Boolean = sub.hasNext
+
+    def next: E = {
+      barrier.await()
+      sub.next()
+    }
+  }
+
+  val c1 = new Callable[R] {
+    override def call(): R = f1(new BarrierIterator(1, data1.iterator))
+  }
+  val c2 = new Callable[R] {
+    override def call(): R = f2(new BarrierIterator(2, data2.iterator))
+  }
+
+  val e: ExecutorService = Executors.newFixedThreadPool(2)
+
+  def run(): (R, R) = {
+    val future1 = e.submit(c1)
+    val future2 = e.submit(c2)
+    val r1 = future1.get()
+    val r2 = future2.get()
+    e.shutdown()
+    (r1, r2)
+  }
 }
 
 object ShuffleSuite {
diff --git a/core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala
index b92a302806f76..d3b1b2b620b4d 100644
--- a/core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala
+++ b/core/src/test/scala/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriterSuite.scala
@@ -68,6 +68,17 @@ class BypassMergeSortShuffleWriterSuite extends SparkFunSuite with BeforeAndAfte
     when(dependency.serializer).thenReturn(Some(new JavaSerializer(conf)))
     when(taskContext.taskMetrics()).thenReturn(taskMetrics)
     when(blockResolver.getDataFile(0, 0)).thenReturn(outputFile)
+    doAnswer(new Answer[Void] {
+      def answer(invocationOnMock: InvocationOnMock): Void = {
+        val tmp: File = invocationOnMock.getArguments()(3).asInstanceOf[File]
+        if (tmp != null) {
+          outputFile.delete
+          tmp.renameTo(outputFile)
+        }
+        null
+      }
+    }).when(blockResolver)
+      .writeIndexFileAndCommit(anyInt, anyInt, any(classOf[Array[Long]]), any(classOf[File]))
     when(blockManager.diskBlockManager).thenReturn(diskBlockManager)
     when(blockManager.getDiskWriter(
       any[BlockId],
@@ -84,7 +95,8 @@ class BypassMergeSortShuffleWriterSuite extends SparkFunSuite with BeforeAndAfte
           args(3).asInstanceOf[Int],
           compressStream = identity,
           syncWrites = false,
-          args(4).asInstanceOf[ShuffleWriteMetrics]
+          args(4).asInstanceOf[ShuffleWriteMetrics],
+          blockId = args(0).asInstanceOf[BlockId]
         )
       }
     })

From ec80c0c2fc63360ee6b5872c24e6c67779ac63f4 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Fri, 13 Nov 2015 00:30:27 -0800
Subject: [PATCH 230/518] [SPARK-11706][STREAMING] Fix the bug that Streaming
 Python tests cannot report failures

This PR just checks the test results and returns 1 if the test fails, so that `run-tests.py` can mark it fail.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9669 from zsxwing/streaming-python-tests.
---
 .../streaming/flume/FlumeTestUtils.scala      |  5 ++--
 .../flume/PollingFlumeTestUtils.scala         |  9 +++---
 .../flume/FlumePollingStreamSuite.scala       |  2 +-
 .../streaming/flume/FlumeStreamSuite.scala    |  2 +-
 python/pyspark/streaming/tests.py             | 30 ++++++++++++-------
 5 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/external/flume/src/main/scala/org/apache/spark/streaming/flume/FlumeTestUtils.scala b/external/flume/src/main/scala/org/apache/spark/streaming/flume/FlumeTestUtils.scala
index 70018c86f92be..fe5dcc8e4b9de 100644
--- a/external/flume/src/main/scala/org/apache/spark/streaming/flume/FlumeTestUtils.scala
+++ b/external/flume/src/main/scala/org/apache/spark/streaming/flume/FlumeTestUtils.scala
@@ -19,6 +19,7 @@ package org.apache.spark.streaming.flume
 
 import java.net.{InetSocketAddress, ServerSocket}
 import java.nio.ByteBuffer
+import java.util.{List => JList}
 import java.util.Collections
 
 import scala.collection.JavaConverters._
@@ -59,10 +60,10 @@ private[flume] class FlumeTestUtils {
   }
 
   /** Send data to the flume receiver */
-  def writeInput(input: Seq[String], enableCompression: Boolean): Unit = {
+  def writeInput(input: JList[String], enableCompression: Boolean): Unit = {
     val testAddress = new InetSocketAddress("localhost", testPort)
 
-    val inputEvents = input.map { item =>
+    val inputEvents = input.asScala.map { item =>
       val event = new AvroFlumeEvent
       event.setBody(ByteBuffer.wrap(item.getBytes(UTF_8)))
       event.setHeaders(Collections.singletonMap("test", "header"))
diff --git a/external/flume/src/main/scala/org/apache/spark/streaming/flume/PollingFlumeTestUtils.scala b/external/flume/src/main/scala/org/apache/spark/streaming/flume/PollingFlumeTestUtils.scala
index a2ab320957db3..bfe7548d4f50e 100644
--- a/external/flume/src/main/scala/org/apache/spark/streaming/flume/PollingFlumeTestUtils.scala
+++ b/external/flume/src/main/scala/org/apache/spark/streaming/flume/PollingFlumeTestUtils.scala
@@ -18,7 +18,7 @@
 package org.apache.spark.streaming.flume
 
 import java.util.concurrent._
-import java.util.{Map => JMap, Collections}
+import java.util.{Collections, List => JList, Map => JMap}
 
 import scala.collection.mutable.ArrayBuffer
 
@@ -137,7 +137,8 @@ private[flume] class PollingFlumeTestUtils {
   /**
    * A Python-friendly method to assert the output
    */
-  def assertOutput(outputHeaders: Seq[JMap[String, String]], outputBodies: Seq[String]): Unit = {
+  def assertOutput(
+      outputHeaders: JList[JMap[String, String]], outputBodies: JList[String]): Unit = {
     require(outputHeaders.size == outputBodies.size)
     val eventSize = outputHeaders.size
     if (eventSize != totalEventsPerChannel * channels.size) {
@@ -151,8 +152,8 @@ private[flume] class PollingFlumeTestUtils {
       var found = false
       var j = 0
       while (j < eventSize && !found) {
-        if (eventBodyToVerify == outputBodies(j) &&
-          eventHeaderToVerify == outputHeaders(j)) {
+        if (eventBodyToVerify == outputBodies.get(j) &&
+          eventHeaderToVerify == outputHeaders.get(j)) {
           found = true
           counter += 1
         }
diff --git a/external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumePollingStreamSuite.scala b/external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumePollingStreamSuite.scala
index ff2fb8eed204c..5fd2711f5f7df 100644
--- a/external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumePollingStreamSuite.scala
+++ b/external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumePollingStreamSuite.scala
@@ -120,7 +120,7 @@ class FlumePollingStreamSuite extends SparkFunSuite with BeforeAndAfter with Log
           case (key, value) => (key.toString, value.toString)
         }).map(_.asJava)
         val bodies = flattenOutputBuffer.map(e => new String(e.event.getBody.array(), UTF_8))
-        utils.assertOutput(headers, bodies)
+        utils.assertOutput(headers.asJava, bodies.asJava)
       }
     } finally {
       ssc.stop()
diff --git a/external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumeStreamSuite.scala b/external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumeStreamSuite.scala
index 5ffb60bd602f9..f315e0a7ca23c 100644
--- a/external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumeStreamSuite.scala
+++ b/external/flume/src/test/scala/org/apache/spark/streaming/flume/FlumeStreamSuite.scala
@@ -54,7 +54,7 @@ class FlumeStreamSuite extends SparkFunSuite with BeforeAndAfter with Matchers w
       val outputBuffer = startContext(utils.getTestPort(), testCompression)
 
       eventually(timeout(10 seconds), interval(100 milliseconds)) {
-        utils.writeInput(input, testCompression)
+        utils.writeInput(input.asJava, testCompression)
       }
 
       eventually(timeout(10 seconds), interval(100 milliseconds)) {
diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index 179479625bca4..6ee864d8d3da6 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -611,12 +611,16 @@ class CheckpointTests(unittest.TestCase):
     @staticmethod
     def tearDownClass():
         # Clean up in the JVM just in case there has been some issues in Python API
-        jStreamingContextOption = StreamingContext._jvm.SparkContext.getActive()
-        if jStreamingContextOption.nonEmpty():
-            jStreamingContextOption.get().stop()
-        jSparkContextOption = SparkContext._jvm.SparkContext.get()
-        if jSparkContextOption.nonEmpty():
-            jSparkContextOption.get().stop()
+        if SparkContext._jvm is not None:
+            jStreamingContextOption = \
+                SparkContext._jvm.org.apache.spark.streaming.StreamingContext.getActive()
+            if jStreamingContextOption.nonEmpty():
+                jStreamingContextOption.get().stop()
+
+    def setUp(self):
+        self.ssc = None
+        self.sc = None
+        self.cpd = None
 
     def tearDown(self):
         if self.ssc is not None:
@@ -626,6 +630,7 @@ def tearDown(self):
         if self.cpd is not None:
             shutil.rmtree(self.cpd)
 
+    @unittest.skip("Enable it when we fix the checkpoint bug")
     def test_get_or_create_and_get_active_or_create(self):
         inputd = tempfile.mkdtemp()
         outputd = tempfile.mkdtemp() + "/"
@@ -648,7 +653,7 @@ def setup():
         self.cpd = tempfile.mkdtemp("test_streaming_cps")
         self.setupCalled = False
         self.ssc = StreamingContext.getOrCreate(self.cpd, setup)
-        self.assertFalse(self.setupCalled)
+        self.assertTrue(self.setupCalled)
 
         self.ssc.start()
 
@@ -1322,11 +1327,16 @@ def search_kinesis_asl_assembly_jar():
             "or 'build/mvn -Pkinesis-asl package' before running this test.")
 
     sys.stderr.write("Running tests: %s \n" % (str(testcases)))
+    failed = False
     for testcase in testcases:
         sys.stderr.write("[Running %s]\n" % (testcase))
         tests = unittest.TestLoader().loadTestsFromTestCase(testcase)
         if xmlrunner:
-            unittest.main(tests, verbosity=3,
-                          testRunner=xmlrunner.XMLTestRunner(output='target/test-reports'))
+            result = xmlrunner.XMLTestRunner(output='target/test-reports', verbosity=3).run(tests)
+            if not result.wasSuccessful():
+                failed = True
         else:
-            unittest.TextTestRunner(verbosity=3).run(tests)
+            result = unittest.TextTestRunner(verbosity=3).run(tests)
+            if not result.wasSuccessful():
+                failed = True
+    sys.exit(failed)

From 7b5d9051cf91c099458d092a6705545899134b3b Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Fri, 13 Nov 2015 18:36:56 +0800
Subject: [PATCH 231/518] [SPARK-11678][SQL] Partition discovery should stop at
 the root path of the table.

https://issues.apache.org/jira/browse/SPARK-11678

The change of this PR is to pass root paths of table to the partition discovery logic. So, the process of partition discovery stops at those root paths instead of going all the way to the root path of the file system.

Author: Yin Huai <yhuai@databricks.com>

Closes #9651 from yhuai/SPARK-11678.
---
 .../datasources/PartitioningUtils.scala       |  68 ++++++---
 .../datasources/json/JSONRelation.scala       |  21 +--
 .../datasources/parquet/ParquetRelation.scala |   2 +-
 .../datasources/text/DefaultSource.scala      |   5 +-
 .../apache/spark/sql/sources/interfaces.scala |  49 ++++++-
 .../parquet/ParquetFilterSuite.scala          |   4 +-
 .../ParquetPartitionDiscoverySuite.scala      | 132 ++++++++++++++++--
 .../spark/sql/hive/orc/OrcRelation.scala      |   2 +-
 .../sql/sources/SimpleTextRelation.scala      |   2 +-
 .../sql/sources/hadoopFsRelationSuites.scala  |   1 +
 10 files changed, 235 insertions(+), 51 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
index 86bc3a1b6dab2..81962f8d63789 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
@@ -75,10 +75,11 @@ private[sql] object PartitioningUtils {
   private[sql] def parsePartitions(
       paths: Seq[Path],
       defaultPartitionName: String,
-      typeInference: Boolean): PartitionSpec = {
+      typeInference: Boolean,
+      basePaths: Set[Path]): PartitionSpec = {
     // First, we need to parse every partition's path and see if we can find partition values.
-    val (partitionValues, optBasePaths) = paths.map { path =>
-      parsePartition(path, defaultPartitionName, typeInference)
+    val (partitionValues, optDiscoveredBasePaths) = paths.map { path =>
+      parsePartition(path, defaultPartitionName, typeInference, basePaths)
     }.unzip
 
     // We create pairs of (path -> path's partition value) here
@@ -101,11 +102,15 @@ private[sql] object PartitioningUtils {
       // It will be recognised as conflicting directory structure:
       //   "hdfs://host:9000/invalidPath"
       //   "hdfs://host:9000/path"
-      val basePaths = optBasePaths.flatMap(x => x)
+      val disvoeredBasePaths = optDiscoveredBasePaths.flatMap(x => x)
       assert(
-        basePaths.distinct.size == 1,
+        disvoeredBasePaths.distinct.size == 1,
         "Conflicting directory structures detected. Suspicious paths:\b" +
-          basePaths.distinct.mkString("\n\t", "\n\t", "\n\n"))
+          disvoeredBasePaths.distinct.mkString("\n\t", "\n\t", "\n\n") +
+          "If provided paths are partition directories, please set " +
+          "\"basePath\" in the options of the data source to specify the " +
+          "root directory of the table. If there are multiple root directories, " +
+          "please load them separately and then union them.")
 
       val resolvedPartitionValues = resolvePartitions(pathsWithPartitionValues)
 
@@ -131,7 +136,7 @@ private[sql] object PartitioningUtils {
 
   /**
    * Parses a single partition, returns column names and values of each partition column, also
-   * the base path.  For example, given:
+   * the path when we stop partition discovery.  For example, given:
    * {{{
    *   path = hdfs://<host>:<port>/path/to/partition/a=42/b=hello/c=3.14
    * }}}
@@ -144,40 +149,63 @@ private[sql] object PartitioningUtils {
    *       Literal.create("hello", StringType),
    *       Literal.create(3.14, FloatType)))
    * }}}
-   * and the base path:
+   * and the path when we stop the discovery is:
    * {{{
-   *   /path/to/partition
+   *   hdfs://<host>:<port>/path/to/partition
    * }}}
    */
   private[sql] def parsePartition(
       path: Path,
       defaultPartitionName: String,
-      typeInference: Boolean): (Option[PartitionValues], Option[Path]) = {
+      typeInference: Boolean,
+      basePaths: Set[Path]): (Option[PartitionValues], Option[Path]) = {
     val columns = ArrayBuffer.empty[(String, Literal)]
     // Old Hadoop versions don't have `Path.isRoot`
     var finished = path.getParent == null
-    var chopped = path
-    var basePath = path
+    // currentPath is the current path that we will use to parse partition column value.
+    var currentPath: Path = path
 
     while (!finished) {
       // Sometimes (e.g., when speculative task is enabled), temporary directories may be left
-      // uncleaned.  Here we simply ignore them.
-      if (chopped.getName.toLowerCase == "_temporary") {
+      // uncleaned. Here we simply ignore them.
+      if (currentPath.getName.toLowerCase == "_temporary") {
         return (None, None)
       }
 
-      val maybeColumn = parsePartitionColumn(chopped.getName, defaultPartitionName, typeInference)
-      maybeColumn.foreach(columns += _)
-      basePath = chopped
-      chopped = chopped.getParent
-      finished = (maybeColumn.isEmpty && !columns.isEmpty) || chopped.getParent == null
+      if (basePaths.contains(currentPath)) {
+        // If the currentPath is one of base paths. We should stop.
+        finished = true
+      } else {
+        // Let's say currentPath is a path of "/table/a=1/", currentPath.getName will give us a=1.
+        // Once we get the string, we try to parse it and find the partition column and value.
+        val maybeColumn =
+          parsePartitionColumn(currentPath.getName, defaultPartitionName, typeInference)
+        maybeColumn.foreach(columns += _)
+
+        // Now, we determine if we should stop.
+        // When we hit any of the following cases, we will stop:
+        //  - In this iteration, we could not parse the value of partition column and value,
+        //    i.e. maybeColumn is None, and columns is not empty. At here we check if columns is
+        //    empty to handle cases like /table/a=1/_temporary/something (we need to find a=1 in
+        //    this case).
+        //  - After we get the new currentPath, this new currentPath represent the top level dir
+        //    i.e. currentPath.getParent == null. For the example of "/table/a=1/",
+        //    the top level dir is "/table".
+        finished =
+          (maybeColumn.isEmpty && !columns.isEmpty) || currentPath.getParent == null
+
+        if (!finished) {
+          // For the above example, currentPath will be "/table/".
+          currentPath = currentPath.getParent
+        }
+      }
     }
 
     if (columns.isEmpty) {
       (None, Some(path))
     } else {
       val (columnNames, values) = columns.reverse.unzip
-      (Some(PartitionValues(columnNames, values)), Some(basePath))
+      (Some(PartitionValues(columnNames, values)), Some(currentPath))
     }
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONRelation.scala
index 85b52f04c8d01..dca638b7f67a5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONRelation.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONRelation.scala
@@ -56,13 +56,14 @@ class DefaultSource extends HadoopFsRelationProvider with DataSourceRegister {
     val primitivesAsString = parameters.get("primitivesAsString").map(_.toBoolean).getOrElse(false)
 
     new JSONRelation(
-      None,
-      samplingRatio,
-      primitivesAsString,
-      dataSchema,
-      None,
-      partitionColumns,
-      paths)(sqlContext)
+      inputRDD = None,
+      samplingRatio = samplingRatio,
+      primitivesAsString = primitivesAsString,
+      maybeDataSchema = dataSchema,
+      maybePartitionSpec = None,
+      userDefinedPartitionColumns = partitionColumns,
+      paths = paths,
+      parameters = parameters)(sqlContext)
   }
 }
 
@@ -73,8 +74,10 @@ private[sql] class JSONRelation(
     val maybeDataSchema: Option[StructType],
     val maybePartitionSpec: Option[PartitionSpec],
     override val userDefinedPartitionColumns: Option[StructType],
-    override val paths: Array[String] = Array.empty[String])(@transient val sqlContext: SQLContext)
-  extends HadoopFsRelation(maybePartitionSpec) {
+    override val paths: Array[String] = Array.empty[String],
+    parameters: Map[String, String] = Map.empty[String, String])
+    (@transient val sqlContext: SQLContext)
+  extends HadoopFsRelation(maybePartitionSpec, parameters) {
 
   /** Constraints to be imposed on schema to be stored. */
   private def checkConstraints(schema: StructType): Unit = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
index 21337b2932aac..cb0aab8cc0d09 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
@@ -109,7 +109,7 @@ private[sql] class ParquetRelation(
     override val userDefinedPartitionColumns: Option[StructType],
     parameters: Map[String, String])(
     val sqlContext: SQLContext)
-  extends HadoopFsRelation(maybePartitionSpec)
+  extends HadoopFsRelation(maybePartitionSpec, parameters)
   with Logging {
 
   private[sql] def this(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/DefaultSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/DefaultSource.scala
index 4b8b8e4e74dad..fbd387bc2ef47 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/DefaultSource.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/DefaultSource.scala
@@ -71,9 +71,10 @@ class DefaultSource extends HadoopFsRelationProvider with DataSourceRegister {
 private[sql] class TextRelation(
     val maybePartitionSpec: Option[PartitionSpec],
     override val userDefinedPartitionColumns: Option[StructType],
-    override val paths: Array[String] = Array.empty[String])
+    override val paths: Array[String] = Array.empty[String],
+    parameters: Map[String, String] = Map.empty[String, String])
     (@transient val sqlContext: SQLContext)
-  extends HadoopFsRelation(maybePartitionSpec) {
+  extends HadoopFsRelation(maybePartitionSpec, parameters) {
 
   /** Data schema is always a single column, named "text". */
   override def dataSchema: StructType = new StructType().add("value", StringType)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index 2be6cd45337f5..b3d3bdf50df63 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -416,12 +416,19 @@ abstract class OutputWriter {
  * @since 1.4.0
  */
 @Experimental
-abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[PartitionSpec])
+abstract class HadoopFsRelation private[sql](
+    maybePartitionSpec: Option[PartitionSpec],
+    parameters: Map[String, String])
   extends BaseRelation with FileRelation with Logging {
 
   override def toString: String = getClass.getSimpleName + paths.mkString("[", ",", "]")
 
-  def this() = this(None)
+  def this() = this(None, Map.empty[String, String])
+
+  def this(parameters: Map[String, String]) = this(None, parameters)
+
+  private[sql] def this(maybePartitionSpec: Option[PartitionSpec]) =
+    this(maybePartitionSpec, Map.empty[String, String])
 
   private val hadoopConf = new Configuration(sqlContext.sparkContext.hadoopConfiguration)
 
@@ -519,13 +526,37 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
   }
 
   /**
-   * Base paths of this relation.  For partitioned relations, it should be either root directories
+   * Paths of this relation.  For partitioned relations, it should be root directories
    * of all partition directories.
    *
    * @since 1.4.0
    */
   def paths: Array[String]
 
+  /**
+   * Contains a set of paths that are considered as the base dirs of the input datasets.
+   * The partitioning discovery logic will make sure it will stop when it reaches any
+   * base path. By default, the paths of the dataset provided by users will be base paths.
+   * For example, if a user uses `sqlContext.read.parquet("/path/something=true/")`, the base path
+   * will be `/path/something=true/`, and the returned DataFrame will not contain a column of
+   * `something`. If users want to override the basePath. They can set `basePath` in the options
+   * to pass the new base path to the data source.
+   * For the above example, if the user-provided base path is `/path/`, the returned
+   * DataFrame will have the column of `something`.
+   */
+  private def basePaths: Set[Path] = {
+    val userDefinedBasePath = parameters.get("basePath").map(basePath => Set(new Path(basePath)))
+    userDefinedBasePath.getOrElse {
+      // If the user does not provide basePath, we will just use paths.
+      val pathSet = paths.toSet
+      pathSet.map(p => new Path(p))
+    }.map { hdfsPath =>
+      // Make the path qualified (consistent with listLeafFiles and listLeafFilesInParallel).
+      val fs = hdfsPath.getFileSystem(hadoopConf)
+      hdfsPath.makeQualified(fs.getUri, fs.getWorkingDirectory)
+    }
+  }
+
   override def inputFiles: Array[String] = cachedLeafStatuses().map(_.getPath.toString).toArray
 
   override def sizeInBytes: Long = cachedLeafStatuses().map(_.getLen).sum
@@ -559,7 +590,10 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
     userDefinedPartitionColumns match {
       case Some(userProvidedSchema) if userProvidedSchema.nonEmpty =>
         val spec = PartitioningUtils.parsePartitions(
-          leafDirs, PartitioningUtils.DEFAULT_PARTITION_NAME, typeInference = false)
+          leafDirs,
+          PartitioningUtils.DEFAULT_PARTITION_NAME,
+          typeInference = false,
+          basePaths = basePaths)
 
         // Without auto inference, all of value in the `row` should be null or in StringType,
         // we need to cast into the data type that user specified.
@@ -577,8 +611,11 @@ abstract class HadoopFsRelation private[sql](maybePartitionSpec: Option[Partitio
 
       case _ =>
         // user did not provide a partitioning schema
-        PartitioningUtils.parsePartitions(leafDirs, PartitioningUtils.DEFAULT_PARTITION_NAME,
-          typeInference = sqlContext.conf.partitionColumnTypeInferenceEnabled())
+        PartitioningUtils.parsePartitions(
+          leafDirs,
+          PartitioningUtils.DEFAULT_PARTITION_NAME,
+          typeInference = sqlContext.conf.partitionColumnTypeInferenceEnabled(),
+          basePaths = basePaths)
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
index 2ac87ad6cd035..458786f77af3f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
@@ -294,7 +294,7 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
         // If the "part = 1" filter gets pushed down, this query will throw an exception since
         // "part" is not a valid column in the actual Parquet file
         checkAnswer(
-          sqlContext.read.parquet(path).filter("part = 1"),
+          sqlContext.read.parquet(dir.getCanonicalPath).filter("part = 1"),
           (1 to 3).map(i => Row(i, i.toString, 1)))
       }
     }
@@ -311,7 +311,7 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
         // If the "part = 1" filter gets pushed down, this query will throw an exception since
         // "part" is not a valid column in the actual Parquet file
         checkAnswer(
-          sqlContext.read.parquet(path).filter("a > 0 and (part = 0 or a > 1)"),
+          sqlContext.read.parquet(dir.getCanonicalPath).filter("a > 0 and (part = 0 or a > 1)"),
           (2 to 3).map(i => Row(i, i.toString, 1)))
       }
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
index 61cc0da50865c..71e9034d97792 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetPartitionDiscoverySuite.scala
@@ -66,7 +66,7 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
       "hdfs://host:9000/path/a=10.5/b=hello")
 
     var exception = intercept[AssertionError] {
-      parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
+      parsePartitions(paths.map(new Path(_)), defaultPartitionName, true, Set.empty[Path])
     }
     assert(exception.getMessage().contains("Conflicting directory structures detected"))
 
@@ -76,7 +76,37 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
       "hdfs://host:9000/path/a=10/b=20",
       "hdfs://host:9000/path/_temporary/path")
 
-    parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
+    parsePartitions(
+      paths.map(new Path(_)),
+      defaultPartitionName,
+      true,
+      Set(new Path("hdfs://host:9000/path/")))
+
+    // Valid
+    paths = Seq(
+      "hdfs://host:9000/path/something=true/table/",
+      "hdfs://host:9000/path/something=true/table/_temporary",
+      "hdfs://host:9000/path/something=true/table/a=10/b=20",
+      "hdfs://host:9000/path/something=true/table/_temporary/path")
+
+    parsePartitions(
+      paths.map(new Path(_)),
+      defaultPartitionName,
+      true,
+      Set(new Path("hdfs://host:9000/path/something=true/table")))
+
+    // Valid
+    paths = Seq(
+      "hdfs://host:9000/path/table=true/",
+      "hdfs://host:9000/path/table=true/_temporary",
+      "hdfs://host:9000/path/table=true/a=10/b=20",
+      "hdfs://host:9000/path/table=true/_temporary/path")
+
+    parsePartitions(
+      paths.map(new Path(_)),
+      defaultPartitionName,
+      true,
+      Set(new Path("hdfs://host:9000/path/table=true")))
 
     // Invalid
     paths = Seq(
@@ -85,7 +115,11 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
       "hdfs://host:9000/path/path1")
 
     exception = intercept[AssertionError] {
-      parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
+      parsePartitions(
+        paths.map(new Path(_)),
+        defaultPartitionName,
+        true,
+        Set(new Path("hdfs://host:9000/path/")))
     }
     assert(exception.getMessage().contains("Conflicting directory structures detected"))
 
@@ -101,19 +135,24 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
       "hdfs://host:9000/tmp/tables/nonPartitionedTable2")
 
     exception = intercept[AssertionError] {
-      parsePartitions(paths.map(new Path(_)), defaultPartitionName, true)
+      parsePartitions(
+        paths.map(new Path(_)),
+        defaultPartitionName,
+        true,
+        Set(new Path("hdfs://host:9000/tmp/tables/")))
     }
     assert(exception.getMessage().contains("Conflicting directory structures detected"))
   }
 
   test("parse partition") {
     def check(path: String, expected: Option[PartitionValues]): Unit = {
-      assert(expected === parsePartition(new Path(path), defaultPartitionName, true)._1)
+      val actual = parsePartition(new Path(path), defaultPartitionName, true, Set.empty[Path])._1
+      assert(expected === actual)
     }
 
     def checkThrows[T <: Throwable: Manifest](path: String, expected: String): Unit = {
       val message = intercept[T] {
-        parsePartition(new Path(path), defaultPartitionName, true)
+        parsePartition(new Path(path), defaultPartitionName, true, Set.empty[Path])
       }.getMessage
 
       assert(message.contains(expected))
@@ -152,8 +191,17 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
   }
 
   test("parse partitions") {
-    def check(paths: Seq[String], spec: PartitionSpec): Unit = {
-      assert(parsePartitions(paths.map(new Path(_)), defaultPartitionName, true) === spec)
+    def check(
+        paths: Seq[String],
+        spec: PartitionSpec,
+        rootPaths: Set[Path] = Set.empty[Path]): Unit = {
+      val actualSpec =
+        parsePartitions(
+          paths.map(new Path(_)),
+          defaultPartitionName,
+          true,
+          rootPaths)
+      assert(actualSpec === spec)
     }
 
     check(Seq(
@@ -232,7 +280,9 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
 
   test("parse partitions with type inference disabled") {
     def check(paths: Seq[String], spec: PartitionSpec): Unit = {
-      assert(parsePartitions(paths.map(new Path(_)), defaultPartitionName, false) === spec)
+      val actualSpec =
+        parsePartitions(paths.map(new Path(_)), defaultPartitionName, false, Set.empty[Path])
+      assert(actualSpec === spec)
     }
 
     check(Seq(
@@ -590,6 +640,70 @@ class ParquetPartitionDiscoverySuite extends QueryTest with ParquetTest with Sha
     }
   }
 
+  test("SPARK-11678: Partition discovery stops at the root path of the dataset") {
+    withTempPath { dir =>
+      val tablePath = new File(dir, "key=value")
+      val df = (1 to 3).map(i => (i, i, i, i)).toDF("a", "b", "c", "d")
+
+      df.write
+        .format("parquet")
+        .partitionBy("b", "c", "d")
+        .save(tablePath.getCanonicalPath)
+
+      Files.touch(new File(s"${tablePath.getCanonicalPath}/", "_SUCCESS"))
+      Files.createParentDirs(new File(s"${dir.getCanonicalPath}/b=1/c=1/.foo/bar"))
+
+      checkAnswer(sqlContext.read.format("parquet").load(tablePath.getCanonicalPath), df)
+    }
+
+    withTempPath { dir =>
+      val path = new File(dir, "key=value")
+      val tablePath = new File(path, "table")
+
+      val df = (1 to 3).map(i => (i, i, i, i)).toDF("a", "b", "c", "d")
+
+      df.write
+        .format("parquet")
+        .partitionBy("b", "c", "d")
+        .save(tablePath.getCanonicalPath)
+
+      Files.touch(new File(s"${tablePath.getCanonicalPath}/", "_SUCCESS"))
+      Files.createParentDirs(new File(s"${dir.getCanonicalPath}/b=1/c=1/.foo/bar"))
+
+      checkAnswer(sqlContext.read.format("parquet").load(tablePath.getCanonicalPath), df)
+    }
+  }
+
+  test("use basePath to specify the root dir of a partitioned table.") {
+    withTempPath { dir =>
+      val tablePath = new File(dir, "table")
+      val df = (1 to 3).map(i => (i, i, i, i)).toDF("a", "b", "c", "d")
+
+      df.write
+        .format("parquet")
+        .partitionBy("b", "c", "d")
+        .save(tablePath.getCanonicalPath)
+
+      val twoPartitionsDF =
+        sqlContext
+          .read
+          .option("basePath", tablePath.getCanonicalPath)
+          .parquet(
+            s"${tablePath.getCanonicalPath}/b=1",
+            s"${tablePath.getCanonicalPath}/b=2")
+
+      checkAnswer(twoPartitionsDF, df.filter("b != 3"))
+
+      intercept[AssertionError] {
+        sqlContext
+          .read
+          .parquet(
+            s"${tablePath.getCanonicalPath}/b=1",
+            s"${tablePath.getCanonicalPath}/b=2")
+      }
+    }
+  }
+
   test("listConflictingPartitionColumns") {
     def makeExpectedMessage(colNameLists: Seq[String], paths: Seq[String]): String = {
       val conflictingColNameLists = colNameLists.zipWithIndex.map { case (list, index) =>
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/orc/OrcRelation.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/orc/OrcRelation.scala
index 45de567039760..1136670b7a0eb 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/orc/OrcRelation.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/orc/OrcRelation.scala
@@ -157,7 +157,7 @@ private[sql] class OrcRelation(
     override val userDefinedPartitionColumns: Option[StructType],
     parameters: Map[String, String])(
     @transient val sqlContext: SQLContext)
-  extends HadoopFsRelation(maybePartitionSpec)
+  extends HadoopFsRelation(maybePartitionSpec, parameters)
   with Logging {
 
   private[sql] def this(
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
index bdc48a383bbbf..01960fd2901b0 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextRelation.scala
@@ -89,7 +89,7 @@ class SimpleTextRelation(
     override val userDefinedPartitionColumns: Option[StructType],
     parameters: Map[String, String])(
     @transient val sqlContext: SQLContext)
-  extends HadoopFsRelation {
+  extends HadoopFsRelation(parameters) {
 
   import sqlContext.sparkContext
 
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/hadoopFsRelationSuites.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/hadoopFsRelationSuites.scala
index 100b97137cff0..665e87e3e3355 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/hadoopFsRelationSuites.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/hadoopFsRelationSuites.scala
@@ -486,6 +486,7 @@ abstract class HadoopFsRelationTest extends QueryTest with SQLTestUtils with Tes
       val df = sqlContext.read
         .format(dataSourceName)
         .option("dataSchema", dataSchema.json)
+        .option("basePath", file.getCanonicalPath)
         .load(s"${file.getCanonicalPath}/p1=*/p2=???")
 
       val expectedPaths = Set(

From 61a28486ccbcdd37461419df958aea222c8b9f09 Mon Sep 17 00:00:00 2001
From: Rishabh Bhardwaj <rbnext29@gmail.com>
Date: Fri, 13 Nov 2015 08:36:46 -0800
Subject: [PATCH 232/518] [SPARK-11445][DOCS] Replaced example code in
 mllib-ensembles.md using include_example

I have made the required changes and tested.
Kindly review the changes.

Author: Rishabh Bhardwaj <rbnext29@gmail.com>

Closes #9407 from rishabhbhardwaj/SPARK-11445.
---
 docs/mllib-ensembles.md                       | 526 +-----------------
 ...GradientBoostingClassificationExample.java |  92 +++
 ...JavaGradientBoostingRegressionExample.java |  96 ++++
 ...JavaRandomForestClassificationExample.java |  89 +++
 .../JavaRandomForestRegressionExample.java    |  95 ++++
 ...radient_boosting_classification_example.py |  57 ++
 .../gradient_boosting_regression_example.py   |  57 ++
 .../random_forest_classification_example.py   |  58 ++
 .../mllib/random_forest_regression_example.py |  59 ++
 ...radientBoostingClassificationExample.scala |  69 +++
 .../GradientBoostingRegressionExample.scala   |  66 +++
 .../RandomForestClassificationExample.scala   |  67 +++
 .../mllib/RandomForestRegressionExample.scala |  68 +++
 13 files changed, 885 insertions(+), 514 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingClassificationExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingRegressionExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestClassificationExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestRegressionExample.java
 create mode 100644 examples/src/main/python/mllib/gradient_boosting_classification_example.py
 create mode 100644 examples/src/main/python/mllib/gradient_boosting_regression_example.py
 create mode 100644 examples/src/main/python/mllib/random_forest_classification_example.py
 create mode 100644 examples/src/main/python/mllib/random_forest_regression_example.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/GradientBoostingClassificationExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/GradientBoostingRegressionExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/RandomForestClassificationExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/RandomForestRegressionExample.scala

diff --git a/docs/mllib-ensembles.md b/docs/mllib-ensembles.md
index fc587298f7d2e..50450e05d2abb 100644
--- a/docs/mllib-ensembles.md
+++ b/docs/mllib-ensembles.md
@@ -98,144 +98,19 @@ The test error is calculated to measure the algorithm accuracy.
 <div data-lang="scala" markdown="1">
 Refer to the [`RandomForest` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.RandomForest) and [`RandomForestModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.model.RandomForestModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.tree.RandomForest
-import org.apache.spark.mllib.tree.model.RandomForestModel
-import org.apache.spark.mllib.util.MLUtils
-
-// Load and parse the data file.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-// Split the data into training and test sets (30% held out for testing)
-val splits = data.randomSplit(Array(0.7, 0.3))
-val (trainingData, testData) = (splits(0), splits(1))
-
-// Train a RandomForest model.
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-val numClasses = 2
-val categoricalFeaturesInfo = Map[Int, Int]()
-val numTrees = 3 // Use more in practice.
-val featureSubsetStrategy = "auto" // Let the algorithm choose.
-val impurity = "gini"
-val maxDepth = 4
-val maxBins = 32
-
-val model = RandomForest.trainClassifier(trainingData, numClasses, categoricalFeaturesInfo,
-  numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins)
-
-// Evaluate model on test instances and compute test error
-val labelAndPreds = testData.map { point =>
-  val prediction = model.predict(point.features)
-  (point.label, prediction)
-}
-val testErr = labelAndPreds.filter(r => r._1 != r._2).count.toDouble / testData.count()
-println("Test Error = " + testErr)
-println("Learned classification forest model:\n" + model.toDebugString)
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = RandomForestModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/RandomForestClassificationExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`RandomForest` Java docs](api/java/org/apache/spark/mllib/tree/RandomForest.html) and [`RandomForestModel` Java docs](api/java/org/apache/spark/mllib/tree/model/RandomForestModel.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-import java.util.HashMap;
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.RandomForest;
-import org.apache.spark.mllib.tree.model.RandomForestModel;
-import org.apache.spark.mllib.util.MLUtils;
-
-SparkConf sparkConf = new SparkConf().setAppName("JavaRandomForestClassification");
-JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-// Load and parse the data file.
-String datapath = "data/mllib/sample_libsvm_data.txt";
-JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD();
-// Split the data into training and test sets (30% held out for testing)
-JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
-JavaRDD<LabeledPoint> trainingData = splits[0];
-JavaRDD<LabeledPoint> testData = splits[1];
-
-// Train a RandomForest model.
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-Integer numClasses = 2;
-HashMap<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-Integer numTrees = 3; // Use more in practice.
-String featureSubsetStrategy = "auto"; // Let the algorithm choose.
-String impurity = "gini";
-Integer maxDepth = 5;
-Integer maxBins = 32;
-Integer seed = 12345;
-
-final RandomForestModel model = RandomForest.trainClassifier(trainingData, numClasses,
-  categoricalFeaturesInfo, numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins,
-  seed);
-
-// Evaluate model on test instances and compute test error
-JavaPairRDD<Double, Double> predictionAndLabel =
-  testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-    @Override
-    public Tuple2<Double, Double> call(LabeledPoint p) {
-      return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-    }
-  });
-Double testErr =
-  1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
-    @Override
-    public Boolean call(Tuple2<Double, Double> pl) {
-      return !pl._1().equals(pl._2());
-    }
-  }).count() / testData.count();
-System.out.println("Test Error: " + testErr);
-System.out.println("Learned classification forest model:\n" + model.toDebugString());
-
-// Save and load model
-model.save(sc.sc(), "myModelPath");
-RandomForestModel sameModel = RandomForestModel.load(sc.sc(), "myModelPath");
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaRandomForestClassificationExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`RandomForest` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.RandomForest) and [`RandomForest` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.RandomForestModel) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.tree import RandomForest, RandomForestModel
-from pyspark.mllib.util import MLUtils
-
-# Load and parse the data file into an RDD of LabeledPoint.
-data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a RandomForest model.
-#  Empty categoricalFeaturesInfo indicates all features are continuous.
-#  Note: Use larger numTrees in practice.
-#  Setting featureSubsetStrategy="auto" lets the algorithm choose.
-model = RandomForest.trainClassifier(trainingData, numClasses=2, categoricalFeaturesInfo={},
-                                     numTrees=3, featureSubsetStrategy="auto",
-                                     impurity='gini', maxDepth=4, maxBins=32)
-
-# Evaluate model on test instances and compute test error
-predictions = model.predict(testData.map(lambda x: x.features))
-labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-testErr = labelsAndPredictions.filter(lambda (v, p): v != p).count() / float(testData.count())
-print('Test Error = ' + str(testErr))
-print('Learned classification forest model:')
-print(model.toDebugString())
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = RandomForestModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example python/mllib/random_forest_classification_example.py %}
 </div>
 
 </div>
@@ -254,147 +129,19 @@ The Mean Squared Error (MSE) is computed at the end to evaluate
 <div data-lang="scala" markdown="1">
 Refer to the [`RandomForest` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.RandomForest) and [`RandomForestModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.model.RandomForestModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.tree.RandomForest
-import org.apache.spark.mllib.tree.model.RandomForestModel
-import org.apache.spark.mllib.util.MLUtils
-
-// Load and parse the data file.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-// Split the data into training and test sets (30% held out for testing)
-val splits = data.randomSplit(Array(0.7, 0.3))
-val (trainingData, testData) = (splits(0), splits(1))
-
-// Train a RandomForest model.
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-val numClasses = 2
-val categoricalFeaturesInfo = Map[Int, Int]()
-val numTrees = 3 // Use more in practice.
-val featureSubsetStrategy = "auto" // Let the algorithm choose.
-val impurity = "variance"
-val maxDepth = 4
-val maxBins = 32
-
-val model = RandomForest.trainRegressor(trainingData, categoricalFeaturesInfo,
-  numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins)
-
-// Evaluate model on test instances and compute test error
-val labelsAndPredictions = testData.map { point =>
-  val prediction = model.predict(point.features)
-  (point.label, prediction)
-}
-val testMSE = labelsAndPredictions.map{ case(v, p) => math.pow((v - p), 2)}.mean()
-println("Test Mean Squared Error = " + testMSE)
-println("Learned regression forest model:\n" + model.toDebugString)
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = RandomForestModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/RandomForestRegressionExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`RandomForest` Java docs](api/java/org/apache/spark/mllib/tree/RandomForest.html) and [`RandomForestModel` Java docs](api/java/org/apache/spark/mllib/tree/model/RandomForestModel.html) for details on the API.
 
-{% highlight java %}
-import java.util.HashMap;
-import scala.Tuple2;
-import org.apache.spark.api.java.function.Function2;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.RandomForest;
-import org.apache.spark.mllib.tree.model.RandomForestModel;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.SparkConf;
-
-SparkConf sparkConf = new SparkConf().setAppName("JavaRandomForest");
-JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-// Load and parse the data file.
-String datapath = "data/mllib/sample_libsvm_data.txt";
-JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD();
-// Split the data into training and test sets (30% held out for testing)
-JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
-JavaRDD<LabeledPoint> trainingData = splits[0];
-JavaRDD<LabeledPoint> testData = splits[1];
-
-// Set parameters.
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-String impurity = "variance";
-Integer maxDepth = 4;
-Integer maxBins = 32;
-
-// Train a RandomForest model.
-final RandomForestModel model = RandomForest.trainRegressor(trainingData,
-  categoricalFeaturesInfo, impurity, maxDepth, maxBins);
-
-// Evaluate model on test instances and compute test error
-JavaPairRDD<Double, Double> predictionAndLabel =
-  testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-    @Override
-    public Tuple2<Double, Double> call(LabeledPoint p) {
-      return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-    }
-  });
-Double testMSE =
-  predictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
-    @Override
-    public Double call(Tuple2<Double, Double> pl) {
-      Double diff = pl._1() - pl._2();
-      return diff * diff;
-    }
-  }).reduce(new Function2<Double, Double, Double>() {
-    @Override
-    public Double call(Double a, Double b) {
-      return a + b;
-    }
-  }) / testData.count();
-System.out.println("Test Mean Squared Error: " + testMSE);
-System.out.println("Learned regression forest model:\n" + model.toDebugString());
-
-// Save and load model
-model.save(sc.sc(), "myModelPath");
-RandomForestModel sameModel = RandomForestModel.load(sc.sc(), "myModelPath");
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaRandomForestRegressionExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`RandomForest` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.RandomForest) and [`RandomForest` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.RandomForestModel) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.tree import RandomForest, RandomForestModel
-from pyspark.mllib.util import MLUtils
-
-# Load and parse the data file into an RDD of LabeledPoint.
-data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a RandomForest model.
-#  Empty categoricalFeaturesInfo indicates all features are continuous.
-#  Note: Use larger numTrees in practice.
-#  Setting featureSubsetStrategy="auto" lets the algorithm choose.
-model = RandomForest.trainRegressor(trainingData, categoricalFeaturesInfo={},
-                                    numTrees=3, featureSubsetStrategy="auto",
-                                    impurity='variance', maxDepth=4, maxBins=32)
-
-# Evaluate model on test instances and compute test error
-predictions = model.predict(testData.map(lambda x: x.features))
-labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-testMSE = labelsAndPredictions.map(lambda (v, p): (v - p) * (v - p)).sum() / float(testData.count())
-print('Test Mean Squared Error = ' + str(testMSE))
-print('Learned regression forest model:')
-print(model.toDebugString())
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = RandomForestModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example python/mllib/random_forest_regression_example.py %}
 </div>
 
 </div>
@@ -492,141 +239,19 @@ The test error is calculated to measure the algorithm accuracy.
 <div data-lang="scala" markdown="1">
 Refer to the [`GradientBoostedTrees` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.GradientBoostedTrees) and [`GradientBoostedTreesModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.model.GradientBoostedTreesModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.tree.GradientBoostedTrees
-import org.apache.spark.mllib.tree.configuration.BoostingStrategy
-import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel
-import org.apache.spark.mllib.util.MLUtils
-
-// Load and parse the data file.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-// Split the data into training and test sets (30% held out for testing)
-val splits = data.randomSplit(Array(0.7, 0.3))
-val (trainingData, testData) = (splits(0), splits(1))
-
-// Train a GradientBoostedTrees model.
-//  The defaultParams for Classification use LogLoss by default.
-val boostingStrategy = BoostingStrategy.defaultParams("Classification")
-boostingStrategy.numIterations = 3 // Note: Use more iterations in practice.
-boostingStrategy.treeStrategy.numClasses = 2
-boostingStrategy.treeStrategy.maxDepth = 5
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-boostingStrategy.treeStrategy.categoricalFeaturesInfo = Map[Int, Int]()
-
-val model = GradientBoostedTrees.train(trainingData, boostingStrategy)
-
-// Evaluate model on test instances and compute test error
-val labelAndPreds = testData.map { point =>
-  val prediction = model.predict(point.features)
-  (point.label, prediction)
-}
-val testErr = labelAndPreds.filter(r => r._1 != r._2).count.toDouble / testData.count()
-println("Test Error = " + testErr)
-println("Learned classification GBT model:\n" + model.toDebugString)
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = GradientBoostedTreesModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/GradientBoostingClassificationExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`GradientBoostedTrees` Java docs](api/java/org/apache/spark/mllib/tree/GradientBoostedTrees.html) and [`GradientBoostedTreesModel` Java docs](api/java/org/apache/spark/mllib/tree/model/GradientBoostedTreesModel.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-import java.util.HashMap;
-import java.util.Map;
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.GradientBoostedTrees;
-import org.apache.spark.mllib.tree.configuration.BoostingStrategy;
-import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel;
-import org.apache.spark.mllib.util.MLUtils;
-
-SparkConf sparkConf = new SparkConf().setAppName("JavaGradientBoostedTrees");
-JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-// Load and parse the data file.
-String datapath = "data/mllib/sample_libsvm_data.txt";
-JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD();
-// Split the data into training and test sets (30% held out for testing)
-JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
-JavaRDD<LabeledPoint> trainingData = splits[0];
-JavaRDD<LabeledPoint> testData = splits[1];
-
-// Train a GradientBoostedTrees model.
-//  The defaultParams for Classification use LogLoss by default.
-BoostingStrategy boostingStrategy = BoostingStrategy.defaultParams("Classification");
-boostingStrategy.setNumIterations(3); // Note: Use more iterations in practice.
-boostingStrategy.getTreeStrategy().setNumClassesForClassification(2);
-boostingStrategy.getTreeStrategy().setMaxDepth(5);
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-boostingStrategy.treeStrategy().setCategoricalFeaturesInfo(categoricalFeaturesInfo);
-
-final GradientBoostedTreesModel model =
-  GradientBoostedTrees.train(trainingData, boostingStrategy);
-
-// Evaluate model on test instances and compute test error
-JavaPairRDD<Double, Double> predictionAndLabel =
-  testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-    @Override
-    public Tuple2<Double, Double> call(LabeledPoint p) {
-      return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-    }
-  });
-Double testErr =
-  1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
-    @Override
-    public Boolean call(Tuple2<Double, Double> pl) {
-      return !pl._1().equals(pl._2());
-    }
-  }).count() / testData.count();
-System.out.println("Test Error: " + testErr);
-System.out.println("Learned classification GBT model:\n" + model.toDebugString());
-
-// Save and load model
-model.save(sc.sc(), "myModelPath");
-GradientBoostedTreesModel sameModel = GradientBoostedTreesModel.load(sc.sc(), "myModelPath");
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaGradientBoostingClassificationExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`GradientBoostedTrees` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.GradientBoostedTrees) and [`GradientBoostedTreesModel` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.GradientBoostedTreesModel) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.tree import GradientBoostedTrees, GradientBoostedTreesModel
-from pyspark.mllib.util import MLUtils
-
-# Load and parse the data file.
-data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a GradientBoostedTrees model.
-#  Notes: (a) Empty categoricalFeaturesInfo indicates all features are continuous.
-#         (b) Use more iterations in practice.
-model = GradientBoostedTrees.trainClassifier(trainingData,
-    categoricalFeaturesInfo={}, numIterations=3)
-
-# Evaluate model on test instances and compute test error
-predictions = model.predict(testData.map(lambda x: x.features))
-labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-testErr = labelsAndPredictions.filter(lambda (v, p): v != p).count() / float(testData.count())
-print('Test Error = ' + str(testErr))
-print('Learned classification GBT model:')
-print(model.toDebugString())
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = GradientBoostedTreesModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example python/mllib/gradient_boosting_classification_example.py %}
 </div>
 
 </div>
@@ -645,146 +270,19 @@ The Mean Squared Error (MSE) is computed at the end to evaluate
 <div data-lang="scala" markdown="1">
 Refer to the [`GradientBoostedTrees` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.GradientBoostedTrees) and [`GradientBoostedTreesModel` Scala docs](api/scala/index.html#org.apache.spark.mllib.tree.model.GradientBoostedTreesModel) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.tree.GradientBoostedTrees
-import org.apache.spark.mllib.tree.configuration.BoostingStrategy
-import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel
-import org.apache.spark.mllib.util.MLUtils
-
-// Load and parse the data file.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-// Split the data into training and test sets (30% held out for testing)
-val splits = data.randomSplit(Array(0.7, 0.3))
-val (trainingData, testData) = (splits(0), splits(1))
-
-// Train a GradientBoostedTrees model.
-//  The defaultParams for Regression use SquaredError by default.
-val boostingStrategy = BoostingStrategy.defaultParams("Regression")
-boostingStrategy.numIterations = 3 // Note: Use more iterations in practice.
-boostingStrategy.treeStrategy.maxDepth = 5
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-boostingStrategy.treeStrategy.categoricalFeaturesInfo = Map[Int, Int]()
-
-val model = GradientBoostedTrees.train(trainingData, boostingStrategy)
-
-// Evaluate model on test instances and compute test error
-val labelsAndPredictions = testData.map { point =>
-  val prediction = model.predict(point.features)
-  (point.label, prediction)
-}
-val testMSE = labelsAndPredictions.map{ case(v, p) => math.pow((v - p), 2)}.mean()
-println("Test Mean Squared Error = " + testMSE)
-println("Learned regression GBT model:\n" + model.toDebugString)
-
-// Save and load model
-model.save(sc, "myModelPath")
-val sameModel = GradientBoostedTreesModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/GradientBoostingRegressionExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`GradientBoostedTrees` Java docs](api/java/org/apache/spark/mllib/tree/GradientBoostedTrees.html) and [`GradientBoostedTreesModel` Java docs](api/java/org/apache/spark/mllib/tree/model/GradientBoostedTreesModel.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-import java.util.HashMap;
-import java.util.Map;
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.function.Function2;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.GradientBoostedTrees;
-import org.apache.spark.mllib.tree.configuration.BoostingStrategy;
-import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel;
-import org.apache.spark.mllib.util.MLUtils;
-
-SparkConf sparkConf = new SparkConf().setAppName("JavaGradientBoostedTrees");
-JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-// Load and parse the data file.
-String datapath = "data/mllib/sample_libsvm_data.txt";
-JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD();
-// Split the data into training and test sets (30% held out for testing)
-JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
-JavaRDD<LabeledPoint> trainingData = splits[0];
-JavaRDD<LabeledPoint> testData = splits[1];
-
-// Train a GradientBoostedTrees model.
-//  The defaultParams for Regression use SquaredError by default.
-BoostingStrategy boostingStrategy = BoostingStrategy.defaultParams("Regression");
-boostingStrategy.setNumIterations(3); // Note: Use more iterations in practice.
-boostingStrategy.getTreeStrategy().setMaxDepth(5);
-//  Empty categoricalFeaturesInfo indicates all features are continuous.
-Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-boostingStrategy.treeStrategy().setCategoricalFeaturesInfo(categoricalFeaturesInfo);
-
-final GradientBoostedTreesModel model =
-  GradientBoostedTrees.train(trainingData, boostingStrategy);
-
-// Evaluate model on test instances and compute test error
-JavaPairRDD<Double, Double> predictionAndLabel =
-  testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-    @Override
-    public Tuple2<Double, Double> call(LabeledPoint p) {
-      return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-    }
-  });
-Double testMSE =
-  predictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
-    @Override
-    public Double call(Tuple2<Double, Double> pl) {
-      Double diff = pl._1() - pl._2();
-      return diff * diff;
-    }
-  }).reduce(new Function2<Double, Double, Double>() {
-    @Override
-    public Double call(Double a, Double b) {
-      return a + b;
-    }
-  }) / data.count();
-System.out.println("Test Mean Squared Error: " + testMSE);
-System.out.println("Learned regression GBT model:\n" + model.toDebugString());
-
-// Save and load model
-model.save(sc.sc(), "myModelPath");
-GradientBoostedTreesModel sameModel = GradientBoostedTreesModel.load(sc.sc(), "myModelPath");
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaGradientBoostingRegressionExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`GradientBoostedTrees` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.GradientBoostedTrees) and [`GradientBoostedTreesModel` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.tree.GradientBoostedTreesModel) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.tree import GradientBoostedTrees, GradientBoostedTreesModel
-from pyspark.mllib.util import MLUtils
-
-# Load and parse the data file.
-data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a GradientBoostedTrees model.
-#  Notes: (a) Empty categoricalFeaturesInfo indicates all features are continuous.
-#         (b) Use more iterations in practice.
-model = GradientBoostedTrees.trainRegressor(trainingData,
-    categoricalFeaturesInfo={}, numIterations=3)
-
-# Evaluate model on test instances and compute test error
-predictions = model.predict(testData.map(lambda x: x.features))
-labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-testMSE = labelsAndPredictions.map(lambda (v, p): (v - p) * (v - p)).sum() / float(testData.count())
-print('Test Mean Squared Error = ' + str(testMSE))
-print('Learned regression GBT model:')
-print(model.toDebugString())
-
-# Save and load model
-model.save(sc, "myModelPath")
-sameModel = GradientBoostedTreesModel.load(sc, "myModelPath")
-{% endhighlight %}
+{% include_example python/mllib/gradient_boosting_regression_example.py %}
 </div>
 
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingClassificationExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingClassificationExample.java
new file mode 100644
index 0000000000000..80faabd2325d0
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingClassificationExample.java
@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.HashMap;
+import java.util.Map;
+
+import scala.Tuple2;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.tree.GradientBoostedTrees;
+import org.apache.spark.mllib.tree.configuration.BoostingStrategy;
+import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel;
+import org.apache.spark.mllib.util.MLUtils;
+// $example off$
+
+public class JavaGradientBoostingClassificationExample {
+  public static void main(String[] args) {
+    // $example on$
+    SparkConf sparkConf = new SparkConf()
+      .setAppName("JavaGradientBoostedTreesClassificationExample");
+    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
+
+    // Load and parse the data file.
+    String datapath = "data/mllib/sample_libsvm_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD();
+    // Split the data into training and test sets (30% held out for testing)
+    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
+    JavaRDD<LabeledPoint> trainingData = splits[0];
+    JavaRDD<LabeledPoint> testData = splits[1];
+
+    // Train a GradientBoostedTrees model.
+    // The defaultParams for Classification use LogLoss by default.
+    BoostingStrategy boostingStrategy = BoostingStrategy.defaultParams("Classification");
+    boostingStrategy.setNumIterations(3); // Note: Use more iterations in practice.
+    boostingStrategy.getTreeStrategy().setNumClasses(2);
+    boostingStrategy.getTreeStrategy().setMaxDepth(5);
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
+    boostingStrategy.treeStrategy().setCategoricalFeaturesInfo(categoricalFeaturesInfo);
+
+    final GradientBoostedTreesModel model =
+      GradientBoostedTrees.train(trainingData, boostingStrategy);
+
+    // Evaluate model on test instances and compute test error
+    JavaPairRDD<Double, Double> predictionAndLabel =
+      testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
+        @Override
+        public Tuple2<Double, Double> call(LabeledPoint p) {
+          return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
+        }
+      });
+    Double testErr =
+      1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
+        @Override
+        public Boolean call(Tuple2<Double, Double> pl) {
+          return !pl._1().equals(pl._2());
+        }
+      }).count() / testData.count();
+    System.out.println("Test Error: " + testErr);
+    System.out.println("Learned classification GBT model:\n" + model.toDebugString());
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myGradientBoostingClassificationModel");
+    GradientBoostedTreesModel sameModel = GradientBoostedTreesModel.load(jsc.sc(),
+      "target/tmp/myGradientBoostingClassificationModel");
+    // $example off$
+  }
+
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingRegressionExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingRegressionExample.java
new file mode 100644
index 0000000000000..216895b368202
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostingRegressionExample.java
@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.HashMap;
+import java.util.Map;
+
+import scala.Tuple2;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.function.Function2;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.tree.GradientBoostedTrees;
+import org.apache.spark.mllib.tree.configuration.BoostingStrategy;
+import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel;
+import org.apache.spark.mllib.util.MLUtils;
+// $example off$
+
+public class JavaGradientBoostingRegressionExample {
+  public static void main(String[] args) {
+    // $example on$
+    SparkConf sparkConf = new SparkConf()
+      .setAppName("JavaGradientBoostedTreesRegressionExample");
+    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
+    // Load and parse the data file.
+    String datapath = "data/mllib/sample_libsvm_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD();
+    // Split the data into training and test sets (30% held out for testing)
+    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
+    JavaRDD<LabeledPoint> trainingData = splits[0];
+    JavaRDD<LabeledPoint> testData = splits[1];
+
+    // Train a GradientBoostedTrees model.
+    // The defaultParams for Regression use SquaredError by default.
+    BoostingStrategy boostingStrategy = BoostingStrategy.defaultParams("Regression");
+    boostingStrategy.setNumIterations(3); // Note: Use more iterations in practice.
+    boostingStrategy.getTreeStrategy().setMaxDepth(5);
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
+    boostingStrategy.treeStrategy().setCategoricalFeaturesInfo(categoricalFeaturesInfo);
+
+    final GradientBoostedTreesModel model =
+      GradientBoostedTrees.train(trainingData, boostingStrategy);
+
+    // Evaluate model on test instances and compute test error
+    JavaPairRDD<Double, Double> predictionAndLabel =
+      testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
+        @Override
+        public Tuple2<Double, Double> call(LabeledPoint p) {
+          return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
+        }
+      });
+    Double testMSE =
+      predictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
+        @Override
+        public Double call(Tuple2<Double, Double> pl) {
+          Double diff = pl._1() - pl._2();
+          return diff * diff;
+        }
+      }).reduce(new Function2<Double, Double, Double>() {
+        @Override
+        public Double call(Double a, Double b) {
+          return a + b;
+        }
+      }) / data.count();
+    System.out.println("Test Mean Squared Error: " + testMSE);
+    System.out.println("Learned regression GBT model:\n" + model.toDebugString());
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myGradientBoostingRegressionModel");
+    GradientBoostedTreesModel sameModel = GradientBoostedTreesModel.load(jsc.sc(),
+      "target/tmp/myGradientBoostingRegressionModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestClassificationExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestClassificationExample.java
new file mode 100644
index 0000000000000..9219eef1ad2d6
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestClassificationExample.java
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.HashMap;
+
+import scala.Tuple2;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.tree.RandomForest;
+import org.apache.spark.mllib.tree.model.RandomForestModel;
+import org.apache.spark.mllib.util.MLUtils;
+// $example off$
+
+public class JavaRandomForestClassificationExample {
+  public static void main(String[] args) {
+    // $example on$
+    SparkConf sparkConf = new SparkConf().setAppName("JavaRandomForestClassificationExample");
+    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
+    // Load and parse the data file.
+    String datapath = "data/mllib/sample_libsvm_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD();
+    // Split the data into training and test sets (30% held out for testing)
+    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
+    JavaRDD<LabeledPoint> trainingData = splits[0];
+    JavaRDD<LabeledPoint> testData = splits[1];
+
+    // Train a RandomForest model.
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    Integer numClasses = 2;
+    HashMap<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
+    Integer numTrees = 3; // Use more in practice.
+    String featureSubsetStrategy = "auto"; // Let the algorithm choose.
+    String impurity = "gini";
+    Integer maxDepth = 5;
+    Integer maxBins = 32;
+    Integer seed = 12345;
+
+    final RandomForestModel model = RandomForest.trainClassifier(trainingData, numClasses,
+      categoricalFeaturesInfo, numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins,
+      seed);
+
+    // Evaluate model on test instances and compute test error
+    JavaPairRDD<Double, Double> predictionAndLabel =
+      testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
+        @Override
+        public Tuple2<Double, Double> call(LabeledPoint p) {
+          return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
+        }
+      });
+    Double testErr =
+      1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
+        @Override
+        public Boolean call(Tuple2<Double, Double> pl) {
+          return !pl._1().equals(pl._2());
+        }
+      }).count() / testData.count();
+    System.out.println("Test Error: " + testErr);
+    System.out.println("Learned classification forest model:\n" + model.toDebugString());
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myRandomForestClassificationModel");
+    RandomForestModel sameModel = RandomForestModel.load(jsc.sc(),
+      "target/tmp/myRandomForestClassificationModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestRegressionExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestRegressionExample.java
new file mode 100644
index 0000000000000..4db926a4218ff
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestRegressionExample.java
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.HashMap;
+import java.util.Map;
+
+import scala.Tuple2;
+
+import org.apache.spark.api.java.function.Function2;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.api.java.function.PairFunction;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.tree.RandomForest;
+import org.apache.spark.mllib.tree.model.RandomForestModel;
+import org.apache.spark.mllib.util.MLUtils;
+import org.apache.spark.SparkConf;
+// $example off$
+
+public class JavaRandomForestRegressionExample {
+  public static void main(String[] args) {
+    // $example on$
+    SparkConf sparkConf = new SparkConf().setAppName("JavaRandomForestRegressionExample");
+    JavaSparkContext jsc = new JavaSparkContext(sparkConf);
+    // Load and parse the data file.
+    String datapath = "data/mllib/sample_libsvm_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD();
+    // Split the data into training and test sets (30% held out for testing)
+    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
+    JavaRDD<LabeledPoint> trainingData = splits[0];
+    JavaRDD<LabeledPoint> testData = splits[1];
+
+    // Set parameters.
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
+    Integer numTrees = 3; // Use more in practice.
+    String featureSubsetStrategy = "auto"; // Let the algorithm choose.
+    String impurity = "variance";
+    Integer maxDepth = 4;
+    Integer maxBins = 32;
+    Integer seed = 12345;
+    // Train a RandomForest model.
+    final RandomForestModel model = RandomForest.trainRegressor(trainingData,
+      categoricalFeaturesInfo, numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins, seed);
+
+    // Evaluate model on test instances and compute test error
+    JavaPairRDD<Double, Double> predictionAndLabel =
+      testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
+        @Override
+        public Tuple2<Double, Double> call(LabeledPoint p) {
+          return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
+        }
+      });
+    Double testMSE =
+      predictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
+        @Override
+        public Double call(Tuple2<Double, Double> pl) {
+          Double diff = pl._1() - pl._2();
+          return diff * diff;
+        }
+      }).reduce(new Function2<Double, Double, Double>() {
+        @Override
+        public Double call(Double a, Double b) {
+          return a + b;
+        }
+      }) / testData.count();
+    System.out.println("Test Mean Squared Error: " + testMSE);
+    System.out.println("Learned regression forest model:\n" + model.toDebugString());
+
+    // Save and load model
+    model.save(jsc.sc(), "target/tmp/myRandomForestRegressionModel");
+    RandomForestModel sameModel = RandomForestModel.load(jsc.sc(),
+      "target/tmp/myRandomForestRegressionModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/python/mllib/gradient_boosting_classification_example.py b/examples/src/main/python/mllib/gradient_boosting_classification_example.py
new file mode 100644
index 0000000000000..a94ea0d582e59
--- /dev/null
+++ b/examples/src/main/python/mllib/gradient_boosting_classification_example.py
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Gradient Boosted Trees Classification Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext
+# $example on$
+from pyspark.mllib.tree import GradientBoostedTrees, GradientBoostedTreesModel
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="PythonGradientBoostedTreesClassificationExample")
+    # $example on$
+    # Load and parse the data file.
+    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a GradientBoostedTrees model.
+    #  Notes: (a) Empty categoricalFeaturesInfo indicates all features are continuous.
+    #         (b) Use more iterations in practice.
+    model = GradientBoostedTrees.trainClassifier(trainingData,
+                                                 categoricalFeaturesInfo={}, numIterations=3)
+
+    # Evaluate model on test instances and compute test error
+    predictions = model.predict(testData.map(lambda x: x.features))
+    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
+    testErr = labelsAndPredictions.filter(lambda (v, p): v != p).count() / float(testData.count())
+    print('Test Error = ' + str(testErr))
+    print('Learned classification GBT model:')
+    print(model.toDebugString())
+
+    # Save and load model
+    model.save(sc, "target/tmp/myGradientBoostingClassificationModel")
+    sameModel = GradientBoostedTreesModel.load(sc,
+                                               "target/tmp/myGradientBoostingClassificationModel")
+    # $example off$
diff --git a/examples/src/main/python/mllib/gradient_boosting_regression_example.py b/examples/src/main/python/mllib/gradient_boosting_regression_example.py
new file mode 100644
index 0000000000000..86040799dc1d9
--- /dev/null
+++ b/examples/src/main/python/mllib/gradient_boosting_regression_example.py
@@ -0,0 +1,57 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Gradient Boosted Trees Regression Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext
+# $example on$
+from pyspark.mllib.tree import GradientBoostedTrees, GradientBoostedTreesModel
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="PythonGradientBoostedTreesRegressionExample")
+    # $example on$
+    # Load and parse the data file.
+    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a GradientBoostedTrees model.
+    #  Notes: (a) Empty categoricalFeaturesInfo indicates all features are continuous.
+    #         (b) Use more iterations in practice.
+    model = GradientBoostedTrees.trainRegressor(trainingData,
+                                                categoricalFeaturesInfo={}, numIterations=3)
+
+    # Evaluate model on test instances and compute test error
+    predictions = model.predict(testData.map(lambda x: x.features))
+    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
+    testMSE = labelsAndPredictions.map(lambda (v, p): (v - p) * (v - p)).sum() /\
+        float(testData.count())
+    print('Test Mean Squared Error = ' + str(testMSE))
+    print('Learned regression GBT model:')
+    print(model.toDebugString())
+
+    # Save and load model
+    model.save(sc, "target/tmp/myGradientBoostingRegressionModel")
+    sameModel = GradientBoostedTreesModel.load(sc, "target/tmp/myGradientBoostingRegressionModel")
+    # $example off$
diff --git a/examples/src/main/python/mllib/random_forest_classification_example.py b/examples/src/main/python/mllib/random_forest_classification_example.py
new file mode 100644
index 0000000000000..324ba50625d25
--- /dev/null
+++ b/examples/src/main/python/mllib/random_forest_classification_example.py
@@ -0,0 +1,58 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Random Forest Classification Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext
+# $example on$
+from pyspark.mllib.tree import RandomForest, RandomForestModel
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="PythonRandomForestClassificationExample")
+    # $example on$
+    # Load and parse the data file into an RDD of LabeledPoint.
+    data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a RandomForest model.
+    #  Empty categoricalFeaturesInfo indicates all features are continuous.
+    #  Note: Use larger numTrees in practice.
+    #  Setting featureSubsetStrategy="auto" lets the algorithm choose.
+    model = RandomForest.trainClassifier(trainingData, numClasses=2, categoricalFeaturesInfo={},
+                                         numTrees=3, featureSubsetStrategy="auto",
+                                         impurity='gini', maxDepth=4, maxBins=32)
+
+    # Evaluate model on test instances and compute test error
+    predictions = model.predict(testData.map(lambda x: x.features))
+    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
+    testErr = labelsAndPredictions.filter(lambda (v, p): v != p).count() / float(testData.count())
+    print('Test Error = ' + str(testErr))
+    print('Learned classification forest model:')
+    print(model.toDebugString())
+
+    # Save and load model
+    model.save(sc, "target/tmp/myRandomForestClassificationModel")
+    sameModel = RandomForestModel.load(sc, "target/tmp/myRandomForestClassificationModel")
+    # $example off$
diff --git a/examples/src/main/python/mllib/random_forest_regression_example.py b/examples/src/main/python/mllib/random_forest_regression_example.py
new file mode 100644
index 0000000000000..f7aa6114eceb3
--- /dev/null
+++ b/examples/src/main/python/mllib/random_forest_regression_example.py
@@ -0,0 +1,59 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Random Forest Regression Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext
+# $example on$
+from pyspark.mllib.tree import RandomForest, RandomForestModel
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="PythonRandomForestRegressionExample")
+    # $example on$
+    # Load and parse the data file into an RDD of LabeledPoint.
+    data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a RandomForest model.
+    #  Empty categoricalFeaturesInfo indicates all features are continuous.
+    #  Note: Use larger numTrees in practice.
+    #  Setting featureSubsetStrategy="auto" lets the algorithm choose.
+    model = RandomForest.trainRegressor(trainingData, categoricalFeaturesInfo={},
+                                        numTrees=3, featureSubsetStrategy="auto",
+                                        impurity='variance', maxDepth=4, maxBins=32)
+
+    # Evaluate model on test instances and compute test error
+    predictions = model.predict(testData.map(lambda x: x.features))
+    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
+    testMSE = labelsAndPredictions.map(lambda (v, p): (v - p) * (v - p)).sum() /\
+        float(testData.count())
+    print('Test Mean Squared Error = ' + str(testMSE))
+    print('Learned regression forest model:')
+    print(model.toDebugString())
+
+    # Save and load model
+    model.save(sc, "target/tmp/myRandomForestRegressionModel")
+    sameModel = RandomForestModel.load(sc, "target/tmp/myRandomForestRegressionModel")
+    # $example off$
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/GradientBoostingClassificationExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/GradientBoostingClassificationExample.scala
new file mode 100644
index 0000000000000..139e1f909bdce
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/GradientBoostingClassificationExample.scala
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
+import org.apache.spark.mllib.tree.GradientBoostedTrees
+import org.apache.spark.mllib.tree.configuration.BoostingStrategy
+import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+
+object GradientBoostingClassificationExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("GradientBoostedTreesClassificationExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    // Load and parse the data file.
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    // Split the data into training and test sets (30% held out for testing)
+    val splits = data.randomSplit(Array(0.7, 0.3))
+    val (trainingData, testData) = (splits(0), splits(1))
+
+    // Train a GradientBoostedTrees model.
+    // The defaultParams for Classification use LogLoss by default.
+    val boostingStrategy = BoostingStrategy.defaultParams("Classification")
+    boostingStrategy.numIterations = 3 // Note: Use more iterations in practice.
+    boostingStrategy.treeStrategy.numClasses = 2
+    boostingStrategy.treeStrategy.maxDepth = 5
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    boostingStrategy.treeStrategy.categoricalFeaturesInfo = Map[Int, Int]()
+
+    val model = GradientBoostedTrees.train(trainingData, boostingStrategy)
+
+    // Evaluate model on test instances and compute test error
+    val labelAndPreds = testData.map { point =>
+      val prediction = model.predict(point.features)
+      (point.label, prediction)
+    }
+    val testErr = labelAndPreds.filter(r => r._1 != r._2).count.toDouble / testData.count()
+    println("Test Error = " + testErr)
+    println("Learned classification GBT model:\n" + model.toDebugString)
+
+    // Save and load model
+    model.save(sc, "target/tmp/myGradientBoostingClassificationModel")
+    val sameModel = GradientBoostedTreesModel.load(sc,
+      "target/tmp/myGradientBoostingClassificationModel")
+    // $example off$
+  }
+}
+// scalastyle:on println
+
+
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/GradientBoostingRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/GradientBoostingRegressionExample.scala
new file mode 100644
index 0000000000000..3dc86da8e4d2b
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/GradientBoostingRegressionExample.scala
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
+import org.apache.spark.mllib.tree.GradientBoostedTrees
+import org.apache.spark.mllib.tree.configuration.BoostingStrategy
+import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+
+object GradientBoostingRegressionExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("GradientBoostedTreesRegressionExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    // Load and parse the data file.
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    // Split the data into training and test sets (30% held out for testing)
+    val splits = data.randomSplit(Array(0.7, 0.3))
+    val (trainingData, testData) = (splits(0), splits(1))
+
+    // Train a GradientBoostedTrees model.
+    // The defaultParams for Regression use SquaredError by default.
+    val boostingStrategy = BoostingStrategy.defaultParams("Regression")
+    boostingStrategy.numIterations = 3 // Note: Use more iterations in practice.
+    boostingStrategy.treeStrategy.maxDepth = 5
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    boostingStrategy.treeStrategy.categoricalFeaturesInfo = Map[Int, Int]()
+
+    val model = GradientBoostedTrees.train(trainingData, boostingStrategy)
+
+    // Evaluate model on test instances and compute test error
+    val labelsAndPredictions = testData.map { point =>
+      val prediction = model.predict(point.features)
+      (point.label, prediction)
+    }
+    val testMSE = labelsAndPredictions.map{ case(v, p) => math.pow((v - p), 2)}.mean()
+    println("Test Mean Squared Error = " + testMSE)
+    println("Learned regression GBT model:\n" + model.toDebugString)
+
+    // Save and load model
+    model.save(sc, "target/tmp/myGradientBoostingRegressionModel")
+    val sameModel = GradientBoostedTreesModel.load(sc,
+      "target/tmp/myGradientBoostingRegressionModel")
+    // $example off$
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/RandomForestClassificationExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/RandomForestClassificationExample.scala
new file mode 100644
index 0000000000000..5e55abd5121c4
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/RandomForestClassificationExample.scala
@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
+import org.apache.spark.mllib.tree.RandomForest
+import org.apache.spark.mllib.tree.model.RandomForestModel
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+
+object RandomForestClassificationExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("RandomForestClassificationExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    // Load and parse the data file.
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    // Split the data into training and test sets (30% held out for testing)
+    val splits = data.randomSplit(Array(0.7, 0.3))
+    val (trainingData, testData) = (splits(0), splits(1))
+
+    // Train a RandomForest model.
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    val numClasses = 2
+    val categoricalFeaturesInfo = Map[Int, Int]()
+    val numTrees = 3 // Use more in practice.
+    val featureSubsetStrategy = "auto" // Let the algorithm choose.
+    val impurity = "gini"
+    val maxDepth = 4
+    val maxBins = 32
+
+    val model = RandomForest.trainClassifier(trainingData, numClasses, categoricalFeaturesInfo,
+      numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins)
+
+    // Evaluate model on test instances and compute test error
+    val labelAndPreds = testData.map { point =>
+      val prediction = model.predict(point.features)
+      (point.label, prediction)
+    }
+    val testErr = labelAndPreds.filter(r => r._1 != r._2).count.toDouble / testData.count()
+    println("Test Error = " + testErr)
+    println("Learned classification forest model:\n" + model.toDebugString)
+
+    // Save and load model
+    model.save(sc, "target/tmp/myRandomForestClassificationModel")
+    val sameModel = RandomForestModel.load(sc, "target/tmp/myRandomForestClassificationModel")
+    // $example off$
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/RandomForestRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/RandomForestRegressionExample.scala
new file mode 100644
index 0000000000000..a54fb3ab7e37a
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/RandomForestRegressionExample.scala
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
+import org.apache.spark.mllib.tree.RandomForest
+import org.apache.spark.mllib.tree.model.RandomForestModel
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+
+object RandomForestRegressionExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("RandomForestRegressionExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    // Load and parse the data file.
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt")
+    // Split the data into training and test sets (30% held out for testing)
+    val splits = data.randomSplit(Array(0.7, 0.3))
+    val (trainingData, testData) = (splits(0), splits(1))
+
+    // Train a RandomForest model.
+    // Empty categoricalFeaturesInfo indicates all features are continuous.
+    val numClasses = 2
+    val categoricalFeaturesInfo = Map[Int, Int]()
+    val numTrees = 3 // Use more in practice.
+    val featureSubsetStrategy = "auto" // Let the algorithm choose.
+    val impurity = "variance"
+    val maxDepth = 4
+    val maxBins = 32
+
+    val model = RandomForest.trainRegressor(trainingData, categoricalFeaturesInfo,
+      numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins)
+
+    // Evaluate model on test instances and compute test error
+    val labelsAndPredictions = testData.map { point =>
+      val prediction = model.predict(point.features)
+      (point.label, prediction)
+    }
+    val testMSE = labelsAndPredictions.map{ case(v, p) => math.pow((v - p), 2)}.mean()
+    println("Test Mean Squared Error = " + testMSE)
+    println("Learned regression forest model:\n" + model.toDebugString)
+
+    // Save and load model
+    model.save(sc, "target/tmp/myRandomForestRegressionModel")
+    val sameModel = RandomForestModel.load(sc, "target/tmp/myRandomForestRegressionModel")
+    // $example off$
+  }
+}
+// scalastyle:on println
+

From 99693fef0a30432d94556154b81872356d921c64 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Fri, 13 Nov 2015 08:43:05 -0800
Subject: [PATCH 233/518] [SPARK-11723][ML][DOC] Use LibSVM data source rather
 than MLUtils.loadLibSVMFile to load DataFrame

Use LibSVM data source rather than MLUtils.loadLibSVMFile to load DataFrame, include:
* Use libSVM data source for all example codes under examples/ml, and remove unused import.
* Use libSVM data source for user guides under ml-*** which were omitted by #8697.
* Fix bug: We should use ```sqlContext.read().format("libsvm").load(path)``` at Java side, but the API doc and user guides misuse as ```sqlContext.read.format("libsvm").load(path)```.
* Code cleanup.

mengxr

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9690 from yanboliang/spark-11723.
---
 docs/ml-ensembles.md                          | 10 ++---
 docs/ml-features.md                           |  8 ++--
 docs/ml-guide.md                              | 10 +----
 docs/ml-linear-methods.md                     |  4 +-
 ...JavaDecisionTreeClassificationExample.java |  8 +---
 .../ml/JavaDecisionTreeRegressionExample.java |  9 ++--
 ...MultilayerPerceptronClassifierExample.java |  6 +--
 .../examples/ml/JavaOneVsRestExample.java     | 23 +++++-----
 .../ml/JavaTrainValidationSplitExample.java   |  6 +--
 .../decision_tree_classification_example.py   |  5 +--
 .../ml/decision_tree_regression_example.py    |  5 +--
 .../main/python/ml/gradient_boosted_trees.py  |  5 +--
 .../src/main/python/ml/logistic_regression.py |  5 +--
 .../multilayer_perceptron_classification.py   |  5 +--
 .../main/python/ml/random_forest_example.py   |  4 +-
 .../DecisionTreeClassificationExample.scala   |  6 +--
 .../examples/ml/DecisionTreeExample.scala     | 42 +++++++------------
 .../ml/DecisionTreeRegressionExample.scala    |  7 ++--
 .../apache/spark/examples/ml/GBTExample.scala |  2 +-
 .../examples/ml/LinearRegressionExample.scala |  2 +-
 .../ml/LogisticRegressionExample.scala        |  4 +-
 ...ultilayerPerceptronClassifierExample.scala |  8 ++--
 .../spark/examples/ml/OneVsRestExample.scala  | 17 ++++----
 .../examples/ml/RandomForestExample.scala     |  2 +-
 .../ml/TrainValidationSplitExample.scala      |  4 +-
 .../ml/source/libsvm/LibSVMRelation.scala     |  2 +-
 26 files changed, 79 insertions(+), 130 deletions(-)

diff --git a/docs/ml-ensembles.md b/docs/ml-ensembles.md
index 58f566c9b4b55..ce15f5e6466ec 100644
--- a/docs/ml-ensembles.md
+++ b/docs/ml-ensembles.md
@@ -195,7 +195,7 @@ import org.apache.spark.ml.feature.*;
 import org.apache.spark.sql.DataFrame;
 
 // Load and parse the data file, converting it to a DataFrame.
-DataFrame data = sqlContext.read.format("libsvm")
+DataFrame data = sqlContext.read().format("libsvm")
   .load("data/mllib/sample_libsvm_data.txt");
 
 // Index labels, adding metadata to the label column.
@@ -384,7 +384,7 @@ import org.apache.spark.ml.regression.RandomForestRegressor;
 import org.apache.spark.sql.DataFrame;
 
 // Load and parse the data file, converting it to a DataFrame.
-DataFrame data = sqlContext.read.format("libsvm")
+DataFrame data = sqlContext.read().format("libsvm")
   .load("data/mllib/sample_libsvm_data.txt");
 
 // Automatically identify categorical features, and index them.
@@ -640,7 +640,7 @@ import org.apache.spark.ml.feature.*;
 import org.apache.spark.sql.DataFrame;
 
 // Load and parse the data file, converting it to a DataFrame.
-DataFrame data sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt");
+DataFrame data sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
 
 // Index labels, adding metadata to the label column.
 // Fit on whole dataset to include all labels in index.
@@ -830,7 +830,7 @@ import org.apache.spark.ml.regression.GBTRegressor;
 import org.apache.spark.sql.DataFrame;
 
 // Load and parse the data file, converting it to a DataFrame.
-DataFrame data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt");
+DataFrame data = sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
 
 // Automatically identify categorical features, and index them.
 // Set maxCategories so features with > 4 distinct values are treated as continuous.
@@ -1000,7 +1000,7 @@ SparkConf conf = new SparkConf().setAppName("JavaOneVsRestExample");
 JavaSparkContext jsc = new JavaSparkContext(conf);
 SQLContext jsql = new SQLContext(jsc);
 
-DataFrame dataFrame = sqlContext.read.format("libsvm")
+DataFrame dataFrame = sqlContext.read().format("libsvm")
   .load("data/mllib/sample_multiclass_classification_data.txt");
 
 DataFrame[] splits = dataFrame.randomSplit(new double[] {0.7, 0.3}, 12345);
diff --git a/docs/ml-features.md b/docs/ml-features.md
index 142afac2f3f95..cd1838d6d2882 100644
--- a/docs/ml-features.md
+++ b/docs/ml-features.md
@@ -1109,7 +1109,7 @@ import org.apache.spark.ml.feature.VectorIndexer;
 import org.apache.spark.ml.feature.VectorIndexerModel;
 import org.apache.spark.sql.DataFrame;
 
-DataFrame data = sqlContext.read.format("libsvm")
+DataFrame data = sqlContext.read().format("libsvm")
   .load("data/mllib/sample_libsvm_data.txt");
 VectorIndexer indexer = new VectorIndexer()
   .setInputCol("features")
@@ -1187,7 +1187,7 @@ for more details on the API.
 import org.apache.spark.ml.feature.Normalizer;
 import org.apache.spark.sql.DataFrame;
 
-DataFrame dataFrame = sqlContext.read.format("libsvm")
+DataFrame dataFrame = sqlContext.read().format("libsvm")
   .load("data/mllib/sample_libsvm_data.txt");
 
 // Normalize each Vector using $L^1$ norm.
@@ -1273,7 +1273,7 @@ import org.apache.spark.ml.feature.StandardScaler;
 import org.apache.spark.ml.feature.StandardScalerModel;
 import org.apache.spark.sql.DataFrame;
 
-DataFrame dataFrame = sqlContext.read.format("libsvm")
+DataFrame dataFrame = sqlContext.read().format("libsvm")
   .load("data/mllib/sample_libsvm_data.txt");
 StandardScaler scaler = new StandardScaler()
   .setInputCol("features")
@@ -1366,7 +1366,7 @@ import org.apache.spark.ml.feature.MinMaxScaler;
 import org.apache.spark.ml.feature.MinMaxScalerModel;
 import org.apache.spark.sql.DataFrame;
 
-DataFrame dataFrame = sqlContext.read.format("libsvm")
+DataFrame dataFrame = sqlContext.read().format("libsvm")
   .load("data/mllib/sample_libsvm_data.txt");
 MinMaxScaler scaler = new MinMaxScaler()
   .setInputCol("features")
diff --git a/docs/ml-guide.md b/docs/ml-guide.md
index c293e71d2870e..be18a05361a17 100644
--- a/docs/ml-guide.md
+++ b/docs/ml-guide.md
@@ -867,10 +867,9 @@ The `ParamMap` which produces the best evaluation metric is selected as the best
 import org.apache.spark.ml.evaluation.RegressionEvaluator
 import org.apache.spark.ml.regression.LinearRegression
 import org.apache.spark.ml.tuning.{ParamGridBuilder, TrainValidationSplit}
-import org.apache.spark.mllib.util.MLUtils
 
 // Prepare training and test data.
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 val Array(training, test) = data.randomSplit(Array(0.9, 0.1), seed = 12345)
 
 val lr = new LinearRegression()
@@ -911,14 +910,9 @@ import org.apache.spark.ml.evaluation.RegressionEvaluator;
 import org.apache.spark.ml.param.ParamMap;
 import org.apache.spark.ml.regression.LinearRegression;
 import org.apache.spark.ml.tuning.*;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.rdd.RDD;
 import org.apache.spark.sql.DataFrame;
 
-DataFrame data = sqlContext.createDataFrame(
-  MLUtils.loadLibSVMFile(jsc.sc(), "data/mllib/sample_libsvm_data.txt"),
-  LabeledPoint.class);
+DataFrame data = jsql.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
 
 // Prepare training and test data.
 DataFrame[] splits = data.randomSplit(new double[] {0.9, 0.1}, 12345);
diff --git a/docs/ml-linear-methods.md b/docs/ml-linear-methods.md
index 16e2ee71293ae..85edfd373465f 100644
--- a/docs/ml-linear-methods.md
+++ b/docs/ml-linear-methods.md
@@ -95,7 +95,7 @@ public class LogisticRegressionWithElasticNetExample {
     String path = "data/mllib/sample_libsvm_data.txt";
 
     // Load training data
-    DataFrame training = sqlContext.read.format("libsvm").load(path);
+    DataFrame training = sqlContext.read().format("libsvm").load(path);
 
     LogisticRegression lr = new LogisticRegression()
       .setMaxIter(10)
@@ -292,7 +292,7 @@ public class LinearRegressionWithElasticNetExample {
     String path = "data/mllib/sample_libsvm_data.txt";
 
     // Load training data
-    DataFrame training = sqlContext.read.format("libsvm").load(path);
+    DataFrame training = sqlContext.read().format("libsvm").load(path);
 
     LinearRegression lr = new LinearRegression()
       .setMaxIter(10)
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java
index 51c1730a8a085..482225e585cf8 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeClassificationExample.java
@@ -26,9 +26,6 @@
 import org.apache.spark.ml.classification.DecisionTreeClassificationModel;
 import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
 import org.apache.spark.ml.feature.*;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.rdd.RDD;
 import org.apache.spark.sql.DataFrame;
 import org.apache.spark.sql.SQLContext;
 // $example off$
@@ -40,9 +37,8 @@ public static void main(String[] args) {
     SQLContext sqlContext = new SQLContext(jsc);
 
     // $example on$
-    // Load and parse the data file, converting it to a DataFrame.
-    RDD<LabeledPoint> rdd = MLUtils.loadLibSVMFile(jsc.sc(), "data/mllib/sample_libsvm_data.txt");
-    DataFrame data = sqlContext.createDataFrame(rdd, LabeledPoint.class);
+    // Load the data stored in LIBSVM format as a DataFrame.
+    DataFrame data = sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
 
     // Index labels, adding metadata to the label column.
     // Fit on whole dataset to include all labels in index.
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java
index a4098a4233ec2..c7f1868dd105a 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaDecisionTreeRegressionExample.java
@@ -27,9 +27,6 @@
 import org.apache.spark.ml.feature.VectorIndexerModel;
 import org.apache.spark.ml.regression.DecisionTreeRegressionModel;
 import org.apache.spark.ml.regression.DecisionTreeRegressor;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.rdd.RDD;
 import org.apache.spark.sql.DataFrame;
 import org.apache.spark.sql.SQLContext;
 // $example off$
@@ -40,9 +37,9 @@ public static void main(String[] args) {
     JavaSparkContext jsc = new JavaSparkContext(conf);
     SQLContext sqlContext = new SQLContext(jsc);
     // $example on$
-    // Load and parse the data file, converting it to a DataFrame.
-    RDD<LabeledPoint> rdd = MLUtils.loadLibSVMFile(jsc.sc(), "data/mllib/sample_libsvm_data.txt");
-    DataFrame data = sqlContext.createDataFrame(rdd, LabeledPoint.class);
+    // Load the data stored in LIBSVM format as a DataFrame.
+    DataFrame data = sqlContext.read().format("libsvm")
+      .load("data/mllib/sample_libsvm_data.txt");
 
     // Automatically identify categorical features, and index them.
     // Set maxCategories so features with > 4 distinct values are treated as continuous.
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java
index f48e1339c5007..84369f6681d04 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaMultilayerPerceptronClassifierExample.java
@@ -21,12 +21,9 @@
 import org.apache.spark.SparkConf;
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.SQLContext;
-import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.ml.classification.MultilayerPerceptronClassificationModel;
 import org.apache.spark.ml.classification.MultilayerPerceptronClassifier;
 import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
 import org.apache.spark.sql.DataFrame;
 // $example off$
 
@@ -43,8 +40,7 @@ public static void main(String[] args) {
     // $example on$
     // Load training data
     String path = "data/mllib/sample_multiclass_classification_data.txt";
-    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), path).toJavaRDD();
-    DataFrame dataFrame = jsql.createDataFrame(data, LabeledPoint.class);
+    DataFrame dataFrame = jsql.read().format("libsvm").load(path);
     // Split the data into train and test
     DataFrame[] splits = dataFrame.randomSplit(new double[]{0.6, 0.4}, 1234L);
     DataFrame train = splits[0];
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java
index e7f2f6f615070..f0d92a56bee73 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java
@@ -27,9 +27,7 @@
 import org.apache.spark.ml.util.MetadataUtils;
 import org.apache.spark.mllib.evaluation.MulticlassMetrics;
 import org.apache.spark.mllib.linalg.Matrix;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.rdd.RDD;
+import org.apache.spark.mllib.linalg.Vector;
 import org.apache.spark.sql.DataFrame;
 import org.apache.spark.sql.SQLContext;
 import org.apache.spark.sql.types.StructField;
@@ -80,31 +78,30 @@ public static void main(String[] args) {
     OneVsRest ovr = new OneVsRest().setClassifier(classifier);
 
     String input = params.input;
-    RDD<LabeledPoint> inputData = MLUtils.loadLibSVMFile(jsc.sc(), input);
-    RDD<LabeledPoint> train;
-    RDD<LabeledPoint> test;
+    DataFrame inputData = jsql.read().format("libsvm").load(input);
+    DataFrame train;
+    DataFrame test;
 
     // compute the train/ test split: if testInput is not provided use part of input
     String testInput = params.testInput;
     if (testInput != null) {
       train = inputData;
       // compute the number of features in the training set.
-      int numFeatures = inputData.first().features().size();
-      test = MLUtils.loadLibSVMFile(jsc.sc(), testInput, numFeatures);
+      int numFeatures = inputData.first().<Vector>getAs(1).size();
+      test = jsql.read().format("libsvm").option("numFeatures",
+        String.valueOf(numFeatures)).load(testInput);
     } else {
       double f = params.fracTest;
-      RDD<LabeledPoint>[] tmp = inputData.randomSplit(new double[]{1 - f, f}, 12345);
+      DataFrame[] tmp = inputData.randomSplit(new double[]{1 - f, f}, 12345);
       train = tmp[0];
       test = tmp[1];
     }
 
     // train the multiclass model
-    DataFrame trainingDataFrame = jsql.createDataFrame(train, LabeledPoint.class);
-    OneVsRestModel ovrModel = ovr.fit(trainingDataFrame.cache());
+    OneVsRestModel ovrModel = ovr.fit(train.cache());
 
     // score the model on test data
-    DataFrame testDataFrame = jsql.createDataFrame(test, LabeledPoint.class);
-    DataFrame predictions = ovrModel.transform(testDataFrame.cache())
+    DataFrame predictions = ovrModel.transform(test.cache())
       .select("prediction", "label");
 
     // obtain metrics
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaTrainValidationSplitExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaTrainValidationSplitExample.java
index 23f834ab4332b..d433905fc8012 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaTrainValidationSplitExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaTrainValidationSplitExample.java
@@ -23,8 +23,6 @@
 import org.apache.spark.ml.param.ParamMap;
 import org.apache.spark.ml.regression.LinearRegression;
 import org.apache.spark.ml.tuning.*;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
 import org.apache.spark.sql.DataFrame;
 import org.apache.spark.sql.SQLContext;
 
@@ -46,9 +44,7 @@ public static void main(String[] args) {
     JavaSparkContext jsc = new JavaSparkContext(conf);
     SQLContext jsql = new SQLContext(jsc);
 
-    DataFrame data = jsql.createDataFrame(
-      MLUtils.loadLibSVMFile(jsc.sc(), "data/mllib/sample_libsvm_data.txt"),
-      LabeledPoint.class);
+    DataFrame data = jsql.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
 
     // Prepare training and test data.
     DataFrame[] splits = data.randomSplit(new double [] {0.9, 0.1}, 12345);
diff --git a/examples/src/main/python/ml/decision_tree_classification_example.py b/examples/src/main/python/ml/decision_tree_classification_example.py
index 0af92050e3e3b..8cda56dbb9bdf 100644
--- a/examples/src/main/python/ml/decision_tree_classification_example.py
+++ b/examples/src/main/python/ml/decision_tree_classification_example.py
@@ -28,7 +28,6 @@
 from pyspark.ml.classification import DecisionTreeClassifier
 from pyspark.ml.feature import StringIndexer, VectorIndexer
 from pyspark.ml.evaluation import MulticlassClassificationEvaluator
-from pyspark.mllib.util import MLUtils
 # $example off$
 
 if __name__ == "__main__":
@@ -36,8 +35,8 @@
     sqlContext = SQLContext(sc)
 
     # $example on$
-    # Load and parse the data file, converting it to a DataFrame.
-    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+    # Load the data stored in LIBSVM format as a DataFrame.
+    data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 
     # Index labels, adding metadata to the label column.
     # Fit on whole dataset to include all labels in index.
diff --git a/examples/src/main/python/ml/decision_tree_regression_example.py b/examples/src/main/python/ml/decision_tree_regression_example.py
index 3857aed538da2..439e398947499 100644
--- a/examples/src/main/python/ml/decision_tree_regression_example.py
+++ b/examples/src/main/python/ml/decision_tree_regression_example.py
@@ -28,7 +28,6 @@
 from pyspark.ml.regression import DecisionTreeRegressor
 from pyspark.ml.feature import VectorIndexer
 from pyspark.ml.evaluation import RegressionEvaluator
-from pyspark.mllib.util import MLUtils
 # $example off$
 
 if __name__ == "__main__":
@@ -36,8 +35,8 @@
     sqlContext = SQLContext(sc)
 
     # $example on$
-    # Load and parse the data file, converting it to a DataFrame.
-    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+    # Load the data stored in LIBSVM format as a DataFrame.
+    data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 
     # Automatically identify categorical features, and index them.
     # We specify maxCategories so features with > 4 distinct values are treated as continuous.
diff --git a/examples/src/main/python/ml/gradient_boosted_trees.py b/examples/src/main/python/ml/gradient_boosted_trees.py
index 6446f0fe5eeab..c3bf8aa2eb1e6 100644
--- a/examples/src/main/python/ml/gradient_boosted_trees.py
+++ b/examples/src/main/python/ml/gradient_boosted_trees.py
@@ -24,7 +24,6 @@
 from pyspark.ml.feature import StringIndexer
 from pyspark.ml.regression import GBTRegressor
 from pyspark.mllib.evaluation import BinaryClassificationMetrics, RegressionMetrics
-from pyspark.mllib.util import MLUtils
 from pyspark.sql import Row, SQLContext
 
 """
@@ -70,8 +69,8 @@ def testRegression(train, test):
     sc = SparkContext(appName="PythonGBTExample")
     sqlContext = SQLContext(sc)
 
-    # Load and parse the data file into a dataframe.
-    df = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+    # Load the data stored in LIBSVM format as a DataFrame.
+    df = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 
     # Map labels into an indexed column of labels in [0, numLabels)
     stringIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel")
diff --git a/examples/src/main/python/ml/logistic_regression.py b/examples/src/main/python/ml/logistic_regression.py
index 55afe1b207fe0..4cd027fdfbe8a 100644
--- a/examples/src/main/python/ml/logistic_regression.py
+++ b/examples/src/main/python/ml/logistic_regression.py
@@ -23,7 +23,6 @@
 from pyspark.ml.classification import LogisticRegression
 from pyspark.mllib.evaluation import MulticlassMetrics
 from pyspark.ml.feature import StringIndexer
-from pyspark.mllib.util import MLUtils
 from pyspark.sql import SQLContext
 
 """
@@ -41,8 +40,8 @@
     sc = SparkContext(appName="PythonLogisticRegressionExample")
     sqlContext = SQLContext(sc)
 
-    # Load and parse the data file into a dataframe.
-    df = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+    # Load the data stored in LIBSVM format as a DataFrame.
+    df = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 
     # Map labels into an indexed column of labels in [0, numLabels)
     stringIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel")
diff --git a/examples/src/main/python/ml/multilayer_perceptron_classification.py b/examples/src/main/python/ml/multilayer_perceptron_classification.py
index d8ef9f39e3fa5..f84588f547fff 100644
--- a/examples/src/main/python/ml/multilayer_perceptron_classification.py
+++ b/examples/src/main/python/ml/multilayer_perceptron_classification.py
@@ -22,7 +22,6 @@
 # $example on$
 from pyspark.ml.classification import MultilayerPerceptronClassifier
 from pyspark.ml.evaluation import MulticlassClassificationEvaluator
-from pyspark.mllib.util import MLUtils
 # $example off$
 
 if __name__ == "__main__":
@@ -32,8 +31,8 @@
 
     # $example on$
     # Load training data
-    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt")\
-        .toDF()
+    data = sqlContext.read.format("libsvm")\
+        .load("data/mllib/sample_multiclass_classification_data.txt")
     # Split the data into train and test
     splits = data.randomSplit([0.6, 0.4], 1234)
     train = splits[0]
diff --git a/examples/src/main/python/ml/random_forest_example.py b/examples/src/main/python/ml/random_forest_example.py
index c7730e1bfacd9..dc6a778670193 100644
--- a/examples/src/main/python/ml/random_forest_example.py
+++ b/examples/src/main/python/ml/random_forest_example.py
@@ -74,8 +74,8 @@ def testRegression(train, test):
     sc = SparkContext(appName="PythonRandomForestExample")
     sqlContext = SQLContext(sc)
 
-    # Load and parse the data file into a dataframe.
-    df = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+    # Load the data stored in LIBSVM format as a DataFrame.
+    df = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 
     # Map labels into an indexed column of labels in [0, numLabels)
     stringIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel")
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
index a24a344f1bcf4..ff8a0a90f1e44 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
@@ -26,7 +26,6 @@ import org.apache.spark.ml.classification.DecisionTreeClassifier
 import org.apache.spark.ml.classification.DecisionTreeClassificationModel
 import org.apache.spark.ml.feature.{StringIndexer, IndexToString, VectorIndexer}
 import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
-import org.apache.spark.mllib.util.MLUtils
 // $example off$
 
 object DecisionTreeClassificationExample {
@@ -34,10 +33,9 @@ object DecisionTreeClassificationExample {
     val conf = new SparkConf().setAppName("DecisionTreeClassificationExample")
     val sc = new SparkContext(conf)
     val sqlContext = new SQLContext(sc)
-    import sqlContext.implicits._
     // $example on$
-    // Load and parse the data file, converting it to a DataFrame.
-    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+    // Load the data stored in LIBSVM format as a DataFrame.
+    val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 
     // Index labels, adding metadata to the label column.
     // Fit on whole dataset to include all labels in index.
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeExample.scala
index f28671f7869fc..c4e98dfaca6c9 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeExample.scala
@@ -32,10 +32,7 @@ import org.apache.spark.ml.regression.{DecisionTreeRegressionModel, DecisionTree
 import org.apache.spark.ml.util.MetadataUtils
 import org.apache.spark.mllib.evaluation.{RegressionMetrics, MulticlassMetrics}
 import org.apache.spark.mllib.linalg.Vector
-import org.apache.spark.mllib.regression.LabeledPoint
 import org.apache.spark.mllib.util.MLUtils
-import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.types.StringType
 import org.apache.spark.sql.{SQLContext, DataFrame}
 
 
@@ -138,15 +135,18 @@ object DecisionTreeExample {
 
   /** Load a dataset from the given path, using the given format */
   private[ml] def loadData(
-      sc: SparkContext,
+      sqlContext: SQLContext,
       path: String,
       format: String,
-      expectedNumFeatures: Option[Int] = None): RDD[LabeledPoint] = {
+      expectedNumFeatures: Option[Int] = None): DataFrame = {
+    import sqlContext.implicits._
+
     format match {
-      case "dense" => MLUtils.loadLabeledPoints(sc, path)
+      case "dense" => MLUtils.loadLabeledPoints(sqlContext.sparkContext, path).toDF()
       case "libsvm" => expectedNumFeatures match {
-        case Some(numFeatures) => MLUtils.loadLibSVMFile(sc, path, numFeatures)
-        case None => MLUtils.loadLibSVMFile(sc, path)
+        case Some(numFeatures) => sqlContext.read.option("numFeatures", numFeatures.toString)
+          .format("libsvm").load(path)
+        case None => sqlContext.read.format("libsvm").load(path)
       }
       case _ => throw new IllegalArgumentException(s"Bad data format: $format")
     }
@@ -169,36 +169,22 @@ object DecisionTreeExample {
       algo: String,
       fracTest: Double): (DataFrame, DataFrame) = {
     val sqlContext = new SQLContext(sc)
-    import sqlContext.implicits._
 
     // Load training data
-    val origExamples: RDD[LabeledPoint] = loadData(sc, input, dataFormat)
+    val origExamples: DataFrame = loadData(sqlContext, input, dataFormat)
 
     // Load or create test set
-    val splits: Array[RDD[LabeledPoint]] = if (testInput != "") {
+    val dataframes: Array[DataFrame] = if (testInput != "") {
       // Load testInput.
-      val numFeatures = origExamples.take(1)(0).features.size
-      val origTestExamples: RDD[LabeledPoint] =
-        loadData(sc, testInput, dataFormat, Some(numFeatures))
+      val numFeatures = origExamples.first().getAs[Vector](1).size
+      val origTestExamples: DataFrame =
+        loadData(sqlContext, testInput, dataFormat, Some(numFeatures))
       Array(origExamples, origTestExamples)
     } else {
       // Split input into training, test.
       origExamples.randomSplit(Array(1.0 - fracTest, fracTest), seed = 12345)
     }
 
-    // For classification, convert labels to Strings since we will index them later with
-    // StringIndexer.
-    def labelsToStrings(data: DataFrame): DataFrame = {
-      algo.toLowerCase match {
-        case "classification" =>
-          data.withColumn("labelString", data("label").cast(StringType))
-        case "regression" =>
-          data
-        case _ =>
-          throw new IllegalArgumentException("Algo ${params.algo} not supported.")
-      }
-    }
-    val dataframes = splits.map(_.toDF()).map(labelsToStrings)
     val training = dataframes(0).cache()
     val test = dataframes(1).cache()
 
@@ -230,7 +216,7 @@ object DecisionTreeExample {
     val labelColName = if (algo == "classification") "indexedLabel" else "label"
     if (algo == "classification") {
       val labelIndexer = new StringIndexer()
-        .setInputCol("labelString")
+        .setInputCol("label")
         .setOutputCol(labelColName)
       stages += labelIndexer
     }
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala
index 64cd986129007..fc402724d2156 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala
@@ -25,17 +25,16 @@ import org.apache.spark.ml.regression.DecisionTreeRegressor
 import org.apache.spark.ml.regression.DecisionTreeRegressionModel
 import org.apache.spark.ml.feature.VectorIndexer
 import org.apache.spark.ml.evaluation.RegressionEvaluator
-import org.apache.spark.mllib.util.MLUtils
 // $example off$
 object DecisionTreeRegressionExample {
   def main(args: Array[String]): Unit = {
     val conf = new SparkConf().setAppName("DecisionTreeRegressionExample")
     val sc = new SparkContext(conf)
     val sqlContext = new SQLContext(sc)
-    import sqlContext.implicits._
+
     // $example on$
-    // Load and parse the data file, converting it to a DataFrame.
-    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+    // Load the data stored in LIBSVM format as a DataFrame.
+    val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 
     // Automatically identify categorical features, and index them.
     // Here, we treat features with > 4 distinct values as continuous.
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/GBTExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/GBTExample.scala
index f4a15f806ea81..6b0be0f34e196 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/GBTExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/GBTExample.scala
@@ -153,7 +153,7 @@ object GBTExample {
     val labelColName = if (algo == "classification") "indexedLabel" else "label"
     if (algo == "classification") {
       val labelIndexer = new StringIndexer()
-        .setInputCol("labelString")
+        .setInputCol("label")
         .setOutputCol(labelColName)
       stages += labelIndexer
     }
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionExample.scala
index b73299fb12d3f..50998c94de3d0 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionExample.scala
@@ -131,7 +131,7 @@ object LinearRegressionExample {
     println(s"Training time: $elapsedTime seconds")
 
     // Print the weights and intercept for linear regression.
-    println(s"Weights: ${lirModel.weights} Intercept: ${lirModel.intercept}")
+    println(s"Weights: ${lirModel.coefficients} Intercept: ${lirModel.intercept}")
 
     println("Training data results:")
     DecisionTreeExample.evaluateRegressionModel(lirModel, training, "label")
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionExample.scala
index 8e3760ddb50a9..a380c90662a50 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionExample.scala
@@ -125,7 +125,7 @@ object LogisticRegressionExample {
     val stages = new mutable.ArrayBuffer[PipelineStage]()
 
     val labelIndexer = new StringIndexer()
-      .setInputCol("labelString")
+      .setInputCol("label")
       .setOutputCol("indexedLabel")
     stages += labelIndexer
 
@@ -149,7 +149,7 @@ object LogisticRegressionExample {
 
     val lorModel = pipelineModel.stages.last.asInstanceOf[LogisticRegressionModel]
     // Print the weights and intercept for logistic regression.
-    println(s"Weights: ${lorModel.weights} Intercept: ${lorModel.intercept}")
+    println(s"Weights: ${lorModel.coefficients} Intercept: ${lorModel.intercept}")
 
     println("Training data results:")
     DecisionTreeExample.evaluateClassificationModel(pipelineModel, training, "indexedLabel")
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala
index 99d5f35b5a565..146b83c8be490 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala
@@ -23,7 +23,6 @@ import org.apache.spark.sql.SQLContext
 // $example on$
 import org.apache.spark.ml.classification.MultilayerPerceptronClassifier
 import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
-import org.apache.spark.mllib.util.MLUtils
 // $example off$
 
 /**
@@ -35,12 +34,11 @@ object MultilayerPerceptronClassifierExample {
     val conf = new SparkConf().setAppName("MultilayerPerceptronClassifierExample")
     val sc = new SparkContext(conf)
     val sqlContext = new SQLContext(sc)
-    import sqlContext.implicits._
 
     // $example on$
-    // Load training data
-    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt")
-      .toDF()
+    // Load the data stored in LIBSVM format as a DataFrame.
+    val data = sqlContext.read.format("libsvm")
+      .load("data/mllib/sample_multiclass_classification_data.txt")
     // Split the data into train and test
     val splits = data.randomSplit(Array(0.6, 0.4), seed = 1234L)
     val train = splits(0)
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/OneVsRestExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/OneVsRestExample.scala
index bab31f585b0ef..8e4f1b09a24b5 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/OneVsRestExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/OneVsRestExample.scala
@@ -27,9 +27,8 @@ import org.apache.spark.examples.mllib.AbstractParams
 import org.apache.spark.ml.classification.{OneVsRest, LogisticRegression}
 import org.apache.spark.ml.util.MetadataUtils
 import org.apache.spark.mllib.evaluation.MulticlassMetrics
-import org.apache.spark.mllib.regression.LabeledPoint
-import org.apache.spark.mllib.util.MLUtils
-import org.apache.spark.rdd.RDD
+import org.apache.spark.mllib.linalg.Vector
+import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.SQLContext
 
 /**
@@ -111,24 +110,24 @@ object OneVsRestExample {
   private def run(params: Params) {
     val conf = new SparkConf().setAppName(s"OneVsRestExample with $params")
     val sc = new SparkContext(conf)
-    val inputData = MLUtils.loadLibSVMFile(sc, params.input)
     val sqlContext = new SQLContext(sc)
-    import sqlContext.implicits._
 
+    val inputData = sqlContext.read.format("libsvm").load(params.input)
     // compute the train/test split: if testInput is not provided use part of input.
     val data = params.testInput match {
       case Some(t) => {
         // compute the number of features in the training set.
-        val numFeatures = inputData.first().features.size
-        val testData = MLUtils.loadLibSVMFile(sc, t, numFeatures)
-        Array[RDD[LabeledPoint]](inputData, testData)
+        val numFeatures = inputData.first().getAs[Vector](1).size
+        val testData = sqlContext.read.option("numFeatures", numFeatures.toString)
+          .format("libsvm").load(t)
+        Array[DataFrame](inputData, testData)
       }
       case None => {
         val f = params.fracTest
         inputData.randomSplit(Array(1 - f, f), seed = 12345)
       }
     }
-    val Array(train, test) = data.map(_.toDF().cache())
+    val Array(train, test) = data.map(_.cache())
 
     // instantiate the base classifier
     val classifier = new LogisticRegression()
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestExample.scala
index 109178f4137b2..7a00d99dfe53d 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestExample.scala
@@ -159,7 +159,7 @@ object RandomForestExample {
     val labelColName = if (algo == "classification") "indexedLabel" else "label"
     if (algo == "classification") {
       val labelIndexer = new StringIndexer()
-        .setInputCol("labelString")
+        .setInputCol("label")
         .setOutputCol(labelColName)
       stages += labelIndexer
     }
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/TrainValidationSplitExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/TrainValidationSplitExample.scala
index 1abdf219b1c00..cd1b0e9358beb 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/TrainValidationSplitExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/TrainValidationSplitExample.scala
@@ -20,7 +20,6 @@ package org.apache.spark.examples.ml
 import org.apache.spark.ml.evaluation.RegressionEvaluator
 import org.apache.spark.ml.regression.LinearRegression
 import org.apache.spark.ml.tuning.{ParamGridBuilder, TrainValidationSplit}
-import org.apache.spark.mllib.util.MLUtils
 import org.apache.spark.sql.SQLContext
 import org.apache.spark.{SparkConf, SparkContext}
 
@@ -39,10 +38,9 @@ object TrainValidationSplitExample {
     val conf = new SparkConf().setAppName("TrainValidationSplitExample")
     val sc = new SparkContext(conf)
     val sqlContext = new SQLContext(sc)
-    import sqlContext.implicits._
 
     // Prepare training and test data.
-    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_libsvm_data.txt").toDF()
+    val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
     val Array(training, test) = data.randomSplit(Array(0.9, 0.1), seed = 12345)
 
     val lr = new LinearRegression()
diff --git a/mllib/src/main/scala/org/apache/spark/ml/source/libsvm/LibSVMRelation.scala b/mllib/src/main/scala/org/apache/spark/ml/source/libsvm/LibSVMRelation.scala
index 1f627777fc68d..11b9815ecc832 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/source/libsvm/LibSVMRelation.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/source/libsvm/LibSVMRelation.scala
@@ -82,7 +82,7 @@ private[libsvm] class LibSVMRelation(val path: String, val numFeatures: Int, val
  *     .load("data/mllib/sample_libsvm_data.txt")
  *
  *   // Java
- *   DataFrame df = sqlContext.read.format("libsvm")
+ *   DataFrame df = sqlContext.read().format("libsvm")
  *     .option("numFeatures, "780")
  *     .load("data/mllib/sample_libsvm_data.txt");
  * }}}

From a24477996e936b0861819ffb420f763f80f0b1da Mon Sep 17 00:00:00 2001
From: Andrew Ray <ray.andrew@gmail.com>
Date: Fri, 13 Nov 2015 10:31:17 -0800
Subject: [PATCH 234/518] [SPARK-11690][PYSPARK] Add pivot to python api

This PR adds pivot to the python api of GroupedData with the same syntax as Scala/Java.

Author: Andrew Ray <ray.andrew@gmail.com>

Closes #9653 from aray/sql-pivot-python.
---
 python/pyspark/sql/group.py | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/python/pyspark/sql/group.py b/python/pyspark/sql/group.py
index 71c0bccc5eeff..227f40bc3cf53 100644
--- a/python/pyspark/sql/group.py
+++ b/python/pyspark/sql/group.py
@@ -17,7 +17,7 @@
 
 from pyspark import since
 from pyspark.rdd import ignore_unicode_prefix
-from pyspark.sql.column import Column, _to_seq
+from pyspark.sql.column import Column, _to_seq, _to_java_column, _create_column_from_literal
 from pyspark.sql.dataframe import DataFrame
 from pyspark.sql.types import *
 
@@ -167,6 +167,23 @@ def sum(self, *cols):
         [Row(sum(age)=7, sum(height)=165)]
         """
 
+    @since(1.6)
+    def pivot(self, pivot_col, *values):
+        """Pivots a column of the current DataFrame and preform the specified aggregation.
+
+        :param pivot_col: Column to pivot
+        :param values: Optional list of values of pivotColumn that will be translated to columns in
+            the output data frame. If values are not provided the method with do an immediate call
+            to .distinct() on the pivot column.
+        >>> df4.groupBy("year").pivot("course", "dotNET", "Java").sum("earnings").collect()
+        [Row(year=2012, dotNET=15000, Java=20000), Row(year=2013, dotNET=48000, Java=30000)]
+        >>> df4.groupBy("year").pivot("course").sum("earnings").collect()
+        [Row(year=2012, Java=20000, dotNET=15000), Row(year=2013, Java=30000, dotNET=48000)]
+        """
+        jgd = self._jdf.pivot(_to_java_column(pivot_col),
+                              _to_seq(self.sql_ctx._sc, values, _create_column_from_literal))
+        return GroupedData(jgd, self.sql_ctx)
+
 
 def _test():
     import doctest
@@ -182,6 +199,11 @@ def _test():
                           StructField('name', StringType())]))
     globs['df3'] = sc.parallelize([Row(name='Alice', age=2, height=80),
                                    Row(name='Bob', age=5, height=85)]).toDF()
+    globs['df4'] = sc.parallelize([Row(course="dotNET", year=2012, earnings=10000),
+                                   Row(course="Java",   year=2012, earnings=20000),
+                                   Row(course="dotNET", year=2012, earnings=5000),
+                                   Row(course="dotNET", year=2013, earnings=48000),
+                                   Row(course="Java",   year=2013, earnings=30000)]).toDF()
 
     (failure_count, test_count) = doctest.testmod(
         pyspark.sql.group, globs=globs,

From 23b8188f75d945ef70fbb1c4dc9720c2c5f8cbc3 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Fri, 13 Nov 2015 11:13:09 -0800
Subject: [PATCH 235/518] [SPARK-11654][SQL][FOLLOW-UP] fix some mistakes and
 clean up

* rename `AppendColumn` to `AppendColumns` to be consistent with the physical plan name.
* clean up stale comments.
* always pass in resolved encoder to `TypedColumn.withInputType`(test added)
* enable a mistakenly disabled java test.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9688 from cloud-fan/follow.
---
 .../sql/catalyst/plans/logical/basicOperators.scala    |  8 ++++----
 .../src/main/scala/org/apache/spark/sql/Column.scala   |  3 ++-
 .../src/main/scala/org/apache/spark/sql/Dataset.scala  | 10 +++-------
 .../scala/org/apache/spark/sql/GroupedDataset.scala    |  4 ++--
 .../apache/spark/sql/execution/SparkStrategies.scala   |  2 +-
 .../test/org/apache/spark/sql/JavaDatasetSuite.java    |  1 +
 .../org/apache/spark/sql/DatasetAggregatorSuite.scala  |  4 ++++
 7 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index d9f046efce0bf..e2b97b27a6c2c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -482,13 +482,13 @@ case class MapPartitions[T, U](
 }
 
 /** Factory for constructing new `AppendColumn` nodes. */
-object AppendColumn {
+object AppendColumns {
   def apply[T, U : Encoder](
       func: T => U,
       tEncoder: ExpressionEncoder[T],
-      child: LogicalPlan): AppendColumn[T, U] = {
+      child: LogicalPlan): AppendColumns[T, U] = {
     val attrs = encoderFor[U].schema.toAttributes
-    new AppendColumn[T, U](func, tEncoder, encoderFor[U], attrs, child)
+    new AppendColumns[T, U](func, tEncoder, encoderFor[U], attrs, child)
   }
 }
 
@@ -497,7 +497,7 @@ object AppendColumn {
  * resulting columns at the end of the input row. tEncoder/uEncoder are used respectively to
  * decode/encode from the JVM object representation expected by `func.`
  */
-case class AppendColumn[T, U](
+case class AppendColumns[T, U](
     func: T => U,
     tEncoder: ExpressionEncoder[T],
     uEncoder: ExpressionEncoder[U],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index 929224460dc09..82e9cd7f50a31 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -58,10 +58,11 @@ class TypedColumn[-T, U](
   private[sql] def withInputType(
       inputEncoder: ExpressionEncoder[_],
       schema: Seq[Attribute]): TypedColumn[T, U] = {
+    val boundEncoder = inputEncoder.bind(schema).asInstanceOf[ExpressionEncoder[Any]]
     new TypedColumn[T, U] (expr transform {
       case ta: TypedAggregateExpression if ta.aEncoder.isEmpty =>
         ta.copy(
-          aEncoder = Some(inputEncoder.asInstanceOf[ExpressionEncoder[Any]]),
+          aEncoder = Some(boundEncoder),
           children = schema)
     }, encoder)
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index b930e4661c1a6..4cc3aa2465f2e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -299,7 +299,7 @@ class Dataset[T] private[sql](
    */
   def groupBy[K : Encoder](func: T => K): GroupedDataset[K, T] = {
     val inputPlan = queryExecution.analyzed
-    val withGroupingKey = AppendColumn(func, resolvedTEncoder, inputPlan)
+    val withGroupingKey = AppendColumns(func, resolvedTEncoder, inputPlan)
     val executed = sqlContext.executePlan(withGroupingKey)
 
     new GroupedDataset(
@@ -364,13 +364,11 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def select[U1: Encoder](c1: TypedColumn[T, U1]): Dataset[U1] = {
-    // We use an unbound encoder since the expression will make up its own schema.
-    // TODO: This probably doesn't work if we are relying on reordering of the input class fields.
     new Dataset[U1](
       sqlContext,
       Project(
         c1.withInputType(
-          resolvedTEncoder.bind(queryExecution.analyzed.output),
+          resolvedTEncoder,
           queryExecution.analyzed.output).named :: Nil,
         logicalPlan))
   }
@@ -382,10 +380,8 @@ class Dataset[T] private[sql](
    */
   protected def selectUntyped(columns: TypedColumn[_, _]*): Dataset[_] = {
     val encoders = columns.map(_.encoder)
-    // We use an unbound encoder since the expression will make up its own schema.
-    // TODO: This probably doesn't work if we are relying on reordering of the input class fields.
     val namedColumns =
-      columns.map(_.withInputType(unresolvedTEncoder, queryExecution.analyzed.output).named)
+      columns.map(_.withInputType(resolvedTEncoder, queryExecution.analyzed.output).named)
     val execution = new QueryExecution(sqlContext, Project(namedColumns, logicalPlan))
 
     new Dataset(sqlContext, execution, ExpressionEncoder.tuple(encoders))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index ae1272ae531fb..9c16940707de9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -89,7 +89,7 @@ class GroupedDataset[K, T] private[sql](
   }
 
   /**
-   * Applies the given function to each group of data.  For each unique group, the function  will
+   * Applies the given function to each group of data.  For each unique group, the function will
    * be passed the group key and an iterator that contains all of the elements in the group. The
    * function can return an iterator containing elements of an arbitrary type which will be returned
    * as a new [[Dataset]].
@@ -162,7 +162,7 @@ class GroupedDataset[K, T] private[sql](
     val encoders = columns.map(_.encoder)
     val namedColumns =
       columns.map(
-        _.withInputType(resolvedTEncoder.bind(dataAttributes), dataAttributes).named)
+        _.withInputType(resolvedTEncoder, dataAttributes).named)
     val aggregate = Aggregate(groupingAttributes, groupingAttributes ++ namedColumns, logicalPlan)
     val execution = new QueryExecution(sqlContext, aggregate)
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index a99ae4674bb12..67201a2c191cd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -321,7 +321,7 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
 
       case logical.MapPartitions(f, tEnc, uEnc, output, child) =>
         execution.MapPartitions(f, tEnc, uEnc, output, planLater(child)) :: Nil
-      case logical.AppendColumn(f, tEnc, uEnc, newCol, child) =>
+      case logical.AppendColumns(f, tEnc, uEnc, newCol, child) =>
         execution.AppendColumns(f, tEnc, uEnc, newCol, planLater(child)) :: Nil
       case logical.MapGroups(f, kEnc, tEnc, uEnc, grouping, output, child) =>
         execution.MapGroups(f, kEnc, tEnc, uEnc, grouping, output, planLater(child)) :: Nil
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index 46169ca07d715..eb6fa1e72e27b 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -157,6 +157,7 @@ public Integer call(Integer v1, Integer v2) throws Exception {
     Assert.assertEquals(6, reduced);
   }
 
+  @Test
   public void testGroupBy() {
     List<String> data = Arrays.asList("a", "foo", "bar");
     Dataset<String> ds = context.createDataset(data, Encoders.STRING());
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
index 20896efdfec16..46f9f077fe7f2 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -161,6 +161,10 @@ class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
       ds.select(ClassInputAgg.toColumn),
       1)
 
+    checkAnswer(
+      ds.select(expr("avg(a)").as[Double], ClassInputAgg.toColumn),
+      (1.0, 1))
+
     checkAnswer(
       ds.groupBy(_.b).agg(ClassInputAgg.toColumn),
       ("one", 1))

From d7b2b97ad67f9700fb8c13422c399f2edb72f770 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Fri, 13 Nov 2015 11:25:33 -0800
Subject: [PATCH 236/518] [SPARK-11727][SQL] Split ExpressionEncoder into
 FlatEncoder and ProductEncoder

also add more tests for encoders, and fix bugs that I found:

* when convert array to catalyst array, we can only skip element conversion for native types(e.g. int, long, boolean), not `AtomicType`(String is AtomicType but we need to convert it)
* we should also handle scala `BigDecimal` when convert from catalyst `Decimal`.
* complex map type should be supported

other issues that still in investigation:

* encode java `BigDecimal` and decode it back, seems we will loss precision info.
* when encode case class that defined inside a object, `ClassNotFound` exception will be thrown.

I'll remove unused code in a follow-up PR.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9693 from cloud-fan/split.
---
 .../spark/sql/catalyst/ScalaReflection.scala  |   2 +-
 .../sql/catalyst/encoders/FlatEncoder.scala   |  50 ++
 .../catalyst/encoders/ProductEncoder.scala    | 452 ++++++++++++++++++
 .../sql/catalyst/encoders/RowEncoder.scala    |  58 +--
 .../sql/catalyst/util/DateTimeUtils.scala     |   2 +-
 .../sql/catalyst/util/GenericArrayData.scala  |   2 +-
 .../encoders/ExpressionEncoderSuite.scala     | 259 ++--------
 .../catalyst/encoders/FlatEncoderSuite.scala  |  74 +++
 .../encoders/ProductEncoderSuite.scala        | 123 +++++
 .../org/apache/spark/sql/GroupedDataset.scala |   7 +-
 .../org/apache/spark/sql/SQLImplicits.scala   |  22 +-
 .../org/apache/spark/sql/functions.scala      |   4 +-
 12 files changed, 766 insertions(+), 289 deletions(-)
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoder.scala
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
 create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
 create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 6d822261b050a..0b3dd351e38e8 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -75,7 +75,7 @@ trait ScalaReflection {
    *
    * @see SPARK-5281
    */
-  private def localTypeOf[T: TypeTag]: `Type` = typeTag[T].in(mirror).tpe
+  def localTypeOf[T: TypeTag]: `Type` = typeTag[T].in(mirror).tpe
 
   /**
    * Returns the Spark SQL DataType for a given scala type.  Where this is not an exact mapping
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoder.scala
new file mode 100644
index 0000000000000..6d307ab13a9fc
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoder.scala
@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.encoders
+
+import scala.reflect.ClassTag
+import scala.reflect.runtime.universe.{typeTag, TypeTag}
+
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.catalyst.expressions.{Literal, CreateNamedStruct, BoundReference}
+import org.apache.spark.sql.catalyst.ScalaReflection
+
+object FlatEncoder {
+  import ScalaReflection.schemaFor
+  import ScalaReflection.dataTypeFor
+
+  def apply[T : TypeTag]: ExpressionEncoder[T] = {
+    // We convert the not-serializable TypeTag into StructType and ClassTag.
+    val tpe = typeTag[T].tpe
+    val mirror = typeTag[T].mirror
+    val cls = mirror.runtimeClass(tpe)
+    assert(!schemaFor(tpe).dataType.isInstanceOf[StructType])
+
+    val input = BoundReference(0, dataTypeFor(tpe), nullable = true)
+    val toRowExpression = CreateNamedStruct(
+      Literal("value") :: ProductEncoder.extractorFor(input, tpe) :: Nil)
+    val fromRowExpression = ProductEncoder.constructorFor(tpe)
+
+    new ExpressionEncoder[T](
+      toRowExpression.dataType,
+      flat = true,
+      toRowExpression.flatten,
+      fromRowExpression,
+      ClassTag[T](cls))
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
new file mode 100644
index 0000000000000..414adb21168ed
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
@@ -0,0 +1,452 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.encoders
+
+import org.apache.spark.util.Utils
+import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.catalyst.ScalaReflectionLock
+import org.apache.spark.sql.catalyst.ScalaReflection
+import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedExtractValue}
+import org.apache.spark.sql.catalyst.expressions
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.util.{DateTimeUtils, ArrayBasedMapData, GenericArrayData}
+
+import scala.reflect.ClassTag
+
+object ProductEncoder {
+  import ScalaReflection.universe._
+  import ScalaReflection.localTypeOf
+  import ScalaReflection.dataTypeFor
+  import ScalaReflection.Schema
+  import ScalaReflection.schemaFor
+  import ScalaReflection.arrayClassFor
+
+  def apply[T <: Product : TypeTag]: ExpressionEncoder[T] = {
+    // We convert the not-serializable TypeTag into StructType and ClassTag.
+    val tpe = typeTag[T].tpe
+    val mirror = typeTag[T].mirror
+    val cls = mirror.runtimeClass(tpe)
+
+    val inputObject = BoundReference(0, ObjectType(cls), nullable = true)
+    val toRowExpression = extractorFor(inputObject, tpe).asInstanceOf[CreateNamedStruct]
+    val fromRowExpression = constructorFor(tpe)
+
+    new ExpressionEncoder[T](
+      toRowExpression.dataType,
+      flat = false,
+      toRowExpression.flatten,
+      fromRowExpression,
+      ClassTag[T](cls))
+  }
+
+  // The Predef.Map is scala.collection.immutable.Map.
+  // Since the map values can be mutable, we explicitly import scala.collection.Map at here.
+  import scala.collection.Map
+
+  def extractorFor(
+      inputObject: Expression,
+      tpe: `Type`): Expression = ScalaReflectionLock.synchronized {
+    if (!inputObject.dataType.isInstanceOf[ObjectType]) {
+      inputObject
+    } else {
+      tpe match {
+        case t if t <:< localTypeOf[Option[_]] =>
+          val TypeRef(_, _, Seq(optType)) = t
+          optType match {
+            // For primitive types we must manually unbox the value of the object.
+            case t if t <:< definitions.IntTpe =>
+              Invoke(
+                UnwrapOption(ObjectType(classOf[java.lang.Integer]), inputObject),
+                "intValue",
+                IntegerType)
+            case t if t <:< definitions.LongTpe =>
+              Invoke(
+                UnwrapOption(ObjectType(classOf[java.lang.Long]), inputObject),
+                "longValue",
+                LongType)
+            case t if t <:< definitions.DoubleTpe =>
+              Invoke(
+                UnwrapOption(ObjectType(classOf[java.lang.Double]), inputObject),
+                "doubleValue",
+                DoubleType)
+            case t if t <:< definitions.FloatTpe =>
+              Invoke(
+                UnwrapOption(ObjectType(classOf[java.lang.Float]), inputObject),
+                "floatValue",
+                FloatType)
+            case t if t <:< definitions.ShortTpe =>
+              Invoke(
+                UnwrapOption(ObjectType(classOf[java.lang.Short]), inputObject),
+                "shortValue",
+                ShortType)
+            case t if t <:< definitions.ByteTpe =>
+              Invoke(
+                UnwrapOption(ObjectType(classOf[java.lang.Byte]), inputObject),
+                "byteValue",
+                ByteType)
+            case t if t <:< definitions.BooleanTpe =>
+              Invoke(
+                UnwrapOption(ObjectType(classOf[java.lang.Boolean]), inputObject),
+                "booleanValue",
+                BooleanType)
+
+            // For non-primitives, we can just extract the object from the Option and then recurse.
+            case other =>
+              val className: String = optType.erasure.typeSymbol.asClass.fullName
+              val classObj = Utils.classForName(className)
+              val optionObjectType = ObjectType(classObj)
+
+              val unwrapped = UnwrapOption(optionObjectType, inputObject)
+              expressions.If(
+                IsNull(unwrapped),
+                expressions.Literal.create(null, schemaFor(optType).dataType),
+                extractorFor(unwrapped, optType))
+          }
+
+        case t if t <:< localTypeOf[Product] =>
+          val formalTypeArgs = t.typeSymbol.asClass.typeParams
+          val TypeRef(_, _, actualTypeArgs) = t
+          val constructorSymbol = t.member(nme.CONSTRUCTOR)
+          val params = if (constructorSymbol.isMethod) {
+            constructorSymbol.asMethod.paramss
+          } else {
+            // Find the primary constructor, and use its parameter ordering.
+            val primaryConstructorSymbol: Option[Symbol] =
+              constructorSymbol.asTerm.alternatives.find(s =>
+                s.isMethod && s.asMethod.isPrimaryConstructor)
+
+            if (primaryConstructorSymbol.isEmpty) {
+              sys.error("Internal SQL error: Product object did not have a primary constructor.")
+            } else {
+              primaryConstructorSymbol.get.asMethod.paramss
+            }
+          }
+
+          CreateNamedStruct(params.head.flatMap { p =>
+            val fieldName = p.name.toString
+            val fieldType = p.typeSignature.substituteTypes(formalTypeArgs, actualTypeArgs)
+            val fieldValue = Invoke(inputObject, fieldName, dataTypeFor(fieldType))
+            expressions.Literal(fieldName) :: extractorFor(fieldValue, fieldType) :: Nil
+          })
+
+        case t if t <:< localTypeOf[Array[_]] =>
+          val TypeRef(_, _, Seq(elementType)) = t
+          toCatalystArray(inputObject, elementType)
+
+        case t if t <:< localTypeOf[Seq[_]] =>
+          val TypeRef(_, _, Seq(elementType)) = t
+          toCatalystArray(inputObject, elementType)
+
+        case t if t <:< localTypeOf[Map[_, _]] =>
+          val TypeRef(_, _, Seq(keyType, valueType)) = t
+
+          val keys =
+            Invoke(
+              Invoke(inputObject, "keysIterator",
+                ObjectType(classOf[scala.collection.Iterator[_]])),
+              "toSeq",
+              ObjectType(classOf[scala.collection.Seq[_]]))
+          val convertedKeys = toCatalystArray(keys, keyType)
+
+          val values =
+            Invoke(
+              Invoke(inputObject, "valuesIterator",
+                ObjectType(classOf[scala.collection.Iterator[_]])),
+              "toSeq",
+              ObjectType(classOf[scala.collection.Seq[_]]))
+          val convertedValues = toCatalystArray(values, valueType)
+
+          val Schema(keyDataType, _) = schemaFor(keyType)
+          val Schema(valueDataType, valueNullable) = schemaFor(valueType)
+          NewInstance(
+            classOf[ArrayBasedMapData],
+            convertedKeys :: convertedValues :: Nil,
+            dataType = MapType(keyDataType, valueDataType, valueNullable))
+
+        case t if t <:< localTypeOf[String] =>
+          StaticInvoke(
+            classOf[UTF8String],
+            StringType,
+            "fromString",
+            inputObject :: Nil)
+
+        case t if t <:< localTypeOf[java.sql.Timestamp] =>
+          StaticInvoke(
+            DateTimeUtils,
+            TimestampType,
+            "fromJavaTimestamp",
+            inputObject :: Nil)
+
+        case t if t <:< localTypeOf[java.sql.Date] =>
+          StaticInvoke(
+            DateTimeUtils,
+            DateType,
+            "fromJavaDate",
+            inputObject :: Nil)
+
+        case t if t <:< localTypeOf[BigDecimal] =>
+          StaticInvoke(
+            Decimal,
+            DecimalType.SYSTEM_DEFAULT,
+            "apply",
+            inputObject :: Nil)
+
+        case t if t <:< localTypeOf[java.math.BigDecimal] =>
+          StaticInvoke(
+            Decimal,
+            DecimalType.SYSTEM_DEFAULT,
+            "apply",
+            inputObject :: Nil)
+
+        case t if t <:< localTypeOf[java.lang.Integer] =>
+          Invoke(inputObject, "intValue", IntegerType)
+        case t if t <:< localTypeOf[java.lang.Long] =>
+          Invoke(inputObject, "longValue", LongType)
+        case t if t <:< localTypeOf[java.lang.Double] =>
+          Invoke(inputObject, "doubleValue", DoubleType)
+        case t if t <:< localTypeOf[java.lang.Float] =>
+          Invoke(inputObject, "floatValue", FloatType)
+        case t if t <:< localTypeOf[java.lang.Short] =>
+          Invoke(inputObject, "shortValue", ShortType)
+        case t if t <:< localTypeOf[java.lang.Byte] =>
+          Invoke(inputObject, "byteValue", ByteType)
+        case t if t <:< localTypeOf[java.lang.Boolean] =>
+          Invoke(inputObject, "booleanValue", BooleanType)
+
+        case other =>
+          throw new UnsupportedOperationException(s"Extractor for type $other is not supported")
+      }
+    }
+  }
+
+  private def toCatalystArray(input: Expression, elementType: `Type`): Expression = {
+    val externalDataType = dataTypeFor(elementType)
+    val Schema(catalystType, nullable) = schemaFor(elementType)
+    if (RowEncoder.isNativeType(catalystType)) {
+      NewInstance(
+        classOf[GenericArrayData],
+        input :: Nil,
+        dataType = ArrayType(catalystType, nullable))
+    } else {
+      MapObjects(extractorFor(_, elementType), input, externalDataType)
+    }
+  }
+
+  def constructorFor(
+      tpe: `Type`,
+      path: Option[Expression] = None): Expression = ScalaReflectionLock.synchronized {
+
+    /** Returns the current path with a sub-field extracted. */
+    def addToPath(part: String): Expression = path
+      .map(p => UnresolvedExtractValue(p, expressions.Literal(part)))
+      .getOrElse(UnresolvedAttribute(part))
+
+    /** Returns the current path with a field at ordinal extracted. */
+    def addToPathOrdinal(ordinal: Int, dataType: DataType): Expression = path
+      .map(p => GetInternalRowField(p, ordinal, dataType))
+      .getOrElse(BoundReference(ordinal, dataType, false))
+
+    /** Returns the current path or `BoundReference`. */
+    def getPath: Expression = path.getOrElse(BoundReference(0, schemaFor(tpe).dataType, true))
+
+    tpe match {
+      case t if !dataTypeFor(t).isInstanceOf[ObjectType] => getPath
+
+      case t if t <:< localTypeOf[Option[_]] =>
+        val TypeRef(_, _, Seq(optType)) = t
+        WrapOption(null, constructorFor(optType, path))
+
+      case t if t <:< localTypeOf[java.lang.Integer] =>
+        val boxedType = classOf[java.lang.Integer]
+        val objectType = ObjectType(boxedType)
+        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
+
+      case t if t <:< localTypeOf[java.lang.Long] =>
+        val boxedType = classOf[java.lang.Long]
+        val objectType = ObjectType(boxedType)
+        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
+
+      case t if t <:< localTypeOf[java.lang.Double] =>
+        val boxedType = classOf[java.lang.Double]
+        val objectType = ObjectType(boxedType)
+        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
+
+      case t if t <:< localTypeOf[java.lang.Float] =>
+        val boxedType = classOf[java.lang.Float]
+        val objectType = ObjectType(boxedType)
+        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
+
+      case t if t <:< localTypeOf[java.lang.Short] =>
+        val boxedType = classOf[java.lang.Short]
+        val objectType = ObjectType(boxedType)
+        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
+
+      case t if t <:< localTypeOf[java.lang.Byte] =>
+        val boxedType = classOf[java.lang.Byte]
+        val objectType = ObjectType(boxedType)
+        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
+
+      case t if t <:< localTypeOf[java.lang.Boolean] =>
+        val boxedType = classOf[java.lang.Boolean]
+        val objectType = ObjectType(boxedType)
+        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
+
+      case t if t <:< localTypeOf[java.sql.Date] =>
+        StaticInvoke(
+          DateTimeUtils,
+          ObjectType(classOf[java.sql.Date]),
+          "toJavaDate",
+          getPath :: Nil,
+          propagateNull = true)
+
+      case t if t <:< localTypeOf[java.sql.Timestamp] =>
+        StaticInvoke(
+          DateTimeUtils,
+          ObjectType(classOf[java.sql.Timestamp]),
+          "toJavaTimestamp",
+          getPath :: Nil,
+          propagateNull = true)
+
+      case t if t <:< localTypeOf[java.lang.String] =>
+        Invoke(getPath, "toString", ObjectType(classOf[String]))
+
+      case t if t <:< localTypeOf[java.math.BigDecimal] =>
+        Invoke(getPath, "toJavaBigDecimal", ObjectType(classOf[java.math.BigDecimal]))
+
+      case t if t <:< localTypeOf[BigDecimal] =>
+        Invoke(getPath, "toBigDecimal", ObjectType(classOf[BigDecimal]))
+
+      case t if t <:< localTypeOf[Array[_]] =>
+        val TypeRef(_, _, Seq(elementType)) = t
+        val primitiveMethod = elementType match {
+          case t if t <:< definitions.IntTpe => Some("toIntArray")
+          case t if t <:< definitions.LongTpe => Some("toLongArray")
+          case t if t <:< definitions.DoubleTpe => Some("toDoubleArray")
+          case t if t <:< definitions.FloatTpe => Some("toFloatArray")
+          case t if t <:< definitions.ShortTpe => Some("toShortArray")
+          case t if t <:< definitions.ByteTpe => Some("toByteArray")
+          case t if t <:< definitions.BooleanTpe => Some("toBooleanArray")
+          case _ => None
+        }
+
+        primitiveMethod.map { method =>
+          Invoke(getPath, method, arrayClassFor(elementType))
+        }.getOrElse {
+          Invoke(
+            MapObjects(
+              p => constructorFor(elementType, Some(p)),
+              getPath,
+              schemaFor(elementType).dataType),
+            "array",
+            arrayClassFor(elementType))
+        }
+
+      case t if t <:< localTypeOf[Seq[_]] =>
+        val TypeRef(_, _, Seq(elementType)) = t
+        val arrayData =
+          Invoke(
+            MapObjects(
+              p => constructorFor(elementType, Some(p)),
+              getPath,
+              schemaFor(elementType).dataType),
+            "array",
+            ObjectType(classOf[Array[Any]]))
+
+        StaticInvoke(
+          scala.collection.mutable.WrappedArray,
+          ObjectType(classOf[Seq[_]]),
+          "make",
+          arrayData :: Nil)
+
+      case t if t <:< localTypeOf[Map[_, _]] =>
+        val TypeRef(_, _, Seq(keyType, valueType)) = t
+
+        val keyData =
+          Invoke(
+            MapObjects(
+              p => constructorFor(keyType, Some(p)),
+              Invoke(getPath, "keyArray", ArrayType(schemaFor(keyType).dataType)),
+              schemaFor(keyType).dataType),
+            "array",
+            ObjectType(classOf[Array[Any]]))
+
+        val valueData =
+          Invoke(
+            MapObjects(
+              p => constructorFor(valueType, Some(p)),
+              Invoke(getPath, "valueArray", ArrayType(schemaFor(valueType).dataType)),
+              schemaFor(valueType).dataType),
+            "array",
+            ObjectType(classOf[Array[Any]]))
+
+        StaticInvoke(
+          ArrayBasedMapData,
+          ObjectType(classOf[Map[_, _]]),
+          "toScalaMap",
+          keyData :: valueData :: Nil)
+
+      case t if t <:< localTypeOf[Product] =>
+        val formalTypeArgs = t.typeSymbol.asClass.typeParams
+        val TypeRef(_, _, actualTypeArgs) = t
+        val constructorSymbol = t.member(nme.CONSTRUCTOR)
+        val params = if (constructorSymbol.isMethod) {
+          constructorSymbol.asMethod.paramss
+        } else {
+          // Find the primary constructor, and use its parameter ordering.
+          val primaryConstructorSymbol: Option[Symbol] =
+            constructorSymbol.asTerm.alternatives.find(s =>
+              s.isMethod && s.asMethod.isPrimaryConstructor)
+
+          if (primaryConstructorSymbol.isEmpty) {
+            sys.error("Internal SQL error: Product object did not have a primary constructor.")
+          } else {
+            primaryConstructorSymbol.get.asMethod.paramss
+          }
+        }
+
+        val className: String = t.erasure.typeSymbol.asClass.fullName
+        val cls = Utils.classForName(className)
+
+        val arguments = params.head.zipWithIndex.map { case (p, i) =>
+          val fieldName = p.name.toString
+          val fieldType = p.typeSignature.substituteTypes(formalTypeArgs, actualTypeArgs)
+          val dataType = schemaFor(fieldType).dataType
+
+          // For tuples, we based grab the inner fields by ordinal instead of name.
+          if (className startsWith "scala.Tuple") {
+            constructorFor(fieldType, Some(addToPathOrdinal(i, dataType)))
+          } else {
+            constructorFor(fieldType, Some(addToPath(fieldName)))
+          }
+        }
+
+        val newInstance = NewInstance(cls, arguments, propagateNull = false, ObjectType(cls))
+
+        if (path.nonEmpty) {
+          expressions.If(
+            IsNull(getPath),
+            expressions.Literal.create(null, ObjectType(cls)),
+            newInstance
+          )
+        } else {
+          newInstance
+        }
+    }
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
index 0b42130a013b2..e0be896bb3548 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
@@ -119,9 +119,17 @@ object RowEncoder {
       CreateStruct(convertedFields)
   }
 
-  private def externalDataTypeFor(dt: DataType): DataType = dt match {
+  /**
+   * Returns true if the value of this data type is same between internal and external.
+   */
+  def isNativeType(dt: DataType): Boolean = dt match {
     case BooleanType | ByteType | ShortType | IntegerType | LongType |
-         FloatType | DoubleType | BinaryType => dt
+         FloatType | DoubleType | BinaryType => true
+    case _ => false
+  }
+
+  private def externalDataTypeFor(dt: DataType): DataType = dt match {
+    case _ if isNativeType(dt) => dt
     case TimestampType => ObjectType(classOf[java.sql.Timestamp])
     case DateType => ObjectType(classOf[java.sql.Date])
     case _: DecimalType => ObjectType(classOf[java.math.BigDecimal])
@@ -137,13 +145,13 @@ object RowEncoder {
       If(
         IsNull(field),
         Literal.create(null, externalDataTypeFor(f.dataType)),
-        constructorFor(BoundReference(i, f.dataType, f.nullable), f.dataType)
+        constructorFor(BoundReference(i, f.dataType, f.nullable))
       )
     }
     CreateExternalRow(fields)
   }
 
-  private def constructorFor(input: Expression, dataType: DataType): Expression = dataType match {
+  private def constructorFor(input: Expression): Expression = input.dataType match {
     case BooleanType | ByteType | ShortType | IntegerType | LongType |
          FloatType | DoubleType | BinaryType => input
 
@@ -170,7 +178,7 @@ object RowEncoder {
     case ArrayType(et, nullable) =>
       val arrayData =
         Invoke(
-          MapObjects(constructorFor(_, et), input, et),
+          MapObjects(constructorFor, input, et),
           "array",
           ObjectType(classOf[Array[_]]))
       StaticInvoke(
@@ -181,10 +189,10 @@ object RowEncoder {
 
     case MapType(kt, vt, valueNullable) =>
       val keyArrayType = ArrayType(kt, false)
-      val keyData = constructorFor(Invoke(input, "keyArray", keyArrayType), keyArrayType)
+      val keyData = constructorFor(Invoke(input, "keyArray", keyArrayType))
 
       val valueArrayType = ArrayType(vt, valueNullable)
-      val valueData = constructorFor(Invoke(input, "valueArray", valueArrayType), valueArrayType)
+      val valueData = constructorFor(Invoke(input, "valueArray", valueArrayType))
 
       StaticInvoke(
         ArrayBasedMapData,
@@ -197,42 +205,8 @@ object RowEncoder {
         If(
           Invoke(input, "isNullAt", BooleanType, Literal(i) :: Nil),
           Literal.create(null, externalDataTypeFor(f.dataType)),
-          constructorFor(getField(input, i, f.dataType), f.dataType))
+          constructorFor(GetInternalRowField(input, i, f.dataType)))
       }
       CreateExternalRow(convertedFields)
   }
-
-  private def getField(
-     row: Expression,
-     ordinal: Int,
-     dataType: DataType): Expression = dataType match {
-    case BooleanType =>
-      Invoke(row, "getBoolean", dataType, Literal(ordinal) :: Nil)
-    case ByteType =>
-      Invoke(row, "getByte", dataType, Literal(ordinal) :: Nil)
-    case ShortType =>
-      Invoke(row, "getShort", dataType, Literal(ordinal) :: Nil)
-    case IntegerType | DateType =>
-      Invoke(row, "getInt", dataType, Literal(ordinal) :: Nil)
-    case LongType | TimestampType =>
-      Invoke(row, "getLong", dataType, Literal(ordinal) :: Nil)
-    case FloatType =>
-      Invoke(row, "getFloat", dataType, Literal(ordinal) :: Nil)
-    case DoubleType =>
-      Invoke(row, "getDouble", dataType, Literal(ordinal) :: Nil)
-    case t: DecimalType =>
-      Invoke(row, "getDecimal", dataType, Seq(ordinal, t.precision, t.scale).map(Literal(_)))
-    case StringType =>
-      Invoke(row, "getUTF8String", dataType, Literal(ordinal) :: Nil)
-    case BinaryType =>
-      Invoke(row, "getBinary", dataType, Literal(ordinal) :: Nil)
-    case CalendarIntervalType =>
-      Invoke(row, "getInterval", dataType, Literal(ordinal) :: Nil)
-    case t: StructType =>
-      Invoke(row, "getStruct", dataType, Literal(ordinal) :: Literal(t.size) :: Nil)
-    case _: ArrayType =>
-      Invoke(row, "getArray", dataType, Literal(ordinal) :: Nil)
-    case _: MapType =>
-      Invoke(row, "getMap", dataType, Literal(ordinal) :: Nil)
-  }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
index f5fff90e5a542..deff8a5378b92 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
@@ -110,7 +110,7 @@ object DateTimeUtils {
   }
 
   def stringToTime(s: String): java.util.Date = {
-    var indexOfGMT = s.indexOf("GMT");
+    val indexOfGMT = s.indexOf("GMT")
     if (indexOfGMT != -1) {
       // ISO8601 with a weird time zone specifier (2000-01-01T00:00GMT+01:00)
       val s0 = s.substring(0, indexOfGMT)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GenericArrayData.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GenericArrayData.scala
index e9bf7b33e35be..96588bb5dc1bc 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GenericArrayData.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GenericArrayData.scala
@@ -23,7 +23,7 @@ import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
 
 class GenericArrayData(val array: Array[Any]) extends ArrayData {
 
-  def this(seq: scala.collection.GenIterable[Any]) = this(seq.toArray)
+  def this(seq: Seq[Any]) = this(seq.toArray)
 
   // TODO: This is boxing.  We should specialize.
   def this(primitiveArray: Array[Int]) = this(primitiveArray.toSeq)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
index b0dacf7f555e0..9fe64b4cf10e4 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
@@ -17,232 +17,27 @@
 
 package org.apache.spark.sql.catalyst.encoders
 
-import scala.collection.mutable.ArrayBuffer
-import scala.reflect.runtime.universe._
+import java.util.Arrays
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.sql.catalyst._
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
 import org.apache.spark.sql.catalyst.util.ArrayData
-import org.apache.spark.sql.types.{StructField, ArrayType}
-
-case class RepeatedStruct(s: Seq[PrimitiveData])
-
-case class NestedArray(a: Array[Array[Int]])
-
-case class BoxedData(
-    intField: java.lang.Integer,
-    longField: java.lang.Long,
-    doubleField: java.lang.Double,
-    floatField: java.lang.Float,
-    shortField: java.lang.Short,
-    byteField: java.lang.Byte,
-    booleanField: java.lang.Boolean)
-
-case class RepeatedData(
-    arrayField: Seq[Int],
-    arrayFieldContainsNull: Seq[java.lang.Integer],
-    mapField: scala.collection.Map[Int, Long],
-    mapFieldNull: scala.collection.Map[Int, java.lang.Long],
-    structField: PrimitiveData)
-
-case class SpecificCollection(l: List[Int])
-
-class ExpressionEncoderSuite extends SparkFunSuite {
-
-  encodeDecodeTest(1)
-  encodeDecodeTest(1L)
-  encodeDecodeTest(1.toDouble)
-  encodeDecodeTest(1.toFloat)
-  encodeDecodeTest(true)
-  encodeDecodeTest(false)
-  encodeDecodeTest(1.toShort)
-  encodeDecodeTest(1.toByte)
-  encodeDecodeTest("hello")
-
-  encodeDecodeTest(PrimitiveData(1, 1, 1, 1, 1, 1, true))
-
-  // TODO: Support creating specific subclasses of Seq.
-  ignore("Specific collection types") { encodeDecodeTest(SpecificCollection(1 :: Nil)) }
-
-  encodeDecodeTest(
-    OptionalData(
-      Some(2), Some(2), Some(2), Some(2), Some(2), Some(2), Some(true),
-      Some(PrimitiveData(1, 1, 1, 1, 1, 1, true))))
-
-  encodeDecodeTest(OptionalData(None, None, None, None, None, None, None, None))
-
-  encodeDecodeTest(
-    BoxedData(1, 1L, 1.0, 1.0f, 1.toShort, 1.toByte, true))
-
-  encodeDecodeTest(
-    BoxedData(null, null, null, null, null, null, null))
-
-  encodeDecodeTest(
-    RepeatedStruct(PrimitiveData(1, 1, 1, 1, 1, 1, true) :: Nil))
-
-  encodeDecodeTest(
-    RepeatedData(
-      Seq(1, 2),
-      Seq(new Integer(1), null, new Integer(2)),
-      Map(1 -> 2L),
-      Map(1 -> null),
-      PrimitiveData(1, 1, 1, 1, 1, 1, true)))
-
-  encodeDecodeTest(("nullable Seq[Integer]", Seq[Integer](1, null)))
-
-  encodeDecodeTest(("Seq[(String, String)]",
-    Seq(("a", "b"))))
-  encodeDecodeTest(("Seq[(Int, Int)]",
-    Seq((1, 2))))
-  encodeDecodeTest(("Seq[(Long, Long)]",
-    Seq((1L, 2L))))
-  encodeDecodeTest(("Seq[(Float, Float)]",
-    Seq((1.toFloat, 2.toFloat))))
-  encodeDecodeTest(("Seq[(Double, Double)]",
-    Seq((1.toDouble, 2.toDouble))))
-  encodeDecodeTest(("Seq[(Short, Short)]",
-    Seq((1.toShort, 2.toShort))))
-  encodeDecodeTest(("Seq[(Byte, Byte)]",
-    Seq((1.toByte, 2.toByte))))
-  encodeDecodeTest(("Seq[(Boolean, Boolean)]",
-    Seq((true, false))))
-
-  // TODO: Decoding/encoding of complex maps.
-  ignore("complex maps") {
-    encodeDecodeTest(("Map[Int, (String, String)]",
-      Map(1 ->("a", "b"))))
-  }
-
-  encodeDecodeTest(("ArrayBuffer[(String, String)]",
-    ArrayBuffer(("a", "b"))))
-  encodeDecodeTest(("ArrayBuffer[(Int, Int)]",
-    ArrayBuffer((1, 2))))
-  encodeDecodeTest(("ArrayBuffer[(Long, Long)]",
-    ArrayBuffer((1L, 2L))))
-  encodeDecodeTest(("ArrayBuffer[(Float, Float)]",
-    ArrayBuffer((1.toFloat, 2.toFloat))))
-  encodeDecodeTest(("ArrayBuffer[(Double, Double)]",
-    ArrayBuffer((1.toDouble, 2.toDouble))))
-  encodeDecodeTest(("ArrayBuffer[(Short, Short)]",
-    ArrayBuffer((1.toShort, 2.toShort))))
-  encodeDecodeTest(("ArrayBuffer[(Byte, Byte)]",
-    ArrayBuffer((1.toByte, 2.toByte))))
-  encodeDecodeTest(("ArrayBuffer[(Boolean, Boolean)]",
-    ArrayBuffer((true, false))))
-
-  encodeDecodeTest(("Seq[Seq[(Int, Int)]]",
-    Seq(Seq((1, 2)))))
-
-  encodeDecodeTestCustom(("Array[Array[(Int, Int)]]",
-    Array(Array((1, 2)))))
-  { (l, r) => l._2(0)(0) == r._2(0)(0) }
-
-  encodeDecodeTestCustom(("Array[Array[(Int, Int)]]",
-    Array(Array(Array((1, 2))))))
-  { (l, r) => l._2(0)(0)(0) == r._2(0)(0)(0) }
-
-  encodeDecodeTestCustom(("Array[Array[Array[(Int, Int)]]]",
-    Array(Array(Array(Array((1, 2)))))))
-  { (l, r) => l._2(0)(0)(0)(0) == r._2(0)(0)(0)(0) }
-
-  encodeDecodeTestCustom(("Array[Array[Array[Array[(Int, Int)]]]]",
-    Array(Array(Array(Array(Array((1, 2))))))))
-  { (l, r) => l._2(0)(0)(0)(0)(0) == r._2(0)(0)(0)(0)(0) }
-
-
-  encodeDecodeTestCustom(("Array[Array[Integer]]",
-    Array(Array[Integer](1))))
-  { (l, r) => l._2(0)(0) == r._2(0)(0) }
-
-  encodeDecodeTestCustom(("Array[Array[Int]]",
-    Array(Array(1))))
-  { (l, r) => l._2(0)(0) == r._2(0)(0) }
-
-  encodeDecodeTestCustom(("Array[Array[Int]]",
-    Array(Array(Array(1)))))
-  { (l, r) => l._2(0)(0)(0) == r._2(0)(0)(0) }
-
-  encodeDecodeTestCustom(("Array[Array[Array[Int]]]",
-    Array(Array(Array(Array(1))))))
-  { (l, r) => l._2(0)(0)(0)(0) == r._2(0)(0)(0)(0) }
-
-  encodeDecodeTestCustom(("Array[Array[Array[Array[Int]]]]",
-    Array(Array(Array(Array(Array(1)))))))
-  { (l, r) => l._2(0)(0)(0)(0)(0) == r._2(0)(0)(0)(0)(0) }
-
-  encodeDecodeTest(("Array[Byte] null",
-    null: Array[Byte]))
-  encodeDecodeTestCustom(("Array[Byte]",
-    Array[Byte](1, 2, 3)))
-    { (l, r) => java.util.Arrays.equals(l._2, r._2) }
-
-  encodeDecodeTest(("Array[Int] null",
-    null: Array[Int]))
-  encodeDecodeTestCustom(("Array[Int]",
-    Array[Int](1, 2, 3)))
-    { (l, r) => java.util.Arrays.equals(l._2, r._2) }
-
-  encodeDecodeTest(("Array[Long] null",
-    null: Array[Long]))
-  encodeDecodeTestCustom(("Array[Long]",
-    Array[Long](1, 2, 3)))
-    { (l, r) => java.util.Arrays.equals(l._2, r._2) }
-
-  encodeDecodeTest(("Array[Double] null",
-    null: Array[Double]))
-  encodeDecodeTestCustom(("Array[Double]",
-    Array[Double](1, 2, 3)))
-    { (l, r) => java.util.Arrays.equals(l._2, r._2) }
-
-  encodeDecodeTest(("Array[Float] null",
-    null: Array[Float]))
-  encodeDecodeTestCustom(("Array[Float]",
-    Array[Float](1, 2, 3)))
-    { (l, r) => java.util.Arrays.equals(l._2, r._2) }
-
-  encodeDecodeTest(("Array[Boolean] null",
-    null: Array[Boolean]))
-  encodeDecodeTestCustom(("Array[Boolean]",
-    Array[Boolean](true, false)))
-    { (l, r) => java.util.Arrays.equals(l._2, r._2) }
-
-  encodeDecodeTest(("Array[Short] null",
-    null: Array[Short]))
-  encodeDecodeTestCustom(("Array[Short]",
-    Array[Short](1, 2, 3)))
-    { (l, r) => java.util.Arrays.equals(l._2, r._2) }
-
-  encodeDecodeTestCustom(("java.sql.Timestamp",
-    new java.sql.Timestamp(1)))
-    { (l, r) => l._2.toString == r._2.toString }
-
-  encodeDecodeTestCustom(("java.sql.Date", new java.sql.Date(1)))
-    { (l, r) => l._2.toString == r._2.toString }
-
-  /** Simplified encodeDecodeTestCustom, where the comparison function can be `Object.equals`. */
-  protected def encodeDecodeTest[T : TypeTag](inputData: T) =
-    encodeDecodeTestCustom[T](inputData)((l, r) => l == r)
-
-  /**
-   * Constructs a test that round-trips `t` through an encoder, checking the results to ensure it
-   * matches the original.
-   */
-  protected def encodeDecodeTestCustom[T : TypeTag](
-      inputData: T)(
-      c: (T, T) => Boolean) = {
-    test(s"encode/decode: $inputData - ${inputData.getClass.getName}") {
-      val encoder = try ExpressionEncoder[T]() catch {
-        case e: Exception =>
-          fail(s"Exception thrown generating encoder", e)
-      }
-      val convertedData = encoder.toRow(inputData)
+import org.apache.spark.sql.types.ArrayType
+
+abstract class ExpressionEncoderSuite extends SparkFunSuite {
+  protected def encodeDecodeTest[T](
+      input: T,
+      encoder: ExpressionEncoder[T],
+      testName: String): Unit = {
+    test(s"encode/decode for $testName: $input") {
+      val row = encoder.toRow(input)
       val schema = encoder.schema.toAttributes
       val boundEncoder = encoder.resolve(schema).bind(schema)
-      val convertedBack = try boundEncoder.fromRow(convertedData) catch {
+      val convertedBack = try boundEncoder.fromRow(row) catch {
         case e: Exception =>
           fail(
            s"""Exception thrown while decoding
-              |Converted: $convertedData
+              |Converted: $row
               |Schema: ${schema.mkString(",")}
               |${encoder.schema.treeString}
               |
@@ -252,18 +47,27 @@ class ExpressionEncoderSuite extends SparkFunSuite {
             """.stripMargin, e)
       }
 
-      if (!c(inputData, convertedBack)) {
+      val isCorrect = (input, convertedBack) match {
+        case (b1: Array[Byte], b2: Array[Byte]) => Arrays.equals(b1, b2)
+        case (b1: Array[Int], b2: Array[Int]) => Arrays.equals(b1, b2)
+        case (b1: Array[Array[_]], b2: Array[Array[_]]) =>
+          Arrays.deepEquals(b1.asInstanceOf[Array[AnyRef]], b2.asInstanceOf[Array[AnyRef]])
+        case (b1: Array[_], b2: Array[_]) =>
+          Arrays.equals(b1.asInstanceOf[Array[AnyRef]], b2.asInstanceOf[Array[AnyRef]])
+        case _ => input == convertedBack
+      }
+
+      if (!isCorrect) {
         val types = convertedBack match {
           case c: Product =>
             c.productIterator.filter(_ != null).map(_.getClass.getName).mkString(",")
           case other => other.getClass.getName
         }
 
-
         val encodedData = try {
-          convertedData.toSeq(encoder.schema).zip(encoder.schema).map {
-            case (a: ArrayData, StructField(_, at: ArrayType, _, _)) =>
-              a.toArray[Any](at.elementType).toSeq
+          row.toSeq(encoder.schema).zip(schema).map {
+            case (a: ArrayData, AttributeReference(_, ArrayType(et, _), _, _)) =>
+              a.toArray[Any](et).toSeq
             case (other, _) =>
               other
           }.mkString("[", ",", "]")
@@ -274,7 +78,7 @@ class ExpressionEncoderSuite extends SparkFunSuite {
         fail(
           s"""Encoded/Decoded data does not match input data
              |
-             |in:  $inputData
+             |in:  $input
              |out: $convertedBack
              |types: $types
              |
@@ -282,11 +86,10 @@ class ExpressionEncoderSuite extends SparkFunSuite {
              |Schema: ${schema.mkString(",")}
              |${encoder.schema.treeString}
              |
-             |Extract Expressions:
-             |$boundEncoder
+             |fromRow Expressions:
+             |${boundEncoder.fromRowExpression.treeString}
          """.stripMargin)
-        }
       }
-
+    }
   }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
new file mode 100644
index 0000000000000..55821c4370684
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
@@ -0,0 +1,74 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.encoders
+
+import java.sql.{Date, Timestamp}
+
+class FlatEncoderSuite extends ExpressionEncoderSuite {
+  encodeDecodeTest(false, FlatEncoder[Boolean], "primitive boolean")
+  encodeDecodeTest(-3.toByte, FlatEncoder[Byte], "primitive byte")
+  encodeDecodeTest(-3.toShort, FlatEncoder[Short], "primitive short")
+  encodeDecodeTest(-3, FlatEncoder[Int], "primitive int")
+  encodeDecodeTest(-3L, FlatEncoder[Long], "primitive long")
+  encodeDecodeTest(-3.7f, FlatEncoder[Float], "primitive float")
+  encodeDecodeTest(-3.7, FlatEncoder[Double], "primitive double")
+
+  encodeDecodeTest(new java.lang.Boolean(false), FlatEncoder[java.lang.Boolean], "boxed boolean")
+  encodeDecodeTest(new java.lang.Byte(-3.toByte), FlatEncoder[java.lang.Byte], "boxed byte")
+  encodeDecodeTest(new java.lang.Short(-3.toShort), FlatEncoder[java.lang.Short], "boxed short")
+  encodeDecodeTest(new java.lang.Integer(-3), FlatEncoder[java.lang.Integer], "boxed int")
+  encodeDecodeTest(new java.lang.Long(-3L), FlatEncoder[java.lang.Long], "boxed long")
+  encodeDecodeTest(new java.lang.Float(-3.7f), FlatEncoder[java.lang.Float], "boxed float")
+  encodeDecodeTest(new java.lang.Double(-3.7), FlatEncoder[java.lang.Double], "boxed double")
+
+  encodeDecodeTest(BigDecimal("32131413.211321313"), FlatEncoder[BigDecimal], "scala decimal")
+  type JDecimal = java.math.BigDecimal
+  // encodeDecodeTest(new JDecimal("231341.23123"), FlatEncoder[JDecimal], "java decimal")
+
+  encodeDecodeTest("hello", FlatEncoder[String], "string")
+  encodeDecodeTest(Date.valueOf("2012-12-23"), FlatEncoder[Date], "date")
+  encodeDecodeTest(Timestamp.valueOf("2016-01-29 10:00:00"), FlatEncoder[Timestamp], "timestamp")
+  encodeDecodeTest(Array[Byte](13, 21, -23), FlatEncoder[Array[Byte]], "binary")
+
+  encodeDecodeTest(Seq(31, -123, 4), FlatEncoder[Seq[Int]], "seq of int")
+  encodeDecodeTest(Seq("abc", "xyz"), FlatEncoder[Seq[String]], "seq of string")
+  encodeDecodeTest(Seq("abc", null, "xyz"), FlatEncoder[Seq[String]], "seq of string with null")
+  encodeDecodeTest(Seq.empty[Int], FlatEncoder[Seq[Int]], "empty seq of int")
+  encodeDecodeTest(Seq.empty[String], FlatEncoder[Seq[String]], "empty seq of string")
+
+  encodeDecodeTest(Seq(Seq(31, -123), null, Seq(4, 67)),
+    FlatEncoder[Seq[Seq[Int]]], "seq of seq of int")
+  encodeDecodeTest(Seq(Seq("abc", "xyz"), Seq[String](null), null, Seq("1", null, "2")),
+    FlatEncoder[Seq[Seq[String]]], "seq of seq of string")
+
+  encodeDecodeTest(Array(31, -123, 4), FlatEncoder[Array[Int]], "array of int")
+  encodeDecodeTest(Array("abc", "xyz"), FlatEncoder[Array[String]], "array of string")
+  encodeDecodeTest(Array("a", null, "x"), FlatEncoder[Array[String]], "array of string with null")
+  encodeDecodeTest(Array.empty[Int], FlatEncoder[Array[Int]], "empty array of int")
+  encodeDecodeTest(Array.empty[String], FlatEncoder[Array[String]], "empty array of string")
+
+  encodeDecodeTest(Array(Array(31, -123), null, Array(4, 67)),
+    FlatEncoder[Array[Array[Int]]], "array of array of int")
+  encodeDecodeTest(Array(Array("abc", "xyz"), Array[String](null), null, Array("1", null, "2")),
+    FlatEncoder[Array[Array[String]]], "array of array of string")
+
+  encodeDecodeTest(Map(1 -> "a", 2 -> "b"), FlatEncoder[Map[Int, String]], "map")
+  encodeDecodeTest(Map(1 -> "a", 2 -> null), FlatEncoder[Map[Int, String]], "map with null")
+  encodeDecodeTest(Map(1 -> Map("a" -> 1), 2 -> Map("b" -> 2)),
+    FlatEncoder[Map[Int, Map[String, Int]]], "map of map")
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
new file mode 100644
index 0000000000000..fda978e7055ea
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.encoders
+
+import scala.collection.mutable.ArrayBuffer
+import scala.reflect.runtime.universe.TypeTag
+
+import org.apache.spark.sql.catalyst.{OptionalData, PrimitiveData}
+
+case class RepeatedStruct(s: Seq[PrimitiveData])
+
+case class NestedArray(a: Array[Array[Int]]) {
+  override def equals(other: Any): Boolean = other match {
+    case NestedArray(otherArray) =>
+      java.util.Arrays.deepEquals(
+        a.asInstanceOf[Array[AnyRef]],
+        otherArray.asInstanceOf[Array[AnyRef]])
+    case _ => false
+  }
+}
+
+case class BoxedData(
+    intField: java.lang.Integer,
+    longField: java.lang.Long,
+    doubleField: java.lang.Double,
+    floatField: java.lang.Float,
+    shortField: java.lang.Short,
+    byteField: java.lang.Byte,
+    booleanField: java.lang.Boolean)
+
+case class RepeatedData(
+    arrayField: Seq[Int],
+    arrayFieldContainsNull: Seq[java.lang.Integer],
+    mapField: scala.collection.Map[Int, Long],
+    mapFieldNull: scala.collection.Map[Int, java.lang.Long],
+    structField: PrimitiveData)
+
+case class SpecificCollection(l: List[Int])
+
+class ProductEncoderSuite extends ExpressionEncoderSuite {
+
+  productTest(PrimitiveData(1, 1, 1, 1, 1, 1, true))
+
+  productTest(
+    OptionalData(Some(2), Some(2), Some(2), Some(2), Some(2), Some(2), Some(true),
+      Some(PrimitiveData(1, 1, 1, 1, 1, 1, true))))
+
+  productTest(OptionalData(None, None, None, None, None, None, None, None))
+
+  productTest(BoxedData(1, 1L, 1.0, 1.0f, 1.toShort, 1.toByte, true))
+
+  productTest(BoxedData(null, null, null, null, null, null, null))
+
+  productTest(RepeatedStruct(PrimitiveData(1, 1, 1, 1, 1, 1, true) :: Nil))
+
+  productTest((1, "test", PrimitiveData(1, 1, 1, 1, 1, 1, true)))
+
+  productTest(
+    RepeatedData(
+      Seq(1, 2),
+      Seq(new Integer(1), null, new Integer(2)),
+      Map(1 -> 2L),
+      Map(1 -> null),
+      PrimitiveData(1, 1, 1, 1, 1, 1, true)))
+
+  productTest(NestedArray(Array(Array(1, -2, 3), null, Array(4, 5, -6))))
+
+  productTest(("Seq[(String, String)]",
+    Seq(("a", "b"))))
+  productTest(("Seq[(Int, Int)]",
+    Seq((1, 2))))
+  productTest(("Seq[(Long, Long)]",
+    Seq((1L, 2L))))
+  productTest(("Seq[(Float, Float)]",
+    Seq((1.toFloat, 2.toFloat))))
+  productTest(("Seq[(Double, Double)]",
+    Seq((1.toDouble, 2.toDouble))))
+  productTest(("Seq[(Short, Short)]",
+    Seq((1.toShort, 2.toShort))))
+  productTest(("Seq[(Byte, Byte)]",
+    Seq((1.toByte, 2.toByte))))
+  productTest(("Seq[(Boolean, Boolean)]",
+    Seq((true, false))))
+
+  productTest(("ArrayBuffer[(String, String)]",
+    ArrayBuffer(("a", "b"))))
+  productTest(("ArrayBuffer[(Int, Int)]",
+    ArrayBuffer((1, 2))))
+  productTest(("ArrayBuffer[(Long, Long)]",
+    ArrayBuffer((1L, 2L))))
+  productTest(("ArrayBuffer[(Float, Float)]",
+    ArrayBuffer((1.toFloat, 2.toFloat))))
+  productTest(("ArrayBuffer[(Double, Double)]",
+    ArrayBuffer((1.toDouble, 2.toDouble))))
+  productTest(("ArrayBuffer[(Short, Short)]",
+    ArrayBuffer((1.toShort, 2.toShort))))
+  productTest(("ArrayBuffer[(Byte, Byte)]",
+    ArrayBuffer((1.toByte, 2.toByte))))
+  productTest(("ArrayBuffer[(Boolean, Boolean)]",
+    ArrayBuffer((true, false))))
+
+  productTest(("Seq[Seq[(Int, Int)]]",
+    Seq(Seq((1, 2)))))
+
+  private def productTest[T <: Product : TypeTag](input: T): Unit = {
+    encodeDecodeTest(input, ProductEncoder[T], input.getClass.getSimpleName)
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 9c16940707de9..ebcf4c8bfe7e6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -22,7 +22,7 @@ import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.function._
-import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor}
+import org.apache.spark.sql.catalyst.encoders.{FlatEncoder, ExpressionEncoder, encoderFor}
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.QueryExecution
@@ -56,9 +56,6 @@ class GroupedDataset[K, T] private[sql](
   private val resolvedKEncoder = unresolvedKEncoder.resolve(groupingAttributes)
   private val resolvedTEncoder = unresolvedTEncoder.resolve(dataAttributes)
 
-  /** Encoders for built in aggregations. */
-  private implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder[Long](flat = true)
-
   private def logicalPlan = queryExecution.analyzed
   private def sqlContext = queryExecution.sqlContext
 
@@ -211,7 +208,7 @@ class GroupedDataset[K, T] private[sql](
    * Returns a [[Dataset]] that contains a tuple with each key and the number of items present
    * for that key.
    */
-  def count(): Dataset[(K, Long)] = agg(functions.count("*").as[Long])
+  def count(): Dataset[(K, Long)] = agg(functions.count("*").as(FlatEncoder[Long]))
 
   /**
    * Applies the given function to each cogrouped data.  For each unique group, the function will
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
index 6da46a5f7ef9a..8471eea1b7d9c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
@@ -37,17 +37,21 @@ import org.apache.spark.unsafe.types.UTF8String
 abstract class SQLImplicits {
   protected def _sqlContext: SQLContext
 
-  implicit def newProductEncoder[T <: Product : TypeTag]: Encoder[T] = ExpressionEncoder[T]()
+  implicit def newProductEncoder[T <: Product : TypeTag]: Encoder[T] = ProductEncoder[T]
 
-  implicit def newIntEncoder: Encoder[Int] = ExpressionEncoder[Int](flat = true)
-  implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder[Long](flat = true)
-  implicit def newDoubleEncoder: Encoder[Double] = ExpressionEncoder[Double](flat = true)
-  implicit def newFloatEncoder: Encoder[Float] = ExpressionEncoder[Float](flat = true)
-  implicit def newByteEncoder: Encoder[Byte] = ExpressionEncoder[Byte](flat = true)
-  implicit def newShortEncoder: Encoder[Short] = ExpressionEncoder[Short](flat = true)
-  implicit def newBooleanEncoder: Encoder[Boolean] = ExpressionEncoder[Boolean](flat = true)
-  implicit def newStringEncoder: Encoder[String] = ExpressionEncoder[String](flat = true)
+  implicit def newIntEncoder: Encoder[Int] = FlatEncoder[Int]
+  implicit def newLongEncoder: Encoder[Long] = FlatEncoder[Long]
+  implicit def newDoubleEncoder: Encoder[Double] = FlatEncoder[Double]
+  implicit def newFloatEncoder: Encoder[Float] = FlatEncoder[Float]
+  implicit def newByteEncoder: Encoder[Byte] = FlatEncoder[Byte]
+  implicit def newShortEncoder: Encoder[Short] = FlatEncoder[Short]
+  implicit def newBooleanEncoder: Encoder[Boolean] = FlatEncoder[Boolean]
+  implicit def newStringEncoder: Encoder[String] = FlatEncoder[String]
 
+  /**
+   * Creates a [[Dataset]] from an RDD.
+   * @since 1.6.0
+   */
   implicit def rddToDatasetHolder[T : Encoder](rdd: RDD[T]): DatasetHolder[T] = {
     DatasetHolder(_sqlContext.createDataset(rdd))
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 53cc6e0cda110..95158de710acf 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -26,7 +26,7 @@ import scala.util.Try
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.sql.catalyst.{SqlParser, ScalaReflection}
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunction, Star}
-import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.sql.catalyst.encoders.FlatEncoder
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.BroadcastHint
@@ -267,7 +267,7 @@ object functions extends LegacyFunctions {
    * @since 1.3.0
    */
   def count(columnName: String): TypedColumn[Any, Long] =
-    count(Column(columnName)).as(ExpressionEncoder[Long](flat = true))
+    count(Column(columnName)).as(FlatEncoder[Long])
 
   /**
    * Aggregate function: returns the number of distinct items in a group.

From 2d2411faa2dd1b7312c4277b2dd9e5678195cfbb Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Fri, 13 Nov 2015 13:09:28 -0800
Subject: [PATCH 237/518] [SPARK-11672][ML] Set active SQLContext in
 MLlibTestSparkContext.beforeAll

Still saw some error messages caused by `SQLContext.getOrCreate`:

https://amplab.cs.berkeley.edu/jenkins/job/Spark-Master-SBT/3997/AMPLAB_JENKINS_BUILD_PROFILE=hadoop2.3,label=spark-test/testReport/junit/org.apache.spark.ml.util/JavaDefaultReadWriteSuite/testDefaultReadWrite/

This PR sets the active SQLContext in beforeAll, which is not automatically set in `new SQLContext`. This makes `SQLContext.getOrCreate` return the right SQLContext.

cc: yhuai

Author: Xiangrui Meng <meng@databricks.com>

Closes #9694 from mengxr/SPARK-11672.3.
---
 .../main/scala/org/apache/spark/ml/util/ReadWrite.scala    | 7 +++++--
 .../apache/spark/mllib/util/MLlibTestSparkContext.scala    | 1 +
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index 85f888c9f2f67..ca896ed6106c4 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -48,8 +48,11 @@ private[util] sealed trait BaseReadWrite {
   /**
    * Returns the user-specified SQL context or the default.
    */
-  protected final def sqlContext: SQLContext = optionSQLContext.getOrElse {
-    SQLContext.getOrCreate(SparkContext.getOrCreate())
+  protected final def sqlContext: SQLContext = {
+    if (optionSQLContext.isEmpty) {
+      optionSQLContext = Some(SQLContext.getOrCreate(SparkContext.getOrCreate()))
+    }
+    optionSQLContext.get
   }
 
   /** Returns the [[SparkContext]] underlying [[sqlContext]] */
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/util/MLlibTestSparkContext.scala b/mllib/src/test/scala/org/apache/spark/mllib/util/MLlibTestSparkContext.scala
index 998ee48186558..378139593b26f 100644
--- a/mllib/src/test/scala/org/apache/spark/mllib/util/MLlibTestSparkContext.scala
+++ b/mllib/src/test/scala/org/apache/spark/mllib/util/MLlibTestSparkContext.scala
@@ -34,6 +34,7 @@ trait MLlibTestSparkContext extends BeforeAndAfterAll { self: Suite =>
     sc = new SparkContext(conf)
     SQLContext.clearActive()
     sqlContext = new SQLContext(sc)
+    SQLContext.setActive(sqlContext)
   }
 
   override def afterAll() {

From 912b94363bb113ab14024a45f17a4d2b82a09e66 Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Fri, 13 Nov 2015 13:14:25 -0800
Subject: [PATCH 238/518] [SPARK-11336] Add links to example codes

https://issues.apache.org/jira/browse/SPARK-11336

mengxr I add a hyperlink of Spark on Github and a hint of their existences in Spark code repo in each code example. I remove the config key for changing the example code dir, since we assume all examples  should be in spark/examples.

The hyperlink, though we cannot use it now, since the Spark v1.6.0 has not been released yet, can be used after the release. So it is not a problem.

I add some screen shots, so you can get an instant feeling.

<img width="949" alt="screen shot 2015-10-27 at 10 47 18 pm" src="https://cloud.githubusercontent.com/assets/2637239/10780634/bd20e072-7cfc-11e5-8960-def4fc62a8ea.png">

<img width="1144" alt="screen shot 2015-10-27 at 10 47 31 pm" src="https://cloud.githubusercontent.com/assets/2637239/10780636/c3f6e180-7cfc-11e5-80b2-233589f4a9a3.png">

Author: Xusen Yin <yinxusen@gmail.com>

Closes #9320 from yinxusen/SPARK-11336.
---
 docs/_plugins/include_example.rb | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/docs/_plugins/include_example.rb b/docs/_plugins/include_example.rb
index 6ee63a5ac69df..549f81fe1b1bc 100644
--- a/docs/_plugins/include_example.rb
+++ b/docs/_plugins/include_example.rb
@@ -28,7 +28,7 @@ def initialize(tag_name, markup, tokens)
  
     def render(context)
       site = context.registers[:site]
-      config_dir = (site.config['code_dir'] || '../examples/src/main').sub(/^\//,'')
+      config_dir = '../examples/src/main'
       @code_dir = File.join(site.source, config_dir)
 
       clean_markup = @markup.strip
@@ -38,7 +38,12 @@ def render(context)
       code = File.open(@file).read.encode("UTF-8")
       code = select_lines(code)
  
-      Pygments.highlight(code, :lexer => @lang)
+      rendered_code = Pygments.highlight(code, :lexer => @lang)
+
+      hint = "<div><small>Find full example code at " \
+        "\"examples/src/main/#{clean_markup}\" in the Spark repo.</small></div>"
+
+      rendered_code + hint
     end
  
     # Trim the code block so as to have the same indention, regardless of their positions in the

From bdfbc1dcaf121a1a1239857adcf54cdfe82c26dc Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Fri, 13 Nov 2015 13:19:04 -0800
Subject: [PATCH 239/518] [MINOR][ML] remove MLlibTestsSparkContext from
 ImpuritySuite

ImpuritySuite doesn't need SparkContext.

Author: Xiangrui Meng <meng@databricks.com>

Closes #9698 from mengxr/remove-mllib-test-context-in-impurity-suite.
---
 .../test/scala/org/apache/spark/mllib/tree/ImpuritySuite.scala | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/tree/ImpuritySuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/tree/ImpuritySuite.scala
index 49aff21fe7914..14152cdd63bc7 100644
--- a/mllib/src/test/scala/org/apache/spark/mllib/tree/ImpuritySuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/mllib/tree/ImpuritySuite.scala
@@ -19,12 +19,11 @@ package org.apache.spark.mllib.tree
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.mllib.tree.impurity.{EntropyAggregator, GiniAggregator}
-import org.apache.spark.mllib.util.MLlibTestSparkContext
 
 /**
  * Test suites for [[GiniAggregator]] and [[EntropyAggregator]].
  */
-class ImpuritySuite extends SparkFunSuite with MLlibTestSparkContext {
+class ImpuritySuite extends SparkFunSuite {
   test("Gini impurity does not support negative labels") {
     val gini = new GiniAggregator(2)
     intercept[IllegalArgumentException] {

From c939c70ac1ab6a26d9fda0a99c4e837f7e5a7935 Mon Sep 17 00:00:00 2001
From: nitin goyal <nitin.goyal@guavus.com>
Date: Fri, 13 Nov 2015 18:09:08 -0800
Subject: [PATCH 240/518] [SPARK-7970] Skip closure cleaning for SQL operations

Also introduces new spark private API in RDD.scala with name 'mapPartitionsInternal' which doesn't closure cleans the RDD elements.

Author: nitin goyal <nitin.goyal@guavus.com>
Author: nitin.goyal <nitin.goyal@guavus.com>

Closes #9253 from nitin2goyal/master.
---
 .../main/scala/org/apache/spark/rdd/RDD.scala  | 18 ++++++++++++++++++
 .../columnar/InMemoryColumnarTableScan.scala   |  4 ++--
 .../apache/spark/sql/execution/Exchange.scala  |  6 +++---
 .../apache/spark/sql/execution/Generate.scala  |  4 ++--
 .../aggregate/SortBasedAggregate.scala         |  2 +-
 .../spark/sql/execution/basicOperators.scala   | 16 ++++++++--------
 .../joins/BroadcastLeftSemiJoinHash.scala      |  4 ++--
 .../sql/execution/joins/CartesianProduct.scala |  2 +-
 .../org/apache/spark/sql/execution/sort.scala  |  2 +-
 9 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index 800ef53cbef07..2aeb5eeaad32c 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -705,6 +705,24 @@ abstract class RDD[T: ClassTag](
       preservesPartitioning)
   }
 
+  /**
+   * [performance] Spark's internal mapPartitions method which skips closure cleaning. It is a
+   * performance API to be used carefully only if we are sure that the RDD elements are
+   * serializable and don't require closure cleaning.
+   *
+   * @param preservesPartitioning indicates whether the input function preserves the partitioner,
+   * which should be `false` unless this is a pair RDD and the input function doesn't modify
+   * the keys.
+   */
+  private[spark] def mapPartitionsInternal[U: ClassTag](
+      f: Iterator[T] => Iterator[U],
+      preservesPartitioning: Boolean = false): RDD[U] = withScope {
+    new MapPartitionsRDD(
+      this,
+      (context: TaskContext, index: Int, iter: Iterator[T]) => f(iter),
+      preservesPartitioning)
+  }
+
   /**
    * Return a new RDD by applying a function to each partition of this RDD, while tracking the index
    * of the original partition.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
index 7eb1ad7cd8198..2cface61e59c4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
@@ -125,7 +125,7 @@ private[sql] case class InMemoryRelation(
 
   private def buildBuffers(): Unit = {
     val output = child.output
-    val cached = child.execute().mapPartitions { rowIterator =>
+    val cached = child.execute().mapPartitionsInternal { rowIterator =>
       new Iterator[CachedBatch] {
         def next(): CachedBatch = {
           val columnBuilders = output.map { attribute =>
@@ -292,7 +292,7 @@ private[sql] case class InMemoryColumnarTableScan(
     val relOutput = relation.output
     val buffers = relation.cachedColumnBuffers
 
-    buffers.mapPartitions { cachedBatchIterator =>
+    buffers.mapPartitionsInternal { cachedBatchIterator =>
       val partitionFilter = newPredicate(
         partitionFilters.reduceOption(And).getOrElse(Literal(true)),
         schema)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
index bc252d98e7144..a161cf0a3185b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
@@ -168,7 +168,7 @@ case class Exchange(
       case RangePartitioning(sortingExpressions, numPartitions) =>
         // Internally, RangePartitioner runs a job on the RDD that samples keys to compute
         // partition bounds. To get accurate samples, we need to copy the mutable keys.
-        val rddForSampling = rdd.mapPartitions { iter =>
+        val rddForSampling = rdd.mapPartitionsInternal { iter =>
           val mutablePair = new MutablePair[InternalRow, Null]()
           iter.map(row => mutablePair.update(row.copy(), null))
         }
@@ -200,12 +200,12 @@ case class Exchange(
     }
     val rddWithPartitionIds: RDD[Product2[Int, InternalRow]] = {
       if (needToCopyObjectsBeforeShuffle(part, serializer)) {
-        rdd.mapPartitions { iter =>
+        rdd.mapPartitionsInternal { iter =>
           val getPartitionKey = getPartitionKeyExtractor()
           iter.map { row => (part.getPartition(getPartitionKey(row)), row.copy()) }
         }
       } else {
-        rdd.mapPartitions { iter =>
+        rdd.mapPartitionsInternal { iter =>
           val getPartitionKey = getPartitionKeyExtractor()
           val mutablePair = new MutablePair[Int, InternalRow]()
           iter.map { row => mutablePair.update(part.getPartition(getPartitionKey(row)), row) }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Generate.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Generate.scala
index 78e33d9f233a6..54b8cb58285c2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Generate.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Generate.scala
@@ -59,7 +59,7 @@ case class Generate(
   protected override def doExecute(): RDD[InternalRow] = {
     // boundGenerator.terminate() should be triggered after all of the rows in the partition
     if (join) {
-      child.execute().mapPartitions { iter =>
+      child.execute().mapPartitionsInternal { iter =>
         val generatorNullRow = InternalRow.fromSeq(Seq.fill[Any](generator.elementTypes.size)(null))
         val joinedRow = new JoinedRow
 
@@ -79,7 +79,7 @@ case class Generate(
         }
       }
     } else {
-      child.execute().mapPartitions { iter =>
+      child.execute().mapPartitionsInternal { iter =>
         iter.flatMap(row => boundGenerator.eval(row)) ++
         LazyIterator(() => boundGenerator.terminate())
       }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
index c8ccbb933df61..ee982453c3287 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortBasedAggregate.scala
@@ -69,7 +69,7 @@ case class SortBasedAggregate(
   protected override def doExecute(): RDD[InternalRow] = attachTree(this, "execute") {
     val numInputRows = longMetric("numInputRows")
     val numOutputRows = longMetric("numOutputRows")
-    child.execute().mapPartitions { iter =>
+    child.execute().mapPartitionsInternal { iter =>
       // Because the constructor of an aggregation iterator will read at least the first row,
       // we need to get the value of iter.hasNext first.
       val hasInput = iter.hasNext
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index ed82c9a6a3770..07925c62cd386 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -43,7 +43,7 @@ case class TungstenProject(projectList: Seq[NamedExpression], child: SparkPlan)
 
   protected override def doExecute(): RDD[InternalRow] = {
     val numRows = longMetric("numRows")
-    child.execute().mapPartitions { iter =>
+    child.execute().mapPartitionsInternal { iter =>
       val project = UnsafeProjection.create(projectList, child.output,
         subexpressionEliminationEnabled)
       iter.map { row =>
@@ -67,7 +67,7 @@ case class Filter(condition: Expression, child: SparkPlan) extends UnaryNode {
   protected override def doExecute(): RDD[InternalRow] = {
     val numInputRows = longMetric("numInputRows")
     val numOutputRows = longMetric("numOutputRows")
-    child.execute().mapPartitions { iter =>
+    child.execute().mapPartitionsInternal { iter =>
       val predicate = newPredicate(condition, child.output)
       iter.filter { row =>
         numInputRows += 1
@@ -161,11 +161,11 @@ case class Limit(limit: Int, child: SparkPlan)
 
   protected override def doExecute(): RDD[InternalRow] = {
     val rdd: RDD[_ <: Product2[Boolean, InternalRow]] = if (sortBasedShuffleOn) {
-      child.execute().mapPartitions { iter =>
+      child.execute().mapPartitionsInternal { iter =>
         iter.take(limit).map(row => (false, row.copy()))
       }
     } else {
-      child.execute().mapPartitions { iter =>
+      child.execute().mapPartitionsInternal { iter =>
         val mutablePair = new MutablePair[Boolean, InternalRow]()
         iter.take(limit).map(row => mutablePair.update(false, row))
       }
@@ -173,7 +173,7 @@ case class Limit(limit: Int, child: SparkPlan)
     val part = new HashPartitioner(1)
     val shuffled = new ShuffledRDD[Boolean, InternalRow, InternalRow](rdd, part)
     shuffled.setSerializer(new SparkSqlSerializer(child.sqlContext.sparkContext.getConf))
-    shuffled.mapPartitions(_.take(limit).map(_._2))
+    shuffled.mapPartitionsInternal(_.take(limit).map(_._2))
   }
 }
 
@@ -294,7 +294,7 @@ case class MapPartitions[T, U](
     child: SparkPlan) extends UnaryNode {
 
   override protected def doExecute(): RDD[InternalRow] = {
-    child.execute().mapPartitions { iter =>
+    child.execute().mapPartitionsInternal { iter =>
       val tBoundEncoder = tEncoder.bind(child.output)
       func(iter.map(tBoundEncoder.fromRow)).map(uEncoder.toRow)
     }
@@ -318,7 +318,7 @@ case class AppendColumns[T, U](
   override def output: Seq[Attribute] = child.output ++ newColumns
 
   override protected def doExecute(): RDD[InternalRow] = {
-    child.execute().mapPartitions { iter =>
+    child.execute().mapPartitionsInternal { iter =>
       val tBoundEncoder = tEncoder.bind(child.output)
       val combiner = GenerateUnsafeRowJoiner.create(tEncoder.schema, uEncoder.schema)
       iter.map { row =>
@@ -350,7 +350,7 @@ case class MapGroups[K, T, U](
     Seq(groupingAttributes.map(SortOrder(_, Ascending)))
 
   override protected def doExecute(): RDD[InternalRow] = {
-    child.execute().mapPartitions { iter =>
+    child.execute().mapPartitionsInternal { iter =>
       val grouped = GroupedIterator(iter, groupingAttributes, child.output)
       val groupKeyEncoder = kEncoder.bind(groupingAttributes)
       val groupDataEncoder = tEncoder.bind(child.output)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastLeftSemiJoinHash.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastLeftSemiJoinHash.scala
index c5cd6a2fd6372..004407b2e6925 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastLeftSemiJoinHash.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastLeftSemiJoinHash.scala
@@ -54,7 +54,7 @@ case class BroadcastLeftSemiJoinHash(
       val hashSet = buildKeyHashSet(input.toIterator, SQLMetrics.nullLongMetric)
       val broadcastedRelation = sparkContext.broadcast(hashSet)
 
-      left.execute().mapPartitions { streamIter =>
+      left.execute().mapPartitionsInternal { streamIter =>
         hashSemiJoin(streamIter, numLeftRows, broadcastedRelation.value, numOutputRows)
       }
     } else {
@@ -62,7 +62,7 @@ case class BroadcastLeftSemiJoinHash(
         HashedRelation(input.toIterator, SQLMetrics.nullLongMetric, rightKeyGenerator, input.size)
       val broadcastedRelation = sparkContext.broadcast(hashRelation)
 
-      left.execute().mapPartitions { streamIter =>
+      left.execute().mapPartitionsInternal { streamIter =>
         val hashedRelation = broadcastedRelation.value
         hashedRelation match {
           case unsafe: UnsafeHashedRelation =>
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/CartesianProduct.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/CartesianProduct.scala
index 0243e196dbc37..f467519b802a7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/CartesianProduct.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/CartesianProduct.scala
@@ -46,7 +46,7 @@ case class CartesianProduct(left: SparkPlan, right: SparkPlan) extends BinaryNod
       row.copy()
     }
 
-    leftResults.cartesian(rightResults).mapPartitions { iter =>
+    leftResults.cartesian(rightResults).mapPartitionsInternal { iter =>
       val joinedRow = new JoinedRow
       iter.map { r =>
         numOutputRows += 1
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala
index 47fe70ab154ec..52ef00ef5b283 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala
@@ -47,7 +47,7 @@ case class Sort(
     if (global) OrderedDistribution(sortOrder) :: Nil else UnspecifiedDistribution :: Nil
 
   protected override def doExecute(): RDD[InternalRow] = attachTree(this, "sort") {
-    child.execute().mapPartitions( { iterator =>
+    child.execute().mapPartitionsInternal( { iterator =>
       val ordering = newOrdering(sortOrder, child.output)
       val sorter = new ExternalSorter[InternalRow, Null, InternalRow](
         TaskContext.get(), ordering = Some(ordering))

From 139c15b624c88b376ffdd05d78795295c8c4fc17 Mon Sep 17 00:00:00 2001
From: hyukjinkwon <gurwls223@gmail.com>
Date: Sat, 14 Nov 2015 18:36:01 +0800
Subject: [PATCH 241/518] [SPARK-11694][SQL] Parquet logical types are not
 being tested properly

All the physical types are properly tested at `ParquetIOSuite` but logical type mapping is not being tested.

Author: hyukjinkwon <gurwls223@gmail.com>
Author: Hyukjin Kwon <gurwls223@gmail.com>

Closes #9660 from HyukjinKwon/SPARK-11694.
---
 .../datasources/parquet/ParquetIOSuite.scala  | 39 ++++++++++++++-----
 .../datasources/parquet/ParquetTest.scala     | 17 ++++++++
 2 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index 82a42d68fedc1..78df363ade5c9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -91,6 +91,33 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
     }
   }
 
+  test("SPARK-11694 Parquet logical types are not being tested properly") {
+    val parquetSchema = MessageTypeParser.parseMessageType(
+      """message root {
+        |  required int32 a(INT_8);
+        |  required int32 b(INT_16);
+        |  required int32 c(DATE);
+        |  required int32 d(DECIMAL(1,0));
+        |  required int64 e(DECIMAL(10,0));
+        |  required binary f(UTF8);
+        |  required binary g(ENUM);
+        |  required binary h(DECIMAL(32,0));
+        |  required fixed_len_byte_array(32) i(DECIMAL(32,0));
+        |}
+      """.stripMargin)
+
+    val expectedSparkTypes = Seq(ByteType, ShortType, DateType, DecimalType(1, 0),
+      DecimalType(10, 0), StringType, StringType, DecimalType(32, 0), DecimalType(32, 0))
+
+    withTempPath { location =>
+      val path = new Path(location.getCanonicalPath)
+      val conf = sparkContext.hadoopConfiguration
+      writeMetadata(parquetSchema, path, conf)
+      val sparkTypes = sqlContext.read.parquet(path.toString).schema.map(_.dataType)
+      assert(sparkTypes === expectedSparkTypes)
+    }
+  }
+
   test("string") {
     val data = (1 to 4).map(i => Tuple1(i.toString))
     // Property spark.sql.parquet.binaryAsString shouldn't affect Parquet files written by Spark SQL
@@ -374,16 +401,10 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
       """.stripMargin)
 
     withTempPath { location =>
-      val extraMetadata = Collections.singletonMap(
-        CatalystReadSupport.SPARK_METADATA_KEY, sparkSchema.toString)
-      val fileMetadata = new FileMetaData(parquetSchema, extraMetadata, "Spark")
+      val extraMetadata = Map(CatalystReadSupport.SPARK_METADATA_KEY -> sparkSchema.toString)
       val path = new Path(location.getCanonicalPath)
-
-      ParquetFileWriter.writeMetadataFile(
-        sparkContext.hadoopConfiguration,
-        path,
-        Collections.singletonList(
-          new Footer(path, new ParquetMetadata(fileMetadata, Collections.emptyList()))))
+      val conf = sparkContext.hadoopConfiguration
+      writeMetadata(parquetSchema, path, conf, extraMetadata)
 
       assertResult(sqlContext.read.parquet(path.toString).schema) {
         StructType(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTest.scala
index 8ffb01fc5b584..fdd7697c91f5e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTest.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTest.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.execution.datasources.parquet
 
 import java.io.File
 
+import org.apache.parquet.schema.MessageType
+
 import scala.collection.JavaConverters._
 import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.TypeTag
@@ -117,6 +119,21 @@ private[sql] trait ParquetTest extends SQLTestUtils {
     ParquetFileWriter.writeMetadataFile(configuration, path, Seq(footer).asJava)
   }
 
+  /**
+   * This is an overloaded version of `writeMetadata` above to allow writing customized
+   * Parquet schema.
+   */
+  protected def writeMetadata(
+      parquetSchema: MessageType, path: Path, configuration: Configuration,
+      extraMetadata: Map[String, String] = Map.empty[String, String]): Unit = {
+    val extraMetadataAsJava = extraMetadata.asJava
+    val createdBy = s"Apache Spark ${org.apache.spark.SPARK_VERSION}"
+    val fileMetadata = new FileMetaData(parquetSchema, extraMetadataAsJava, createdBy)
+    val parquetMetadata = new ParquetMetadata(fileMetadata, Seq.empty[BlockMetaData].asJava)
+    val footer = new Footer(path, parquetMetadata)
+    ParquetFileWriter.writeMetadataFile(configuration, path, Seq(footer).asJava)
+  }
+
   protected def readAllFootersWithoutSummaryFiles(
       path: Path, configuration: Configuration): Seq[Footer] = {
     val fs = path.getFileSystem(configuration)

From 9a73b33a9a440d7312b92df9f6a9b9e17917b582 Mon Sep 17 00:00:00 2001
From: Kai Jiang <jiangkai@gmail.com>
Date: Sat, 14 Nov 2015 11:59:37 +0000
Subject: [PATCH 242/518] [MINOR][DOCS] typo in docs/configuration.md

`<\code>` end tag missing backslash in
docs/configuration.md{L308-L339}

ref #8795

Author: Kai Jiang <jiangkai@gmail.com>

Closes #9715 from vectorijk/minor-typo-docs.
---
 docs/configuration.md | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/docs/configuration.md b/docs/configuration.md
index c276e8e90decf..d961f43acf4ab 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -305,7 +305,7 @@ Apart from these, the following properties are also available, and may be useful
   <td>daily</td>
   <td>
     Set the time interval by which the executor logs will be rolled over.
-    Rolling is disabled by default. Valid values are <code>daily</code>, <code>hourly<code>, <code>minutely<code> or
+    Rolling is disabled by default. Valid values are <code>daily</code>, <code>hourly</code>, <code>minutely</code> or
     any interval in seconds. See <code>spark.executor.logs.rolling.maxRetainedFiles</code>
     for automatic cleaning of old logs.
   </td>
@@ -330,13 +330,13 @@ Apart from these, the following properties are also available, and may be useful
   <td><code>spark.python.profile</code></td>
   <td>false</td>
   <td>
-    Enable profiling in Python worker, the profile result will show up by <code>sc.show_profiles()<code>,
+    Enable profiling in Python worker, the profile result will show up by <code>sc.show_profiles()</code>,
     or it will be displayed before the driver exiting. It also can be dumped into disk by
-    <code>sc.dump_profiles(path)<code>. If some of the profile results had been displayed manually,
+    <code>sc.dump_profiles(path)</code>. If some of the profile results had been displayed manually,
     they will not be displayed automatically before driver exiting.
 
-    By default the <code>pyspark.profiler.BasicProfiler<code> will be used, but this can be overridden by
-    passing a profiler class in as a parameter to the <code>SparkContext<code> constructor.
+    By default the <code>pyspark.profiler.BasicProfiler</code> will be used, but this can be overridden by
+    passing a profiler class in as a parameter to the <code>SparkContext</code> constructor.
   </td>
 </tr>
 <tr>

From 9461f5ee80e51fd709d6f573c333936cb3c2acc5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=A1bor=20Lipt=C3=A1k?= <gliptak@gmail.com>
Date: Sat, 14 Nov 2015 12:02:02 +0000
Subject: [PATCH 243/518] =?UTF-8?q?[SPARK-11573]=20Correct=20'reflective?=
 =?UTF-8?q?=20access=20of=20structural=20type=20member=20meth=E2=80=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…od should be enabled' Scala warnings

Author: Gábor Lipták <gliptak@gmail.com>

Closes #9550 from gliptak/SPARK-11573.
---
 .../apache/spark/streaming/receiver/BlockGeneratorSuite.scala    | 1 +
 1 file changed, 1 insertion(+)

diff --git a/streaming/src/test/scala/org/apache/spark/streaming/receiver/BlockGeneratorSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/receiver/BlockGeneratorSuite.scala
index 2f11b255f1104..92ad9fe52b777 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/receiver/BlockGeneratorSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/receiver/BlockGeneratorSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.streaming.receiver
 
 import scala.collection.mutable
+import scala.language.reflectiveCalls
 
 import org.scalatest.BeforeAndAfter
 import org.scalatest.Matchers._

From 22e96b87fb0a0eb4f2f1a8fc29a742ceabff952a Mon Sep 17 00:00:00 2001
From: Rohan Bhanderi <rohan.bhanderi@sjsu.edu>
Date: Sat, 14 Nov 2015 13:38:53 +0000
Subject: [PATCH 244/518] Typo in comment: use 2 seconds instead of 1

Use 2 seconds batch size as duration specified in JavaStreamingContext constructor is 2000 ms

Author: Rohan Bhanderi <rohan.bhanderi@sjsu.edu>

Closes #9714 from RohanBhanderi/patch-2.
---
 .../org/apache/spark/examples/streaming/JavaKafkaWordCount.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/src/main/java/org/apache/spark/examples/streaming/JavaKafkaWordCount.java b/examples/src/main/java/org/apache/spark/examples/streaming/JavaKafkaWordCount.java
index 16ae9a3319ee2..337f8ffb5bfb0 100644
--- a/examples/src/main/java/org/apache/spark/examples/streaming/JavaKafkaWordCount.java
+++ b/examples/src/main/java/org/apache/spark/examples/streaming/JavaKafkaWordCount.java
@@ -66,7 +66,7 @@ public static void main(String[] args) {
 
     StreamingExamples.setStreamingLogLevels();
     SparkConf sparkConf = new SparkConf().setAppName("JavaKafkaWordCount");
-    // Create the context with a 1 second batch size
+    // Create the context with 2 seconds batch size
     JavaStreamingContext jssc = new JavaStreamingContext(sparkConf, new Duration(2000));
 
     int numThreads = Integer.parseInt(args[3]);

From d83c2f9f0b08d6d5d369d9fae04cdb15448e7f0d Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Sat, 14 Nov 2015 21:04:18 -0800
Subject: [PATCH 245/518] [SPARK-11736][SQL] Add monotonically_increasing_id to
 function registry.

https://issues.apache.org/jira/browse/SPARK-11736

Author: Yin Huai <yhuai@databricks.com>

Closes #9703 from yhuai/MonotonicallyIncreasingID.
---
 .../apache/spark/sql/catalyst/analysis/FunctionRegistry.scala | 3 ++-
 .../scala/org/apache/spark/sql/ColumnExpressionSuite.scala    | 4 ++++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index 870808aa560e5..a8f4d257acd0a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -281,7 +281,8 @@ object FunctionRegistry {
     expression[Sha1]("sha1"),
     expression[Sha2]("sha2"),
     expression[SparkPartitionID]("spark_partition_id"),
-    expression[InputFileName]("input_file_name")
+    expression[InputFileName]("input_file_name"),
+    expression[MonotonicallyIncreasingID]("monotonically_increasing_id")
   )
 
   val builtin: SimpleFunctionRegistry = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
index 010df2a341589..8674da7a79c81 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
@@ -563,6 +563,10 @@ class ColumnExpressionSuite extends QueryTest with SharedSQLContext {
       df.select(monotonicallyIncreasingId()),
       Row(0L) :: Row(1L) :: Row((1L << 33) + 0L) :: Row((1L << 33) + 1L) :: Nil
     )
+    checkAnswer(
+      df.select(expr("monotonically_increasing_id()")),
+      Row(0L) :: Row(1L) :: Row((1L << 33) + 0L) :: Row((1L << 33) + 1L) :: Nil
+    )
   }
 
   test("sparkPartitionId") {

From d22fc10887fdc6a86f6122648a823d0d37d4d795 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Sun, 15 Nov 2015 10:33:53 -0800
Subject: [PATCH 246/518] [SPARK-11734][SQL] Rename TungstenProject -> Project,
 TungstenSort -> Sort

I didn't remove the old Sort operator, since we still use it in randomized tests. I moved it into test module and renamed it ReferenceSort.

Author: Reynold Xin <rxin@databricks.com>

Closes #9700 from rxin/SPARK-11734.
---
 .../apache/spark/sql/execution/Exchange.scala |  7 +-
 .../sql/execution/{sort.scala => Sort.scala}  | 55 +-----------
 .../spark/sql/execution/SparkPlanner.scala    |  2 +-
 .../spark/sql/execution/SparkStrategies.scala | 16 +---
 .../spark/sql/execution/basicOperators.scala  |  2 +-
 .../datasources/DataSourceStrategy.scala      |  2 +-
 .../spark/sql/ColumnExpressionSuite.scala     |  4 +-
 .../spark/sql/execution/PlannerSuite.scala    |  6 +-
 .../spark/sql/execution/ReferenceSort.scala   | 61 +++++++++++++
 .../execution/RowFormatConvertersSuite.scala  |  4 +-
 .../spark/sql/execution/SortSuite.scala       | 69 +++++++++++++--
 .../sql/execution/TungstenSortSuite.scala     | 86 -------------------
 .../execution/metric/SQLMetricsSuite.scala    | 12 +--
 .../execution/HiveTypeCoercionSuite.scala     |  4 +-
 .../ParquetHadoopFsRelationSuite.scala        |  2 +-
 15 files changed, 148 insertions(+), 184 deletions(-)
 rename sql/core/src/main/scala/org/apache/spark/sql/execution/{sort.scala => Sort.scala} (65%)
 create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/ReferenceSort.scala
 delete mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
index a161cf0a3185b..62cbc518e02af 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Exchange.scala
@@ -51,7 +51,7 @@ case class Exchange(
     }
 
     val simpleNodeName = if (tungstenMode) "TungstenExchange" else "Exchange"
-    s"${simpleNodeName}${extraInfo}"
+    s"$simpleNodeName$extraInfo"
   }
 
   /**
@@ -475,10 +475,7 @@ private[sql] case class EnsureRequirements(sqlContext: SQLContext) extends Rule[
       if (requiredOrdering.nonEmpty) {
         // If child.outputOrdering is [a, b] and requiredOrdering is [a], we do not need to sort.
         if (requiredOrdering != child.outputOrdering.take(requiredOrdering.length)) {
-          sqlContext.planner.BasicOperators.getSortOperator(
-            requiredOrdering,
-            global = false,
-            child)
+          Sort(requiredOrdering, global = false, child = child)
         } else {
           child
         }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Sort.scala
similarity index 65%
rename from sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/Sort.scala
index 52ef00ef5b283..24207cb46fd29 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/sort.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Sort.scala
@@ -17,68 +17,22 @@
 
 package org.apache.spark.sql.execution
 
+import org.apache.spark.{InternalAccumulator, SparkEnv, TaskContext}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.errors._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.physical.{Distribution, OrderedDistribution, UnspecifiedDistribution}
 import org.apache.spark.sql.execution.metric.SQLMetrics
-import org.apache.spark.sql.types.StructType
-import org.apache.spark.util.CompletionIterator
-import org.apache.spark.util.collection.ExternalSorter
-import org.apache.spark.{InternalAccumulator, SparkEnv, TaskContext}
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-// This file defines various sort operators.
-////////////////////////////////////////////////////////////////////////////////////////////////////
 
 /**
- * Performs a sort, spilling to disk as needed.
- * @param global when true performs a global sort of all partitions by shuffling the data first
- *               if necessary.
- */
-case class Sort(
-    sortOrder: Seq[SortOrder],
-    global: Boolean,
-    child: SparkPlan)
-  extends UnaryNode {
-
-  override def requiredChildDistribution: Seq[Distribution] =
-    if (global) OrderedDistribution(sortOrder) :: Nil else UnspecifiedDistribution :: Nil
-
-  protected override def doExecute(): RDD[InternalRow] = attachTree(this, "sort") {
-    child.execute().mapPartitionsInternal( { iterator =>
-      val ordering = newOrdering(sortOrder, child.output)
-      val sorter = new ExternalSorter[InternalRow, Null, InternalRow](
-        TaskContext.get(), ordering = Some(ordering))
-      sorter.insertAll(iterator.map(r => (r.copy(), null)))
-      val baseIterator = sorter.iterator.map(_._1)
-      val context = TaskContext.get()
-      context.taskMetrics().incDiskBytesSpilled(sorter.diskBytesSpilled)
-      context.taskMetrics().incMemoryBytesSpilled(sorter.memoryBytesSpilled)
-      context.internalMetricsToAccumulators(
-        InternalAccumulator.PEAK_EXECUTION_MEMORY).add(sorter.peakMemoryUsedBytes)
-      // TODO(marmbrus): The complex type signature below thwarts inference for no reason.
-      CompletionIterator[InternalRow, Iterator[InternalRow]](baseIterator, sorter.stop())
-    }, preservesPartitioning = true)
-  }
-
-  override def output: Seq[Attribute] = child.output
-
-  override def outputOrdering: Seq[SortOrder] = sortOrder
-}
-
-/**
- * Optimized version of [[Sort]] that operates on binary data (implemented as part of
- * Project Tungsten).
+ * Performs (external) sorting.
  *
  * @param global when true performs a global sort of all partitions by shuffling the data first
  *               if necessary.
  * @param testSpillFrequency Method for configuring periodic spilling in unit tests. If set, will
  *                           spill every `frequency` records.
  */
-
-case class TungstenSort(
+case class Sort(
     sortOrder: Seq[SortOrder],
     global: Boolean,
     child: SparkPlan,
@@ -107,7 +61,7 @@ case class TungstenSort(
     val dataSize = longMetric("dataSize")
     val spillSize = longMetric("spillSize")
 
-    child.execute().mapPartitions { iter =>
+    child.execute().mapPartitionsInternal { iter =>
       val ordering = newOrdering(sortOrder, childOutput)
 
       // The comparator for comparing prefix
@@ -143,5 +97,4 @@ case class TungstenSort(
       sortedIterator
     }
   }
-
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
index b7c5476346b2a..6e9a4df828246 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanner.scala
@@ -80,7 +80,7 @@ class SparkPlanner(val sqlContext: SQLContext) extends SparkStrategies {
       filterCondition.map(Filter(_, scan)).getOrElse(scan)
     } else {
       val scan = scanBuilder((projectSet ++ filterSet).toSeq)
-      TungstenProject(projectList, filterCondition.map(Filter(_, scan)).getOrElse(scan))
+      Project(projectList, filterCondition.map(Filter(_, scan)).getOrElse(scan))
     }
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index 67201a2c191cd..3d4ce633c07c9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -302,16 +302,6 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
   object BasicOperators extends Strategy {
     def numPartitions: Int = self.numPartitions
 
-    /**
-     * Picks an appropriate sort operator.
-     *
-     * @param global when true performs a global sort of all partitions by shuffling the data first
-     *               if necessary.
-     */
-    def getSortOperator(sortExprs: Seq[SortOrder], global: Boolean, child: SparkPlan): SparkPlan = {
-      execution.TungstenSort(sortExprs, global, child)
-    }
-
     def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case r: RunnableCommand => ExecutedCommand(r) :: Nil
 
@@ -339,11 +329,11 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
       case logical.SortPartitions(sortExprs, child) =>
         // This sort only sorts tuples within a partition. Its requiredDistribution will be
         // an UnspecifiedDistribution.
-        getSortOperator(sortExprs, global = false, planLater(child)) :: Nil
+        execution.Sort(sortExprs, global = false, child = planLater(child)) :: Nil
       case logical.Sort(sortExprs, global, child) =>
-        getSortOperator(sortExprs, global, planLater(child)):: Nil
+        execution.Sort(sortExprs, global, planLater(child)) :: Nil
       case logical.Project(projectList, child) =>
-        execution.TungstenProject(projectList, planLater(child)) :: Nil
+        execution.Project(projectList, planLater(child)) :: Nil
       case logical.Filter(condition, child) =>
         execution.Filter(condition, planLater(child)) :: Nil
       case e @ logical.Expand(_, _, child) =>
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index 07925c62cd386..e79092efdaa3e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -30,7 +30,7 @@ import org.apache.spark.util.random.PoissonSampler
 import org.apache.spark.{HashPartitioner, SparkEnv}
 
 
-case class TungstenProject(projectList: Seq[NamedExpression], child: SparkPlan) extends UnaryNode {
+case class Project(projectList: Seq[NamedExpression], child: SparkPlan) extends UnaryNode {
 
   override private[sql] lazy val metrics = Map(
     "numRows" -> SQLMetrics.createLongMetric(sparkContext, "number of rows"))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
index 824c89a90eb8a..9bbbfa7c77cba 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
@@ -343,7 +343,7 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
         requestedColumns,
         scanBuilder(requestedColumns, candidatePredicates, pushedFilters),
         relation.relation)
-      execution.TungstenProject(
+      execution.Project(
         projects, filterCondition.map(execution.Filter(_, scan)).getOrElse(scan))
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
index 8674da7a79c81..3eae3f6d85066 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
@@ -20,7 +20,7 @@ package org.apache.spark.sql
 import org.apache.spark.sql.catalyst.expressions.NamedExpression
 import org.scalatest.Matchers._
 
-import org.apache.spark.sql.execution.TungstenProject
+import org.apache.spark.sql.execution.Project
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
@@ -619,7 +619,7 @@ class ColumnExpressionSuite extends QueryTest with SharedSQLContext {
 
     def checkNumProjects(df: DataFrame, expectedNumProjects: Int): Unit = {
       val projects = df.queryExecution.executedPlan.collect {
-        case tungstenProject: TungstenProject => tungstenProject
+        case tungstenProject: Project => tungstenProject
       }
       assert(projects.size === expectedNumProjects)
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
index 8c41d79dae817..be53ec3e271c6 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -365,7 +365,7 @@ class PlannerSuite extends SharedSQLContext {
     )
     val outputPlan = EnsureRequirements(sqlContext).apply(inputPlan)
     assertDistributionRequirementsAreSatisfied(outputPlan)
-    if (outputPlan.collect { case s: TungstenSort => true; case s: Sort => true }.isEmpty) {
+    if (outputPlan.collect { case s: Sort => true }.isEmpty) {
       fail(s"Sort should have been added:\n$outputPlan")
     }
   }
@@ -381,7 +381,7 @@ class PlannerSuite extends SharedSQLContext {
     )
     val outputPlan = EnsureRequirements(sqlContext).apply(inputPlan)
     assertDistributionRequirementsAreSatisfied(outputPlan)
-    if (outputPlan.collect { case s: TungstenSort => true; case s: Sort => true }.nonEmpty) {
+    if (outputPlan.collect { case s: Sort => true }.nonEmpty) {
       fail(s"No sorts should have been added:\n$outputPlan")
     }
   }
@@ -398,7 +398,7 @@ class PlannerSuite extends SharedSQLContext {
     )
     val outputPlan = EnsureRequirements(sqlContext).apply(inputPlan)
     assertDistributionRequirementsAreSatisfied(outputPlan)
-    if (outputPlan.collect { case s: TungstenSort => true; case s: Sort => true }.isEmpty) {
+    if (outputPlan.collect { case s: Sort => true }.isEmpty) {
       fail(s"Sort should have been added:\n$outputPlan")
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ReferenceSort.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReferenceSort.scala
new file mode 100644
index 0000000000000..9575d26fd123f
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ReferenceSort.scala
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.{InternalAccumulator, TaskContext}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.errors._
+import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder}
+import org.apache.spark.sql.catalyst.plans.physical._
+import org.apache.spark.util.CompletionIterator
+import org.apache.spark.util.collection.ExternalSorter
+
+
+/**
+ * A reference sort implementation used to compare against our normal sort.
+ */
+case class ReferenceSort(
+    sortOrder: Seq[SortOrder],
+    global: Boolean,
+    child: SparkPlan)
+  extends UnaryNode {
+
+  override def requiredChildDistribution: Seq[Distribution] =
+    if (global) OrderedDistribution(sortOrder) :: Nil else UnspecifiedDistribution :: Nil
+
+  protected override def doExecute(): RDD[InternalRow] = attachTree(this, "sort") {
+    child.execute().mapPartitions( { iterator =>
+      val ordering = newOrdering(sortOrder, child.output)
+      val sorter = new ExternalSorter[InternalRow, Null, InternalRow](
+        TaskContext.get(), ordering = Some(ordering))
+      sorter.insertAll(iterator.map(r => (r.copy(), null)))
+      val baseIterator = sorter.iterator.map(_._1)
+      val context = TaskContext.get()
+      context.taskMetrics().incDiskBytesSpilled(sorter.diskBytesSpilled)
+      context.taskMetrics().incMemoryBytesSpilled(sorter.memoryBytesSpilled)
+      context.internalMetricsToAccumulators(
+        InternalAccumulator.PEAK_EXECUTION_MEMORY).add(sorter.peakMemoryUsedBytes)
+      CompletionIterator[InternalRow, Iterator[InternalRow]](baseIterator, sorter.stop())
+    }, preservesPartitioning = true)
+  }
+
+  override def output: Seq[Attribute] = child.output
+
+  override def outputOrdering: Seq[SortOrder] = sortOrder
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RowFormatConvertersSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RowFormatConvertersSuite.scala
index b3fceeab64cfe..6876ab0f02b10 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RowFormatConvertersSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RowFormatConvertersSuite.scala
@@ -33,9 +33,9 @@ class RowFormatConvertersSuite extends SparkPlanTest with SharedSQLContext {
     case c: ConvertToSafe => c
   }
 
-  private val outputsSafe = Sort(Nil, false, PhysicalRDD(Seq.empty, null, "name"))
+  private val outputsSafe = ReferenceSort(Nil, false, PhysicalRDD(Seq.empty, null, "name"))
   assert(!outputsSafe.outputsUnsafeRows)
-  private val outputsUnsafe = TungstenSort(Nil, false, PhysicalRDD(Seq.empty, null, "name"))
+  private val outputsUnsafe = Sort(Nil, false, PhysicalRDD(Seq.empty, null, "name"))
   assert(outputsUnsafe.outputsUnsafeRows)
 
   test("planner should insert unsafe->safe conversions when required") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SortSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SortSuite.scala
index 847c188a30333..e5d34be4c65e8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SortSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SortSuite.scala
@@ -17,15 +17,22 @@
 
 package org.apache.spark.sql.execution
 
-import org.apache.spark.sql.Row
+import scala.util.Random
+
+import org.apache.spark.AccumulatorSuite
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.test.SharedSQLContext
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.{RandomDataGenerator, Row}
+
 
+/**
+ * Test sorting. Many of the test cases generate random data and compares the sorted result with one
+ * sorted by a reference implementation ([[ReferenceSort]]).
+ */
 class SortSuite extends SparkPlanTest with SharedSQLContext {
   import testImplicits.localSeqToDataFrameHolder
 
-  // This test was originally added as an example of how to use [[SparkPlanTest]];
-  // it's not designed to be a comprehensive test of ExternalSort.
   test("basic sorting using ExternalSort") {
 
     val input = Seq(
@@ -36,14 +43,66 @@ class SortSuite extends SparkPlanTest with SharedSQLContext {
 
     checkAnswer(
       input.toDF("a", "b", "c"),
-      Sort('a.asc :: 'b.asc :: Nil, global = true, _: SparkPlan),
+      (child: SparkPlan) => Sort('a.asc :: 'b.asc :: Nil, global = true, child = child),
       input.sortBy(t => (t._1, t._2)).map(Row.fromTuple),
       sortAnswers = false)
 
     checkAnswer(
       input.toDF("a", "b", "c"),
-      Sort('b.asc :: 'a.asc :: Nil, global = true, _: SparkPlan),
+      (child: SparkPlan) => Sort('b.asc :: 'a.asc :: Nil, global = true, child = child),
       input.sortBy(t => (t._2, t._1)).map(Row.fromTuple),
       sortAnswers = false)
   }
+
+  test("sort followed by limit") {
+    checkThatPlansAgree(
+      (1 to 100).map(v => Tuple1(v)).toDF("a"),
+      (child: SparkPlan) => Limit(10, Sort('a.asc :: Nil, global = true, child = child)),
+      (child: SparkPlan) => Limit(10, ReferenceSort('a.asc :: Nil, global = true, child)),
+      sortAnswers = false
+    )
+  }
+
+  test("sorting does not crash for large inputs") {
+    val sortOrder = 'a.asc :: Nil
+    val stringLength = 1024 * 1024 * 2
+    checkThatPlansAgree(
+      Seq(Tuple1("a" * stringLength), Tuple1("b" * stringLength)).toDF("a").repartition(1),
+      Sort(sortOrder, global = true, _: SparkPlan, testSpillFrequency = 1),
+      ReferenceSort(sortOrder, global = true, _: SparkPlan),
+      sortAnswers = false
+    )
+  }
+
+  test("sorting updates peak execution memory") {
+    AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "unsafe external sort") {
+      checkThatPlansAgree(
+        (1 to 100).map(v => Tuple1(v)).toDF("a"),
+        (child: SparkPlan) => Sort('a.asc :: Nil, global = true, child = child),
+        (child: SparkPlan) => ReferenceSort('a.asc :: Nil, global = true, child),
+        sortAnswers = false)
+    }
+  }
+
+  // Test sorting on different data types
+  for (
+    dataType <- DataTypeTestUtils.atomicTypes ++ Set(NullType);
+    nullable <- Seq(true, false);
+    sortOrder <- Seq('a.asc :: Nil, 'a.desc :: Nil);
+    randomDataGenerator <- RandomDataGenerator.forType(dataType, nullable)
+  ) {
+    test(s"sorting on $dataType with nullable=$nullable, sortOrder=$sortOrder") {
+      val inputData = Seq.fill(1000)(randomDataGenerator())
+      val inputDf = sqlContext.createDataFrame(
+        sparkContext.parallelize(Random.shuffle(inputData).map(v => Row(v))),
+        StructType(StructField("a", dataType, nullable = true) :: Nil)
+      )
+      checkThatPlansAgree(
+        inputDf,
+        p => ConvertToSafe(Sort(sortOrder, global = true, p: SparkPlan, testSpillFrequency = 23)),
+        ReferenceSort(sortOrder, global = true, _: SparkPlan),
+        sortAnswers = false
+      )
+    }
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala
deleted file mode 100644
index 7c860d1d58d5a..0000000000000
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/TungstenSortSuite.scala
+++ /dev/null
@@ -1,86 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.execution
-
-import scala.util.Random
-
-import org.apache.spark.AccumulatorSuite
-import org.apache.spark.sql.{RandomDataGenerator, Row, SQLConf}
-import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.test.SharedSQLContext
-import org.apache.spark.sql.types._
-
-/**
- * A test suite that generates randomized data to test the [[TungstenSort]] operator.
- */
-class TungstenSortSuite extends SparkPlanTest with SharedSQLContext {
-  import testImplicits.localSeqToDataFrameHolder
-
-  test("sort followed by limit") {
-    checkThatPlansAgree(
-      (1 to 100).map(v => Tuple1(v)).toDF("a"),
-      (child: SparkPlan) => Limit(10, TungstenSort('a.asc :: Nil, true, child)),
-      (child: SparkPlan) => Limit(10, Sort('a.asc :: Nil, global = true, child)),
-      sortAnswers = false
-    )
-  }
-
-  test("sorting does not crash for large inputs") {
-    val sortOrder = 'a.asc :: Nil
-    val stringLength = 1024 * 1024 * 2
-    checkThatPlansAgree(
-      Seq(Tuple1("a" * stringLength), Tuple1("b" * stringLength)).toDF("a").repartition(1),
-      TungstenSort(sortOrder, global = true, _: SparkPlan, testSpillFrequency = 1),
-      Sort(sortOrder, global = true, _: SparkPlan),
-      sortAnswers = false
-    )
-  }
-
-  test("sorting updates peak execution memory") {
-    AccumulatorSuite.verifyPeakExecutionMemorySet(sparkContext, "unsafe external sort") {
-      checkThatPlansAgree(
-        (1 to 100).map(v => Tuple1(v)).toDF("a"),
-        (child: SparkPlan) => TungstenSort('a.asc :: Nil, true, child),
-        (child: SparkPlan) => Sort('a.asc :: Nil, global = true, child),
-        sortAnswers = false)
-    }
-  }
-
-  // Test sorting on different data types
-  for (
-    dataType <- DataTypeTestUtils.atomicTypes ++ Set(NullType);
-    nullable <- Seq(true, false);
-    sortOrder <- Seq('a.asc :: Nil, 'a.desc :: Nil);
-    randomDataGenerator <- RandomDataGenerator.forType(dataType, nullable)
-  ) {
-    test(s"sorting on $dataType with nullable=$nullable, sortOrder=$sortOrder") {
-      val inputData = Seq.fill(1000)(randomDataGenerator())
-      val inputDf = sqlContext.createDataFrame(
-        sparkContext.parallelize(Random.shuffle(inputData).map(v => Row(v))),
-        StructType(StructField("a", dataType, nullable = true) :: Nil)
-      )
-      checkThatPlansAgree(
-        inputDf,
-        plan => ConvertToSafe(
-          TungstenSort(sortOrder, global = true, plan: SparkPlan, testSpillFrequency = 23)),
-        Sort(sortOrder, global = true, _: SparkPlan),
-        sortAnswers = false
-      )
-    }
-  }
-}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index 486bfbbd70887..5e2b4154dd7ce 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -114,17 +114,7 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
     // PhysicalRDD(nodeId = 1) -> Project(nodeId = 0)
     val df = person.select('name)
     testSparkPlanMetrics(df, 1, Map(
-      0L ->("TungstenProject", Map(
-        "number of rows" -> 2L)))
-    )
-  }
-
-  test("TungstenProject metrics") {
-    // Assume the execution plan is
-    // PhysicalRDD(nodeId = 1) -> TungstenProject(nodeId = 0)
-    val df = person.select('name)
-    testSparkPlanMetrics(df, 1, Map(
-      0L ->("TungstenProject", Map(
+      0L ->("Project", Map(
         "number of rows" -> 2L)))
     )
   }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveTypeCoercionSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveTypeCoercionSuite.scala
index 4cf4e13890294..5bd323ea096a4 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveTypeCoercionSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveTypeCoercionSuite.scala
@@ -18,7 +18,7 @@
 package org.apache.spark.sql.hive.execution
 
 import org.apache.spark.sql.catalyst.expressions.{Cast, EqualTo}
-import org.apache.spark.sql.execution.TungstenProject
+import org.apache.spark.sql.execution.Project
 import org.apache.spark.sql.hive.test.TestHive
 
 /**
@@ -44,7 +44,7 @@ class HiveTypeCoercionSuite extends HiveComparisonTest {
   test("[SPARK-2210] boolean cast on boolean value should be removed") {
     val q = "select cast(cast(key=0 as boolean) as boolean) from src"
     val project = TestHive.sql(q).queryExecution.executedPlan.collect {
-      case e: TungstenProject => e
+      case e: Project => e
     }.head
 
     // No cast expression introduced
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
index b6db6225805a1..e866493ee6c96 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/ParquetHadoopFsRelationSuite.scala
@@ -151,7 +151,7 @@ class ParquetHadoopFsRelationSuite extends HadoopFsRelationTest {
       val df = sqlContext.read.parquet(path).filter('a === 0).select('b)
       val physicalPlan = df.queryExecution.executedPlan
 
-      assert(physicalPlan.collect { case p: execution.TungstenProject => p }.length === 1)
+      assert(physicalPlan.collect { case p: execution.Project => p }.length === 1)
       assert(physicalPlan.collect { case p: execution.Filter => p }.length === 1)
     }
   }

From 64e55511033afb6ef42be142eb371bfbc31f5230 Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Sun, 15 Nov 2015 13:23:05 -0800
Subject: [PATCH 247/518] [SPARK-11672][ML] set active SQLContext in
 JavaDefaultReadWriteSuite

The same as #9694, but for Java test suite. yhuai

Author: Xiangrui Meng <meng@databricks.com>

Closes #9719 from mengxr/SPARK-11672.4.
---
 .../apache/spark/ml/util/JavaDefaultReadWriteSuite.java    | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java b/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
index c39538014be81..01ff1ea658610 100644
--- a/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
+++ b/mllib/src/test/java/org/apache/spark/ml/util/JavaDefaultReadWriteSuite.java
@@ -32,17 +32,23 @@
 public class JavaDefaultReadWriteSuite {
 
   JavaSparkContext jsc = null;
+  SQLContext sqlContext = null;
   File tempDir = null;
 
   @Before
   public void setUp() {
     jsc = new JavaSparkContext("local[2]", "JavaDefaultReadWriteSuite");
+    SQLContext.clearActive();
+    sqlContext = new SQLContext(jsc);
+    SQLContext.setActive(sqlContext);
     tempDir = Utils.createTempDir(
       System.getProperty("java.io.tmpdir"), "JavaDefaultReadWriteSuite");
   }
 
   @After
   public void tearDown() {
+    sqlContext = null;
+    SQLContext.clearActive();
     if (jsc != null) {
       jsc.stop();
       jsc = null;
@@ -64,7 +70,6 @@ public void testDefaultReadWrite() throws IOException {
     } catch (IOException e) {
       // expected
     }
-    SQLContext sqlContext = new SQLContext(jsc);
     instance.write().context(sqlContext).overwrite().save(outputPath);
     MyParams newInstance = MyParams.load(outputPath);
     Assert.assertEquals("UID should match.", instance.uid(), newInstance.uid());

From 3e2e1873b2762d07e49de8f9ea709bf3fa2d171c Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Sun, 15 Nov 2015 13:59:59 -0800
Subject: [PATCH 248/518] [SPARK-11738] [SQL] Making ArrayType orderable

https://issues.apache.org/jira/browse/SPARK-11738

Author: Yin Huai <yhuai@databricks.com>

Closes #9718 from yhuai/makingArrayOrderable.
---
 .../sql/catalyst/analysis/CheckAnalysis.scala |  32 +----
 .../expressions/codegen/CodeGenerator.scala   |  43 ++++++
 .../expressions/collectionOperations.scala    |   2 +
 .../sql/catalyst/expressions/ordering.scala   |   6 +
 .../spark/sql/catalyst/util/TypeUtils.scala   |   1 +
 .../spark/sql/types/AbstractDataType.scala    |   1 +
 .../apache/spark/sql/types/ArrayType.scala    |  48 +++++++
 .../analysis/AnalysisErrorSuite.scala         |  37 ++++--
 .../ExpressionTypeCheckingSuite.scala         |  32 ++---
 .../sql/catalyst/analysis/TestRelations.scala |   3 +
 .../expressions/CodeGenerationSuite.scala     |  36 -----
 .../catalyst/expressions/OrderingSuite.scala  | 124 ++++++++++++++++++
 .../spark/sql/DataFrameFunctionsSuite.scala   |  12 +-
 .../execution/AggregationQuerySuite.scala     |  52 ++++++++
 14 files changed, 335 insertions(+), 94 deletions(-)
 create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/OrderingSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
index 5a4b0c1e39ce1..7b2c93d63d673 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala
@@ -137,32 +137,14 @@ trait CheckAnalysis {
               case e => e.children.foreach(checkValidAggregateExpression)
             }
 
-            def checkSupportedGroupingDataType(
-                expressionString: String,
-                dataType: DataType): Unit = dataType match {
-              case BinaryType =>
-                failAnalysis(s"expression $expressionString cannot be used in " +
-                  s"grouping expression because it is in binary type or its inner field is " +
-                  s"in binary type")
-              case a: ArrayType =>
-                failAnalysis(s"expression $expressionString cannot be used in " +
-                  s"grouping expression because it is in array type or its inner field is " +
-                  s"in array type")
-              case m: MapType =>
-                failAnalysis(s"expression $expressionString cannot be used in " +
-                  s"grouping expression because it is in map type or its inner field is " +
-                  s"in map type")
-              case s: StructType =>
-                s.fields.foreach { f =>
-                  checkSupportedGroupingDataType(expressionString, f.dataType)
-                }
-              case udt: UserDefinedType[_] =>
-                checkSupportedGroupingDataType(expressionString, udt.sqlType)
-              case _ => // OK
-            }
-
             def checkValidGroupingExprs(expr: Expression): Unit = {
-              checkSupportedGroupingDataType(expr.prettyString, expr.dataType)
+              // Check if the data type of expr is orderable.
+              if (!RowOrdering.isOrderable(expr.dataType)) {
+                failAnalysis(
+                  s"expression ${expr.prettyString} cannot be used as a grouping expression " +
+                    s"because its data type ${expr.dataType.simpleString} is not a orderable " +
+                    s"data type.")
+              }
 
               if (!expr.deterministic) {
                 // This is just a sanity check, our analysis rule PullOutNondeterministic should
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
index ccd91d3549b53..1718cfbd35332 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
@@ -267,6 +267,49 @@ class CodeGenContext {
     case dt: DataType if isPrimitiveType(dt) => s"($c1 > $c2 ? 1 : $c1 < $c2 ? -1 : 0)"
     case BinaryType => s"org.apache.spark.sql.catalyst.util.TypeUtils.compareBinary($c1, $c2)"
     case NullType => "0"
+    case array: ArrayType =>
+      val elementType = array.elementType
+      val elementA = freshName("elementA")
+      val isNullA = freshName("isNullA")
+      val elementB = freshName("elementB")
+      val isNullB = freshName("isNullB")
+      val compareFunc = freshName("compareArray")
+      val minLength = freshName("minLength")
+      val funcCode: String =
+        s"""
+          public int $compareFunc(ArrayData a, ArrayData b) {
+            int lengthA = a.numElements();
+            int lengthB = b.numElements();
+            int $minLength = (lengthA > lengthB) ? lengthB : lengthA;
+            for (int i = 0; i < $minLength; i++) {
+              boolean $isNullA = a.isNullAt(i);
+              boolean $isNullB = b.isNullAt(i);
+              if ($isNullA && $isNullB) {
+                // Nothing
+              } else if ($isNullA) {
+                return -1;
+              } else if ($isNullB) {
+                return 1;
+              } else {
+                ${javaType(elementType)} $elementA = ${getValue("a", elementType, "i")};
+                ${javaType(elementType)} $elementB = ${getValue("b", elementType, "i")};
+                int comp = ${genComp(elementType, elementA, elementB)};
+                if (comp != 0) {
+                  return comp;
+                }
+              }
+            }
+
+            if (lengthA < lengthB) {
+              return -1;
+            } else if (lengthA > lengthB) {
+              return 1;
+            }
+            return 0;
+          }
+        """
+      addNewFunction(compareFunc, funcCode)
+      s"this.$compareFunc($c1, $c2)"
     case schema: StructType =>
       val comparisons = GenerateOrdering.genComparisons(this, schema)
       val compareFunc = freshName("compareStruct")
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
index 2cf19b939f734..741ad1f3efd8a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
@@ -68,6 +68,7 @@ case class SortArray(base: Expression, ascendingOrder: Expression)
   private lazy val lt: Comparator[Any] = {
     val ordering = base.dataType match {
       case _ @ ArrayType(n: AtomicType, _) => n.ordering.asInstanceOf[Ordering[Any]]
+      case _ @ ArrayType(a: ArrayType, _) => a.interpretedOrdering.asInstanceOf[Ordering[Any]]
       case _ @ ArrayType(s: StructType, _) => s.interpretedOrdering.asInstanceOf[Ordering[Any]]
     }
 
@@ -90,6 +91,7 @@ case class SortArray(base: Expression, ascendingOrder: Expression)
   private lazy val gt: Comparator[Any] = {
     val ordering = base.dataType match {
       case _ @ ArrayType(n: AtomicType, _) => n.ordering.asInstanceOf[Ordering[Any]]
+      case _ @ ArrayType(a: ArrayType, _) => a.interpretedOrdering.asInstanceOf[Ordering[Any]]
       case _ @ ArrayType(s: StructType, _) => s.interpretedOrdering.asInstanceOf[Ordering[Any]]
     }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ordering.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ordering.scala
index 6407c73bc97d9..6112259fed619 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ordering.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ordering.scala
@@ -48,6 +48,10 @@ class InterpretedOrdering(ordering: Seq[SortOrder]) extends Ordering[InternalRow
             dt.ordering.asInstanceOf[Ordering[Any]].compare(left, right)
           case dt: AtomicType if order.direction == Descending =>
             dt.ordering.asInstanceOf[Ordering[Any]].reverse.compare(left, right)
+          case a: ArrayType if order.direction == Ascending =>
+            a.interpretedOrdering.asInstanceOf[Ordering[Any]].compare(left, right)
+          case a: ArrayType if order.direction == Descending =>
+            a.interpretedOrdering.asInstanceOf[Ordering[Any]].reverse.compare(left, right)
           case s: StructType if order.direction == Ascending =>
             s.interpretedOrdering.asInstanceOf[Ordering[Any]].compare(left, right)
           case s: StructType if order.direction == Descending =>
@@ -86,6 +90,8 @@ object RowOrdering {
     case NullType => true
     case dt: AtomicType => true
     case struct: StructType => struct.fields.forall(f => isOrderable(f.dataType))
+    case array: ArrayType => isOrderable(array.elementType)
+    case udt: UserDefinedType[_] => isOrderable(udt.sqlType)
     case _ => false
   }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala
index bcf4d78fb9371..f603cbfb0cc21 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala
@@ -57,6 +57,7 @@ object TypeUtils {
   def getInterpretedOrdering(t: DataType): Ordering[Any] = {
     t match {
       case i: AtomicType => i.ordering.asInstanceOf[Ordering[Any]]
+      case a: ArrayType => a.interpretedOrdering.asInstanceOf[Ordering[Any]]
       case s: StructType => s.interpretedOrdering.asInstanceOf[Ordering[Any]]
     }
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/AbstractDataType.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/AbstractDataType.scala
index 1d2d007c2b4d2..a5ae8bb0e5eb6 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/AbstractDataType.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/AbstractDataType.scala
@@ -84,6 +84,7 @@ private[sql] object TypeCollection {
    * Types that can be ordered/compared. In the long run we should probably make this a trait
    * that can be mixed into each data type, and perhaps create an [[AbstractDataType]].
    */
+  // TODO: Should we consolidate this with RowOrdering.isOrderable?
   val Ordered = TypeCollection(
     BooleanType,
     ByteType, ShortType, IntegerType, LongType,
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/ArrayType.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/ArrayType.scala
index 5770f59b53077..a001eadcc61d0 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/ArrayType.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/ArrayType.scala
@@ -17,10 +17,13 @@
 
 package org.apache.spark.sql.types
 
+import org.apache.spark.sql.catalyst.util.ArrayData
 import org.json4s.JsonDSL._
 
 import org.apache.spark.annotation.DeveloperApi
 
+import scala.math.Ordering
+
 
 object ArrayType extends AbstractDataType {
   /** Construct a [[ArrayType]] object with the given element type. The `containsNull` is true. */
@@ -81,4 +84,49 @@ case class ArrayType(elementType: DataType, containsNull: Boolean) extends DataT
   override private[spark] def existsRecursively(f: (DataType) => Boolean): Boolean = {
     f(this) || elementType.existsRecursively(f)
   }
+
+  @transient
+  private[sql] lazy val interpretedOrdering: Ordering[ArrayData] = new Ordering[ArrayData] {
+    private[this] val elementOrdering: Ordering[Any] = elementType match {
+      case dt: AtomicType => dt.ordering.asInstanceOf[Ordering[Any]]
+      case a : ArrayType => a.interpretedOrdering.asInstanceOf[Ordering[Any]]
+      case s: StructType => s.interpretedOrdering.asInstanceOf[Ordering[Any]]
+      case other =>
+        throw new IllegalArgumentException(s"Type $other does not support ordered operations")
+    }
+
+    def compare(x: ArrayData, y: ArrayData): Int = {
+      val leftArray = x
+      val rightArray = y
+      val minLength = scala.math.min(leftArray.numElements(), rightArray.numElements())
+      var i = 0
+      while (i < minLength) {
+        val isNullLeft = leftArray.isNullAt(i)
+        val isNullRight = rightArray.isNullAt(i)
+        if (isNullLeft && isNullRight) {
+          // Do nothing.
+        } else if (isNullLeft) {
+          return -1
+        } else if (isNullRight) {
+          return 1
+        } else {
+          val comp =
+            elementOrdering.compare(
+              leftArray.get(i, elementType),
+              rightArray.get(i, elementType))
+          if (comp != 0) {
+            return comp
+          }
+        }
+        i += 1
+      }
+      if (leftArray.numElements() < rightArray.numElements()) {
+        return -1
+      } else if (leftArray.numElements() > rightArray.numElements()) {
+        return 1
+      } else {
+        return 0
+      }
+    }
+  }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
index 2e7c3bd67b554..ee435578743fc 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisErrorSuite.scala
@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.Inner
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.util.{GenericArrayData, ArrayData}
+import org.apache.spark.sql.catalyst.util.{MapData, ArrayBasedMapData, GenericArrayData, ArrayData}
 import org.apache.spark.sql.types._
 
 import scala.beans.{BeanProperty, BeanInfo}
@@ -53,21 +53,29 @@ private[sql] class GroupableUDT extends UserDefinedType[GroupableData] {
 }
 
 @BeanInfo
-private[sql] case class UngroupableData(@BeanProperty data: Array[Int])
+private[sql] case class UngroupableData(@BeanProperty data: Map[Int, Int])
 
 private[sql] class UngroupableUDT extends UserDefinedType[UngroupableData] {
 
-  override def sqlType: DataType = ArrayType(IntegerType)
+  override def sqlType: DataType = MapType(IntegerType, IntegerType)
 
-  override def serialize(obj: Any): ArrayData = {
+  override def serialize(obj: Any): MapData = {
     obj match {
-      case groupableData: UngroupableData => new GenericArrayData(groupableData.data)
+      case groupableData: UngroupableData =>
+        val keyArray = new GenericArrayData(groupableData.data.keys.toSeq)
+        val valueArray = new GenericArrayData(groupableData.data.values.toSeq)
+        new ArrayBasedMapData(keyArray, valueArray)
     }
   }
 
   override def deserialize(datum: Any): UngroupableData = {
     datum match {
-      case data: Array[Int] => UngroupableData(data)
+      case data: MapData =>
+        val keyArray = data.keyArray().array
+        val valueArray = data.valueArray().array
+        assert(keyArray.length == valueArray.length)
+        val mapData = keyArray.zip(valueArray).toMap.asInstanceOf[Map[Int, Int]]
+        UngroupableData(mapData)
     }
   }
 
@@ -154,8 +162,8 @@ class AnalysisErrorSuite extends AnalysisTest {
 
   errorTest(
     "sorting by unsupported column types",
-    listRelation.orderBy('list.asc),
-    "sort" :: "type" :: "array<int>" :: Nil)
+    mapRelation.orderBy('map.asc),
+    "sort" :: "type" :: "map<int,int>" :: Nil)
 
   errorTest(
     "non-boolean filters",
@@ -259,32 +267,33 @@ class AnalysisErrorSuite extends AnalysisTest {
         case true =>
           assertAnalysisSuccess(plan, true)
         case false =>
-          assertAnalysisError(plan, "expression a cannot be used in grouping expression" :: Nil)
+          assertAnalysisError(plan, "expression a cannot be used as a grouping expression" :: Nil)
       }
-
     }
 
     val supportedDataTypes = Seq(
-      StringType,
+      StringType, BinaryType,
       NullType, BooleanType,
       ByteType, ShortType, IntegerType, LongType,
       FloatType, DoubleType, DecimalType(25, 5), DecimalType(6, 5),
       DateType, TimestampType,
+      ArrayType(IntegerType),
       new StructType()
         .add("f1", FloatType, nullable = true)
         .add("f2", StringType, nullable = true),
+      new StructType()
+        .add("f1", FloatType, nullable = true)
+        .add("f2", ArrayType(BooleanType, containsNull = true), nullable = true),
       new GroupableUDT())
     supportedDataTypes.foreach { dataType =>
       checkDataType(dataType, shouldSuccess = true)
     }
 
     val unsupportedDataTypes = Seq(
-      BinaryType,
-      ArrayType(IntegerType),
       MapType(StringType, LongType),
       new StructType()
         .add("f1", FloatType, nullable = true)
-        .add("f2", ArrayType(BooleanType, containsNull = true), nullable = true),
+        .add("f2", MapType(StringType, LongType), nullable = true),
       new UngroupableUDT())
     unsupportedDataTypes.foreach { dataType =>
       checkDataType(dataType, shouldSuccess = false)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala
index b902982add8ff..ba1866efc84e1 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ExpressionTypeCheckingSuite.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
-import org.apache.spark.sql.types.{TypeCollection, StringType}
+import org.apache.spark.sql.types.{LongType, TypeCollection, StringType}
 
 class ExpressionTypeCheckingSuite extends SparkFunSuite {
 
@@ -32,7 +32,8 @@ class ExpressionTypeCheckingSuite extends SparkFunSuite {
     'intField.int,
     'stringField.string,
     'booleanField.boolean,
-    'complexField.array(StringType))
+    'arrayField.array(StringType),
+    'mapField.map(StringType, LongType))
 
   def assertError(expr: Expression, errorMessage: String): Unit = {
     val e = intercept[AnalysisException] {
@@ -90,9 +91,9 @@ class ExpressionTypeCheckingSuite extends SparkFunSuite {
     assertError(BitwiseOr('booleanField, 'booleanField), "requires integral type")
     assertError(BitwiseXor('booleanField, 'booleanField), "requires integral type")
 
-    assertError(MaxOf('complexField, 'complexField),
+    assertError(MaxOf('mapField, 'mapField),
       s"requires ${TypeCollection.Ordered.simpleString} type")
-    assertError(MinOf('complexField, 'complexField),
+    assertError(MinOf('mapField, 'mapField),
       s"requires ${TypeCollection.Ordered.simpleString} type")
   }
 
@@ -109,20 +110,20 @@ class ExpressionTypeCheckingSuite extends SparkFunSuite {
     assertSuccess(EqualTo('intField, 'booleanField))
     assertSuccess(EqualNullSafe('intField, 'booleanField))
 
-    assertErrorForDifferingTypes(EqualTo('intField, 'complexField))
-    assertErrorForDifferingTypes(EqualNullSafe('intField, 'complexField))
+    assertErrorForDifferingTypes(EqualTo('intField, 'mapField))
+    assertErrorForDifferingTypes(EqualNullSafe('intField, 'mapField))
     assertErrorForDifferingTypes(LessThan('intField, 'booleanField))
     assertErrorForDifferingTypes(LessThanOrEqual('intField, 'booleanField))
     assertErrorForDifferingTypes(GreaterThan('intField, 'booleanField))
     assertErrorForDifferingTypes(GreaterThanOrEqual('intField, 'booleanField))
 
-    assertError(LessThan('complexField, 'complexField),
+    assertError(LessThan('mapField, 'mapField),
       s"requires ${TypeCollection.Ordered.simpleString} type")
-    assertError(LessThanOrEqual('complexField, 'complexField),
+    assertError(LessThanOrEqual('mapField, 'mapField),
       s"requires ${TypeCollection.Ordered.simpleString} type")
-    assertError(GreaterThan('complexField, 'complexField),
+    assertError(GreaterThan('mapField, 'mapField),
       s"requires ${TypeCollection.Ordered.simpleString} type")
-    assertError(GreaterThanOrEqual('complexField, 'complexField),
+    assertError(GreaterThanOrEqual('mapField, 'mapField),
       s"requires ${TypeCollection.Ordered.simpleString} type")
 
     assertError(If('intField, 'stringField, 'stringField),
@@ -130,10 +131,10 @@ class ExpressionTypeCheckingSuite extends SparkFunSuite {
     assertErrorForDifferingTypes(If('booleanField, 'intField, 'booleanField))
 
     assertError(
-      CaseWhen(Seq('booleanField, 'intField, 'booleanField, 'complexField)),
+      CaseWhen(Seq('booleanField, 'intField, 'booleanField, 'mapField)),
       "THEN and ELSE expressions should all be same type or coercible to a common type")
     assertError(
-      CaseKeyWhen('intField, Seq('intField, 'stringField, 'intField, 'complexField)),
+      CaseKeyWhen('intField, Seq('intField, 'stringField, 'intField, 'mapField)),
       "THEN and ELSE expressions should all be same type or coercible to a common type")
     assertError(
       CaseWhen(Seq('booleanField, 'intField, 'intField, 'intField)),
@@ -147,9 +148,10 @@ class ExpressionTypeCheckingSuite extends SparkFunSuite {
     // We will cast String to Double for sum and average
     assertSuccess(Sum('stringField))
     assertSuccess(Average('stringField))
+    assertSuccess(Min('arrayField))
 
-    assertError(Min('complexField), "min does not support ordering on type")
-    assertError(Max('complexField), "max does not support ordering on type")
+    assertError(Min('mapField), "min does not support ordering on type")
+    assertError(Max('mapField), "max does not support ordering on type")
     assertError(Sum('booleanField), "function sum requires numeric type")
     assertError(Average('booleanField), "function average requires numeric type")
   }
@@ -184,7 +186,7 @@ class ExpressionTypeCheckingSuite extends SparkFunSuite {
 
     assertError(Round('intField, 'intField), "Only foldable Expression is allowed")
     assertError(Round('intField, 'booleanField), "requires int type")
-    assertError(Round('intField, 'complexField), "requires int type")
+    assertError(Round('intField, 'mapField), "requires int type")
     assertError(Round('booleanField, 'intField), "requires numeric type")
   }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TestRelations.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TestRelations.scala
index 05b870705e7ea..bc07b609a3413 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TestRelations.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TestRelations.scala
@@ -48,4 +48,7 @@ object TestRelations {
 
   val listRelation = LocalRelation(
     AttributeReference("list", ArrayType(IntegerType))())
+
+  val mapRelation = LocalRelation(
+    AttributeReference("map", MapType(IntegerType, IntegerType))())
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala
index e323467af5f4a..002ed16dcfe7e 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala
@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
-import scala.math._
-
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.{Row, RandomDataGenerator}
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
@@ -49,40 +47,6 @@ class CodeGenerationSuite extends SparkFunSuite with ExpressionEvalHelper {
     futures.foreach(Await.result(_, 10.seconds))
   }
 
-  // Test GenerateOrdering for all common types. For each type, we construct random input rows that
-  // contain two columns of that type, then for pairs of randomly-generated rows we check that
-  // GenerateOrdering agrees with RowOrdering.
-  (DataTypeTestUtils.atomicTypes ++ Set(NullType)).foreach { dataType =>
-    test(s"GenerateOrdering with $dataType") {
-      val rowOrdering = InterpretedOrdering.forSchema(Seq(dataType, dataType))
-      val genOrdering = GenerateOrdering.generate(
-        BoundReference(0, dataType, nullable = true).asc ::
-          BoundReference(1, dataType, nullable = true).asc :: Nil)
-      val rowType = StructType(
-        StructField("a", dataType, nullable = true) ::
-          StructField("b", dataType, nullable = true) :: Nil)
-      val maybeDataGenerator = RandomDataGenerator.forType(rowType, nullable = false)
-      assume(maybeDataGenerator.isDefined)
-      val randGenerator = maybeDataGenerator.get
-      val toCatalyst = CatalystTypeConverters.createToCatalystConverter(rowType)
-      for (_ <- 1 to 50) {
-        val a = toCatalyst(randGenerator()).asInstanceOf[InternalRow]
-        val b = toCatalyst(randGenerator()).asInstanceOf[InternalRow]
-        withClue(s"a = $a, b = $b") {
-          assert(genOrdering.compare(a, a) === 0)
-          assert(genOrdering.compare(b, b) === 0)
-          assert(rowOrdering.compare(a, a) === 0)
-          assert(rowOrdering.compare(b, b) === 0)
-          assert(signum(genOrdering.compare(a, b)) === -1 * signum(genOrdering.compare(b, a)))
-          assert(signum(rowOrdering.compare(a, b)) === -1 * signum(rowOrdering.compare(b, a)))
-          assert(
-            signum(rowOrdering.compare(a, b)) === signum(genOrdering.compare(a, b)),
-            "Generated and non-generated orderings should agree")
-        }
-      }
-    }
-  }
-
   test("SPARK-8443: split wide projections into blocks due to JVM code size limit") {
     val length = 5000
     val expressions = List.fill(length)(EqualTo(Literal(1), Literal(1)))
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/OrderingSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/OrderingSuite.scala
new file mode 100644
index 0000000000000..7ad8657bde128
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/OrderingSuite.scala
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions
+
+import scala.math._
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.{Row, RandomDataGenerator}
+import org.apache.spark.sql.catalyst.{InternalRow, CatalystTypeConverters}
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions.codegen.GenerateOrdering
+import org.apache.spark.sql.types._
+
+class OrderingSuite extends SparkFunSuite with ExpressionEvalHelper {
+
+  def compareArrays(a: Seq[Any], b: Seq[Any], expected: Int): Unit = {
+    test(s"compare two arrays: a = $a, b = $b") {
+      val dataType = ArrayType(IntegerType)
+      val rowType = StructType(StructField("array", dataType, nullable = true) :: Nil)
+      val toCatalyst = CatalystTypeConverters.createToCatalystConverter(rowType)
+      val rowA = toCatalyst(Row(a)).asInstanceOf[InternalRow]
+      val rowB = toCatalyst(Row(b)).asInstanceOf[InternalRow]
+      Seq(Ascending, Descending).foreach { direction =>
+        val sortOrder = direction match {
+          case Ascending => BoundReference(0, dataType, nullable = true).asc
+          case Descending => BoundReference(0, dataType, nullable = true).desc
+        }
+        val expectedCompareResult = direction match {
+          case Ascending => signum(expected)
+          case Descending => -1 * signum(expected)
+        }
+        val intOrdering = new InterpretedOrdering(sortOrder :: Nil)
+        val genOrdering = GenerateOrdering.generate(sortOrder :: Nil)
+        Seq(intOrdering, genOrdering).foreach { ordering =>
+          assert(ordering.compare(rowA, rowA) === 0)
+          assert(ordering.compare(rowB, rowB) === 0)
+          assert(signum(ordering.compare(rowA, rowB)) === expectedCompareResult)
+          assert(signum(ordering.compare(rowB, rowA)) === -1 * expectedCompareResult)
+        }
+      }
+    }
+  }
+
+  // Two arrays have the same size.
+  compareArrays(Seq[Any](), Seq[Any](), 0)
+  compareArrays(Seq[Any](1), Seq[Any](1), 0)
+  compareArrays(Seq[Any](1, 2), Seq[Any](1, 2), 0)
+  compareArrays(Seq[Any](1, 2, 2), Seq[Any](1, 2, 3), -1)
+
+  // Two arrays have different sizes.
+  compareArrays(Seq[Any](), Seq[Any](1), -1)
+  compareArrays(Seq[Any](1, 2, 3), Seq[Any](1, 2, 3, 4), -1)
+  compareArrays(Seq[Any](1, 2, 3), Seq[Any](1, 2, 3, 2), -1)
+  compareArrays(Seq[Any](1, 2, 3), Seq[Any](1, 2, 2, 2), 1)
+
+  // Arrays having nulls.
+  compareArrays(Seq[Any](1, 2, 3), Seq[Any](1, 2, 3, null), -1)
+  compareArrays(Seq[Any](), Seq[Any](null), -1)
+  compareArrays(Seq[Any](null), Seq[Any](null), 0)
+  compareArrays(Seq[Any](null, null), Seq[Any](null, null), 0)
+  compareArrays(Seq[Any](null), Seq[Any](null, null), -1)
+  compareArrays(Seq[Any](null), Seq[Any](1), -1)
+  compareArrays(Seq[Any](null), Seq[Any](null, 1), -1)
+  compareArrays(Seq[Any](null, 1), Seq[Any](1, 1), -1)
+  compareArrays(Seq[Any](1, null, 1), Seq[Any](1, null, 1), 0)
+  compareArrays(Seq[Any](1, null, 1), Seq[Any](1, null, 2), -1)
+
+  // Test GenerateOrdering for all common types. For each type, we construct random input rows that
+  // contain two columns of that type, then for pairs of randomly-generated rows we check that
+  // GenerateOrdering agrees with RowOrdering.
+  {
+    val structType =
+      new StructType()
+        .add("f1", FloatType, nullable = true)
+        .add("f2", ArrayType(BooleanType, containsNull = true), nullable = true)
+    val arrayOfStructType = ArrayType(structType)
+    val complexTypes = ArrayType(IntegerType) :: structType :: arrayOfStructType :: Nil
+    (DataTypeTestUtils.atomicTypes ++ complexTypes ++ Set(NullType)).foreach { dataType =>
+      test(s"GenerateOrdering with $dataType") {
+        val rowOrdering = InterpretedOrdering.forSchema(Seq(dataType, dataType))
+        val genOrdering = GenerateOrdering.generate(
+          BoundReference(0, dataType, nullable = true).asc ::
+            BoundReference(1, dataType, nullable = true).asc :: Nil)
+        val rowType = StructType(
+          StructField("a", dataType, nullable = true) ::
+            StructField("b", dataType, nullable = true) :: Nil)
+        val maybeDataGenerator = RandomDataGenerator.forType(rowType, nullable = false)
+        assume(maybeDataGenerator.isDefined)
+        val randGenerator = maybeDataGenerator.get
+        val toCatalyst = CatalystTypeConverters.createToCatalystConverter(rowType)
+        for (_ <- 1 to 50) {
+          val a = toCatalyst(randGenerator()).asInstanceOf[InternalRow]
+          val b = toCatalyst(randGenerator()).asInstanceOf[InternalRow]
+          withClue(s"a = $a, b = $b") {
+            assert(genOrdering.compare(a, a) === 0)
+            assert(genOrdering.compare(b, b) === 0)
+            assert(rowOrdering.compare(a, a) === 0)
+            assert(rowOrdering.compare(b, b) === 0)
+            assert(signum(genOrdering.compare(a, b)) === -1 * signum(genOrdering.compare(b, a)))
+            assert(signum(rowOrdering.compare(a, b)) === -1 * signum(rowOrdering.compare(b, a)))
+            assert(
+              signum(rowOrdering.compare(a, b)) === signum(genOrdering.compare(a, b)),
+              "Generated and non-generated orderings should agree")
+          }
+        }
+      }
+    }
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
index 3a3f19af1473b..aff9efe4b2b16 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
@@ -308,10 +308,14 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
         Row(null, null))
     )
 
-    val df2 = Seq((Array[Array[Int]](Array(2)), "x")).toDF("a", "b")
-    assert(intercept[AnalysisException] {
-      df2.selectExpr("sort_array(a)").collect()
-    }.getMessage().contains("does not support sorting array of type array<int>"))
+    val df2 = Seq((Array[Array[Int]](Array(2), Array(1), Array(2, 4), null), "x")).toDF("a", "b")
+    checkAnswer(
+      df2.selectExpr("sort_array(a, true)", "sort_array(a, false)"),
+      Seq(
+        Row(
+          Seq[Seq[Int]](null, Seq(1), Seq(2), Seq(2, 4)),
+          Seq[Seq[Int]](Seq(2, 4), Seq(2), Seq(1), null)))
+    )
 
     val df3 = Seq(("xxx", "x")).toDF("a", "b")
     assert(intercept[AnalysisException] {
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 61e3e913c23ea..6dde79f74d3d8 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -132,6 +132,22 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
       (3, null, null)).toDF("key", "value1", "value2")
     data2.write.saveAsTable("agg2")
 
+    val data3 = Seq[(Seq[Integer], Integer, Integer)](
+      (Seq[Integer](1, 1), 10, -10),
+      (Seq[Integer](null), -60, 60),
+      (Seq[Integer](1, 1), 30, -30),
+      (Seq[Integer](1), 30, 30),
+      (Seq[Integer](2), 1, 1),
+      (null, -10, 10),
+      (Seq[Integer](2, 3), -1, null),
+      (Seq[Integer](2, 3), 1, 1),
+      (Seq[Integer](2, 3, 4), null, 1),
+      (Seq[Integer](null), 100, -10),
+      (Seq[Integer](3), null, 3),
+      (null, null, null),
+      (Seq[Integer](3), null, null)).toDF("key", "value1", "value2")
+    data3.write.saveAsTable("agg3")
+
     val emptyDF = sqlContext.createDataFrame(
       sparkContext.emptyRDD[Row],
       StructType(StructField("key", StringType) :: StructField("value", IntegerType) :: Nil))
@@ -146,6 +162,7 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
   override def afterAll(): Unit = {
     sqlContext.sql("DROP TABLE IF EXISTS agg1")
     sqlContext.sql("DROP TABLE IF EXISTS agg2")
+    sqlContext.sql("DROP TABLE IF EXISTS agg3")
     sqlContext.dropTempTable("emptyTable")
   }
 
@@ -266,6 +283,41 @@ abstract class AggregationQuerySuite extends QueryTest with SQLTestUtils with Te
         Row(100, null) ::
         Row(null, 3) ::
         Row(null, null) :: Nil)
+
+    checkAnswer(
+      sqlContext.sql(
+        """
+          |SELECT DISTINCT key
+          |FROM agg3
+        """.stripMargin),
+      Row(Seq[Integer](1, 1)) ::
+        Row(Seq[Integer](null)) ::
+        Row(Seq[Integer](1)) ::
+        Row(Seq[Integer](2)) ::
+        Row(null) ::
+        Row(Seq[Integer](2, 3)) ::
+        Row(Seq[Integer](2, 3, 4)) ::
+        Row(Seq[Integer](3)) :: Nil)
+
+    checkAnswer(
+      sqlContext.sql(
+        """
+          |SELECT value1, key
+          |FROM agg3
+          |GROUP BY value1, key
+        """.stripMargin),
+      Row(10, Seq[Integer](1, 1)) ::
+        Row(-60, Seq[Integer](null)) ::
+        Row(30, Seq[Integer](1, 1)) ::
+        Row(30, Seq[Integer](1)) ::
+        Row(1, Seq[Integer](2)) ::
+        Row(-10, null) ::
+        Row(-1, Seq[Integer](2, 3)) ::
+        Row(1, Seq[Integer](2, 3)) ::
+        Row(null, Seq[Integer](2, 3, 4)) ::
+        Row(100, Seq[Integer](null)) ::
+        Row(null, Seq[Integer](3)) ::
+        Row(null, null) :: Nil)
   }
 
   test("case in-sensitive resolution") {

From 72c1d68b4ab6acb3f85971e10947caabb4bd846d Mon Sep 17 00:00:00 2001
From: Yu Gao <ygao@us.ibm.com>
Date: Sun, 15 Nov 2015 14:53:59 -0800
Subject: [PATCH 249/518] [SPARK-10181][SQL] Do kerberos login for credentials
 during hive client initialization

On driver process start up, UserGroupInformation.loginUserFromKeytab is called with the principal and keytab passed in, and therefore static var UserGroupInfomation,loginUser is set to that principal with kerberos credentials saved in its private credential set, and all threads within the driver process are supposed to see and use this login credentials to authenticate with Hive and Hadoop. However, because of IsolatedClientLoader, UserGroupInformation class is not shared for hive metastore clients, and instead it is loaded separately and of course not able to see the prepared kerberos login credentials in the main thread.

The first proposed fix would cause other classloader conflict errors, and is not an appropriate solution. This new change does kerberos login during hive client initialization, which will make credentials ready for the particular hive client instance.

 yhuai Please take a look and let me know. If you are not the right person to talk to, could you point me to someone responsible for this?

Author: Yu Gao <ygao@us.ibm.com>
Author: gaoyu <gaoyu@gaoyu-macbookpro.roam.corp.google.com>
Author: Yu Gao <crystalgaoyu@gmail.com>

Closes #9272 from yolandagao/master.
---
 .../org/apache/spark/deploy/SparkSubmit.scala | 17 ++++++++++---
 .../spark/sql/hive/client/ClientWrapper.scala | 24 ++++++++++++++++++-
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
index 84ae122f44370..09d2ec90c9333 100644
--- a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
@@ -39,7 +39,7 @@ import org.apache.ivy.plugins.matcher.GlobPatternMatcher
 import org.apache.ivy.plugins.repository.file.FileRepository
 import org.apache.ivy.plugins.resolver.{FileSystemResolver, ChainResolver, IBiblioResolver}
 
-import org.apache.spark.{SparkUserAppException, SPARK_VERSION}
+import org.apache.spark.{SparkException, SparkUserAppException, SPARK_VERSION}
 import org.apache.spark.api.r.RUtils
 import org.apache.spark.deploy.rest._
 import org.apache.spark.util.{ChildFirstURLClassLoader, MutableURLClassLoader, Utils}
@@ -521,8 +521,19 @@ object SparkSubmit {
         sysProps.put("spark.yarn.isPython", "true")
       }
       if (args.principal != null) {
-        require(args.keytab != null, "Keytab must be specified when the keytab is specified")
-        UserGroupInformation.loginUserFromKeytab(args.principal, args.keytab)
+        require(args.keytab != null, "Keytab must be specified when principal is specified")
+        if (!new File(args.keytab).exists()) {
+          throw new SparkException(s"Keytab file: ${args.keytab} does not exist")
+        } else {
+          // Add keytab and principal configurations in sysProps to make them available
+          // for later use; e.g. in spark sql, the isolated class loader used to talk
+          // to HiveMetastore will use these settings. They will be set as Java system
+          // properties and then loaded by SparkConf
+          sysProps.put("spark.yarn.keytab", args.keytab)
+          sysProps.put("spark.yarn.principal", args.principal)
+
+          UserGroupInformation.loginUserFromKeytab(args.principal, args.keytab)
+        }
       }
     }
 
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/ClientWrapper.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/ClientWrapper.scala
index f1c2489b38271..598ccdeee4ad2 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/ClientWrapper.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/ClientWrapper.scala
@@ -32,9 +32,10 @@ import org.apache.hadoop.hive.ql.processors._
 import org.apache.hadoop.hive.ql.session.SessionState
 import org.apache.hadoop.hive.ql.{Driver, metadata}
 import org.apache.hadoop.hive.shims.{HadoopShims, ShimLoader}
+import org.apache.hadoop.security.UserGroupInformation
 import org.apache.hadoop.util.VersionInfo
 
-import org.apache.spark.Logging
+import org.apache.spark.{SparkConf, SparkException, Logging}
 import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.execution.QueryExecutionException
 import org.apache.spark.util.{CircularBuffer, Utils}
@@ -149,6 +150,27 @@ private[hive] class ClientWrapper(
     val original = Thread.currentThread().getContextClassLoader
     // Switch to the initClassLoader.
     Thread.currentThread().setContextClassLoader(initClassLoader)
+
+    // Set up kerberos credentials for UserGroupInformation.loginUser within
+    // current class loader
+    // Instead of using the spark conf of the current spark context, a new
+    // instance of SparkConf is needed for the original value of spark.yarn.keytab
+    // and spark.yarn.principal set in SparkSubmit, as yarn.Client resets the
+    // keytab configuration for the link name in distributed cache
+    val sparkConf = new SparkConf
+    if (sparkConf.contains("spark.yarn.principal") && sparkConf.contains("spark.yarn.keytab")) {
+      val principalName = sparkConf.get("spark.yarn.principal")
+      val keytabFileName = sparkConf.get("spark.yarn.keytab")
+      if (!new File(keytabFileName).exists()) {
+        throw new SparkException(s"Keytab file: ${keytabFileName}" +
+          " specified in spark.yarn.keytab does not exist")
+      } else {
+        logInfo("Attempting to login to Kerberos" +
+          s" using principal: ${principalName} and keytab: ${keytabFileName}")
+        UserGroupInformation.loginUserFromKeytab(principalName, keytabFileName)
+      }
+    }
+
     val ret = try {
       val initialConf = new HiveConf(classOf[SessionState])
       // HiveConf is a Hadoop Configuration, which has a field of classLoader and

From d7d9fa0b8750166f8b74f9bc321df26908683a8b Mon Sep 17 00:00:00 2001
From: zero323 <matthew.szymkiewicz@gmail.com>
Date: Sun, 15 Nov 2015 19:15:27 -0800
Subject: [PATCH 250/518] [SPARK-11086][SPARKR] Use dropFactors column-wise
 instead of nested loop when createDataFrame

Use `dropFactors` column-wise instead of nested loop when `createDataFrame` from a `data.frame`

At this moment SparkR createDataFrame is using nested loop to convert factors to character when called on a local data.frame.  It works but is incredibly slow especially with data.table (~ 2 orders of magnitude compared to PySpark / Pandas version on a DateFrame of size 1M rows x 2 columns).

A simple improvement is to apply `dropFactor `column-wise and then reshape output list.

It should at least partially address [SPARK-8277](https://issues.apache.org/jira/browse/SPARK-8277).

Author: zero323 <matthew.szymkiewicz@gmail.com>

Closes #9099 from zero323/SPARK-11086.
---
 R/pkg/R/SQLContext.R             | 54 +++++++++++++++++++-------------
 R/pkg/inst/tests/test_sparkSQL.R | 16 ++++++++++
 2 files changed, 49 insertions(+), 21 deletions(-)

diff --git a/R/pkg/R/SQLContext.R b/R/pkg/R/SQLContext.R
index fd013fdb304df..a62b25fde926d 100644
--- a/R/pkg/R/SQLContext.R
+++ b/R/pkg/R/SQLContext.R
@@ -17,27 +17,33 @@
 
 # SQLcontext.R: SQLContext-driven functions
 
+
+# Map top level R type to SQL type
+getInternalType <- function(x) {
+  # class of POSIXlt is c("POSIXlt" "POSIXt")
+  switch(class(x)[[1]],
+         integer = "integer",
+         character = "string",
+         logical = "boolean",
+         double = "double",
+         numeric = "double",
+         raw = "binary",
+         list = "array",
+         struct = "struct",
+         environment = "map",
+         Date = "date",
+         POSIXlt = "timestamp",
+         POSIXct = "timestamp",
+         stop(paste("Unsupported type for DataFrame:", class(x))))
+}
+
 #' infer the SQL type
 infer_type <- function(x) {
   if (is.null(x)) {
     stop("can not infer type from NULL")
   }
 
-  # class of POSIXlt is c("POSIXlt" "POSIXt")
-  type <- switch(class(x)[[1]],
-                 integer = "integer",
-                 character = "string",
-                 logical = "boolean",
-                 double = "double",
-                 numeric = "double",
-                 raw = "binary",
-                 list = "array",
-                 struct = "struct",
-                 environment = "map",
-                 Date = "date",
-                 POSIXlt = "timestamp",
-                 POSIXct = "timestamp",
-                 stop(paste("Unsupported type for DataFrame:", class(x))))
+  type <- getInternalType(x)
 
   if (type == "map") {
     stopifnot(length(x) > 0)
@@ -90,19 +96,25 @@ createDataFrame <- function(sqlContext, data, schema = NULL, samplingRatio = 1.0
       if (is.null(schema)) {
         schema <- names(data)
       }
-      n <- nrow(data)
-      m <- ncol(data)
+
       # get rid of factor type
-      dropFactor <- function(x) {
+      cleanCols <- function(x) {
         if (is.factor(x)) {
           as.character(x)
         } else {
           x
         }
       }
-      data <- lapply(1:n, function(i) {
-        lapply(1:m, function(j) { dropFactor(data[i,j]) })
-      })
+
+      # drop factors and wrap lists
+      data <- setNames(lapply(data, cleanCols), NULL)
+
+      # check if all columns have supported type
+      lapply(data, getInternalType)
+
+      # convert to rows
+      args <- list(FUN = list, SIMPLIFY = FALSE, USE.NAMES = FALSE)
+      data <- do.call(mapply, append(args, data))
   }
   if (is.list(data)) {
     sc <- callJStatic("org.apache.spark.sql.api.r.SQLUtils", "getJavaSparkContext", sqlContext)
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index af024e6183a37..8ff06276599e2 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -242,6 +242,14 @@ test_that("create DataFrame from list or data.frame", {
   expect_equal(count(df), 3)
   ldf2 <- collect(df)
   expect_equal(ldf$a, ldf2$a)
+
+  irisdf <- createDataFrame(sqlContext, iris)
+  iris_collected <- collect(irisdf)
+  expect_equivalent(iris_collected[,-5], iris[,-5])
+  expect_equal(iris_collected$Species, as.character(iris$Species))
+
+  mtcarsdf <- createDataFrame(sqlContext, mtcars)
+  expect_equivalent(collect(mtcarsdf), mtcars)
 })
 
 test_that("create DataFrame with different data types", {
@@ -283,6 +291,14 @@ test_that("create DataFrame with complex types", {
   expect_equal(s$b, 3L)
 })
 
+test_that("create DataFrame from a data.frame with complex types", {
+  ldf <- data.frame(row.names=1:2)
+  ldf$a_list <- list(list(1, 2), list(3, 4))
+  sdf <- createDataFrame(sqlContext, ldf)
+
+  expect_equivalent(ldf, collect(sdf))
+})
+
 # For test map type and struct type in DataFrame
 mockLinesMapType <- c("{\"name\":\"Bob\",\"info\":{\"age\":16,\"height\":176.5}}",
                       "{\"name\":\"Alice\",\"info\":{\"age\":20,\"height\":164.3}}",

From 835a79d78ee879a3c36dde85e5b3591243bf3957 Mon Sep 17 00:00:00 2001
From: Sun Rui <rui.sun@intel.com>
Date: Sun, 15 Nov 2015 19:29:09 -0800
Subject: [PATCH 251/518] [SPARK-10500][SPARKR] sparkr.zip cannot be created if
 /R/lib is unwritable

The basic idea is that:
The archive of the SparkR package itself, that is sparkr.zip, is created during build process and is contained in the Spark binary distribution. No change to it after the distribution is installed as the directory it resides ($SPARK_HOME/R/lib) may not be writable.

When there is R source code contained in jars or Spark packages specified with "--jars" or "--packages" command line option, a temporary directory is created by calling Utils.createTempDir() where the R packages built from the R source code will be installed. The temporary directory is writable, and won't interfere with each other when there are multiple SparkR sessions, and will be deleted when this SparkR session ends. The R binary packages installed in the temporary directory then are packed into an archive named rpkg.zip.

sparkr.zip and rpkg.zip are distributed to the cluster in YARN modes.

The distribution of rpkg.zip in Standalone modes is not supported in this PR, and will be address in another PR.

Various R files are updated to accept multiple lib paths (one is for SparkR package, the other is for other R packages)  so that these package can be accessed in R.

Author: Sun Rui <rui.sun@intel.com>

Closes #9390 from sun-rui/SPARK-10500.
---
 R/install-dev.bat                             |  6 +++
 R/install-dev.sh                              |  4 ++
 R/pkg/R/sparkR.R                              | 14 +++++-
 R/pkg/inst/profile/general.R                  |  3 +-
 R/pkg/inst/worker/daemon.R                    |  5 ++-
 R/pkg/inst/worker/worker.R                    |  3 +-
 .../org/apache/spark/api/r/RBackend.scala     |  1 +
 .../scala/org/apache/spark/api/r/RRDD.scala   |  4 +-
 .../scala/org/apache/spark/api/r/RUtils.scala | 37 +++++++++++++---
 .../apache/spark/deploy/RPackageUtils.scala   | 26 ++++++++---
 .../org/apache/spark/deploy/RRunner.scala     |  5 ++-
 .../org/apache/spark/deploy/SparkSubmit.scala | 43 +++++++++++++++----
 .../spark/deploy/SparkSubmitSuite.scala       |  5 +--
 make-distribution.sh                          |  1 +
 14 files changed, 121 insertions(+), 36 deletions(-)

diff --git a/R/install-dev.bat b/R/install-dev.bat
index 008a5c668bc45..ed1c91ae3a0ff 100644
--- a/R/install-dev.bat
+++ b/R/install-dev.bat
@@ -25,3 +25,9 @@ set SPARK_HOME=%~dp0..
 MKDIR %SPARK_HOME%\R\lib
 
 R.exe CMD INSTALL --library="%SPARK_HOME%\R\lib"  %SPARK_HOME%\R\pkg\
+
+rem Zip the SparkR package so that it can be distributed to worker nodes on YARN
+pushd %SPARK_HOME%\R\lib
+%JAVA_HOME%\bin\jar.exe cfM "%SPARK_HOME%\R\lib\sparkr.zip" SparkR
+popd
+
diff --git a/R/install-dev.sh b/R/install-dev.sh
index 59d98c9c7a646..4972bb9217072 100755
--- a/R/install-dev.sh
+++ b/R/install-dev.sh
@@ -42,4 +42,8 @@ Rscript -e ' if("devtools" %in% rownames(installed.packages())) { library(devtoo
 # Install SparkR to $LIB_DIR
 R CMD INSTALL --library=$LIB_DIR $FWDIR/pkg/
 
+# Zip the SparkR package so that it can be distributed to worker nodes on YARN
+cd $LIB_DIR
+jar cfM "$LIB_DIR/sparkr.zip" SparkR
+
 popd > /dev/null
diff --git a/R/pkg/R/sparkR.R b/R/pkg/R/sparkR.R
index ebe2b2b8dc1d0..7ff3fa628b9ca 100644
--- a/R/pkg/R/sparkR.R
+++ b/R/pkg/R/sparkR.R
@@ -48,6 +48,12 @@ sparkR.stop <- function() {
       }
     }
 
+    # Remove the R package lib path from .libPaths()
+    if (exists(".libPath", envir = env)) {
+      libPath <- get(".libPath", envir = env)
+      .libPaths(.libPaths()[.libPaths() != libPath])
+    }
+
     if (exists(".backendLaunched", envir = env)) {
       callJStatic("SparkRHandler", "stopBackend")
     }
@@ -155,14 +161,20 @@ sparkR.init <- function(
     f <- file(path, open="rb")
     backendPort <- readInt(f)
     monitorPort <- readInt(f)
+    rLibPath <- readString(f)
     close(f)
     file.remove(path)
     if (length(backendPort) == 0 || backendPort == 0 ||
-        length(monitorPort) == 0 || monitorPort == 0) {
+        length(monitorPort) == 0 || monitorPort == 0 ||
+        length(rLibPath) != 1) {
       stop("JVM failed to launch")
     }
     assign(".monitorConn", socketConnection(port = monitorPort), envir = .sparkREnv)
     assign(".backendLaunched", 1, envir = .sparkREnv)
+    if (rLibPath != "") {
+      assign(".libPath", rLibPath, envir = .sparkREnv)
+      .libPaths(c(rLibPath, .libPaths()))
+    }
   }
 
   .sparkREnv$backendPort <- backendPort
diff --git a/R/pkg/inst/profile/general.R b/R/pkg/inst/profile/general.R
index 2a8a8213d0849..c55fe9ba7af7a 100644
--- a/R/pkg/inst/profile/general.R
+++ b/R/pkg/inst/profile/general.R
@@ -17,6 +17,7 @@
 
 .First <- function() {
   packageDir <- Sys.getenv("SPARKR_PACKAGE_DIR")
-  .libPaths(c(packageDir, .libPaths()))
+  dirs <- strsplit(packageDir, ",")[[1]]
+  .libPaths(c(dirs, .libPaths()))
   Sys.setenv(NOAWT=1)
 }
diff --git a/R/pkg/inst/worker/daemon.R b/R/pkg/inst/worker/daemon.R
index 3584b418a71a9..f55beac6c8c07 100644
--- a/R/pkg/inst/worker/daemon.R
+++ b/R/pkg/inst/worker/daemon.R
@@ -18,10 +18,11 @@
 # Worker daemon
 
 rLibDir <- Sys.getenv("SPARKR_RLIBDIR")
-script <- paste(rLibDir, "SparkR/worker/worker.R", sep = "/")
+dirs <- strsplit(rLibDir, ",")[[1]]
+script <- file.path(dirs[[1]], "SparkR", "worker", "worker.R")
 
 # preload SparkR package, speedup worker
-.libPaths(c(rLibDir, .libPaths()))
+.libPaths(c(dirs, .libPaths()))
 suppressPackageStartupMessages(library(SparkR))
 
 port <- as.integer(Sys.getenv("SPARKR_WORKER_PORT"))
diff --git a/R/pkg/inst/worker/worker.R b/R/pkg/inst/worker/worker.R
index 0c3b0d1f4be20..3ae072beca11b 100644
--- a/R/pkg/inst/worker/worker.R
+++ b/R/pkg/inst/worker/worker.R
@@ -35,10 +35,11 @@ bootTime <- currentTimeSecs()
 bootElap <- elapsedSecs()
 
 rLibDir <- Sys.getenv("SPARKR_RLIBDIR")
+dirs <- strsplit(rLibDir, ",")[[1]]
 # Set libPaths to include SparkR package as loadNamespace needs this
 # TODO: Figure out if we can avoid this by not loading any objects that require
 # SparkR namespace
-.libPaths(c(rLibDir, .libPaths()))
+.libPaths(c(dirs, .libPaths()))
 suppressPackageStartupMessages(library(SparkR))
 
 port <- as.integer(Sys.getenv("SPARKR_WORKER_PORT"))
diff --git a/core/src/main/scala/org/apache/spark/api/r/RBackend.scala b/core/src/main/scala/org/apache/spark/api/r/RBackend.scala
index b7e72d4d0ed0b..8b3be0da2c8c4 100644
--- a/core/src/main/scala/org/apache/spark/api/r/RBackend.scala
+++ b/core/src/main/scala/org/apache/spark/api/r/RBackend.scala
@@ -113,6 +113,7 @@ private[spark] object RBackend extends Logging {
       val dos = new DataOutputStream(new FileOutputStream(f))
       dos.writeInt(boundPort)
       dos.writeInt(listenPort)
+      SerDe.writeString(dos, RUtils.rPackages.getOrElse(""))
       dos.close()
       f.renameTo(new File(path))
 
diff --git a/core/src/main/scala/org/apache/spark/api/r/RRDD.scala b/core/src/main/scala/org/apache/spark/api/r/RRDD.scala
index 6b418e908cb53..7509b3d3f44bb 100644
--- a/core/src/main/scala/org/apache/spark/api/r/RRDD.scala
+++ b/core/src/main/scala/org/apache/spark/api/r/RRDD.scala
@@ -400,14 +400,14 @@ private[r] object RRDD {
 
     val rOptions = "--vanilla"
     val rLibDir = RUtils.sparkRPackagePath(isDriver = false)
-    val rExecScript = rLibDir + "/SparkR/worker/" + script
+    val rExecScript = rLibDir(0) + "/SparkR/worker/" + script
     val pb = new ProcessBuilder(Arrays.asList(rCommand, rOptions, rExecScript))
     // Unset the R_TESTS environment variable for workers.
     // This is set by R CMD check as startup.Rs
     // (http://svn.r-project.org/R/trunk/src/library/tools/R/testing.R)
     // and confuses worker script which tries to load a non-existent file
     pb.environment().put("R_TESTS", "")
-    pb.environment().put("SPARKR_RLIBDIR", rLibDir)
+    pb.environment().put("SPARKR_RLIBDIR", rLibDir.mkString(","))
     pb.environment().put("SPARKR_WORKER_PORT", port.toString)
     pb.redirectErrorStream(true)  // redirect stderr into stdout
     val proc = pb.start()
diff --git a/core/src/main/scala/org/apache/spark/api/r/RUtils.scala b/core/src/main/scala/org/apache/spark/api/r/RUtils.scala
index fd5646b5b6372..16157414fd120 100644
--- a/core/src/main/scala/org/apache/spark/api/r/RUtils.scala
+++ b/core/src/main/scala/org/apache/spark/api/r/RUtils.scala
@@ -23,6 +23,10 @@ import java.util.Arrays
 import org.apache.spark.{SparkEnv, SparkException}
 
 private[spark] object RUtils {
+  // Local path where R binary packages built from R source code contained in the spark
+  // packages specified with "--packages" or "--jars" command line option reside.
+  var rPackages: Option[String] = None
+
   /**
    * Get the SparkR package path in the local spark distribution.
    */
@@ -34,11 +38,15 @@ private[spark] object RUtils {
   }
 
   /**
-   * Get the SparkR package path in various deployment modes.
+   * Get the list of paths for R packages in various deployment modes, of which the first
+   * path is for the SparkR package itself. The second path is for R packages built as
+   * part of Spark Packages, if any exist. Spark Packages can be provided through the
+   *  "--packages" or "--jars" command line options.
+   *
    * This assumes that Spark properties `spark.master` and `spark.submit.deployMode`
    * and environment variable `SPARK_HOME` are set.
    */
-  def sparkRPackagePath(isDriver: Boolean): String = {
+  def sparkRPackagePath(isDriver: Boolean): Seq[String] = {
     val (master, deployMode) =
       if (isDriver) {
         (sys.props("spark.master"), sys.props("spark.submit.deployMode"))
@@ -51,15 +59,30 @@ private[spark] object RUtils {
     val isYarnClient = master != null && master.contains("yarn") && deployMode == "client"
 
     // In YARN mode, the SparkR package is distributed as an archive symbolically
-    // linked to the "sparkr" file in the current directory. Note that this does not apply
-    // to the driver in client mode because it is run outside of the cluster.
+    // linked to the "sparkr" file in the current directory and additional R packages
+    // are distributed as an archive symbolically linked to the "rpkg" file in the
+    // current directory.
+    //
+    // Note that this does not apply to the driver in client mode because it is run
+    // outside of the cluster.
     if (isYarnCluster || (isYarnClient && !isDriver)) {
-      new File("sparkr").getAbsolutePath
+      val sparkRPkgPath = new File("sparkr").getAbsolutePath
+      val rPkgPath = new File("rpkg")
+      if (rPkgPath.exists()) {
+        Seq(sparkRPkgPath, rPkgPath.getAbsolutePath)
+      } else {
+        Seq(sparkRPkgPath)
+      }
     } else {
       // Otherwise, assume the package is local
       // TODO: support this for Mesos
-      localSparkRPackagePath.getOrElse {
-        throw new SparkException("SPARK_HOME not set. Can't locate SparkR package.")
+      val sparkRPkgPath = localSparkRPackagePath.getOrElse {
+          throw new SparkException("SPARK_HOME not set. Can't locate SparkR package.")
+      }
+      if (!rPackages.isEmpty) {
+        Seq(sparkRPkgPath, rPackages.get)
+      } else {
+        Seq(sparkRPkgPath)
       }
     }
   }
diff --git a/core/src/main/scala/org/apache/spark/deploy/RPackageUtils.scala b/core/src/main/scala/org/apache/spark/deploy/RPackageUtils.scala
index 7d160b6790eaa..d46dc87a92c97 100644
--- a/core/src/main/scala/org/apache/spark/deploy/RPackageUtils.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/RPackageUtils.scala
@@ -100,20 +100,29 @@ private[deploy] object RPackageUtils extends Logging {
    * Runs the standard R package installation code to build the R package from source.
    * Multiple runs don't cause problems.
    */
-  private def rPackageBuilder(dir: File, printStream: PrintStream, verbose: Boolean): Boolean = {
+  private def rPackageBuilder(
+      dir: File,
+      printStream: PrintStream,
+      verbose: Boolean,
+      libDir: String): Boolean = {
     // this code should be always running on the driver.
-    val pathToSparkR = RUtils.localSparkRPackagePath.getOrElse(
-      throw new SparkException("SPARK_HOME not set. Can't locate SparkR package."))
     val pathToPkg = Seq(dir, "R", "pkg").mkString(File.separator)
-    val installCmd = baseInstallCmd ++ Seq(pathToSparkR, pathToPkg)
+    val installCmd = baseInstallCmd ++ Seq(libDir, pathToPkg)
     if (verbose) {
       print(s"Building R package with the command: $installCmd", printStream)
     }
     try {
       val builder = new ProcessBuilder(installCmd.asJava)
       builder.redirectErrorStream(true)
+
+      // Put the SparkR package directory into R library search paths in case this R package
+      // may depend on SparkR.
       val env = builder.environment()
-      env.clear()
+      val rPackageDir = RUtils.sparkRPackagePath(isDriver = true)
+      env.put("SPARKR_PACKAGE_DIR", rPackageDir.mkString(","))
+      env.put("R_PROFILE_USER",
+        Seq(rPackageDir(0), "SparkR", "profile", "general.R").mkString(File.separator))
+
       val process = builder.start()
       new RedirectThread(process.getInputStream, printStream, "redirect R packaging").start()
       process.waitFor() == 0
@@ -170,8 +179,11 @@ private[deploy] object RPackageUtils extends Logging {
         if (checkManifestForR(jar)) {
           print(s"$file contains R source code. Now installing package.", printStream, Level.INFO)
           val rSource = extractRFolder(jar, printStream, verbose)
+          if (RUtils.rPackages.isEmpty) {
+            RUtils.rPackages = Some(Utils.createTempDir().getAbsolutePath)
+          }
           try {
-            if (!rPackageBuilder(rSource, printStream, verbose)) {
+            if (!rPackageBuilder(rSource, printStream, verbose, RUtils.rPackages.get)) {
               print(s"ERROR: Failed to build R package in $file.", printStream)
               print(RJarDoc, printStream)
             }
@@ -208,7 +220,7 @@ private[deploy] object RPackageUtils extends Logging {
     }
   }
 
-  /** Zips all the libraries found with SparkR in the R/lib directory for distribution with Yarn. */
+  /** Zips all the R libraries built for distribution to the cluster. */
   private[deploy] def zipRLibraries(dir: File, name: String): File = {
     val filesToBundle = listFilesRecursively(dir, Seq(".zip"))
     // create a zip file from scratch, do not append to existing file.
diff --git a/core/src/main/scala/org/apache/spark/deploy/RRunner.scala b/core/src/main/scala/org/apache/spark/deploy/RRunner.scala
index ed183cf16a9cb..661f7317c674b 100644
--- a/core/src/main/scala/org/apache/spark/deploy/RRunner.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/RRunner.scala
@@ -82,9 +82,10 @@ object RRunner {
         val env = builder.environment()
         env.put("EXISTING_SPARKR_BACKEND_PORT", sparkRBackendPort.toString)
         val rPackageDir = RUtils.sparkRPackagePath(isDriver = true)
-        env.put("SPARKR_PACKAGE_DIR", rPackageDir)
+        // Put the R package directories into an env variable of comma-separated paths
+        env.put("SPARKR_PACKAGE_DIR", rPackageDir.mkString(","))
         env.put("R_PROFILE_USER",
-          Seq(rPackageDir, "SparkR", "profile", "general.R").mkString(File.separator))
+          Seq(rPackageDir(0), "SparkR", "profile", "general.R").mkString(File.separator))
         builder.redirectErrorStream(true) // Ugly but needed for stdout and stderr to synchronize
         val process = builder.start()
 
diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
index 09d2ec90c9333..2e912b59afdb8 100644
--- a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
@@ -83,6 +83,7 @@ object SparkSubmit {
   private val PYSPARK_SHELL = "pyspark-shell"
   private val SPARKR_SHELL = "sparkr-shell"
   private val SPARKR_PACKAGE_ARCHIVE = "sparkr.zip"
+  private val R_PACKAGE_ARCHIVE = "rpkg.zip"
 
   private val CLASS_NOT_FOUND_EXIT_STATUS = 101
 
@@ -362,22 +363,46 @@ object SparkSubmit {
       }
     }
 
-    // In YARN mode for an R app, add the SparkR package archive to archives
-    // that can be distributed with the job
+    // In YARN mode for an R app, add the SparkR package archive and the R package
+    // archive containing all of the built R libraries to archives so that they can
+    // be distributed with the job
     if (args.isR && clusterManager == YARN) {
-      val rPackagePath = RUtils.localSparkRPackagePath
-      if (rPackagePath.isEmpty) {
+      val sparkRPackagePath = RUtils.localSparkRPackagePath
+      if (sparkRPackagePath.isEmpty) {
         printErrorAndExit("SPARK_HOME does not exist for R application in YARN mode.")
       }
-      val rPackageFile =
-        RPackageUtils.zipRLibraries(new File(rPackagePath.get), SPARKR_PACKAGE_ARCHIVE)
-      if (!rPackageFile.exists()) {
+      val sparkRPackageFile = new File(sparkRPackagePath.get, SPARKR_PACKAGE_ARCHIVE)
+      if (!sparkRPackageFile.exists()) {
         printErrorAndExit(s"$SPARKR_PACKAGE_ARCHIVE does not exist for R application in YARN mode.")
       }
-      val localURI = Utils.resolveURI(rPackageFile.getAbsolutePath)
+      val sparkRPackageURI = Utils.resolveURI(sparkRPackageFile.getAbsolutePath).toString
 
+      // Distribute the SparkR package.
       // Assigns a symbol link name "sparkr" to the shipped package.
-      args.archives = mergeFileLists(args.archives, localURI.toString + "#sparkr")
+      args.archives = mergeFileLists(args.archives, sparkRPackageURI + "#sparkr")
+
+      // Distribute the R package archive containing all the built R packages.
+      if (!RUtils.rPackages.isEmpty) {
+        val rPackageFile =
+          RPackageUtils.zipRLibraries(new File(RUtils.rPackages.get), R_PACKAGE_ARCHIVE)
+        if (!rPackageFile.exists()) {
+          printErrorAndExit("Failed to zip all the built R packages.")
+        }
+
+        val rPackageURI = Utils.resolveURI(rPackageFile.getAbsolutePath).toString
+        // Assigns a symbol link name "rpkg" to the shipped package.
+        args.archives = mergeFileLists(args.archives, rPackageURI + "#rpkg")
+      }
+    }
+
+    // TODO: Support distributing R packages with standalone cluster
+    if (args.isR && clusterManager == STANDALONE && !RUtils.rPackages.isEmpty) {
+      printErrorAndExit("Distributing R packages with standalone cluster is not supported.")
+    }
+
+    // TODO: Support SparkR with mesos cluster
+    if (args.isR && clusterManager == MESOS) {
+      printErrorAndExit("SparkR is not supported for Mesos cluster.")
     }
 
     // If we're running a R app, set the main class to our specific R runner
diff --git a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
index 66a50512003dc..42e748ec6d528 100644
--- a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
@@ -28,6 +28,7 @@ import org.scalatest.concurrent.Timeouts
 import org.scalatest.time.SpanSugar._
 
 import org.apache.spark._
+import org.apache.spark.api.r.RUtils
 import org.apache.spark.deploy.SparkSubmit._
 import org.apache.spark.deploy.SparkSubmitUtils.MavenCoordinate
 import org.apache.spark.util.{ResetSystemProperties, Utils}
@@ -369,9 +370,6 @@ class SparkSubmitSuite
   }
 
   test("correctly builds R packages included in a jar with --packages") {
-    // TODO(SPARK-9603): Building a package to $SPARK_HOME/R/lib is unavailable on Jenkins.
-    // It's hard to write the test in SparkR (because we can't create the repository dynamically)
-    /*
     assume(RUtils.isRInstalled, "R isn't installed on this machine.")
     val main = MavenCoordinate("my.great.lib", "mylib", "0.1")
     val sparkHome = sys.props.getOrElse("spark.test.home", fail("spark.test.home is not set!"))
@@ -389,7 +387,6 @@ class SparkSubmitSuite
         rScriptDir)
       runSparkSubmit(args)
     }
-    */
   }
 
   test("resolves command line argument paths correctly") {
diff --git a/make-distribution.sh b/make-distribution.sh
index e1c2afdbc6d87..d7d27e253f721 100755
--- a/make-distribution.sh
+++ b/make-distribution.sh
@@ -220,6 +220,7 @@ cp -r "$SPARK_HOME/ec2" "$DISTDIR"
 if [ -d "$SPARK_HOME"/R/lib/SparkR ]; then
   mkdir -p "$DISTDIR"/R/lib
   cp -r "$SPARK_HOME/R/lib/SparkR" "$DISTDIR"/R/lib
+  cp "$SPARK_HOME/R/lib/sparkr.zip" "$DISTDIR"/R/lib
 fi
 
 # Download and copy in tachyon, if requested

From b58765caa6d7e6933050565c5d423c45e7e70ba6 Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Sun, 15 Nov 2015 21:10:46 -0800
Subject: [PATCH 252/518] [SPARK-9928][SQL] Removal of LogicalLocalTable

LogicalLocalTable in ExistingRDD.scala is replaced by localRelation in LocalRelation.scala?

Do you know any reason why we still keep this class?

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9717 from gatorsmile/LogicalLocalTable.
---
 .../spark/sql/execution/ExistingRDD.scala     | 22 -------------------
 1 file changed, 22 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
index 7a466cf6a0a94..8b41d3d3d892e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
@@ -110,25 +110,3 @@ private[sql] object PhysicalRDD {
     PhysicalRDD(output, rdd, relation.toString, relation.isInstanceOf[HadoopFsRelation])
   }
 }
-
-/** Logical plan node for scanning data from a local collection. */
-private[sql]
-case class LogicalLocalTable(output: Seq[Attribute], rows: Seq[InternalRow])(sqlContext: SQLContext)
-   extends LogicalPlan with MultiInstanceRelation {
-
-  override def children: Seq[LogicalPlan] = Nil
-
-  override def newInstance(): this.type =
-    LogicalLocalTable(output.map(_.newInstance()), rows)(sqlContext).asInstanceOf[this.type]
-
-  override def sameResult(plan: LogicalPlan): Boolean = plan match {
-    case LogicalRDD(_, otherRDD) => rows == rows
-    case _ => false
-  }
-
-  @transient override lazy val statistics: Statistics = Statistics(
-    // TODO: Improve the statistics estimation.
-    // This is made small enough so it can be broadcasted.
-    sizeInBytes = sqlContext.conf.autoBroadcastJoinThreshold - 1
-  )
-}

From fd50fa4c3eff42e8adeeabe399ddba0edac930c8 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Sun, 15 Nov 2015 22:38:30 -0800
Subject: [PATCH 253/518] Revert "[SPARK-11572] Exit AsynchronousListenerBus
 thread when stop() is called"

This reverts commit 3e0a6cf1e02a19b37c68d3026415d53bb57a576b.
---
 .../org/apache/spark/util/AsynchronousListenerBus.scala  | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala b/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
index b3b54af972cb4..c20627b056bef 100644
--- a/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
+++ b/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
@@ -66,12 +66,15 @@ private[spark] abstract class AsynchronousListenerBus[L <: AnyRef, E](name: Stri
           processingEvent = true
         }
         try {
-          if (stopped.get()) {
+          val event = eventQueue.poll
+          if (event == null) {
             // Get out of the while loop and shutdown the daemon thread
+            if (!stopped.get) {
+              throw new IllegalStateException("Polling `null` from eventQueue means" +
+                " the listener bus has been stopped. So `stopped` must be true")
+            }
             return
           }
-          val event = eventQueue.poll
-          assert(event != null, "event queue was empty but the listener bus was not stopped")
           postToAll(event)
         } finally {
           self.synchronized {

From 42de5253f327bd7ee258b0efb5024f3847fa3b51 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Mon, 16 Nov 2015 00:06:14 -0800
Subject: [PATCH 254/518] [SPARK-11745][SQL] Enable more JSON parsing options

This patch adds the following options to the JSON data source, for dealing with non-standard JSON files:
* `allowComments` (default `false`): ignores Java/C++ style comment in JSON records
* `allowUnquotedFieldNames` (default `false`): allows unquoted JSON field names
* `allowSingleQuotes` (default `true`): allows single quotes in addition to double quotes
* `allowNumericLeadingZeros` (default `false`): allows leading zeros in numbers (e.g. 00012)

To avoid passing a lot of options throughout the json package, I introduced a new JSONOptions case class to define all JSON config options.

Also updated documentation to explain these options.

Scala

![screen shot 2015-11-15 at 6 12 12 pm](https://cloud.githubusercontent.com/assets/323388/11172965/e3ace6ec-8bc4-11e5-805e-2d78f80d0ed6.png)

Python

![screen shot 2015-11-15 at 6 11 28 pm](https://cloud.githubusercontent.com/assets/323388/11172964/e23ed6ee-8bc4-11e5-8216-312f5983acd5.png)

Author: Reynold Xin <rxin@databricks.com>

Closes #9724 from rxin/SPARK-11745.
---
 python/pyspark/sql/readwriter.py              |  10 ++
 .../apache/spark/sql/DataFrameReader.scala    |  22 ++--
 .../spark/sql/execution/SparkPlan.scala       |  17 +--
 .../datasources/json/InferSchema.scala        |  34 +++---
 .../datasources/json/JSONOptions.scala        |  64 ++++++++++
 .../datasources/json/JSONRelation.scala       |  20 ++-
 .../datasources/json/JacksonParser.scala      |  82 +++++++------
 .../json/JsonParsingOptionsSuite.scala        | 114 ++++++++++++++++++
 .../datasources/json/JsonSuite.scala          |  29 ++---
 9 files changed, 286 insertions(+), 106 deletions(-)
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONOptions.scala
 create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonParsingOptionsSuite.scala

diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py
index 927f4077424dc..7b8ddb9feba34 100644
--- a/python/pyspark/sql/readwriter.py
+++ b/python/pyspark/sql/readwriter.py
@@ -153,6 +153,16 @@ def json(self, path, schema=None):
                      or RDD of Strings storing JSON objects.
         :param schema: an optional :class:`StructType` for the input schema.
 
+        You can set the following JSON-specific options to deal with non-standard JSON files:
+            * ``primitivesAsString`` (default ``false``): infers all primitive values as a string \
+                type
+            * ``allowComments`` (default ``false``): ignores Java/C++ style comment in JSON records
+            * ``allowUnquotedFieldNames`` (default ``false``): allows unquoted JSON field names
+            * ``allowSingleQuotes`` (default ``true``): allows single quotes in addition to double \
+                quotes
+            * ``allowNumericLeadingZeros`` (default ``false``): allows leading zeros in numbers \
+                (e.g. 00012)
+
         >>> df1 = sqlContext.read.json('python/test_support/sql/people.json')
         >>> df1.dtypes
         [('age', 'bigint'), ('name', 'string')]
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
index 6a194a443ab17..5872fbded3833 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
@@ -29,7 +29,7 @@ import org.apache.spark.api.java.JavaRDD
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.execution.datasources.jdbc.{JDBCPartition, JDBCPartitioningInfo, JDBCRelation}
-import org.apache.spark.sql.execution.datasources.json.JSONRelation
+import org.apache.spark.sql.execution.datasources.json.{JSONOptions, JSONRelation}
 import org.apache.spark.sql.execution.datasources.parquet.ParquetRelation
 import org.apache.spark.sql.execution.datasources.{LogicalRelation, ResolvedDataSource}
 import org.apache.spark.sql.types.StructType
@@ -227,6 +227,15 @@ class DataFrameReader private[sql](sqlContext: SQLContext) extends Logging {
    * This function goes through the input once to determine the input schema. If you know the
    * schema in advance, use the version that specifies the schema to avoid the extra scan.
    *
+   * You can set the following JSON-specific options to deal with non-standard JSON files:
+   * <li>`primitivesAsString` (default `false`): infers all primitive values as a string type</li>
+   * <li>`allowComments` (default `false`): ignores Java/C++ style comment in JSON records</li>
+   * <li>`allowUnquotedFieldNames` (default `false`): allows unquoted JSON field names</li>
+   * <li>`allowSingleQuotes` (default `true`): allows single quotes in addition to double quotes
+   * </li>
+   * <li>`allowNumericLeadingZeros` (default `false`): allows leading zeros in numbers
+   * (e.g. 00012)</li>
+   *
    * @param path input path
    * @since 1.4.0
    */
@@ -255,16 +264,13 @@ class DataFrameReader private[sql](sqlContext: SQLContext) extends Logging {
    * @since 1.4.0
    */
   def json(jsonRDD: RDD[String]): DataFrame = {
-    val samplingRatio = extraOptions.getOrElse("samplingRatio", "1.0").toDouble
-    val primitivesAsString = extraOptions.getOrElse("primitivesAsString", "false").toBoolean
     sqlContext.baseRelationToDataFrame(
       new JSONRelation(
         Some(jsonRDD),
-        samplingRatio,
-        primitivesAsString,
-        userSpecifiedSchema,
-        None,
-        None)(sqlContext)
+        maybeDataSchema = userSpecifiedSchema,
+        maybePartitionSpec = None,
+        userDefinedPartitionColumns = None,
+        parameters = extraOptions.toMap)(sqlContext)
     )
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
index 1b833002f434c..534a3bcb8364d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
@@ -221,22 +221,6 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
 
   private[this] def isTesting: Boolean = sys.props.contains("spark.testing")
 
-  protected def newProjection(
-      expressions: Seq[Expression], inputSchema: Seq[Attribute]): Projection = {
-    log.debug(s"Creating Projection: $expressions, inputSchema: $inputSchema")
-    try {
-      GenerateProjection.generate(expressions, inputSchema)
-    } catch {
-      case e: Exception =>
-        if (isTesting) {
-          throw e
-        } else {
-          log.error("Failed to generate projection, fallback to interpret", e)
-          new InterpretedProjection(expressions, inputSchema)
-        }
-    }
-  }
-
   protected def newMutableProjection(
       expressions: Seq[Expression], inputSchema: Seq[Attribute]): () => MutableProjection = {
     log.debug(s"Creating MutableProj: $expressions, inputSchema: $inputSchema")
@@ -282,6 +266,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
         }
     }
   }
+
   /**
    * Creates a row ordering for the given schema, in natural ascending order.
    */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/InferSchema.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/InferSchema.scala
index b9914c581a657..922fd5b21167b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/InferSchema.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/InferSchema.scala
@@ -25,33 +25,36 @@ import org.apache.spark.sql.execution.datasources.json.JacksonUtils.nextUntil
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
 
-private[sql] object InferSchema {
+
+private[json] object InferSchema {
+
   /**
    * Infer the type of a collection of json records in three stages:
    *   1. Infer the type of each record
    *   2. Merge types by choosing the lowest type necessary to cover equal keys
    *   3. Replace any remaining null fields with string, the top type
    */
-  def apply(
+  def infer(
       json: RDD[String],
-      samplingRatio: Double = 1.0,
       columnNameOfCorruptRecords: String,
-      primitivesAsString: Boolean = false): StructType = {
-    require(samplingRatio > 0, s"samplingRatio ($samplingRatio) should be greater than 0")
-    val schemaData = if (samplingRatio > 0.99) {
+      configOptions: JSONOptions): StructType = {
+    require(configOptions.samplingRatio > 0,
+      s"samplingRatio (${configOptions.samplingRatio}) should be greater than 0")
+    val schemaData = if (configOptions.samplingRatio > 0.99) {
       json
     } else {
-      json.sample(withReplacement = false, samplingRatio, 1)
+      json.sample(withReplacement = false, configOptions.samplingRatio, 1)
     }
 
     // perform schema inference on each row and merge afterwards
     val rootType = schemaData.mapPartitions { iter =>
       val factory = new JsonFactory()
+      configOptions.setJacksonOptions(factory)
       iter.map { row =>
         try {
           Utils.tryWithResource(factory.createParser(row)) { parser =>
             parser.nextToken()
-            inferField(parser, primitivesAsString)
+            inferField(parser, configOptions)
           }
         } catch {
           case _: JsonParseException =>
@@ -71,14 +74,14 @@ private[sql] object InferSchema {
   /**
    * Infer the type of a json document from the parser's token stream
    */
-  private def inferField(parser: JsonParser, primitivesAsString: Boolean): DataType = {
+  private def inferField(parser: JsonParser, configOptions: JSONOptions): DataType = {
     import com.fasterxml.jackson.core.JsonToken._
     parser.getCurrentToken match {
       case null | VALUE_NULL => NullType
 
       case FIELD_NAME =>
         parser.nextToken()
-        inferField(parser, primitivesAsString)
+        inferField(parser, configOptions)
 
       case VALUE_STRING if parser.getTextLength < 1 =>
         // Zero length strings and nulls have special handling to deal
@@ -95,7 +98,7 @@ private[sql] object InferSchema {
         while (nextUntil(parser, END_OBJECT)) {
           builder += StructField(
             parser.getCurrentName,
-            inferField(parser, primitivesAsString),
+            inferField(parser, configOptions),
             nullable = true)
         }
 
@@ -107,14 +110,15 @@ private[sql] object InferSchema {
         // the type as we pass through all JSON objects.
         var elementType: DataType = NullType
         while (nextUntil(parser, END_ARRAY)) {
-          elementType = compatibleType(elementType, inferField(parser, primitivesAsString))
+          elementType = compatibleType(
+            elementType, inferField(parser, configOptions))
         }
 
         ArrayType(elementType)
 
-      case (VALUE_NUMBER_INT | VALUE_NUMBER_FLOAT) if primitivesAsString => StringType
+      case (VALUE_NUMBER_INT | VALUE_NUMBER_FLOAT) if configOptions.primitivesAsString => StringType
 
-      case (VALUE_TRUE | VALUE_FALSE) if primitivesAsString => StringType
+      case (VALUE_TRUE | VALUE_FALSE) if configOptions.primitivesAsString => StringType
 
       case VALUE_NUMBER_INT | VALUE_NUMBER_FLOAT =>
         import JsonParser.NumberType._
@@ -178,7 +182,7 @@ private[sql] object InferSchema {
   /**
    * Returns the most general data type for two given data types.
    */
-  private[json] def compatibleType(t1: DataType, t2: DataType): DataType = {
+  def compatibleType(t1: DataType, t2: DataType): DataType = {
     HiveTypeCoercion.findTightestCommonTypeOfTwo(t1, t2).getOrElse {
       // t1 or t2 is a StructType, ArrayType, or an unexpected type.
       (t1, t2) match {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONOptions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONOptions.scala
new file mode 100644
index 0000000000000..c132ead20e7d6
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONOptions.scala
@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.json
+
+import com.fasterxml.jackson.core.{JsonParser, JsonFactory}
+
+/**
+ * Options for the JSON data source.
+ *
+ * Most of these map directly to Jackson's internal options, specified in [[JsonParser.Feature]].
+ */
+case class JSONOptions(
+    samplingRatio: Double = 1.0,
+    primitivesAsString: Boolean = false,
+    allowComments: Boolean = false,
+    allowUnquotedFieldNames: Boolean = false,
+    allowSingleQuotes: Boolean = true,
+    allowNumericLeadingZeros: Boolean = false,
+    allowNonNumericNumbers: Boolean = false) {
+
+  /** Sets config options on a Jackson [[JsonFactory]]. */
+  def setJacksonOptions(factory: JsonFactory): Unit = {
+    factory.configure(JsonParser.Feature.ALLOW_COMMENTS, allowComments)
+    factory.configure(JsonParser.Feature.ALLOW_UNQUOTED_FIELD_NAMES, allowUnquotedFieldNames)
+    factory.configure(JsonParser.Feature.ALLOW_SINGLE_QUOTES, allowSingleQuotes)
+    factory.configure(JsonParser.Feature.ALLOW_NUMERIC_LEADING_ZEROS, allowNumericLeadingZeros)
+    factory.configure(JsonParser.Feature.ALLOW_NON_NUMERIC_NUMBERS, allowNonNumericNumbers)
+  }
+}
+
+
+object JSONOptions {
+  def createFromConfigMap(parameters: Map[String, String]): JSONOptions = JSONOptions(
+    samplingRatio =
+      parameters.get("samplingRatio").map(_.toDouble).getOrElse(1.0),
+    primitivesAsString =
+      parameters.get("primitivesAsString").map(_.toBoolean).getOrElse(false),
+    allowComments =
+      parameters.get("allowComments").map(_.toBoolean).getOrElse(false),
+    allowUnquotedFieldNames =
+      parameters.get("allowUnquotedFieldNames").map(_.toBoolean).getOrElse(false),
+    allowSingleQuotes =
+      parameters.get("allowSingleQuotes").map(_.toBoolean).getOrElse(true),
+    allowNumericLeadingZeros =
+      parameters.get("allowNumericLeadingZeros").map(_.toBoolean).getOrElse(false),
+    allowNonNumericNumbers =
+      parameters.get("allowNonNumericNumbers").map(_.toBoolean).getOrElse(true)
+  )
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONRelation.scala
index dca638b7f67a5..3e61ba35bea8e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONRelation.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JSONRelation.scala
@@ -52,13 +52,9 @@ class DefaultSource extends HadoopFsRelationProvider with DataSourceRegister {
       dataSchema: Option[StructType],
       partitionColumns: Option[StructType],
       parameters: Map[String, String]): HadoopFsRelation = {
-    val samplingRatio = parameters.get("samplingRatio").map(_.toDouble).getOrElse(1.0)
-    val primitivesAsString = parameters.get("primitivesAsString").map(_.toBoolean).getOrElse(false)
 
     new JSONRelation(
       inputRDD = None,
-      samplingRatio = samplingRatio,
-      primitivesAsString = primitivesAsString,
       maybeDataSchema = dataSchema,
       maybePartitionSpec = None,
       userDefinedPartitionColumns = partitionColumns,
@@ -69,8 +65,6 @@ class DefaultSource extends HadoopFsRelationProvider with DataSourceRegister {
 
 private[sql] class JSONRelation(
     val inputRDD: Option[RDD[String]],
-    val samplingRatio: Double,
-    val primitivesAsString: Boolean,
     val maybeDataSchema: Option[StructType],
     val maybePartitionSpec: Option[PartitionSpec],
     override val userDefinedPartitionColumns: Option[StructType],
@@ -79,6 +73,8 @@ private[sql] class JSONRelation(
     (@transient val sqlContext: SQLContext)
   extends HadoopFsRelation(maybePartitionSpec, parameters) {
 
+  val options: JSONOptions = JSONOptions.createFromConfigMap(parameters)
+
   /** Constraints to be imposed on schema to be stored. */
   private def checkConstraints(schema: StructType): Unit = {
     if (schema.fieldNames.length != schema.fieldNames.distinct.length) {
@@ -109,17 +105,16 @@ private[sql] class JSONRelation(
       classOf[Text]).map(_._2.toString) // get the text line
   }
 
-  override lazy val dataSchema = {
+  override lazy val dataSchema: StructType = {
     val jsonSchema = maybeDataSchema.getOrElse {
       val files = cachedLeafStatuses().filterNot { status =>
         val name = status.getPath.getName
         name.startsWith("_") || name.startsWith(".")
       }.toArray
-      InferSchema(
+      InferSchema.infer(
         inputRDD.getOrElse(createBaseRdd(files)),
-        samplingRatio,
         sqlContext.conf.columnNameOfCorruptRecord,
-        primitivesAsString)
+        options)
     }
     checkConstraints(jsonSchema)
 
@@ -132,10 +127,11 @@ private[sql] class JSONRelation(
       inputPaths: Array[FileStatus],
       broadcastedConf: Broadcast[SerializableConfiguration]): RDD[InternalRow] = {
     val requiredDataSchema = StructType(requiredColumns.map(dataSchema(_)))
-    val rows = JacksonParser(
+    val rows = JacksonParser.parse(
       inputRDD.getOrElse(createBaseRdd(inputPaths)),
       requiredDataSchema,
-      sqlContext.conf.columnNameOfCorruptRecord)
+      sqlContext.conf.columnNameOfCorruptRecord,
+      options)
 
     rows.mapPartitions { iterator =>
       val unsafeProjection = UnsafeProjection.create(requiredDataSchema)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JacksonParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JacksonParser.scala
index 4f53eeb081b93..bfa1405041058 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JacksonParser.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JacksonParser.scala
@@ -18,11 +18,10 @@
 package org.apache.spark.sql.execution.datasources.json
 
 import java.io.ByteArrayOutputStream
+import scala.collection.mutable.ArrayBuffer
 
 import com.fasterxml.jackson.core._
 
-import scala.collection.mutable.ArrayBuffer
-
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
@@ -32,18 +31,23 @@ import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 import org.apache.spark.util.Utils
 
-private[sql] object JacksonParser {
-  def apply(
-      json: RDD[String],
+object JacksonParser {
+
+  def parse(
+      input: RDD[String],
       schema: StructType,
-      columnNameOfCorruptRecords: String): RDD[InternalRow] = {
-    parseJson(json, schema, columnNameOfCorruptRecords)
+      columnNameOfCorruptRecords: String,
+      configOptions: JSONOptions): RDD[InternalRow] = {
+
+    input.mapPartitions { iter =>
+      parseJson(iter, schema, columnNameOfCorruptRecords, configOptions)
+    }
   }
 
   /**
    * Parse the current token (and related children) according to a desired schema
    */
-  private[sql] def convertField(
+  def convertField(
       factory: JsonFactory,
       parser: JsonParser,
       schema: DataType): Any = {
@@ -226,9 +230,10 @@ private[sql] object JacksonParser {
   }
 
   private def parseJson(
-      json: RDD[String],
+      input: Iterator[String],
       schema: StructType,
-      columnNameOfCorruptRecords: String): RDD[InternalRow] = {
+      columnNameOfCorruptRecords: String,
+      configOptions: JSONOptions): Iterator[InternalRow] = {
 
     def failedRecord(record: String): Seq[InternalRow] = {
       // create a row even if no corrupt record column is present
@@ -241,37 +246,36 @@ private[sql] object JacksonParser {
       Seq(row)
     }
 
-    json.mapPartitions { iter =>
-      val factory = new JsonFactory()
-
-      iter.flatMap { record =>
-        if (record.trim.isEmpty) {
-          Nil
-        } else {
-          try {
-            Utils.tryWithResource(factory.createParser(record)) { parser =>
-              parser.nextToken()
-
-              convertField(factory, parser, schema) match {
-                case null => failedRecord(record)
-                case row: InternalRow => row :: Nil
-                case array: ArrayData =>
-                  if (array.numElements() == 0) {
-                    Nil
-                  } else {
-                    array.toArray[InternalRow](schema)
-                  }
-                case _ =>
-                  sys.error(
-                    s"Failed to parse record $record. Please make sure that each line of " +
-                      "the file (or each string in the RDD) is a valid JSON object or " +
-                      "an array of JSON objects.")
-              }
+    val factory = new JsonFactory()
+    configOptions.setJacksonOptions(factory)
+
+    input.flatMap { record =>
+      if (record.trim.isEmpty) {
+        Nil
+      } else {
+        try {
+          Utils.tryWithResource(factory.createParser(record)) { parser =>
+            parser.nextToken()
+
+            convertField(factory, parser, schema) match {
+              case null => failedRecord(record)
+              case row: InternalRow => row :: Nil
+              case array: ArrayData =>
+                if (array.numElements() == 0) {
+                  Nil
+                } else {
+                  array.toArray[InternalRow](schema)
+                }
+              case _ =>
+                sys.error(
+                  s"Failed to parse record $record. Please make sure that each line of " +
+                    "the file (or each string in the RDD) is a valid JSON object or " +
+                    "an array of JSON objects.")
             }
-          } catch {
-            case _: JsonProcessingException =>
-              failedRecord(record)
           }
+        } catch {
+          case _: JsonProcessingException =>
+            failedRecord(record)
         }
       }
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonParsingOptionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonParsingOptionsSuite.scala
new file mode 100644
index 0000000000000..4cc0a3a9585d9
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonParsingOptionsSuite.scala
@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.json
+
+import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.test.SharedSQLContext
+
+/**
+ * Test cases for various [[JSONOptions]].
+ */
+class JsonParsingOptionsSuite extends QueryTest with SharedSQLContext {
+
+  test("allowComments off") {
+    val str = """{'name': /* hello */ 'Reynold Xin'}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.json(rdd)
+
+    assert(df.schema.head.name == "_corrupt_record")
+  }
+
+  test("allowComments on") {
+    val str = """{'name': /* hello */ 'Reynold Xin'}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.option("allowComments", "true").json(rdd)
+
+    assert(df.schema.head.name == "name")
+    assert(df.first().getString(0) == "Reynold Xin")
+  }
+
+  test("allowSingleQuotes off") {
+    val str = """{'name': 'Reynold Xin'}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.option("allowSingleQuotes", "false").json(rdd)
+
+    assert(df.schema.head.name == "_corrupt_record")
+  }
+
+  test("allowSingleQuotes on") {
+    val str = """{'name': 'Reynold Xin'}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.json(rdd)
+
+    assert(df.schema.head.name == "name")
+    assert(df.first().getString(0) == "Reynold Xin")
+  }
+
+  test("allowUnquotedFieldNames off") {
+    val str = """{name: 'Reynold Xin'}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.json(rdd)
+
+    assert(df.schema.head.name == "_corrupt_record")
+  }
+
+  test("allowUnquotedFieldNames on") {
+    val str = """{name: 'Reynold Xin'}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.option("allowUnquotedFieldNames", "true").json(rdd)
+
+    assert(df.schema.head.name == "name")
+    assert(df.first().getString(0) == "Reynold Xin")
+  }
+
+  test("allowNumericLeadingZeros off") {
+    val str = """{"age": 0018}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.json(rdd)
+
+    assert(df.schema.head.name == "_corrupt_record")
+  }
+
+  test("allowNumericLeadingZeros on") {
+    val str = """{"age": 0018}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.option("allowNumericLeadingZeros", "true").json(rdd)
+
+    assert(df.schema.head.name == "age")
+    assert(df.first().getLong(0) == 18)
+  }
+
+  // The following two tests are not really working - need to look into Jackson's
+  // JsonParser.Feature.ALLOW_NON_NUMERIC_NUMBERS.
+  ignore("allowNonNumericNumbers off") {
+    val str = """{"age": NaN}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.json(rdd)
+
+    assert(df.schema.head.name == "_corrupt_record")
+  }
+
+  ignore("allowNonNumericNumbers on") {
+    val str = """{"age": NaN}"""
+    val rdd = sqlContext.sparkContext.parallelize(Seq(str))
+    val df = sqlContext.read.option("allowNonNumericNumbers", "true").json(rdd)
+
+    assert(df.schema.head.name == "age")
+    assert(df.first().getDouble(0).isNaN)
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
index 28b8f02bdf87f..6042b1178affe 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
@@ -588,7 +588,7 @@ class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
       relation.isInstanceOf[JSONRelation],
       "The DataFrame returned by jsonFile should be based on JSONRelation.")
     assert(relation.asInstanceOf[JSONRelation].paths === Array(path))
-    assert(relation.asInstanceOf[JSONRelation].samplingRatio === (0.49 +- 0.001))
+    assert(relation.asInstanceOf[JSONRelation].options.samplingRatio === (0.49 +- 0.001))
 
     val schema = StructType(StructField("a", LongType, true) :: Nil)
     val logicalRelation =
@@ -597,7 +597,7 @@ class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
     val relationWithSchema = logicalRelation.relation.asInstanceOf[JSONRelation]
     assert(relationWithSchema.paths === Array(path))
     assert(relationWithSchema.schema === schema)
-    assert(relationWithSchema.samplingRatio > 0.99)
+    assert(relationWithSchema.options.samplingRatio > 0.99)
   }
 
   test("Loading a JSON dataset from a text file") {
@@ -1165,31 +1165,28 @@ class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
   test("JSONRelation equality test") {
     val relation0 = new JSONRelation(
       Some(empty),
-      1.0,
-      false,
       Some(StructType(StructField("a", IntegerType, true) :: Nil)),
-      None, None)(sqlContext)
+      None,
+      None)(sqlContext)
     val logicalRelation0 = LogicalRelation(relation0)
     val relation1 = new JSONRelation(
       Some(singleRow),
-      1.0,
-      false,
       Some(StructType(StructField("a", IntegerType, true) :: Nil)),
-      None, None)(sqlContext)
+      None,
+      None)(sqlContext)
     val logicalRelation1 = LogicalRelation(relation1)
     val relation2 = new JSONRelation(
       Some(singleRow),
-      0.5,
-      false,
       Some(StructType(StructField("a", IntegerType, true) :: Nil)),
-      None, None)(sqlContext)
+      None,
+      None,
+      parameters = Map("samplingRatio" -> "0.5"))(sqlContext)
     val logicalRelation2 = LogicalRelation(relation2)
     val relation3 = new JSONRelation(
       Some(singleRow),
-      1.0,
-      false,
       Some(StructType(StructField("b", IntegerType, true) :: Nil)),
-      None, None)(sqlContext)
+      None,
+      None)(sqlContext)
     val logicalRelation3 = LogicalRelation(relation3)
 
     assert(relation0 !== relation1)
@@ -1232,7 +1229,7 @@ class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
 
   test("SPARK-6245 JsonRDD.inferSchema on empty RDD") {
     // This is really a test that it doesn't throw an exception
-    val emptySchema = InferSchema(empty, 1.0, "")
+    val emptySchema = InferSchema.infer(empty, "", JSONOptions())
     assert(StructType(Seq()) === emptySchema)
   }
 
@@ -1256,7 +1253,7 @@ class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
   }
 
   test("SPARK-8093 Erase empty structs") {
-    val emptySchema = InferSchema(emptyRecords, 1.0, "")
+    val emptySchema = InferSchema.infer(emptyRecords, "", JSONOptions())
     assert(StructType(Seq()) === emptySchema)
   }
 

From 7f8eb3bf6ed64eefc5472f5c5fb02e2db1e3f618 Mon Sep 17 00:00:00 2001
From: hyukjinkwon <gurwls223@gmail.com>
Date: Mon, 16 Nov 2015 21:30:10 +0800
Subject: [PATCH 255/518] [SPARK-11044][SQL] Parquet writer version fixed as
 version1

https://issues.apache.org/jira/browse/SPARK-11044

Spark writes a parquet file only with writer version1 ignoring the writer version given by user.

So, in this PR, it keeps the writer version if given or sets version1 as default.

Author: hyukjinkwon <gurwls223@gmail.com>
Author: HyukjinKwon <gurwls223@gmail.com>

Closes #9060 from HyukjinKwon/SPARK-11044.
---
 .../parquet/CatalystWriteSupport.scala        |  2 +-
 .../datasources/parquet/ParquetIOSuite.scala  | 34 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystWriteSupport.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystWriteSupport.scala
index 483363d2c1a21..6862dea5e6c3b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystWriteSupport.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystWriteSupport.scala
@@ -429,7 +429,7 @@ private[parquet] object CatalystWriteSupport {
   def setSchema(schema: StructType, configuration: Configuration): Unit = {
     schema.map(_.name).foreach(CatalystSchemaConverter.checkFieldName)
     configuration.set(SPARK_ROW_SCHEMA, schema.json)
-    configuration.set(
+    configuration.setIfUnset(
       ParquetOutputFormat.WRITER_VERSION,
       ParquetProperties.WriterVersion.PARQUET_1_0.toString)
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index 78df363ade5c9..2aa5dca847c8f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.execution.datasources.parquet
 
 import java.util.Collections
 
+import org.apache.parquet.column.{Encoding, ParquetProperties}
+
 import scala.collection.JavaConverters._
 import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.TypeTag
@@ -534,6 +536,38 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
     }
   }
 
+  test("SPARK-11044 Parquet writer version fixed as version1 ") {
+    // For dictionary encoding, Parquet changes the encoding types according to its writer
+    // version. So, this test checks one of the encoding types in order to ensure that
+    // the file is written with writer version2.
+    withTempPath { dir =>
+      val clonedConf = new Configuration(hadoopConfiguration)
+      try {
+        // Write a Parquet file with writer version2.
+        hadoopConfiguration.set(ParquetOutputFormat.WRITER_VERSION,
+          ParquetProperties.WriterVersion.PARQUET_2_0.toString)
+
+        // By default, dictionary encoding is enabled from Parquet 1.2.0 but
+        // it is enabled just in case.
+        hadoopConfiguration.setBoolean(ParquetOutputFormat.ENABLE_DICTIONARY, true)
+        val path = s"${dir.getCanonicalPath}/part-r-0.parquet"
+        sqlContext.range(1 << 16).selectExpr("(id % 4) AS i")
+          .coalesce(1).write.mode("overwrite").parquet(path)
+
+        val blockMetadata = readFooter(new Path(path), hadoopConfiguration).getBlocks.asScala.head
+        val columnChunkMetadata = blockMetadata.getColumns.asScala.head
+
+        // If the file is written with version2, this should include
+        // Encoding.RLE_DICTIONARY type. For version1, it is Encoding.PLAIN_DICTIONARY
+        assert(columnChunkMetadata.getEncodings.contains(Encoding.RLE_DICTIONARY))
+      } finally {
+        // Manually clear the hadoop configuration for other tests.
+        hadoopConfiguration.clear()
+        clonedConf.asScala.foreach(entry => hadoopConfiguration.set(entry.getKey, entry.getValue))
+      }
+    }
+  }
+
   test("read dictionary encoded decimals written as INT32") {
     checkAnswer(
       // Decimal column in this file is encoded using plain dictionary

From e388b39d10fc269cdd3d630ea7d4ae80fd0efa97 Mon Sep 17 00:00:00 2001
From: hyukjinkwon <gurwls223@gmail.com>
Date: Mon, 16 Nov 2015 21:59:33 +0800
Subject: [PATCH 256/518] [SPARK-11692][SQL] Support for Parquet logical types,
 JSON and  BSON (embedded types)

Parquet supports some JSON and BSON datatypes. They are represented as binary for BSON and string (UTF-8) for JSON internally.

I searched a bit and found Apache drill also supports both in this way, [link](https://drill.apache.org/docs/parquet-format/).

Author: hyukjinkwon <gurwls223@gmail.com>
Author: Hyukjin Kwon <gurwls223@gmail.com>

Closes #9658 from HyukjinKwon/SPARK-11692.
---
 .../parquet/CatalystSchemaConverter.scala     |  3 ++-
 .../datasources/parquet/ParquetIOSuite.scala  | 25 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystSchemaConverter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystSchemaConverter.scala
index f28a18e2756e4..5f9f9083098a7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystSchemaConverter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystSchemaConverter.scala
@@ -170,9 +170,10 @@ private[parquet] class CatalystSchemaConverter(
 
       case BINARY =>
         originalType match {
-          case UTF8 | ENUM => StringType
+          case UTF8 | ENUM | JSON => StringType
           case null if assumeBinaryIsString => StringType
           case null => BinaryType
+          case BSON => BinaryType
           case DECIMAL => makeDecimalType()
           case _ => illegalType()
         }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index 2aa5dca847c8f..a148facd056a5 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -259,6 +259,31 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
     }
   }
 
+  test("SPARK-11692 Support for Parquet logical types, JSON and BSON (embedded types)") {
+    val parquetSchema = MessageTypeParser.parseMessageType(
+      """message root {
+        |  required binary a(JSON);
+        |  required binary b(BSON);
+        |}
+      """.stripMargin)
+
+    withTempPath { location =>
+      val extraMetadata = Map.empty[String, String].asJava
+      val fileMetadata = new FileMetaData(parquetSchema, extraMetadata, "Spark")
+      val path = new Path(location.getCanonicalPath)
+      val footer = List(
+        new Footer(path, new ParquetMetadata(fileMetadata, Collections.emptyList()))
+      ).asJava
+
+      ParquetFileWriter.writeMetadataFile(sparkContext.hadoopConfiguration, path, footer)
+
+      val jsonDataType = sqlContext.read.parquet(path.toString).schema(0).dataType
+      assert(jsonDataType === StringType)
+      val bsonDataType = sqlContext.read.parquet(path.toString).schema(1).dataType
+      assert(bsonDataType === BinaryType)
+    }
+  }
+
   test("compression codec") {
     def compressionCodecFor(path: String, codecName: String): String = {
       val codecs = for {

From 0e79604aed116bdcb40e03553a2d103b5b1cdbae Mon Sep 17 00:00:00 2001
From: xin Wu <xinwu@us.ibm.com>
Date: Mon, 16 Nov 2015 08:10:48 -0800
Subject: [PATCH 257/518] [SPARK-11522][SQL] input_file_name() returns "" for
 external tables

When computing partition for non-parquet relation, `HadoopRDD.compute` is used. but it does not set the thread local variable `inputFileName` in `NewSqlHadoopRDD`, like `NewSqlHadoopRDD.compute` does.. Yet, when getting the `inputFileName`, `NewSqlHadoopRDD.inputFileName` is exptected, which is empty now.
Adding the setting inputFileName in HadoopRDD.compute resolves this issue.

Author: xin Wu <xinwu@us.ibm.com>

Closes #9542 from xwu0226/SPARK-11522.
---
 .../org/apache/spark/rdd/HadoopRDD.scala      |  7 ++
 .../sql/hive/execution/HiveUDFSuite.scala     | 93 ++++++++++++++++++-
 2 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala b/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
index 0453614f6a1d3..7db583468792e 100644
--- a/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
@@ -213,6 +213,12 @@ class HadoopRDD[K, V](
 
       val inputMetrics = context.taskMetrics.getInputMetricsForReadMethod(DataReadMethod.Hadoop)
 
+      // Sets the thread local variable for the file's name
+      split.inputSplit.value match {
+        case fs: FileSplit => SqlNewHadoopRDD.setInputFileName(fs.getPath.toString)
+        case _ => SqlNewHadoopRDD.unsetInputFileName()
+      }
+
       // Find a function that will return the FileSystem bytes read by this thread. Do this before
       // creating RecordReader, because RecordReader's constructor might read some bytes
       val bytesReadCallback = inputMetrics.bytesReadCallback.orElse {
@@ -250,6 +256,7 @@ class HadoopRDD[K, V](
 
       override def close() {
         if (reader != null) {
+          SqlNewHadoopRDD.unsetInputFileName()
           // Close the reader and release it. Note: it's very important that we don't close the
           // reader more than once, since that exposes us to MAPREDUCE-5918 when running against
           // Hadoop 1.x and older Hadoop 2.x releases. That bug can lead to non-deterministic
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala
index 5ab477efc4ee0..9deb1a6db15ad 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.hive.execution
 
-import java.io.{DataInput, DataOutput}
+import java.io.{PrintWriter, File, DataInput, DataOutput}
 import java.util.{ArrayList, Arrays, Properties}
 
 import org.apache.hadoop.conf.Configuration
@@ -28,6 +28,7 @@ import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectIn
 import org.apache.hadoop.hive.serde2.objectinspector.{ObjectInspector, ObjectInspectorFactory}
 import org.apache.hadoop.hive.serde2.{AbstractSerDe, SerDeStats}
 import org.apache.hadoop.io.Writable
+import org.apache.spark.sql.test.SQLTestUtils
 import org.apache.spark.sql.{AnalysisException, QueryTest, Row}
 import org.apache.spark.sql.hive.test.TestHiveSingleton
 import org.apache.spark.util.Utils
@@ -44,7 +45,7 @@ case class ListStringCaseClass(l: Seq[String])
 /**
  * A test suite for Hive custom UDFs.
  */
-class HiveUDFSuite extends QueryTest with TestHiveSingleton {
+class HiveUDFSuite extends QueryTest with TestHiveSingleton with SQLTestUtils {
 
   import hiveContext.{udf, sql}
   import hiveContext.implicits._
@@ -348,6 +349,94 @@ class HiveUDFSuite extends QueryTest with TestHiveSingleton {
 
     sqlContext.dropTempTable("testUDF")
   }
+
+  test("SPARK-11522 select input_file_name from non-parquet table"){
+
+    withTempDir { tempDir =>
+
+      // EXTERNAL OpenCSVSerde table pointing to LOCATION
+
+      val file1 = new File(tempDir + "/data1")
+      val writer1 = new PrintWriter(file1)
+      writer1.write("1,2")
+      writer1.close()
+
+      val file2 = new File(tempDir + "/data2")
+      val writer2 = new PrintWriter(file2)
+      writer2.write("1,2")
+      writer2.close()
+
+      sql(
+        s"""CREATE EXTERNAL TABLE csv_table(page_id INT, impressions INT)
+        ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'
+        WITH SERDEPROPERTIES (
+          \"separatorChar\" = \",\",
+          \"quoteChar\"     = \"\\\"\",
+          \"escapeChar\"    = \"\\\\\")
+        LOCATION '$tempDir'
+      """)
+
+      val answer1 =
+        sql("SELECT input_file_name() FROM csv_table").head().getString(0)
+      assert(answer1.contains("data1") || answer1.contains("data2"))
+
+      val count1 = sql("SELECT input_file_name() FROM csv_table").distinct().count()
+      assert(count1 == 2)
+      sql("DROP TABLE csv_table")
+
+      // EXTERNAL pointing to LOCATION
+
+      sql(
+        s"""CREATE EXTERNAL TABLE external_t5 (c1 int, c2 int)
+        ROW FORMAT DELIMITED FIELDS TERMINATED BY ','
+        LOCATION '$tempDir'
+      """)
+
+      val answer2 =
+        sql("SELECT input_file_name() as file FROM external_t5").head().getString(0)
+      assert(answer1.contains("data1") || answer1.contains("data2"))
+
+      val count2 = sql("SELECT input_file_name() as file FROM external_t5").distinct().count
+      assert(count2 == 2)
+      sql("DROP TABLE external_t5")
+    }
+
+    withTempDir { tempDir =>
+
+      // External parquet pointing to LOCATION
+
+      val parquetLocation = tempDir + "/external_parquet"
+      sql("SELECT 1, 2").write.parquet(parquetLocation)
+
+      sql(
+        s"""CREATE EXTERNAL TABLE external_parquet(c1 int, c2 int)
+        STORED AS PARQUET
+        LOCATION '$parquetLocation'
+      """)
+
+      val answer3 =
+        sql("SELECT input_file_name() as file FROM external_parquet").head().getString(0)
+      assert(answer3.contains("external_parquet"))
+
+      val count3 = sql("SELECT input_file_name() as file FROM external_parquet").distinct().count
+      assert(count3 == 1)
+      sql("DROP TABLE external_parquet")
+    }
+
+    // Non-External parquet pointing to /tmp/...
+
+    sql("CREATE TABLE parquet_tmp(c1 int, c2 int) " +
+      " STORED AS parquet " +
+      " AS SELECT 1, 2")
+
+    val answer4 =
+      sql("SELECT input_file_name() as file FROM parquet_tmp").head().getString(0)
+    assert(answer4.contains("parquet_tmp"))
+
+    val count4 = sql("SELECT input_file_name() as file FROM parquet_tmp").distinct().count
+    assert(count4 == 1)
+    sql("DROP TABLE parquet_tmp")
+  }
 }
 
 class TestPair(x: Int, y: Int) extends Writable with Serializable {

From 06f1fdba6d1425afddfc1d45a20dbe9bede15e7a Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Mon, 16 Nov 2015 08:58:40 -0800
Subject: [PATCH 258/518] [SPARK-11752] [SQL] fix timezone problem for
 DateTimeUtils.getSeconds

code snippet to reproduce it:
```
TimeZone.setDefault(TimeZone.getTimeZone("Asia/Shanghai"))
val t = Timestamp.valueOf("1900-06-11 12:14:50.789")
val us = fromJavaTimestamp(t)
assert(getSeconds(us) === t.getSeconds)
```

it will be good to add a regression test for it, but the reproducing code need to change the default timezone, and even we change it back, the `lazy val defaultTimeZone` in `DataTimeUtils` is fixed.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9728 from cloud-fan/seconds.
---
 .../spark/sql/catalyst/util/DateTimeUtils.scala    | 14 ++++++++------
 .../sql/catalyst/util/DateTimeUtilsSuite.scala     |  2 +-
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
index deff8a5378b92..8fb3f41f1bd6a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
@@ -395,16 +395,19 @@ object DateTimeUtils {
   /**
    * Returns the microseconds since year zero (-17999) from microseconds since epoch.
    */
-  def absoluteMicroSecond(microsec: SQLTimestamp): SQLTimestamp = {
+  private def absoluteMicroSecond(microsec: SQLTimestamp): SQLTimestamp = {
     microsec + toYearZero * MICROS_PER_DAY
   }
 
+  private def localTimestamp(microsec: SQLTimestamp): SQLTimestamp = {
+    absoluteMicroSecond(microsec) + defaultTimeZone.getOffset(microsec / 1000) * 1000L
+  }
+
   /**
    * Returns the hour value of a given timestamp value. The timestamp is expressed in microseconds.
    */
   def getHours(microsec: SQLTimestamp): Int = {
-    val localTs = absoluteMicroSecond(microsec) + defaultTimeZone.getOffset(microsec / 1000) * 1000L
-    ((localTs / MICROS_PER_SECOND / 3600) % 24).toInt
+    ((localTimestamp(microsec) / MICROS_PER_SECOND / 3600) % 24).toInt
   }
 
   /**
@@ -412,8 +415,7 @@ object DateTimeUtils {
    * microseconds.
    */
   def getMinutes(microsec: SQLTimestamp): Int = {
-    val localTs = absoluteMicroSecond(microsec) + defaultTimeZone.getOffset(microsec / 1000) * 1000L
-    ((localTs / MICROS_PER_SECOND / 60) % 60).toInt
+    ((localTimestamp(microsec) / MICROS_PER_SECOND / 60) % 60).toInt
   }
 
   /**
@@ -421,7 +423,7 @@ object DateTimeUtils {
    * microseconds.
    */
   def getSeconds(microsec: SQLTimestamp): Int = {
-    ((absoluteMicroSecond(microsec) / MICROS_PER_SECOND) % 60).toInt
+    ((localTimestamp(microsec) / MICROS_PER_SECOND) % 60).toInt
   }
 
   private[this] def isLeapYear(year: Int): Boolean = {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
index 64d15e6b910c1..60d45422bc9b8 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
@@ -358,7 +358,7 @@ class DateTimeUtilsSuite extends SparkFunSuite {
     assert(getSeconds(c.getTimeInMillis * 1000) === 9)
   }
 
-  test("hours / miniute / seconds") {
+  test("hours / minutes / seconds") {
     Seq(Timestamp.valueOf("2015-06-11 10:12:35.789"),
       Timestamp.valueOf("2015-06-11 20:13:40.789"),
       Timestamp.valueOf("1900-06-11 12:14:50.789"),

From b0c3fd34e4cfa3f0472d83e71ffe774430cfdc87 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Mon, 16 Nov 2015 09:03:42 -0800
Subject: [PATCH 259/518] [SPARK-11743] [SQL] Add UserDefinedType support to
 RowEncoder

JIRA: https://issues.apache.org/jira/browse/SPARK-11743

RowEncoder doesn't support UserDefinedType now. We should add the support for it.

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #9712 from viirya/rowencoder-udt.
---
 .../main/scala/org/apache/spark/sql/Row.scala | 14 +++-
 .../sql/catalyst/encoders/RowEncoder.scala    | 24 +++++-
 .../sql/catalyst/expressions/objects.scala    | 48 +++++------
 .../catalyst/encoders/RowEncoderSuite.scala   | 82 ++++++++++++++++++-
 4 files changed, 139 insertions(+), 29 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Row.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Row.scala
index ed2fdf9f2f7cf..0f0f200122c34 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Row.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Row.scala
@@ -152,7 +152,7 @@ trait Row extends Serializable {
    *   BinaryType -> byte array
    *   ArrayType -> scala.collection.Seq (use getList for java.util.List)
    *   MapType -> scala.collection.Map (use getJavaMap for java.util.Map)
-   *   StructType -> org.apache.spark.sql.Row
+   *   StructType -> org.apache.spark.sql.Row (or Product)
    * }}}
    */
   def apply(i: Int): Any = get(i)
@@ -177,7 +177,7 @@ trait Row extends Serializable {
    *   BinaryType -> byte array
    *   ArrayType -> scala.collection.Seq (use getList for java.util.List)
    *   MapType -> scala.collection.Map (use getJavaMap for java.util.Map)
-   *   StructType -> org.apache.spark.sql.Row
+   *   StructType -> org.apache.spark.sql.Row (or Product)
    * }}}
    */
   def get(i: Int): Any
@@ -306,7 +306,15 @@ trait Row extends Serializable {
    *
    * @throws ClassCastException when data type does not match.
    */
-  def getStruct(i: Int): Row = getAs[Row](i)
+  def getStruct(i: Int): Row = {
+    // Product and Row both are recoginized as StructType in a Row
+    val t = get(i)
+    if (t.isInstanceOf[Product]) {
+      Row.fromTuple(t.asInstanceOf[Product])
+    } else {
+      t.asInstanceOf[Row]
+    }
+  }
 
   /**
    * Returns the value at position i.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
index e0be896bb3548..9bb1602494b68 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
@@ -50,6 +50,14 @@ object RowEncoder {
     case BooleanType | ByteType | ShortType | IntegerType | LongType |
          FloatType | DoubleType | BinaryType => inputObject
 
+    case udt: UserDefinedType[_] =>
+      val obj = NewInstance(
+        udt.userClass.getAnnotation(classOf[SQLUserDefinedType]).udt(),
+        Nil,
+        false,
+        dataType = ObjectType(udt.userClass.getAnnotation(classOf[SQLUserDefinedType]).udt()))
+      Invoke(obj, "serialize", udt.sqlType, inputObject :: Nil)
+
     case TimestampType =>
       StaticInvoke(
         DateTimeUtils,
@@ -109,11 +117,16 @@ object RowEncoder {
 
     case StructType(fields) =>
       val convertedFields = fields.zipWithIndex.map { case (f, i) =>
+        val method = if (f.dataType.isInstanceOf[StructType]) {
+          "getStruct"
+        } else {
+          "get"
+        }
         If(
           Invoke(inputObject, "isNullAt", BooleanType, Literal(i) :: Nil),
           Literal.create(null, f.dataType),
           extractorsFor(
-            Invoke(inputObject, "get", externalDataTypeFor(f.dataType), Literal(i) :: Nil),
+            Invoke(inputObject, method, externalDataTypeFor(f.dataType), Literal(i) :: Nil),
             f.dataType))
       }
       CreateStruct(convertedFields)
@@ -137,6 +150,7 @@ object RowEncoder {
     case _: ArrayType => ObjectType(classOf[scala.collection.Seq[_]])
     case _: MapType => ObjectType(classOf[scala.collection.Map[_, _]])
     case _: StructType => ObjectType(classOf[Row])
+    case udt: UserDefinedType[_] => ObjectType(udt.userClass)
   }
 
   private def constructorFor(schema: StructType): Expression = {
@@ -155,6 +169,14 @@ object RowEncoder {
     case BooleanType | ByteType | ShortType | IntegerType | LongType |
          FloatType | DoubleType | BinaryType => input
 
+    case udt: UserDefinedType[_] =>
+      val obj = NewInstance(
+        udt.userClass.getAnnotation(classOf[SQLUserDefinedType]).udt(),
+        Nil,
+        false,
+        dataType = ObjectType(udt.userClass.getAnnotation(classOf[SQLUserDefinedType]).udt()))
+      Invoke(obj, "deserialize", ObjectType(udt.userClass), input :: Nil)
+
     case TimestampType =>
       StaticInvoke(
         DateTimeUtils,
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index 4f58464221b4b..5cd19de68391c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -113,7 +113,7 @@ case class Invoke(
     arguments: Seq[Expression] = Nil) extends Expression {
 
   override def nullable: Boolean = true
-  override def children: Seq[Expression] = targetObject :: Nil
+  override def children: Seq[Expression] = arguments.+:(targetObject)
 
   override def eval(input: InternalRow): Any =
     throw new UnsupportedOperationException("Only code-generated evaluation is supported.")
@@ -343,33 +343,35 @@ case class MapObjects(
   private lazy val loopAttribute = AttributeReference("loopVar", elementType)()
   private lazy val completeFunction = function(loopAttribute)
 
+  private def itemAccessorMethod(dataType: DataType): String => String = dataType match {
+    case IntegerType => (i: String) => s".getInt($i)"
+    case LongType => (i: String) => s".getLong($i)"
+    case FloatType => (i: String) => s".getFloat($i)"
+    case DoubleType => (i: String) => s".getDouble($i)"
+    case ByteType => (i: String) => s".getByte($i)"
+    case ShortType => (i: String) => s".getShort($i)"
+    case BooleanType => (i: String) => s".getBoolean($i)"
+    case StringType => (i: String) => s".getUTF8String($i)"
+    case s: StructType => (i: String) => s".getStruct($i, ${s.size})"
+    case a: ArrayType => (i: String) => s".getArray($i)"
+    case _: MapType => (i: String) => s".getMap($i)"
+    case udt: UserDefinedType[_] => itemAccessorMethod(udt.sqlType)
+  }
+
   private lazy val (lengthFunction, itemAccessor, primitiveElement) = inputData.dataType match {
     case ObjectType(cls) if classOf[Seq[_]].isAssignableFrom(cls) =>
       (".size()", (i: String) => s".apply($i)", false)
     case ObjectType(cls) if cls.isArray =>
       (".length", (i: String) => s"[$i]", false)
-    case ArrayType(s: StructType, _) =>
-      (".numElements()", (i: String) => s".getStruct($i, ${s.size})", false)
-    case ArrayType(a: ArrayType, _) =>
-      (".numElements()", (i: String) => s".getArray($i)", true)
-    case ArrayType(IntegerType, _) =>
-      (".numElements()", (i: String) => s".getInt($i)", true)
-    case ArrayType(LongType, _) =>
-      (".numElements()", (i: String) => s".getLong($i)", true)
-    case ArrayType(FloatType, _) =>
-      (".numElements()", (i: String) => s".getFloat($i)", true)
-    case ArrayType(DoubleType, _) =>
-      (".numElements()", (i: String) => s".getDouble($i)", true)
-    case ArrayType(ByteType, _) =>
-      (".numElements()", (i: String) => s".getByte($i)", true)
-    case ArrayType(ShortType, _) =>
-      (".numElements()", (i: String) => s".getShort($i)", true)
-    case ArrayType(BooleanType, _) =>
-      (".numElements()", (i: String) => s".getBoolean($i)", true)
-    case ArrayType(StringType, _) =>
-      (".numElements()", (i: String) => s".getUTF8String($i)", false)
-    case ArrayType(_: MapType, _) =>
-      (".numElements()", (i: String) => s".getMap($i)", false)
+    case ArrayType(t, _) =>
+      val (sqlType, primitiveElement) = t match {
+        case m: MapType => (m, false)
+        case s: StructType => (s, false)
+        case s: StringType => (s, false)
+        case udt: UserDefinedType[_] => (udt.sqlType, false)
+        case o => (o, true)
+      }
+      (".numElements()", itemAccessorMethod(sqlType), primitiveElement)
   }
 
   override def nullable: Boolean = true
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
index e8301e8e06b52..c868ddec1bab2 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
@@ -19,14 +19,62 @@ package org.apache.spark.sql.catalyst.encoders
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.{RandomDataGenerator, Row}
+import org.apache.spark.sql.catalyst.util.{GenericArrayData, ArrayData}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
+@SQLUserDefinedType(udt = classOf[ExamplePointUDT])
+class ExamplePoint(val x: Double, val y: Double) extends Serializable {
+  override def hashCode: Int = 41 * (41 + x.toInt) + y.toInt
+  override def equals(that: Any): Boolean = {
+    if (that.isInstanceOf[ExamplePoint]) {
+      val e = that.asInstanceOf[ExamplePoint]
+      (this.x == e.x || (this.x.isNaN && e.x.isNaN) || (this.x.isInfinity && e.x.isInfinity)) &&
+        (this.y == e.y || (this.y.isNaN && e.y.isNaN) || (this.y.isInfinity && e.y.isInfinity))
+    } else {
+      false
+    }
+  }
+}
+
+/**
+ * User-defined type for [[ExamplePoint]].
+ */
+class ExamplePointUDT extends UserDefinedType[ExamplePoint] {
+
+  override def sqlType: DataType = ArrayType(DoubleType, false)
+
+  override def pyUDT: String = "pyspark.sql.tests.ExamplePointUDT"
+
+  override def serialize(obj: Any): GenericArrayData = {
+    obj match {
+      case p: ExamplePoint =>
+        val output = new Array[Any](2)
+        output(0) = p.x
+        output(1) = p.y
+        new GenericArrayData(output)
+    }
+  }
+
+  override def deserialize(datum: Any): ExamplePoint = {
+    datum match {
+      case values: ArrayData =>
+        new ExamplePoint(values.getDouble(0), values.getDouble(1))
+    }
+  }
+
+  override def userClass: Class[ExamplePoint] = classOf[ExamplePoint]
+
+  private[spark] override def asNullable: ExamplePointUDT = this
+}
+
 class RowEncoderSuite extends SparkFunSuite {
 
   private val structOfString = new StructType().add("str", StringType)
+  private val structOfUDT = new StructType().add("udt", new ExamplePointUDT, false)
   private val arrayOfString = ArrayType(StringType)
   private val mapOfString = MapType(StringType, StringType)
+  private val arrayOfUDT = ArrayType(new ExamplePointUDT, false)
 
   encodeDecodeTest(
     new StructType()
@@ -41,7 +89,8 @@ class RowEncoderSuite extends SparkFunSuite {
       .add("string", StringType)
       .add("binary", BinaryType)
       .add("date", DateType)
-      .add("timestamp", TimestampType))
+      .add("timestamp", TimestampType)
+      .add("udt", new ExamplePointUDT, false))
 
   encodeDecodeTest(
     new StructType()
@@ -68,7 +117,36 @@ class RowEncoderSuite extends SparkFunSuite {
       .add("structOfArray", new StructType().add("array", arrayOfString))
       .add("structOfMap", new StructType().add("map", mapOfString))
       .add("structOfArrayAndMap",
-        new StructType().add("array", arrayOfString).add("map", mapOfString)))
+        new StructType().add("array", arrayOfString).add("map", mapOfString))
+      .add("structOfUDT", structOfUDT))
+
+  test(s"encode/decode: arrayOfUDT") {
+    val schema = new StructType()
+      .add("arrayOfUDT", arrayOfUDT)
+
+    val encoder = RowEncoder(schema)
+
+    val input: Row = Row(Seq(new ExamplePoint(0.1, 0.2), new ExamplePoint(0.3, 0.4)))
+    val row = encoder.toRow(input)
+    val convertedBack = encoder.fromRow(row)
+    assert(input.getSeq[ExamplePoint](0) == convertedBack.getSeq[ExamplePoint](0))
+  }
+
+  test(s"encode/decode: Product") {
+    val schema = new StructType()
+      .add("structAsProduct",
+        new StructType()
+          .add("int", IntegerType)
+          .add("string", StringType)
+          .add("double", DoubleType))
+
+    val encoder = RowEncoder(schema)
+
+    val input: Row = Row((100, "test", 0.123))
+    val row = encoder.toRow(input)
+    val convertedBack = encoder.fromRow(row)
+    assert(input.getStruct(0) == convertedBack.getStruct(0))
+  }
 
   private def encodeDecodeTest(schema: StructType): Unit = {
     test(s"encode/decode: ${schema.simpleString}") {

From de5e531d337075fd849437e88846873bca8685e6 Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Mon, 16 Nov 2015 11:21:17 -0800
Subject: [PATCH 260/518] [SPARK-11731][STREAMING] Enable batching on Driver
 WriteAheadLog by default

Using batching on the driver for the WriteAheadLog should be an improvement for all environments and use cases. Users will be able to scale to much higher number of receivers with the BatchedWriteAheadLog. Therefore we should turn it on by default, and QA it in the QA period.

I've also added some tests to make sure the default configurations are correct regarding recent additions:
 - batching on by default
 - closeFileAfterWrite off by default
 - parallelRecovery off by default

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9695 from brkyvz/enable-batch-wal.
---
 .../streaming/util/WriteAheadLogUtils.scala   |  2 +-
 .../streaming/JavaWriteAheadLogSuite.java     |  1 +
 .../streaming/ReceivedBlockTrackerSuite.scala |  9 +++++--
 .../streaming/util/WriteAheadLogSuite.scala   | 24 ++++++++++++++++++-
 .../util/WriteAheadLogUtilsSuite.scala        | 19 ++++++++++++---
 5 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/WriteAheadLogUtils.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/WriteAheadLogUtils.scala
index 731a369fc92c0..7f9e2c9734970 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/WriteAheadLogUtils.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/WriteAheadLogUtils.scala
@@ -67,7 +67,7 @@ private[streaming] object WriteAheadLogUtils extends Logging {
   }
 
   def isBatchingEnabled(conf: SparkConf, isDriver: Boolean): Boolean = {
-    isDriver && conf.getBoolean(DRIVER_WAL_BATCHING_CONF_KEY, defaultValue = false)
+    isDriver && conf.getBoolean(DRIVER_WAL_BATCHING_CONF_KEY, defaultValue = true)
   }
 
   /**
diff --git a/streaming/src/test/java/org/apache/spark/streaming/JavaWriteAheadLogSuite.java b/streaming/src/test/java/org/apache/spark/streaming/JavaWriteAheadLogSuite.java
index 175b8a496b4e5..09b5f8ed03279 100644
--- a/streaming/src/test/java/org/apache/spark/streaming/JavaWriteAheadLogSuite.java
+++ b/streaming/src/test/java/org/apache/spark/streaming/JavaWriteAheadLogSuite.java
@@ -108,6 +108,7 @@ public void close() {
   public void testCustomWAL() {
     SparkConf conf = new SparkConf();
     conf.set("spark.streaming.driver.writeAheadLog.class", JavaWriteAheadLogSuite.class.getName());
+    conf.set("spark.streaming.driver.writeAheadLog.allowBatching", "false");
     WriteAheadLog wal = WriteAheadLogUtils.createLogForDriver(conf, null, null);
 
     String data1 = "data1";
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala
index 7db17abb7947c..081f5a1c93e6e 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockTrackerSuite.scala
@@ -330,8 +330,13 @@ class ReceivedBlockTrackerSuite
     : Seq[ReceivedBlockTrackerLogEvent] = {
     logFiles.flatMap {
       file => new FileBasedWriteAheadLogReader(file, hadoopConf).toSeq
-    }.map { byteBuffer =>
-      Utils.deserialize[ReceivedBlockTrackerLogEvent](byteBuffer.array)
+    }.flatMap { byteBuffer =>
+      val validBuffer = if (WriteAheadLogUtils.isBatchingEnabled(conf, isDriver = true)) {
+        Utils.deserialize[Array[Array[Byte]]](byteBuffer.array()).map(ByteBuffer.wrap)
+      } else {
+        Array(byteBuffer)
+      }
+      validBuffer.map(b => Utils.deserialize[ReceivedBlockTrackerLogEvent](b.array()))
     }.toList
   }
 
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
index 4273fd7dda8be..7f80d6ecdbbb5 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
@@ -20,7 +20,7 @@ import java.io._
 import java.nio.ByteBuffer
 import java.util.{Iterator => JIterator}
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.concurrent.{TimeUnit, CountDownLatch, ThreadPoolExecutor}
+import java.util.concurrent.{RejectedExecutionException, TimeUnit, CountDownLatch, ThreadPoolExecutor}
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
@@ -190,6 +190,28 @@ abstract class CommonWriteAheadLogTests(
     }
     assert(!nonexistentTempPath.exists(), "Directory created just by attempting to read segment")
   }
+
+  test(testPrefix + "parallel recovery not enabled if closeFileAfterWrite = false") {
+    // write some data
+    val writtenData = (1 to 10).map { i =>
+      val data = generateRandomData()
+      val file = testDir + s"/log-$i-$i"
+      writeDataManually(data, file, allowBatching)
+      data
+    }.flatten
+
+    val wal = createWriteAheadLog(testDir, closeFileAfterWrite, allowBatching)
+    // create iterator but don't materialize it
+    val readData = wal.readAll().asScala.map(byteBufferToString)
+    wal.close()
+    if (closeFileAfterWrite) {
+      // the threadpool is shutdown by the wal.close call above, therefore we shouldn't be able
+      // to materialize the iterator with parallel recovery
+      intercept[RejectedExecutionException](readData.toArray)
+    } else {
+      assert(readData.toSeq === writtenData)
+    }
+  }
 }
 
 class FileBasedWriteAheadLogSuite
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogUtilsSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogUtilsSuite.scala
index 9152728191ea1..bfc5b0cf60fb1 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogUtilsSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogUtilsSuite.scala
@@ -56,19 +56,19 @@ class WriteAheadLogUtilsSuite extends SparkFunSuite {
   test("log selection and creation") {
 
     val emptyConf = new SparkConf()  // no log configuration
-    assertDriverLogClass[FileBasedWriteAheadLog](emptyConf)
+    assertDriverLogClass[FileBasedWriteAheadLog](emptyConf, isBatched = true)
     assertReceiverLogClass[FileBasedWriteAheadLog](emptyConf)
 
     // Verify setting driver WAL class
     val driverWALConf = new SparkConf().set("spark.streaming.driver.writeAheadLog.class",
       classOf[MockWriteAheadLog0].getName())
-    assertDriverLogClass[MockWriteAheadLog0](driverWALConf)
+    assertDriverLogClass[MockWriteAheadLog0](driverWALConf, isBatched = true)
     assertReceiverLogClass[FileBasedWriteAheadLog](driverWALConf)
 
     // Verify setting receiver WAL class
     val receiverWALConf = new SparkConf().set("spark.streaming.receiver.writeAheadLog.class",
       classOf[MockWriteAheadLog0].getName())
-    assertDriverLogClass[FileBasedWriteAheadLog](receiverWALConf)
+    assertDriverLogClass[FileBasedWriteAheadLog](receiverWALConf, isBatched = true)
     assertReceiverLogClass[MockWriteAheadLog0](receiverWALConf)
 
     // Verify setting receiver WAL class with 1-arg constructor
@@ -104,6 +104,19 @@ class WriteAheadLogUtilsSuite extends SparkFunSuite {
     assertDriverLogClass[FileBasedWriteAheadLog](receiverWALConf, isBatched = true)
     assertReceiverLogClass[MockWriteAheadLog0](receiverWALConf)
   }
+
+  test("batching is enabled by default in WriteAheadLog") {
+    val conf = new SparkConf()
+    assert(WriteAheadLogUtils.isBatchingEnabled(conf, isDriver = true))
+    // batching is not valid for receiver WALs
+    assert(!WriteAheadLogUtils.isBatchingEnabled(conf, isDriver = false))
+  }
+
+  test("closeFileAfterWrite is disabled by default in WriteAheadLog") {
+    val conf = new SparkConf()
+    assert(!WriteAheadLogUtils.shouldCloseFileAfterWrite(conf, isDriver = true))
+    assert(!WriteAheadLogUtils.shouldCloseFileAfterWrite(conf, isDriver = false))
+  }
 }
 
 object WriteAheadLogUtilsSuite {

From ace0db47141ffd457c2091751038fc291f6d5a8b Mon Sep 17 00:00:00 2001
From: Daniel Jalova <djalova@us.ibm.com>
Date: Mon, 16 Nov 2015 11:29:27 -0800
Subject: [PATCH 261/518] [SPARK-6328][PYTHON] Python API for StreamingListener

Author: Daniel Jalova <djalova@us.ibm.com>

Closes #9186 from djalova/SPARK-6328.
---
 python/pyspark/streaming/__init__.py          |   3 +-
 python/pyspark/streaming/context.py           |   8 ++
 python/pyspark/streaming/listener.py          |  75 +++++++++++
 python/pyspark/streaming/tests.py             | 126 +++++++++++++++++-
 .../api/java/JavaStreamingListener.scala      |  76 +++++++++++
 5 files changed, 286 insertions(+), 2 deletions(-)
 create mode 100644 python/pyspark/streaming/listener.py

diff --git a/python/pyspark/streaming/__init__.py b/python/pyspark/streaming/__init__.py
index d2644a1d4ffab..66e8f8ef001e3 100644
--- a/python/pyspark/streaming/__init__.py
+++ b/python/pyspark/streaming/__init__.py
@@ -17,5 +17,6 @@
 
 from pyspark.streaming.context import StreamingContext
 from pyspark.streaming.dstream import DStream
+from pyspark.streaming.listener import StreamingListener
 
-__all__ = ['StreamingContext', 'DStream']
+__all__ = ['StreamingContext', 'DStream', 'StreamingListener']
diff --git a/python/pyspark/streaming/context.py b/python/pyspark/streaming/context.py
index 8be56c9915265..1388b6d044e04 100644
--- a/python/pyspark/streaming/context.py
+++ b/python/pyspark/streaming/context.py
@@ -363,3 +363,11 @@ def union(self, *dstreams):
         first = dstreams[0]
         jrest = [d._jdstream for d in dstreams[1:]]
         return DStream(self._jssc.union(first._jdstream, jrest), self, first._jrdd_deserializer)
+
+    def addStreamingListener(self, streamingListener):
+        """
+        Add a [[org.apache.spark.streaming.scheduler.StreamingListener]] object for
+        receiving system events related to streaming.
+        """
+        self._jssc.addStreamingListener(self._jvm.JavaStreamingListenerWrapper(
+            self._jvm.PythonStreamingListenerWrapper(streamingListener)))
diff --git a/python/pyspark/streaming/listener.py b/python/pyspark/streaming/listener.py
new file mode 100644
index 0000000000000..b830797f5c0a0
--- /dev/null
+++ b/python/pyspark/streaming/listener.py
@@ -0,0 +1,75 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+__all__ = ["StreamingListener"]
+
+
+class StreamingListener(object):
+
+    def __init__(self):
+        pass
+
+    def onReceiverStarted(self, receiverStarted):
+        """
+        Called when a receiver has been started
+        """
+        pass
+
+    def onReceiverError(self, receiverError):
+        """
+        Called when a receiver has reported an error
+        """
+        pass
+
+    def onReceiverStopped(self, receiverStopped):
+        """
+        Called when a receiver has been stopped
+        """
+        pass
+
+    def onBatchSubmitted(self, batchSubmitted):
+        """
+        Called when a batch of jobs has been submitted for processing.
+        """
+        pass
+
+    def onBatchStarted(self, batchStarted):
+        """
+        Called when processing of a batch of jobs has started.
+        """
+        pass
+
+    def onBatchCompleted(self, batchCompleted):
+        """
+        Called when processing of a batch of jobs has completed.
+        """
+        pass
+
+    def onOutputOperationStarted(self, outputOperationStarted):
+        """
+        Called when processing of a job of a batch has started.
+        """
+        pass
+
+    def onOutputOperationCompleted(self, outputOperationCompleted):
+        """
+        Called when processing of a job of a batch has completed
+        """
+        pass
+
+    class Java:
+        implements = ["org.apache.spark.streaming.api.java.PythonStreamingListener"]
diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index 6ee864d8d3da6..2983028413bb8 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -48,6 +48,7 @@
 from pyspark.streaming.flume import FlumeUtils
 from pyspark.streaming.mqtt import MQTTUtils
 from pyspark.streaming.kinesis import KinesisUtils, InitialPositionInStream
+from pyspark.streaming.listener import StreamingListener
 
 
 class PySparkStreamingTestCase(unittest.TestCase):
@@ -403,6 +404,128 @@ def func(dstream):
         self._test_func(input, func, expected)
 
 
+class StreamingListenerTests(PySparkStreamingTestCase):
+
+    duration = .5
+
+    class BatchInfoCollector(StreamingListener):
+
+        def __init__(self):
+            super(StreamingListener, self).__init__()
+            self.batchInfosCompleted = []
+            self.batchInfosStarted = []
+            self.batchInfosSubmitted = []
+
+        def onBatchSubmitted(self, batchSubmitted):
+            self.batchInfosSubmitted.append(batchSubmitted.batchInfo())
+
+        def onBatchStarted(self, batchStarted):
+            self.batchInfosStarted.append(batchStarted.batchInfo())
+
+        def onBatchCompleted(self, batchCompleted):
+            self.batchInfosCompleted.append(batchCompleted.batchInfo())
+
+    def test_batch_info_reports(self):
+        batch_collector = self.BatchInfoCollector()
+        self.ssc.addStreamingListener(batch_collector)
+        input = [[1], [2], [3], [4]]
+
+        def func(dstream):
+            return dstream.map(int)
+        expected = [[1], [2], [3], [4]]
+        self._test_func(input, func, expected)
+
+        batchInfosSubmitted = batch_collector.batchInfosSubmitted
+        batchInfosStarted = batch_collector.batchInfosStarted
+        batchInfosCompleted = batch_collector.batchInfosCompleted
+
+        self.wait_for(batchInfosCompleted, 4)
+
+        self.assertGreaterEqual(len(batchInfosSubmitted), 4)
+        for info in batchInfosSubmitted:
+            self.assertGreaterEqual(info.batchTime().milliseconds(), 0)
+            self.assertGreaterEqual(info.submissionTime(), 0)
+
+            for streamId in info.streamIdToInputInfo():
+                streamInputInfo = info.streamIdToInputInfo()[streamId]
+                self.assertGreaterEqual(streamInputInfo.inputStreamId(), 0)
+                self.assertGreaterEqual(streamInputInfo.numRecords, 0)
+                for key in streamInputInfo.metadata():
+                    self.assertIsNotNone(streamInputInfo.metadata()[key])
+                self.assertIsNotNone(streamInputInfo.metadataDescription())
+
+            for outputOpId in info.outputOperationInfos():
+                outputInfo = info.outputOperationInfos()[outputOpId]
+                self.assertGreaterEqual(outputInfo.batchTime().milliseconds(), 0)
+                self.assertGreaterEqual(outputInfo.id(), 0)
+                self.assertIsNotNone(outputInfo.name())
+                self.assertIsNotNone(outputInfo.description())
+                self.assertGreaterEqual(outputInfo.startTime(), -1)
+                self.assertGreaterEqual(outputInfo.endTime(), -1)
+                self.assertIsNone(outputInfo.failureReason())
+
+            self.assertEqual(info.schedulingDelay(), -1)
+            self.assertEqual(info.processingDelay(), -1)
+            self.assertEqual(info.totalDelay(), -1)
+            self.assertEqual(info.numRecords(), 0)
+
+        self.assertGreaterEqual(len(batchInfosStarted), 4)
+        for info in batchInfosStarted:
+            self.assertGreaterEqual(info.batchTime().milliseconds(), 0)
+            self.assertGreaterEqual(info.submissionTime(), 0)
+
+            for streamId in info.streamIdToInputInfo():
+                streamInputInfo = info.streamIdToInputInfo()[streamId]
+                self.assertGreaterEqual(streamInputInfo.inputStreamId(), 0)
+                self.assertGreaterEqual(streamInputInfo.numRecords, 0)
+                for key in streamInputInfo.metadata():
+                    self.assertIsNotNone(streamInputInfo.metadata()[key])
+                self.assertIsNotNone(streamInputInfo.metadataDescription())
+
+            for outputOpId in info.outputOperationInfos():
+                outputInfo = info.outputOperationInfos()[outputOpId]
+                self.assertGreaterEqual(outputInfo.batchTime().milliseconds(), 0)
+                self.assertGreaterEqual(outputInfo.id(), 0)
+                self.assertIsNotNone(outputInfo.name())
+                self.assertIsNotNone(outputInfo.description())
+                self.assertGreaterEqual(outputInfo.startTime(), -1)
+                self.assertGreaterEqual(outputInfo.endTime(), -1)
+                self.assertIsNone(outputInfo.failureReason())
+
+            self.assertGreaterEqual(info.schedulingDelay(), 0)
+            self.assertEqual(info.processingDelay(), -1)
+            self.assertEqual(info.totalDelay(), -1)
+            self.assertEqual(info.numRecords(), 0)
+
+        self.assertGreaterEqual(len(batchInfosCompleted), 4)
+        for info in batchInfosCompleted:
+            self.assertGreaterEqual(info.batchTime().milliseconds(), 0)
+            self.assertGreaterEqual(info.submissionTime(), 0)
+
+            for streamId in info.streamIdToInputInfo():
+                streamInputInfo = info.streamIdToInputInfo()[streamId]
+                self.assertGreaterEqual(streamInputInfo.inputStreamId(), 0)
+                self.assertGreaterEqual(streamInputInfo.numRecords, 0)
+                for key in streamInputInfo.metadata():
+                    self.assertIsNotNone(streamInputInfo.metadata()[key])
+                self.assertIsNotNone(streamInputInfo.metadataDescription())
+
+            for outputOpId in info.outputOperationInfos():
+                outputInfo = info.outputOperationInfos()[outputOpId]
+                self.assertGreaterEqual(outputInfo.batchTime().milliseconds(), 0)
+                self.assertGreaterEqual(outputInfo.id(), 0)
+                self.assertIsNotNone(outputInfo.name())
+                self.assertIsNotNone(outputInfo.description())
+                self.assertGreaterEqual(outputInfo.startTime(), 0)
+                self.assertGreaterEqual(outputInfo.endTime(), 0)
+                self.assertIsNone(outputInfo.failureReason())
+
+            self.assertGreaterEqual(info.schedulingDelay(), 0)
+            self.assertGreaterEqual(info.processingDelay(), 0)
+            self.assertGreaterEqual(info.totalDelay(), 0)
+            self.assertEqual(info.numRecords(), 0)
+
+
 class WindowFunctionTests(PySparkStreamingTestCase):
 
     timeout = 15
@@ -1308,7 +1431,8 @@ def search_kinesis_asl_assembly_jar():
 
     os.environ["PYSPARK_SUBMIT_ARGS"] = "--jars %s pyspark-shell" % jars
     testcases = [BasicOperationTests, WindowFunctionTests, StreamingContextTests, CheckpointTests,
-                 KafkaStreamTests, FlumeStreamTests, FlumePollingStreamTests, MQTTStreamTests]
+                 KafkaStreamTests, FlumeStreamTests, FlumePollingStreamTests, MQTTStreamTests,
+                 StreamingListenerTests]
 
     if kinesis_jar_present is True:
         testcases.append(KinesisStreamTests)
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
index 34429074fe804..7bfd6bd5af759 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaStreamingListener.scala
@@ -18,6 +18,82 @@
 package org.apache.spark.streaming.api.java
 
 import org.apache.spark.streaming.Time
+import org.apache.spark.streaming.scheduler.StreamingListener
+
+private[streaming] trait PythonStreamingListener{
+
+  /** Called when a receiver has been started */
+  def onReceiverStarted(receiverStarted: JavaStreamingListenerReceiverStarted) { }
+
+  /** Called when a receiver has reported an error */
+  def onReceiverError(receiverError: JavaStreamingListenerReceiverError) { }
+
+  /** Called when a receiver has been stopped */
+  def onReceiverStopped(receiverStopped: JavaStreamingListenerReceiverStopped) { }
+
+  /** Called when a batch of jobs has been submitted for processing. */
+  def onBatchSubmitted(batchSubmitted: JavaStreamingListenerBatchSubmitted) { }
+
+  /** Called when processing of a batch of jobs has started.  */
+  def onBatchStarted(batchStarted: JavaStreamingListenerBatchStarted) { }
+
+  /** Called when processing of a batch of jobs has completed. */
+  def onBatchCompleted(batchCompleted: JavaStreamingListenerBatchCompleted) { }
+
+  /** Called when processing of a job of a batch has started. */
+  def onOutputOperationStarted(
+      outputOperationStarted: JavaStreamingListenerOutputOperationStarted) { }
+
+  /** Called when processing of a job of a batch has completed. */
+  def onOutputOperationCompleted(
+      outputOperationCompleted: JavaStreamingListenerOutputOperationCompleted) { }
+}
+
+private[streaming] class PythonStreamingListenerWrapper(listener: PythonStreamingListener)
+  extends JavaStreamingListener {
+
+  /** Called when a receiver has been started */
+  override def onReceiverStarted(receiverStarted: JavaStreamingListenerReceiverStarted): Unit = {
+    listener.onReceiverStarted(receiverStarted)
+  }
+
+  /** Called when a receiver has reported an error */
+  override def onReceiverError(receiverError: JavaStreamingListenerReceiverError): Unit = {
+    listener.onReceiverError(receiverError)
+  }
+
+  /** Called when a receiver has been stopped */
+  override def onReceiverStopped(receiverStopped: JavaStreamingListenerReceiverStopped): Unit = {
+    listener.onReceiverStopped(receiverStopped)
+  }
+
+  /** Called when a batch of jobs has been submitted for processing. */
+  override def onBatchSubmitted(batchSubmitted: JavaStreamingListenerBatchSubmitted): Unit = {
+    listener.onBatchSubmitted(batchSubmitted)
+  }
+
+  /** Called when processing of a batch of jobs has started.  */
+  override def onBatchStarted(batchStarted: JavaStreamingListenerBatchStarted): Unit = {
+    listener.onBatchStarted(batchStarted)
+  }
+
+  /** Called when processing of a batch of jobs has completed. */
+  override def onBatchCompleted(batchCompleted: JavaStreamingListenerBatchCompleted): Unit = {
+    listener.onBatchCompleted(batchCompleted)
+  }
+
+  /** Called when processing of a job of a batch has started. */
+  override def onOutputOperationStarted(
+    outputOperationStarted: JavaStreamingListenerOutputOperationStarted): Unit = {
+      listener.onOutputOperationStarted(outputOperationStarted)
+  }
+
+  /** Called when processing of a job of a batch has completed. */
+  override def onOutputOperationCompleted(
+    outputOperationCompleted: JavaStreamingListenerOutputOperationCompleted): Unit = {
+      listener.onOutputOperationCompleted(outputOperationCompleted)
+  }
+}
 
 /**
  * A listener interface for receiving information about an ongoing streaming  computation.

From 24477d2705bcf2a851acc241deb8376c5450dc73 Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Mon, 16 Nov 2015 11:43:18 -0800
Subject: [PATCH 262/518] [SPARK-11718][YARN][CORE] Fix explicitly killed
 executor dies silently issue

Currently if dynamic allocation is enabled, explicitly killing executor will not get response, so the executor metadata is wrong in driver side. Which will make dynamic allocation on Yarn fail to work.

The problem is  `disableExecutor` returns false for pending killing executors when `onDisconnect` is detected, so no further implementation is done.

One solution is to bypass these explicitly killed executors to use `super.onDisconnect` to remove executor. This is simple.

Another solution is still querying the loss reason for these explicitly kill executors. Since executor may get killed and informed in the same AM-RM communication, so current way of adding pending loss reason request is not worked (container complete is already processed), here we should store this loss reason for later query.

Here this PR chooses solution 2.

Please help to review. vanzin I think this part is changed by you previously, would you please help to review? Thanks a lot.

Author: jerryshao <sshao@hortonworks.com>

Closes #9684 from jerryshao/SPARK-11718.
---
 .../spark/scheduler/TaskSchedulerImpl.scala   |  1 +
 .../CoarseGrainedSchedulerBackend.scala       |  6 ++--
 .../spark/deploy/yarn/YarnAllocator.scala     | 30 +++++++++++++++----
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index 43d7d80b7aae1..5f136690f456c 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -473,6 +473,7 @@ private[spark] class TaskSchedulerImpl(
              // If the host mapping still exists, it means we don't know the loss reason for the
              // executor. So call removeExecutor() to update tasks running on that executor when
              // the real loss reason is finally known.
+             logError(s"Actual reason for lost executor $executorId: ${reason.message}")
              removeExecutor(executorId, reason)
 
            case None =>
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
index f71d98feac050..3373caf0d15eb 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
@@ -269,7 +269,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
      * Stop making resource offers for the given executor. The executor is marked as lost with
      * the loss reason still pending.
      *
-     * @return Whether executor was alive.
+     * @return Whether executor should be disabled
      */
     protected def disableExecutor(executorId: String): Boolean = {
       val shouldDisable = CoarseGrainedSchedulerBackend.this.synchronized {
@@ -277,7 +277,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
           executorsPendingLossReason += executorId
           true
         } else {
-          false
+          // Returns true for explicitly killed executors, we also need to get pending loss reasons;
+          // For others return false.
+          executorsPendingToRemove.contains(executorId)
         }
       }
 
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index 4d9e777cb4134..7e39c3ea56af3 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -35,7 +35,7 @@ import org.apache.hadoop.yarn.util.RackResolver
 
 import org.apache.log4j.{Level, Logger}
 
-import org.apache.spark.{Logging, SecurityManager, SparkConf}
+import org.apache.spark.{Logging, SecurityManager, SparkConf, SparkException}
 import org.apache.spark.deploy.yarn.YarnSparkHadoopUtil._
 import org.apache.spark.rpc.{RpcCallContext, RpcEndpointRef}
 import org.apache.spark.scheduler.{ExecutorExited, ExecutorLossReason}
@@ -96,6 +96,10 @@ private[yarn] class YarnAllocator(
   // was lost.
   private val pendingLossReasonRequests = new HashMap[String, mutable.Buffer[RpcCallContext]]
 
+  // Maintain loss reasons for already released executors, it will be added when executor loss
+  // reason is got from AM-RM call, and be removed after querying this loss reason.
+  private val releasedExecutorLossReasons = new HashMap[String, ExecutorLossReason]
+
   // Keep track of which container is running which executor to remove the executors later
   // Visible for testing.
   private[yarn] val executorIdToContainer = new HashMap[String, Container]
@@ -202,8 +206,7 @@ private[yarn] class YarnAllocator(
    */
   def killExecutor(executorId: String): Unit = synchronized {
     if (executorIdToContainer.contains(executorId)) {
-      val container = executorIdToContainer.remove(executorId).get
-      containerIdToExecutorId.remove(container.getId)
+      val container = executorIdToContainer.get(executorId).get
       internalReleaseContainer(container)
       numExecutorsRunning -= 1
     } else {
@@ -514,9 +517,18 @@ private[yarn] class YarnAllocator(
 
       containerIdToExecutorId.remove(containerId).foreach { eid =>
         executorIdToContainer.remove(eid)
-        pendingLossReasonRequests.remove(eid).foreach { pendingRequests =>
-          // Notify application of executor loss reasons so it can decide whether it should abort
-          pendingRequests.foreach(_.reply(exitReason))
+        pendingLossReasonRequests.remove(eid) match {
+          case Some(pendingRequests) =>
+            // Notify application of executor loss reasons so it can decide whether it should abort
+            pendingRequests.foreach(_.reply(exitReason))
+
+          case None =>
+            // We cannot find executor for pending reasons. This is because completed container
+            // is processed before querying pending result. We should store it for later query.
+            // This is usually happened when explicitly killing a container, the result will be
+            // returned in one AM-RM communication. So query RPC will be later than this completed
+            // container process.
+            releasedExecutorLossReasons.put(eid, exitReason)
         }
         if (!alreadyReleased) {
           // The executor could have gone away (like no route to host, node failure, etc)
@@ -538,8 +550,14 @@ private[yarn] class YarnAllocator(
     if (executorIdToContainer.contains(eid)) {
       pendingLossReasonRequests
         .getOrElseUpdate(eid, new ArrayBuffer[RpcCallContext]) += context
+    } else if (releasedExecutorLossReasons.contains(eid)) {
+      // Executor is already released explicitly before getting the loss reason, so directly send
+      // the pre-stored lost reason
+      context.reply(releasedExecutorLossReasons.remove(eid).get)
     } else {
       logWarning(s"Tried to get the loss reason for non-existent executor $eid")
+      context.sendFailure(
+        new SparkException(s"Fail to find loss reason for non-existent executor $eid"))
     }
   }
 

From b1a9662623951079e80bd7498e064c4cae4977e9 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Mon, 16 Nov 2015 12:45:34 -0800
Subject: [PATCH 263/518] [SPARK-11754][SQL] consolidate
 `ExpressionEncoder.tuple` and `Encoders.tuple`

These 2 are very similar, we can consolidate them into one.

Also add tests for it and fix a bug.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9729 from cloud-fan/tuple.
---
 .../scala/org/apache/spark/sql/Encoder.scala  |  95 ++++------------
 .../catalyst/encoders/ExpressionEncoder.scala | 104 ++++++++++--------
 .../encoders/ProductEncoderSuite.scala        |  29 +++++
 3 files changed, 108 insertions(+), 120 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index 5f619d6c339e3..c8b017e251637 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -19,10 +19,8 @@ package org.apache.spark.sql
 
 import scala.reflect.ClassTag
 
-import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
-import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.types.{ObjectType, StructField, StructType}
-import org.apache.spark.util.Utils
+import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor}
+import org.apache.spark.sql.types.StructType
 
 /**
  * Used to convert a JVM object of type `T` to and from the internal Spark SQL representation.
@@ -49,83 +47,34 @@ object Encoders {
   def DOUBLE: Encoder[java.lang.Double] = ExpressionEncoder(flat = true)
   def STRING: Encoder[java.lang.String] = ExpressionEncoder(flat = true)
 
-  def tuple[T1, T2](enc1: Encoder[T1], enc2: Encoder[T2]): Encoder[(T1, T2)] = {
-    tuple(Seq(enc1, enc2).map(_.asInstanceOf[ExpressionEncoder[_]]))
-      .asInstanceOf[ExpressionEncoder[(T1, T2)]]
+  def tuple[T1, T2](
+      e1: Encoder[T1],
+      e2: Encoder[T2]): Encoder[(T1, T2)] = {
+    ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2))
   }
 
   def tuple[T1, T2, T3](
-      enc1: Encoder[T1],
-      enc2: Encoder[T2],
-      enc3: Encoder[T3]): Encoder[(T1, T2, T3)] = {
-    tuple(Seq(enc1, enc2, enc3).map(_.asInstanceOf[ExpressionEncoder[_]]))
-      .asInstanceOf[ExpressionEncoder[(T1, T2, T3)]]
+      e1: Encoder[T1],
+      e2: Encoder[T2],
+      e3: Encoder[T3]): Encoder[(T1, T2, T3)] = {
+    ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2), encoderFor(e3))
   }
 
   def tuple[T1, T2, T3, T4](
-      enc1: Encoder[T1],
-      enc2: Encoder[T2],
-      enc3: Encoder[T3],
-      enc4: Encoder[T4]): Encoder[(T1, T2, T3, T4)] = {
-    tuple(Seq(enc1, enc2, enc3, enc4).map(_.asInstanceOf[ExpressionEncoder[_]]))
-      .asInstanceOf[ExpressionEncoder[(T1, T2, T3, T4)]]
+      e1: Encoder[T1],
+      e2: Encoder[T2],
+      e3: Encoder[T3],
+      e4: Encoder[T4]): Encoder[(T1, T2, T3, T4)] = {
+    ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2), encoderFor(e3), encoderFor(e4))
   }
 
   def tuple[T1, T2, T3, T4, T5](
-      enc1: Encoder[T1],
-      enc2: Encoder[T2],
-      enc3: Encoder[T3],
-      enc4: Encoder[T4],
-      enc5: Encoder[T5]): Encoder[(T1, T2, T3, T4, T5)] = {
-    tuple(Seq(enc1, enc2, enc3, enc4, enc5).map(_.asInstanceOf[ExpressionEncoder[_]]))
-      .asInstanceOf[ExpressionEncoder[(T1, T2, T3, T4, T5)]]
-  }
-
-  private def tuple(encoders: Seq[ExpressionEncoder[_]]): ExpressionEncoder[_] = {
-    assert(encoders.length > 1)
-    // make sure all encoders are resolved, i.e. `Attribute` has been resolved to `BoundReference`.
-    assert(encoders.forall(_.fromRowExpression.find(_.isInstanceOf[Attribute]).isEmpty))
-
-    val schema = StructType(encoders.zipWithIndex.map {
-      case (e, i) => StructField(s"_${i + 1}", if (e.flat) e.schema.head.dataType else e.schema)
-    })
-
-    val cls = Utils.getContextOrSparkClassLoader.loadClass(s"scala.Tuple${encoders.size}")
-
-    val extractExpressions = encoders.map {
-      case e if e.flat => e.toRowExpressions.head
-      case other => CreateStruct(other.toRowExpressions)
-    }.zipWithIndex.map { case (expr, index) =>
-      expr.transformUp {
-        case BoundReference(0, t: ObjectType, _) =>
-          Invoke(
-            BoundReference(0, ObjectType(cls), nullable = true),
-            s"_${index + 1}",
-            t)
-      }
-    }
-
-    val constructExpressions = encoders.zipWithIndex.map { case (enc, index) =>
-      if (enc.flat) {
-        enc.fromRowExpression.transform {
-          case b: BoundReference => b.copy(ordinal = index)
-        }
-      } else {
-        enc.fromRowExpression.transformUp {
-          case BoundReference(ordinal, dt, _) =>
-            GetInternalRowField(BoundReference(index, enc.schema, nullable = true), ordinal, dt)
-        }
-      }
-    }
-
-    val constructExpression =
-      NewInstance(cls, constructExpressions, propagateNull = false, ObjectType(cls))
-
-    new ExpressionEncoder[Any](
-      schema,
-      flat = false,
-      extractExpressions,
-      constructExpression,
-      ClassTag(cls))
+      e1: Encoder[T1],
+      e2: Encoder[T2],
+      e3: Encoder[T3],
+      e4: Encoder[T4],
+      e5: Encoder[T5]): Encoder[(T1, T2, T3, T4, T5)] = {
+    ExpressionEncoder.tuple(
+      encoderFor(e1), encoderFor(e2), encoderFor(e3), encoderFor(e4), encoderFor(e5))
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index 0d3e4aafb0af4..9a1a8f5cbbdc3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -67,47 +67,77 @@ object ExpressionEncoder {
   def tuple(encoders: Seq[ExpressionEncoder[_]]): ExpressionEncoder[_] = {
     encoders.foreach(_.assertUnresolved())
 
-    val schema =
-      StructType(
-        encoders.zipWithIndex.map {
-          case (e, i) => StructField(s"_${i + 1}", if (e.flat) e.schema.head.dataType else e.schema)
-        })
+    val schema = StructType(encoders.zipWithIndex.map {
+      case (e, i) => StructField(s"_${i + 1}", if (e.flat) e.schema.head.dataType else e.schema)
+    })
+
     val cls = Utils.getContextOrSparkClassLoader.loadClass(s"scala.Tuple${encoders.size}")
 
-    // Rebind the encoders to the nested schema.
-    val newConstructExpressions = encoders.zipWithIndex.map {
-      case (e, i) if !e.flat => e.nested(i).fromRowExpression
-      case (e, i) => e.shift(i).fromRowExpression
+    val toRowExpressions = encoders.map {
+      case e if e.flat => e.toRowExpressions.head
+      case other => CreateStruct(other.toRowExpressions)
+    }.zipWithIndex.map { case (expr, index) =>
+      expr.transformUp {
+        case BoundReference(0, t, _) =>
+          Invoke(
+            BoundReference(0, ObjectType(cls), nullable = true),
+            s"_${index + 1}",
+            t)
+      }
     }
 
-    val constructExpression =
-      NewInstance(cls, newConstructExpressions, false, ObjectType(cls))
-
-    val input = BoundReference(0, ObjectType(cls), false)
-    val extractExpressions = encoders.zipWithIndex.map {
-      case (e, i) if !e.flat => CreateStruct(e.toRowExpressions.map(_ transformUp {
-        case b: BoundReference =>
-          Invoke(input, s"_${i + 1}", b.dataType, Nil)
-      }))
-      case (e, i) => e.toRowExpressions.head transformUp {
-        case b: BoundReference =>
-          Invoke(input, s"_${i + 1}", b.dataType, Nil)
+    val fromRowExpressions = encoders.zipWithIndex.map { case (enc, index) =>
+      if (enc.flat) {
+        enc.fromRowExpression.transform {
+          case b: BoundReference => b.copy(ordinal = index)
+        }
+      } else {
+        val input = BoundReference(index, enc.schema, nullable = true)
+        enc.fromRowExpression.transformUp {
+          case UnresolvedAttribute(nameParts) =>
+            assert(nameParts.length == 1)
+            UnresolvedExtractValue(input, Literal(nameParts.head))
+          case BoundReference(ordinal, dt, _) => GetInternalRowField(input, ordinal, dt)
+        }
       }
     }
 
+    val fromRowExpression =
+      NewInstance(cls, fromRowExpressions, propagateNull = false, ObjectType(cls))
+
     new ExpressionEncoder[Any](
       schema,
-      false,
-      extractExpressions,
-      constructExpression,
-      ClassTag.apply(cls))
+      flat = false,
+      toRowExpressions,
+      fromRowExpression,
+      ClassTag(cls))
   }
 
-  /** A helper for producing encoders of Tuple2 from other encoders. */
   def tuple[T1, T2](
       e1: ExpressionEncoder[T1],
       e2: ExpressionEncoder[T2]): ExpressionEncoder[(T1, T2)] =
-    tuple(e1 :: e2 :: Nil).asInstanceOf[ExpressionEncoder[(T1, T2)]]
+    tuple(Seq(e1, e2)).asInstanceOf[ExpressionEncoder[(T1, T2)]]
+
+  def tuple[T1, T2, T3](
+      e1: ExpressionEncoder[T1],
+      e2: ExpressionEncoder[T2],
+      e3: ExpressionEncoder[T3]): ExpressionEncoder[(T1, T2, T3)] =
+    tuple(Seq(e1, e2, e3)).asInstanceOf[ExpressionEncoder[(T1, T2, T3)]]
+
+  def tuple[T1, T2, T3, T4](
+      e1: ExpressionEncoder[T1],
+      e2: ExpressionEncoder[T2],
+      e3: ExpressionEncoder[T3],
+      e4: ExpressionEncoder[T4]): ExpressionEncoder[(T1, T2, T3, T4)] =
+    tuple(Seq(e1, e2, e3, e4)).asInstanceOf[ExpressionEncoder[(T1, T2, T3, T4)]]
+
+  def tuple[T1, T2, T3, T4, T5](
+      e1: ExpressionEncoder[T1],
+      e2: ExpressionEncoder[T2],
+      e3: ExpressionEncoder[T3],
+      e4: ExpressionEncoder[T4],
+      e5: ExpressionEncoder[T5]): ExpressionEncoder[(T1, T2, T3, T4, T5)] =
+    tuple(Seq(e1, e2, e3, e4, e5)).asInstanceOf[ExpressionEncoder[(T1, T2, T3, T4, T5)]]
 }
 
 /**
@@ -208,26 +238,6 @@ case class ExpressionEncoder[T](
     })
   }
 
-  /**
-   * Returns a copy of this encoder where the expressions used to create an object given an
-   * input row have been modified to pull the object out from a nested struct, instead of the
-   * top level fields.
-   */
-  private def nested(i: Int): ExpressionEncoder[T] = {
-    // We don't always know our input type at this point since it might be unresolved.
-    // We fill in null and it will get unbound to the actual attribute at this position.
-    val input = BoundReference(i, NullType, nullable = true)
-    copy(fromRowExpression = fromRowExpression transformUp {
-      case u: Attribute =>
-        UnresolvedExtractValue(input, Literal(u.name))
-      case b: BoundReference =>
-        GetStructField(
-          input,
-          StructField(s"i[${b.ordinal}]", b.dataType),
-          b.ordinal)
-    })
-  }
-
   protected val attrs = toRowExpressions.flatMap(_.collect {
     case _: UnresolvedAttribute => ""
     case a: Attribute => s"#${a.exprId}"
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
index fda978e7055ea..bc539d62c537d 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
@@ -117,6 +117,35 @@ class ProductEncoderSuite extends ExpressionEncoderSuite {
   productTest(("Seq[Seq[(Int, Int)]]",
     Seq(Seq((1, 2)))))
 
+  encodeDecodeTest(
+    1 -> 10L,
+    ExpressionEncoder.tuple(FlatEncoder[Int], FlatEncoder[Long]),
+    "tuple with 2 flat encoders")
+
+  encodeDecodeTest(
+    (PrimitiveData(1, 1, 1, 1, 1, 1, true), (3, 30L)),
+    ExpressionEncoder.tuple(ProductEncoder[PrimitiveData], ProductEncoder[(Int, Long)]),
+    "tuple with 2 product encoders")
+
+  encodeDecodeTest(
+    (PrimitiveData(1, 1, 1, 1, 1, 1, true), 3),
+    ExpressionEncoder.tuple(ProductEncoder[PrimitiveData], FlatEncoder[Int]),
+    "tuple with flat encoder and product encoder")
+
+  encodeDecodeTest(
+    (3, PrimitiveData(1, 1, 1, 1, 1, 1, true)),
+    ExpressionEncoder.tuple(FlatEncoder[Int], ProductEncoder[PrimitiveData]),
+    "tuple with product encoder and flat encoder")
+
+  encodeDecodeTest(
+    (1, (10, 100L)),
+    {
+      val intEnc = FlatEncoder[Int]
+      val longEnc = FlatEncoder[Long]
+      ExpressionEncoder.tuple(intEnc, ExpressionEncoder.tuple(intEnc, longEnc))
+    },
+    "nested tuple encoder")
+
   private def productTest[T <: Product : TypeTag](input: T): Unit = {
     encodeDecodeTest(input, ProductEncoder[T], input.getClass.getSimpleName)
   }

From 985b38dd2fa5d8f1e23f1c420ce6262e7e3ed181 Mon Sep 17 00:00:00 2001
From: Zee Chen <zeechen@us.ibm.com>
Date: Mon, 16 Nov 2015 14:21:28 -0800
Subject: [PATCH 264/518] [SPARK-11390][SQL] Query plan with/without
 filterPushdown indistinguishable
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…ishable

Propagate pushed filters to PhyicalRDD in DataSourceStrategy.apply

Author: Zee Chen <zeechen@us.ibm.com>

Closes #9679 from zeocio/spark-11390.
---
 .../apache/spark/sql/execution/ExistingRDD.scala   |  6 ++++--
 .../execution/datasources/DataSourceStrategy.scala |  6 ++++--
 .../apache/spark/sql/execution/PlannerSuite.scala  | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
index 8b41d3d3d892e..62620ec642c78 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
@@ -106,7 +106,9 @@ private[sql] object PhysicalRDD {
   def createFromDataSource(
       output: Seq[Attribute],
       rdd: RDD[InternalRow],
-      relation: BaseRelation): PhysicalRDD = {
-    PhysicalRDD(output, rdd, relation.toString, relation.isInstanceOf[HadoopFsRelation])
+      relation: BaseRelation,
+      extraInformation: String = ""): PhysicalRDD = {
+    PhysicalRDD(output, rdd, relation.toString + extraInformation,
+      relation.isInstanceOf[HadoopFsRelation])
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
index 9bbbfa7c77cba..544d5eccec037 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSourceStrategy.scala
@@ -315,6 +315,8 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
     // `Filter`s or cannot be handled by `relation`.
     val filterCondition = unhandledPredicates.reduceLeftOption(expressions.And)
 
+    val pushedFiltersString = pushedFilters.mkString(" PushedFilter: [", ",", "] ")
+
     if (projects.map(_.toAttribute) == projects &&
         projectSet.size == projects.size &&
         filterSet.subsetOf(projectSet)) {
@@ -332,7 +334,7 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
       val scan = execution.PhysicalRDD.createFromDataSource(
         projects.map(_.toAttribute),
         scanBuilder(requestedColumns, candidatePredicates, pushedFilters),
-        relation.relation)
+        relation.relation, pushedFiltersString)
       filterCondition.map(execution.Filter(_, scan)).getOrElse(scan)
     } else {
       // Don't request columns that are only referenced by pushed filters.
@@ -342,7 +344,7 @@ private[sql] object DataSourceStrategy extends Strategy with Logging {
       val scan = execution.PhysicalRDD.createFromDataSource(
         requestedColumns,
         scanBuilder(requestedColumns, candidatePredicates, pushedFilters),
-        relation.relation)
+        relation.relation, pushedFiltersString)
       execution.Project(
         projects, filterCondition.map(execution.Filter(_, scan)).getOrElse(scan))
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
index be53ec3e271c6..dfec139985f73 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -160,6 +160,20 @@ class PlannerSuite extends SharedSQLContext {
     }
   }
 
+  test("SPARK-11390 explain should print PushedFilters of PhysicalRDD") {
+    withTempPath { file =>
+      val path = file.getCanonicalPath
+      testData.write.parquet(path)
+      val df = sqlContext.read.parquet(path)
+      sqlContext.registerDataFrameAsTable(df, "testPushed")
+
+      withTempTable("testPushed") {
+        val exp = sql("select * from testPushed where key = 15").queryExecution.executedPlan
+        assert(exp.toString.contains("PushedFilter: [EqualTo(key,15)]"))
+      }
+    }
+  }
+
   test("efficient limit -> project -> sort") {
     {
       val query =

From 3c025087b58f475a9bcb5c8f4b2b2df804915b2b Mon Sep 17 00:00:00 2001
From: Davies Liu <davies.liu@gmail.com>
Date: Mon, 16 Nov 2015 14:50:38 -0800
Subject: [PATCH 265/518] Revert "[SPARK-11271][SPARK-11016][CORE] Use Spark
 BitSet instead of RoaringBitmap to reduce memory usage"

This reverts commit e209fa271ae57dc8849f8b1241bf1ea7d6d3d62c.
---
 core/pom.xml                                  |  4 ++
 .../apache/spark/scheduler/MapStatus.scala    | 13 ++---
 .../spark/serializer/KryoSerializer.scala     | 10 +++-
 .../apache/spark/util/collection/BitSet.scala | 28 ++---------
 .../serializer/KryoSerializerSuite.scala      |  6 +++
 .../spark/util/collection/BitSetSuite.scala   | 49 -------------------
 pom.xml                                       |  5 ++
 7 files changed, 33 insertions(+), 82 deletions(-)

diff --git a/core/pom.xml b/core/pom.xml
index 7e1205a076f26..37e3f168ab374 100644
--- a/core/pom.xml
+++ b/core/pom.xml
@@ -177,6 +177,10 @@
       <groupId>net.jpountz.lz4</groupId>
       <artifactId>lz4</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.roaringbitmap</groupId>
+      <artifactId>RoaringBitmap</artifactId>
+    </dependency>
     <dependency>
       <groupId>commons-net</groupId>
       <artifactId>commons-net</artifactId>
diff --git a/core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala b/core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala
index 180c8d1827e13..1efce124c0a6b 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala
@@ -19,8 +19,9 @@ package org.apache.spark.scheduler
 
 import java.io.{Externalizable, ObjectInput, ObjectOutput}
 
+import org.roaringbitmap.RoaringBitmap
+
 import org.apache.spark.storage.BlockManagerId
-import org.apache.spark.util.collection.BitSet
 import org.apache.spark.util.Utils
 
 /**
@@ -132,7 +133,7 @@ private[spark] class CompressedMapStatus(
 private[spark] class HighlyCompressedMapStatus private (
     private[this] var loc: BlockManagerId,
     private[this] var numNonEmptyBlocks: Int,
-    private[this] var emptyBlocks: BitSet,
+    private[this] var emptyBlocks: RoaringBitmap,
     private[this] var avgSize: Long)
   extends MapStatus with Externalizable {
 
@@ -145,7 +146,7 @@ private[spark] class HighlyCompressedMapStatus private (
   override def location: BlockManagerId = loc
 
   override def getSizeForBlock(reduceId: Int): Long = {
-    if (emptyBlocks.get(reduceId)) {
+    if (emptyBlocks.contains(reduceId)) {
       0
     } else {
       avgSize
@@ -160,7 +161,7 @@ private[spark] class HighlyCompressedMapStatus private (
 
   override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException {
     loc = BlockManagerId(in)
-    emptyBlocks = new BitSet
+    emptyBlocks = new RoaringBitmap()
     emptyBlocks.readExternal(in)
     avgSize = in.readLong()
   }
@@ -176,15 +177,15 @@ private[spark] object HighlyCompressedMapStatus {
     // From a compression standpoint, it shouldn't matter whether we track empty or non-empty
     // blocks. From a performance standpoint, we benefit from tracking empty blocks because
     // we expect that there will be far fewer of them, so we will perform fewer bitmap insertions.
+    val emptyBlocks = new RoaringBitmap()
     val totalNumBlocks = uncompressedSizes.length
-    val emptyBlocks = new BitSet(totalNumBlocks)
     while (i < totalNumBlocks) {
       var size = uncompressedSizes(i)
       if (size > 0) {
         numNonEmptyBlocks += 1
         totalSize += size
       } else {
-        emptyBlocks.set(i)
+        emptyBlocks.add(i)
       }
       i += 1
     }
diff --git a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
index bc51d4f2820c8..c5195c1143a8f 100644
--- a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
+++ b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
@@ -30,6 +30,7 @@ import com.esotericsoftware.kryo.io.{Input => KryoInput, Output => KryoOutput}
 import com.esotericsoftware.kryo.serializers.{JavaSerializer => KryoJavaSerializer}
 import com.twitter.chill.{AllScalaRegistrar, EmptyScalaKryoInstantiator}
 import org.apache.avro.generic.{GenericData, GenericRecord}
+import org.roaringbitmap.{ArrayContainer, BitmapContainer, RoaringArray, RoaringBitmap}
 
 import org.apache.spark._
 import org.apache.spark.api.python.PythonBroadcast
@@ -38,7 +39,7 @@ import org.apache.spark.network.util.ByteUnit
 import org.apache.spark.scheduler.{CompressedMapStatus, HighlyCompressedMapStatus}
 import org.apache.spark.storage._
 import org.apache.spark.util.{Utils, BoundedPriorityQueue, SerializableConfiguration, SerializableJobConf}
-import org.apache.spark.util.collection.{BitSet, CompactBuffer}
+import org.apache.spark.util.collection.CompactBuffer
 
 /**
  * A Spark serializer that uses the [[https://code.google.com/p/kryo/ Kryo serialization library]].
@@ -362,7 +363,12 @@ private[serializer] object KryoSerializer {
     classOf[StorageLevel],
     classOf[CompressedMapStatus],
     classOf[HighlyCompressedMapStatus],
-    classOf[BitSet],
+    classOf[RoaringBitmap],
+    classOf[RoaringArray],
+    classOf[RoaringArray.Element],
+    classOf[Array[RoaringArray.Element]],
+    classOf[ArrayContainer],
+    classOf[BitmapContainer],
     classOf[CompactBuffer[_]],
     classOf[BlockManagerId],
     classOf[Array[Byte]],
diff --git a/core/src/main/scala/org/apache/spark/util/collection/BitSet.scala b/core/src/main/scala/org/apache/spark/util/collection/BitSet.scala
index 85c5bdbfcebc0..7ab67fc3a2de9 100644
--- a/core/src/main/scala/org/apache/spark/util/collection/BitSet.scala
+++ b/core/src/main/scala/org/apache/spark/util/collection/BitSet.scala
@@ -17,21 +17,14 @@
 
 package org.apache.spark.util.collection
 
-import java.io.{Externalizable, ObjectInput, ObjectOutput}
-
-import org.apache.spark.util.{Utils => UUtils}
-
-
 /**
  * A simple, fixed-size bit set implementation. This implementation is fast because it avoids
  * safety/bound checking.
  */
-class BitSet(private[this] var numBits: Int) extends Externalizable {
+class BitSet(numBits: Int) extends Serializable {
 
-  private var words = new Array[Long](bit2words(numBits))
-  private def numWords = words.length
-
-  def this() = this(0)
+  private val words = new Array[Long](bit2words(numBits))
+  private val numWords = words.length
 
   /**
    * Compute the capacity (number of bits) that can be represented
@@ -237,19 +230,4 @@ class BitSet(private[this] var numBits: Int) extends Externalizable {
 
   /** Return the number of longs it would take to hold numBits. */
   private def bit2words(numBits: Int) = ((numBits - 1) >> 6) + 1
-
-  override def writeExternal(out: ObjectOutput): Unit = UUtils.tryOrIOException {
-    out.writeInt(numBits)
-    words.foreach(out.writeLong(_))
-  }
-
-  override def readExternal(in: ObjectInput): Unit = UUtils.tryOrIOException {
-    numBits = in.readInt()
-    words = new Array[Long](bit2words(numBits))
-    var index = 0
-    while (index < words.length) {
-      words(index) = in.readLong()
-      index += 1
-    }
-  }
 }
diff --git a/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala b/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala
index afe2e80358ca0..e428414cf6e85 100644
--- a/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala
@@ -322,6 +322,12 @@ class KryoSerializerSuite extends SparkFunSuite with SharedSparkContext {
     val conf = new SparkConf(false)
     conf.set("spark.kryo.registrationRequired", "true")
 
+    // these cases require knowing the internals of RoaringBitmap a little.  Blocks span 2^16
+    // values, and they use a bitmap (dense) if they have more than 4096 values, and an
+    // array (sparse) if they use less.  So we just create two cases, one sparse and one dense.
+    // and we use a roaring bitmap for the empty blocks, so we trigger the dense case w/ mostly
+    // empty blocks
+
     val ser = new KryoSerializer(conf).newInstance()
     val denseBlockSizes = new Array[Long](5000)
     val sparseBlockSizes = Array[Long](0L, 1L, 0L, 2L)
diff --git a/core/src/test/scala/org/apache/spark/util/collection/BitSetSuite.scala b/core/src/test/scala/org/apache/spark/util/collection/BitSetSuite.scala
index b0db0988eeaab..69dbfa9cd7141 100644
--- a/core/src/test/scala/org/apache/spark/util/collection/BitSetSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/collection/BitSetSuite.scala
@@ -17,10 +17,7 @@
 
 package org.apache.spark.util.collection
 
-import java.io.{File, FileInputStream, FileOutputStream, ObjectInputStream, ObjectOutputStream}
-
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.util.{Utils => UUtils}
 
 class BitSetSuite extends SparkFunSuite {
 
@@ -155,50 +152,4 @@ class BitSetSuite extends SparkFunSuite {
     assert(bitsetDiff.nextSetBit(85) === 85)
     assert(bitsetDiff.nextSetBit(86) === -1)
   }
-
-  test("read and write externally") {
-    val tempDir = UUtils.createTempDir()
-    val outputFile = File.createTempFile("bits", null, tempDir)
-
-    val fos = new FileOutputStream(outputFile)
-    val oos = new ObjectOutputStream(fos)
-
-    // Create BitSet
-    val setBits = Seq(0, 9, 1, 10, 90, 96)
-    val bitset = new BitSet(100)
-
-    for (i <- 0 until 100) {
-      assert(!bitset.get(i))
-    }
-
-    setBits.foreach(i => bitset.set(i))
-
-    for (i <- 0 until 100) {
-      if (setBits.contains(i)) {
-        assert(bitset.get(i))
-      } else {
-        assert(!bitset.get(i))
-      }
-    }
-    assert(bitset.cardinality() === setBits.size)
-
-    bitset.writeExternal(oos)
-    oos.close()
-
-    val fis = new FileInputStream(outputFile)
-    val ois = new ObjectInputStream(fis)
-
-    // Read BitSet from the file
-    val bitset2 = new BitSet(0)
-    bitset2.readExternal(ois)
-
-    for (i <- 0 until 100) {
-      if (setBits.contains(i)) {
-        assert(bitset2.get(i))
-      } else {
-        assert(!bitset2.get(i))
-      }
-    }
-    assert(bitset2.cardinality() === setBits.size)
-  }
 }
diff --git a/pom.xml b/pom.xml
index 01afa80617891..2a8a445057174 100644
--- a/pom.xml
+++ b/pom.xml
@@ -634,6 +634,11 @@
           </exclusion>
         </exclusions>
       </dependency>
+      <dependency>
+        <groupId>org.roaringbitmap</groupId>
+        <artifactId>RoaringBitmap</artifactId>
+        <version>0.4.5</version>
+      </dependency>
       <dependency>
         <groupId>commons-net</groupId>
         <artifactId>commons-net</artifactId>

From bcea0bfda66a30ee86790b048de5cb47b4d0b32f Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Mon, 16 Nov 2015 15:06:06 -0800
Subject: [PATCH 266/518] [SPARK-11742][STREAMING] Add the failure info to the
 batch lists

<img width="1365" alt="screen shot 2015-11-13 at 9 57 43 pm" src="https://cloud.githubusercontent.com/assets/1000778/11162322/9b88e204-8a51-11e5-8c57-a44889cab713.png">

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9711 from zsxwing/failure-info.
---
 .../spark/streaming/ui/AllBatchesTable.scala  | 61 +++++++++++++++++--
 .../apache/spark/streaming/ui/BatchPage.scala | 49 ++-------------
 .../apache/spark/streaming/ui/UIUtils.scala   | 60 ++++++++++++++++++
 3 files changed, 120 insertions(+), 50 deletions(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/ui/AllBatchesTable.scala b/streaming/src/main/scala/org/apache/spark/streaming/ui/AllBatchesTable.scala
index 125cafd41b8af..d33972342731d 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/ui/AllBatchesTable.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/ui/AllBatchesTable.scala
@@ -33,6 +33,22 @@ private[ui] abstract class BatchTableBase(tableId: String, batchInterval: Long)
         {SparkUIUtils.tooltip("Time taken to process all jobs of a batch", "top")}</th>
   }
 
+  /**
+   * Return the first failure reason if finding in the batches.
+   */
+  protected def getFirstFailureReason(batches: Seq[BatchUIData]): Option[String] = {
+    batches.flatMap(_.outputOperations.flatMap(_._2.failureReason)).headOption
+  }
+
+  protected def getFirstFailureTableCell(batch: BatchUIData): Seq[Node] = {
+    val firstFailureReason = batch.outputOperations.flatMap(_._2.failureReason).headOption
+    firstFailureReason.map { failureReason =>
+      val failureReasonForUI = UIUtils.createOutputOperationFailureForUI(failureReason)
+      UIUtils.failureReasonCell(
+        failureReasonForUI, rowspan = 1, includeFirstLineInExpandDetails = false)
+    }.getOrElse(<td>-</td>)
+  }
+
   protected def baseRow(batch: BatchUIData): Seq[Node] = {
     val batchTime = batch.batchTime.milliseconds
     val formattedBatchTime = UIUtils.formatBatchTime(batchTime, batchInterval)
@@ -97,9 +113,17 @@ private[ui] class ActiveBatchTable(
     waitingBatches: Seq[BatchUIData],
     batchInterval: Long) extends BatchTableBase("active-batches-table", batchInterval) {
 
+  private val firstFailureReason = getFirstFailureReason(runningBatches)
+
   override protected def columns: Seq[Node] = super.columns ++ {
     <th>Output Ops: Succeeded/Total</th>
-      <th>Status</th>
+      <th>Status</th> ++ {
+      if (firstFailureReason.nonEmpty) {
+        <th>Error</th>
+      } else {
+        Nil
+      }
+    }
   }
 
   override protected def renderRows: Seq[Node] = {
@@ -110,20 +134,41 @@ private[ui] class ActiveBatchTable(
   }
 
   private def runningBatchRow(batch: BatchUIData): Seq[Node] = {
-    baseRow(batch) ++ createOutputOperationProgressBar(batch) ++ <td>processing</td>
+    baseRow(batch) ++ createOutputOperationProgressBar(batch) ++ <td>processing</td> ++ {
+      if (firstFailureReason.nonEmpty) {
+        getFirstFailureTableCell(batch)
+      } else {
+        Nil
+      }
+    }
   }
 
   private def waitingBatchRow(batch: BatchUIData): Seq[Node] = {
-    baseRow(batch) ++ createOutputOperationProgressBar(batch) ++ <td>queued</td>
+    baseRow(batch) ++ createOutputOperationProgressBar(batch) ++ <td>queued</td>++ {
+      if (firstFailureReason.nonEmpty) {
+        // Waiting batches have not run yet, so must have no failure reasons.
+        <td>-</td>
+      } else {
+        Nil
+      }
+    }
   }
 }
 
 private[ui] class CompletedBatchTable(batches: Seq[BatchUIData], batchInterval: Long)
   extends BatchTableBase("completed-batches-table", batchInterval) {
 
+  private val firstFailureReason = getFirstFailureReason(batches)
+
   override protected def columns: Seq[Node] = super.columns ++ {
     <th>Total Delay {SparkUIUtils.tooltip("Total time taken to handle a batch", "top")}</th>
-      <th>Output Ops: Succeeded/Total</th>
+      <th>Output Ops: Succeeded/Total</th> ++ {
+      if (firstFailureReason.nonEmpty) {
+        <th>Error</th>
+      } else {
+        Nil
+      }
+    }
   }
 
   override protected def renderRows: Seq[Node] = {
@@ -138,6 +183,12 @@ private[ui] class CompletedBatchTable(batches: Seq[BatchUIData], batchInterval:
       <td sorttable_customkey={totalDelay.getOrElse(Long.MaxValue).toString}>
         {formattedTotalDelay}
       </td>
-    } ++ createOutputOperationProgressBar(batch)
+    } ++ createOutputOperationProgressBar(batch)++ {
+      if (firstFailureReason.nonEmpty) {
+        getFirstFailureTableCell(batch)
+      } else {
+        Nil
+      }
+    }
   }
 }
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/ui/BatchPage.scala b/streaming/src/main/scala/org/apache/spark/streaming/ui/BatchPage.scala
index 2ed925572826e..bc1711930d3ac 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/ui/BatchPage.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/ui/BatchPage.scala
@@ -149,7 +149,7 @@ private[ui] class BatchPage(parent: StreamingTab) extends WebUIPage("batch") {
             total = sparkJob.numTasks - sparkJob.numSkippedTasks)
         }
       </td>
-      {failureReasonCell(lastFailureReason, rowspan = 1)}
+      {UIUtils.failureReasonCell(lastFailureReason)}
     </tr>
   }
 
@@ -245,48 +245,6 @@ private[ui] class BatchPage(parent: StreamingTab) extends WebUIPage("batch") {
     </div>
   }
 
-  private def failureReasonCell(
-      failureReason: String,
-      rowspan: Int,
-      includeFirstLineInExpandDetails: Boolean = true): Seq[Node] = {
-    val isMultiline = failureReason.indexOf('\n') >= 0
-    // Display the first line by default
-    val failureReasonSummary = StringEscapeUtils.escapeHtml4(
-      if (isMultiline) {
-        failureReason.substring(0, failureReason.indexOf('\n'))
-      } else {
-        failureReason
-      })
-    val failureDetails =
-      if (isMultiline && !includeFirstLineInExpandDetails) {
-        // Skip the first line
-        failureReason.substring(failureReason.indexOf('\n') + 1)
-      } else {
-        failureReason
-      }
-    val details = if (isMultiline) {
-      // scalastyle:off
-      <span onclick="this.parentNode.querySelector('.stacktrace-details').classList.toggle('collapsed')"
-            class="expand-details">
-        +details
-      </span> ++
-        <div class="stacktrace-details collapsed">
-          <pre>{failureDetails}</pre>
-        </div>
-      // scalastyle:on
-    } else {
-      ""
-    }
-
-    if (rowspan == 1) {
-      <td valign="middle" style="max-width: 300px">{failureReasonSummary}{details}</td>
-    } else {
-      <td valign="middle" style="max-width: 300px" rowspan={rowspan.toString}>
-        {failureReasonSummary}{details}
-      </td>
-    }
-  }
-
   private def getJobData(sparkJobId: SparkJobId): Option[JobUIData] = {
     sparkListener.activeJobs.get(sparkJobId).orElse {
       sparkListener.completedJobs.find(_.jobId == sparkJobId).orElse {
@@ -434,8 +392,9 @@ private[ui] class BatchPage(parent: StreamingTab) extends WebUIPage("batch") {
   private def outputOpStatusCell(outputOp: OutputOperationUIData, rowspan: Int): Seq[Node] = {
     outputOp.failureReason match {
       case Some(failureReason) =>
-        val failureReasonForUI = generateOutputOperationStatusForUI(failureReason)
-        failureReasonCell(failureReasonForUI, rowspan, includeFirstLineInExpandDetails = false)
+        val failureReasonForUI = UIUtils.createOutputOperationFailureForUI(failureReason)
+        UIUtils.failureReasonCell(
+          failureReasonForUI, rowspan, includeFirstLineInExpandDetails = false)
       case None =>
         if (outputOp.endTime.isEmpty) {
           <td rowspan={rowspan.toString}>-</td>
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/ui/UIUtils.scala b/streaming/src/main/scala/org/apache/spark/streaming/ui/UIUtils.scala
index 86cfb1fa47370..d89f7ad3e16b7 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/ui/UIUtils.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/ui/UIUtils.scala
@@ -17,6 +17,10 @@
 
 package org.apache.spark.streaming.ui
 
+import scala.xml.Node
+
+import org.apache.commons.lang3.StringEscapeUtils
+
 import java.text.SimpleDateFormat
 import java.util.TimeZone
 import java.util.concurrent.TimeUnit
@@ -124,4 +128,60 @@ private[streaming] object UIUtils {
       }
     }
   }
+
+  def createOutputOperationFailureForUI(failure: String): String = {
+    if (failure.startsWith("org.apache.spark.Spark")) {
+      // SparkException or SparkDriverExecutionException
+      "Failed due to Spark job error\n" + failure
+    } else {
+      var nextLineIndex = failure.indexOf("\n")
+      if (nextLineIndex < 0) {
+        nextLineIndex = failure.size
+      }
+      val firstLine = failure.substring(0, nextLineIndex)
+      s"Failed due to error: $firstLine\n$failure"
+    }
+  }
+
+  def failureReasonCell(
+      failureReason: String,
+      rowspan: Int = 1,
+      includeFirstLineInExpandDetails: Boolean = true): Seq[Node] = {
+    val isMultiline = failureReason.indexOf('\n') >= 0
+    // Display the first line by default
+    val failureReasonSummary = StringEscapeUtils.escapeHtml4(
+      if (isMultiline) {
+        failureReason.substring(0, failureReason.indexOf('\n'))
+      } else {
+        failureReason
+      })
+    val failureDetails =
+      if (isMultiline && !includeFirstLineInExpandDetails) {
+        // Skip the first line
+        failureReason.substring(failureReason.indexOf('\n') + 1)
+      } else {
+        failureReason
+      }
+    val details = if (isMultiline) {
+      // scalastyle:off
+      <span onclick="this.parentNode.querySelector('.stacktrace-details').classList.toggle('collapsed')"
+            class="expand-details">
+        +details
+      </span> ++
+        <div class="stacktrace-details collapsed">
+          <pre>{failureDetails}</pre>
+        </div>
+      // scalastyle:on
+    } else {
+      ""
+    }
+
+    if (rowspan == 1) {
+      <td valign="middle" style="max-width: 300px">{failureReasonSummary}{details}</td>
+    } else {
+      <td valign="middle" style="max-width: 300px" rowspan={rowspan.toString}>
+        {failureReasonSummary}{details}
+      </td>
+    }
+  }
 }

From 31296628ac7cd7be71e0edca335dc8604f62bb47 Mon Sep 17 00:00:00 2001
From: Bartlomiej Alberski <bartlomiej.alberski@allegrogroup.com>
Date: Mon, 16 Nov 2015 15:14:38 -0800
Subject: [PATCH 267/518] [SPARK-11553][SQL] Primitive Row accessors should not
 convert null to default value

Invocation of getters for type extending AnyVal returns default value (if field value is null) instead of throwing NPE. Please check comments for SPARK-11553 issue for more details.

Author: Bartlomiej Alberski <bartlomiej.alberski@allegrogroup.com>

Closes #9642 from alberskib/bugfix/SPARK-11553.
---
 .../main/scala/org/apache/spark/sql/Row.scala | 32 ++++++++++++-----
 .../scala/org/apache/spark/sql/RowTest.scala  | 20 +++++++++++
 .../local/NestedLoopJoinNodeSuite.scala       | 36 +++++++++++--------
 3 files changed, 65 insertions(+), 23 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Row.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Row.scala
index 0f0f200122c34..b14c66cc5ac88 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Row.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Row.scala
@@ -191,7 +191,7 @@ trait Row extends Serializable {
    * @throws ClassCastException when data type does not match.
    * @throws NullPointerException when value is null.
    */
-  def getBoolean(i: Int): Boolean = getAs[Boolean](i)
+  def getBoolean(i: Int): Boolean = getAnyValAs[Boolean](i)
 
   /**
    * Returns the value at position i as a primitive byte.
@@ -199,7 +199,7 @@ trait Row extends Serializable {
    * @throws ClassCastException when data type does not match.
    * @throws NullPointerException when value is null.
    */
-  def getByte(i: Int): Byte = getAs[Byte](i)
+  def getByte(i: Int): Byte = getAnyValAs[Byte](i)
 
   /**
    * Returns the value at position i as a primitive short.
@@ -207,7 +207,7 @@ trait Row extends Serializable {
    * @throws ClassCastException when data type does not match.
    * @throws NullPointerException when value is null.
    */
-  def getShort(i: Int): Short = getAs[Short](i)
+  def getShort(i: Int): Short = getAnyValAs[Short](i)
 
   /**
    * Returns the value at position i as a primitive int.
@@ -215,7 +215,7 @@ trait Row extends Serializable {
    * @throws ClassCastException when data type does not match.
    * @throws NullPointerException when value is null.
    */
-  def getInt(i: Int): Int = getAs[Int](i)
+  def getInt(i: Int): Int = getAnyValAs[Int](i)
 
   /**
    * Returns the value at position i as a primitive long.
@@ -223,7 +223,7 @@ trait Row extends Serializable {
    * @throws ClassCastException when data type does not match.
    * @throws NullPointerException when value is null.
    */
-  def getLong(i: Int): Long = getAs[Long](i)
+  def getLong(i: Int): Long = getAnyValAs[Long](i)
 
   /**
    * Returns the value at position i as a primitive float.
@@ -232,7 +232,7 @@ trait Row extends Serializable {
    * @throws ClassCastException when data type does not match.
    * @throws NullPointerException when value is null.
    */
-  def getFloat(i: Int): Float = getAs[Float](i)
+  def getFloat(i: Int): Float = getAnyValAs[Float](i)
 
   /**
    * Returns the value at position i as a primitive double.
@@ -240,13 +240,12 @@ trait Row extends Serializable {
    * @throws ClassCastException when data type does not match.
    * @throws NullPointerException when value is null.
    */
-  def getDouble(i: Int): Double = getAs[Double](i)
+  def getDouble(i: Int): Double = getAnyValAs[Double](i)
 
   /**
    * Returns the value at position i as a String object.
    *
    * @throws ClassCastException when data type does not match.
-   * @throws NullPointerException when value is null.
    */
   def getString(i: Int): String = getAs[String](i)
 
@@ -318,6 +317,8 @@ trait Row extends Serializable {
 
   /**
    * Returns the value at position i.
+   * For primitive types if value is null it returns 'zero value' specific for primitive
+   * ie. 0 for Int - use isNullAt to ensure that value is not null
    *
    * @throws ClassCastException when data type does not match.
    */
@@ -325,6 +326,8 @@ trait Row extends Serializable {
 
   /**
    * Returns the value of a given fieldName.
+   * For primitive types if value is null it returns 'zero value' specific for primitive
+   * ie. 0 for Int - use isNullAt to ensure that value is not null
    *
    * @throws UnsupportedOperationException when schema is not defined.
    * @throws IllegalArgumentException when fieldName do not exist.
@@ -344,6 +347,8 @@ trait Row extends Serializable {
 
   /**
    * Returns a Map(name -> value) for the requested fieldNames
+   * For primitive types if value is null it returns 'zero value' specific for primitive
+   * ie. 0 for Int - use isNullAt to ensure that value is not null
    *
    * @throws UnsupportedOperationException when schema is not defined.
    * @throws IllegalArgumentException when fieldName do not exist.
@@ -458,4 +463,15 @@ trait Row extends Serializable {
    * start, end, and separator strings.
    */
   def mkString(start: String, sep: String, end: String): String = toSeq.mkString(start, sep, end)
+
+  /**
+   * Returns the value of a given fieldName.
+   *
+   * @throws UnsupportedOperationException when schema is not defined.
+   * @throws ClassCastException when data type does not match.
+   * @throws NullPointerException when value is null.
+   */
+  private def getAnyValAs[T <: AnyVal](i: Int): T =
+    if (isNullAt(i)) throw new NullPointerException(s"Value at index $i in null")
+    else getAs[T](i)
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/RowTest.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/RowTest.scala
index 01ff84cb56054..5c22a72192541 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/RowTest.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/RowTest.scala
@@ -29,8 +29,10 @@ class RowTest extends FunSpec with Matchers {
     StructField("col2", StringType) ::
     StructField("col3", IntegerType) :: Nil)
   val values = Array("value1", "value2", 1)
+  val valuesWithoutCol3 = Array[Any](null, "value2", null)
 
   val sampleRow: Row = new GenericRowWithSchema(values, schema)
+  val sampleRowWithoutCol3: Row = new GenericRowWithSchema(valuesWithoutCol3, schema)
   val noSchemaRow: Row = new GenericRow(values)
 
   describe("Row (without schema)") {
@@ -68,6 +70,24 @@ class RowTest extends FunSpec with Matchers {
       )
       sampleRow.getValuesMap(List("col1", "col2")) shouldBe expected
     }
+
+    it("getValuesMap() retrieves null value on non AnyVal Type") {
+      val expected = Map(
+        "col1" -> null,
+        "col2" -> "value2"
+      )
+      sampleRowWithoutCol3.getValuesMap[String](List("col1", "col2")) shouldBe expected
+    }
+
+    it("getAs() on type extending AnyVal throws an exception when accessing field that is null") {
+      intercept[NullPointerException] {
+        sampleRowWithoutCol3.getInt(sampleRowWithoutCol3.fieldIndex("col3"))
+      }
+    }
+
+    it("getAs() on type extending AnyVal does not throw exception when value is null"){
+      sampleRowWithoutCol3.getAs[String](sampleRowWithoutCol3.fieldIndex("col1")) shouldBe null
+    }
   }
 
   describe("row equals") {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/local/NestedLoopJoinNodeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/local/NestedLoopJoinNodeSuite.scala
index 252f7cc8971f4..45df2ea6552d8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/local/NestedLoopJoinNodeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/local/NestedLoopJoinNodeSuite.scala
@@ -58,8 +58,14 @@ class NestedLoopJoinNodeSuite extends LocalNodeTest {
       val hashJoinNode = makeUnsafeNode(leftNode, rightNode)
       val expectedOutput = generateExpectedOutput(leftInput, rightInput, joinType)
       val actualOutput = hashJoinNode.collect().map { row =>
-        // (id, name, id, nickname)
-        (row.getInt(0), row.getString(1), row.getInt(2), row.getString(3))
+        // (
+        //   id, name,
+        //   id, nickname
+        // )
+        (
+          Option(row.get(0)).map(_.asInstanceOf[Int]), Option(row.getString(1)),
+          Option(row.get(2)).map(_.asInstanceOf[Int]), Option(row.getString(3))
+        )
       }
       assert(actualOutput.toSet === expectedOutput.toSet)
     }
@@ -95,36 +101,36 @@ class NestedLoopJoinNodeSuite extends LocalNodeTest {
   private def generateExpectedOutput(
       leftInput: Array[(Int, String)],
       rightInput: Array[(Int, String)],
-      joinType: JoinType): Array[(Int, String, Int, String)] = {
+      joinType: JoinType): Array[(Option[Int], Option[String], Option[Int], Option[String])] = {
     joinType match {
       case LeftOuter =>
         val rightInputMap = rightInput.toMap
         leftInput.map { case (k, v) =>
-          val rightKey = rightInputMap.get(k).map { _ => k }.getOrElse(0)
-          val rightValue = rightInputMap.getOrElse(k, null)
-          (k, v, rightKey, rightValue)
+          val rightKey = rightInputMap.get(k).map { _ => k }
+          val rightValue = rightInputMap.get(k)
+          (Some(k), Some(v), rightKey, rightValue)
         }
 
       case RightOuter =>
         val leftInputMap = leftInput.toMap
         rightInput.map { case (k, v) =>
-          val leftKey = leftInputMap.get(k).map { _ => k }.getOrElse(0)
-          val leftValue = leftInputMap.getOrElse(k, null)
-          (leftKey, leftValue, k, v)
+          val leftKey = leftInputMap.get(k).map { _ => k }
+          val leftValue = leftInputMap.get(k)
+          (leftKey, leftValue, Some(k), Some(v))
         }
 
       case FullOuter =>
         val leftInputMap = leftInput.toMap
         val rightInputMap = rightInput.toMap
         val leftOutput = leftInput.map { case (k, v) =>
-          val rightKey = rightInputMap.get(k).map { _ => k }.getOrElse(0)
-          val rightValue = rightInputMap.getOrElse(k, null)
-          (k, v, rightKey, rightValue)
+          val rightKey = rightInputMap.get(k).map { _ => k }
+          val rightValue = rightInputMap.get(k)
+          (Some(k), Some(v), rightKey, rightValue)
         }
         val rightOutput = rightInput.map { case (k, v) =>
-          val leftKey = leftInputMap.get(k).map { _ => k }.getOrElse(0)
-          val leftValue = leftInputMap.getOrElse(k, null)
-          (leftKey, leftValue, k, v)
+          val leftKey = leftInputMap.get(k).map { _ => k }
+          val leftValue = leftInputMap.get(k)
+          (leftKey, leftValue, Some(k), Some(v))
         }
         (leftOutput ++ rightOutput).distinct
 

From 75ee12f09c2645c1ad682764d512965f641eb5c2 Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Mon, 16 Nov 2015 15:22:12 -0800
Subject: [PATCH 268/518] [SPARK-8658][SQL] AttributeReference's equals method
 compares all the members

This fix is to change the equals method to check all of the specified fields for equality of AttributeReference.

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9216 from gatorsmile/namedExpressEqual.
---
 .../sql/catalyst/expressions/namedExpressions.scala  |  4 +++-
 .../sql/catalyst/plans/logical/basicOperators.scala  | 10 +++++-----
 .../sql/catalyst/plans/physical/partitioning.scala   | 12 ++++++------
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
index f80bcfcb0b0bf..e3daddace241d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
@@ -194,7 +194,9 @@ case class AttributeReference(
   def sameRef(other: AttributeReference): Boolean = this.exprId == other.exprId
 
   override def equals(other: Any): Boolean = other match {
-    case ar: AttributeReference => name == ar.name && exprId == ar.exprId && dataType == ar.dataType
+    case ar: AttributeReference =>
+      name == ar.name && dataType == ar.dataType && nullable == ar.nullable &&
+        metadata == ar.metadata && exprId == ar.exprId && qualifiers == ar.qualifiers
     case _ => false
   }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index e2b97b27a6c2c..45630a591d349 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.types._
-import org.apache.spark.util.collection.OpenHashSet
+import scala.collection.mutable.ArrayBuffer
 
 case class Project(projectList: Seq[NamedExpression], child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = projectList.map(_.toAttribute)
@@ -244,12 +244,12 @@ private[sql] object Expand {
    */
   private def buildNonSelectExprSet(
       bitmask: Int,
-      exprs: Seq[Expression]): OpenHashSet[Expression] = {
-    val set = new OpenHashSet[Expression](2)
+      exprs: Seq[Expression]): ArrayBuffer[Expression] = {
+    val set = new ArrayBuffer[Expression](2)
 
     var bit = exprs.length - 1
     while (bit >= 0) {
-      if (((bitmask >> bit) & 1) == 0) set.add(exprs(bit))
+      if (((bitmask >> bit) & 1) == 0) set += exprs(bit)
       bit -= 1
     }
 
@@ -279,7 +279,7 @@ private[sql] object Expand {
 
       (child.output :+ gid).map(expr => expr transformDown {
         // TODO this causes a problem when a column is used both for grouping and aggregation.
-        case x: Expression if nonSelectedGroupExprSet.contains(x) =>
+        case x: Expression if nonSelectedGroupExprSet.exists(_.semanticEquals(x)) =>
           // if the input attribute in the Invalid Grouping Expression set of for this group
           // replace it with constant null
           Literal.create(null, expr.dataType)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index 86b9417477ba3..f6fb31a2af594 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -235,17 +235,17 @@ case class HashPartitioning(expressions: Seq[Expression], numPartitions: Int)
   override def satisfies(required: Distribution): Boolean = required match {
     case UnspecifiedDistribution => true
     case ClusteredDistribution(requiredClustering) =>
-      expressions.toSet.subsetOf(requiredClustering.toSet)
+      expressions.forall(x => requiredClustering.exists(_.semanticEquals(x)))
     case _ => false
   }
 
   override def compatibleWith(other: Partitioning): Boolean = other match {
-    case o: HashPartitioning => this == o
+    case o: HashPartitioning => this.semanticEquals(o)
     case _ => false
   }
 
   override def guarantees(other: Partitioning): Boolean = other match {
-    case o: HashPartitioning => this == o
+    case o: HashPartitioning => this.semanticEquals(o)
     case _ => false
   }
 
@@ -276,17 +276,17 @@ case class RangePartitioning(ordering: Seq[SortOrder], numPartitions: Int)
       val minSize = Seq(requiredOrdering.size, ordering.size).min
       requiredOrdering.take(minSize) == ordering.take(minSize)
     case ClusteredDistribution(requiredClustering) =>
-      ordering.map(_.child).toSet.subsetOf(requiredClustering.toSet)
+      ordering.map(_.child).forall(x => requiredClustering.exists(_.semanticEquals(x)))
     case _ => false
   }
 
   override def compatibleWith(other: Partitioning): Boolean = other match {
-    case o: RangePartitioning => this == o
+    case o: RangePartitioning => this.semanticEquals(o)
     case _ => false
   }
 
   override def guarantees(other: Partitioning): Boolean = other match {
-    case o: RangePartitioning => this == o
+    case o: RangePartitioning => this.semanticEquals(o)
     case _ => false
   }
 }

From fd14936be7beff543dbbcf270f2f9749f7a803c4 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Mon, 16 Nov 2015 15:32:49 -0800
Subject: [PATCH 269/518] [SPARK-11625][SQL] add java test for typed aggregate

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9591 from cloud-fan/agg-test.
---
 .../spark/api/java/function/Function.java     |  2 +-
 .../org/apache/spark/sql/GroupedDataset.scala | 34 ++++++++++-
 .../spark/sql/expressions/Aggregator.scala    |  2 +-
 .../apache/spark/sql/JavaDatasetSuite.java    | 56 +++++++++++++++++++
 .../spark/sql/DatasetAggregatorSuite.scala    |  7 +--
 5 files changed, 92 insertions(+), 9 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/api/java/function/Function.java b/core/src/main/java/org/apache/spark/api/java/function/Function.java
index d00551bb0add6..b9d9777a75651 100644
--- a/core/src/main/java/org/apache/spark/api/java/function/Function.java
+++ b/core/src/main/java/org/apache/spark/api/java/function/Function.java
@@ -25,5 +25,5 @@
  * when mapping RDDs of other types.
  */
 public interface Function<T1, R> extends Serializable {
-  public R call(T1 v1) throws Exception;
+  R call(T1 v1) throws Exception;
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index ebcf4c8bfe7e6..467cd42b9b8dc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -145,9 +145,37 @@ class GroupedDataset[K, T] private[sql](
     reduce(f.call _)
   }
 
-  // To ensure valid overloading.
-  protected def agg(expr: Column, exprs: Column*): DataFrame =
-    groupedData.agg(expr, exprs: _*)
+  /**
+   * Compute aggregates by specifying a series of aggregate columns, and return a [[DataFrame]].
+   * We can call `as[T : Encoder]` to turn the returned [[DataFrame]] to [[Dataset]] again.
+   *
+   * The available aggregate methods are defined in [[org.apache.spark.sql.functions]].
+   *
+   * {{{
+   *   // Selects the age of the oldest employee and the aggregate expense for each department
+   *
+   *   // Scala:
+   *   import org.apache.spark.sql.functions._
+   *   df.groupBy("department").agg(max("age"), sum("expense"))
+   *
+   *   // Java:
+   *   import static org.apache.spark.sql.functions.*;
+   *   df.groupBy("department").agg(max("age"), sum("expense"));
+   * }}}
+   *
+   * We can also use `Aggregator.toColumn` to pass in typed aggregate functions.
+   *
+   * @since 1.6.0
+   */
+  @scala.annotation.varargs
+  def agg(expr: Column, exprs: Column*): DataFrame =
+    groupedData.agg(withEncoder(expr), exprs.map(withEncoder): _*)
+
+  private def withEncoder(c: Column): Column = c match {
+    case tc: TypedColumn[_, _] =>
+      tc.withInputType(resolvedTEncoder.bind(dataAttributes), dataAttributes)
+    case _ => c
+  }
 
   /**
    * Internal helper function for building typed aggregations that return tuples.  For simplicity
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
index 360c9a5bc15e7..72610e735f782 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
@@ -47,7 +47,7 @@ import org.apache.spark.sql.{Dataset, DataFrame, TypedColumn}
  * @tparam B The type of the intermediate value of the reduction.
  * @tparam C The type of the final result.
  */
-abstract class Aggregator[-A, B, C] {
+abstract class Aggregator[-A, B, C] extends Serializable {
 
   /** A zero value for this aggregation. Should satisfy the property that any b + zero = b */
   def zero: B
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index eb6fa1e72e27b..d9b22506fbd3b 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -34,6 +34,7 @@
 import org.apache.spark.sql.Encoders;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.GroupedDataset;
+import org.apache.spark.sql.expressions.Aggregator;
 import org.apache.spark.sql.test.TestSQLContext;
 
 import static org.apache.spark.sql.functions.*;
@@ -381,4 +382,59 @@ public void testNestedTupleEncoder() {
       context.createDataset(data3, encoder3);
     Assert.assertEquals(data3, ds3.collectAsList());
   }
+
+  @Test
+  public void testTypedAggregation() {
+    Encoder<Tuple2<String, Integer>> encoder = Encoders.tuple(Encoders.STRING(), Encoders.INT());
+    List<Tuple2<String, Integer>> data =
+      Arrays.asList(tuple2("a", 1), tuple2("a", 2), tuple2("b", 3));
+    Dataset<Tuple2<String, Integer>> ds = context.createDataset(data, encoder);
+
+    GroupedDataset<String, Tuple2<String, Integer>> grouped = ds.groupBy(
+      new MapFunction<Tuple2<String, Integer>, String>() {
+        @Override
+        public String call(Tuple2<String, Integer> value) throws Exception {
+          return value._1();
+        }
+      },
+      Encoders.STRING());
+
+    Dataset<Tuple2<String, Integer>> agged =
+      grouped.agg(new IntSumOf().toColumn(Encoders.INT(), Encoders.INT()));
+    Assert.assertEquals(Arrays.asList(tuple2("a", 3), tuple2("b", 3)), agged.collectAsList());
+
+    Dataset<Tuple4<String, Integer, Long, Long>> agged2 = grouped.agg(
+      new IntSumOf().toColumn(Encoders.INT(), Encoders.INT()),
+      expr("sum(_2)"),
+      count("*"))
+      .as(Encoders.tuple(Encoders.STRING(), Encoders.INT(), Encoders.LONG(), Encoders.LONG()));
+    Assert.assertEquals(
+      Arrays.asList(
+        new Tuple4<String, Integer, Long, Long>("a", 3, 3L, 2L),
+        new Tuple4<String, Integer, Long, Long>("b", 3, 3L, 1L)),
+      agged2.collectAsList());
+  }
+
+  static class IntSumOf extends Aggregator<Tuple2<String, Integer>, Integer, Integer> {
+
+    @Override
+    public Integer zero() {
+      return 0;
+    }
+
+    @Override
+    public Integer reduce(Integer l, Tuple2<String, Integer> t) {
+      return l + t._2();
+    }
+
+    @Override
+    public Integer merge(Integer b1, Integer b2) {
+      return b1 + b2;
+    }
+
+    @Override
+    public Integer finish(Integer reduction) {
+      return reduction;
+    }
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
index 46f9f077fe7f2..9377589790011 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -25,7 +25,7 @@ import org.apache.spark.sql.functions._
 import org.apache.spark.sql.expressions.Aggregator
 
 /** An `Aggregator` that adds up any numeric type returned by the given function. */
-class SumOf[I, N : Numeric](f: I => N) extends Aggregator[I, N, N] with Serializable {
+class SumOf[I, N : Numeric](f: I => N) extends Aggregator[I, N, N] {
   val numeric = implicitly[Numeric[N]]
 
   override def zero: N = numeric.zero
@@ -37,7 +37,7 @@ class SumOf[I, N : Numeric](f: I => N) extends Aggregator[I, N, N] with Serializ
   override def finish(reduction: N): N = reduction
 }
 
-object TypedAverage extends Aggregator[(String, Int), (Long, Long), Double] with Serializable {
+object TypedAverage extends Aggregator[(String, Int), (Long, Long), Double] {
   override def zero: (Long, Long) = (0, 0)
 
   override def reduce(countAndSum: (Long, Long), input: (String, Int)): (Long, Long) = {
@@ -51,8 +51,7 @@ object TypedAverage extends Aggregator[(String, Int), (Long, Long), Double] with
   override def finish(countAndSum: (Long, Long)): Double = countAndSum._2 / countAndSum._1
 }
 
-object ComplexResultAgg extends Aggregator[(String, Int), (Long, Long), (Long, Long)]
-  with Serializable {
+object ComplexResultAgg extends Aggregator[(String, Int), (Long, Long), (Long, Long)] {
 
   override def zero: (Long, Long) = (0, 0)
 

From ea6f53e48a911b49dc175ccaac8c80e0a1d97a09 Mon Sep 17 00:00:00 2001
From: Shivaram Venkataraman <shivaram@cs.berkeley.edu>
Date: Mon, 16 Nov 2015 16:57:50 -0800
Subject: [PATCH 270/518] [SPARKR][HOTFIX] Disable flaky SparkR package build
 test

See https://github.com/apache/spark/pull/9390#issuecomment-157160063 and https://gist.github.com/shivaram/3a2fecce60768a603dac for more information

Author: Shivaram Venkataraman <shivaram@cs.berkeley.edu>

Closes #9744 from shivaram/sparkr-package-test-disable.
---
 .../test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
index 42e748ec6d528..d494b0caab85f 100644
--- a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
@@ -369,7 +369,9 @@ class SparkSubmitSuite
     }
   }
 
-  test("correctly builds R packages included in a jar with --packages") {
+  // TODO(SPARK-9603): Building a package is flaky on Jenkins Maven builds.
+  // See https://gist.github.com/shivaram/3a2fecce60768a603dac for a error log
+  ignore("correctly builds R packages included in a jar with --packages") {
     assume(RUtils.isRInstalled, "R isn't installed on this machine.")
     val main = MavenCoordinate("my.great.lib", "mylib", "0.1")
     val sparkHome = sys.props.getOrElse("spark.test.home", fail("spark.test.home is not set!"))

From 30f3cfda1cce8760f15c0a48a8c47f09a5167fca Mon Sep 17 00:00:00 2001
From: Kousuke Saruta <sarutak@oss.nttdata.co.jp>
Date: Mon, 16 Nov 2015 16:59:16 -0800
Subject: [PATCH 271/518] [SPARK-11480][CORE][WEBUI] Wrong callsite is
 displayed when using AsyncRDDActions#takeAsync

When we call AsyncRDDActions#takeAsync, actually another DAGScheduler#runJob is called from another thread so we cannot get proper callsite infomation.

Following screenshots are before this patch applied and after.

Before:
<img width="1268" alt="2015-11-04 1 26 40" src="https://cloud.githubusercontent.com/assets/4736016/10914069/0ffc1306-8294-11e5-8e89-c4fadf58dd12.png">
<img width="1258" alt="2015-11-04 1 26 52" src="https://cloud.githubusercontent.com/assets/4736016/10914070/0ffe84ce-8294-11e5-8b2a-69d36276bedb.png">

After:
<img width="1268" alt="2015-11-04 0 48 07" src="https://cloud.githubusercontent.com/assets/4736016/10914080/1d8cfb7a-8294-11e5-9e09-ede25c2563e8.png">
<img width="1269" alt="2015-11-04 0 48 26" src="https://cloud.githubusercontent.com/assets/4736016/10914081/1d934e3a-8294-11e5-8b5e-e3dc37aaced3.png">

Author: Kousuke Saruta <sarutak@oss.nttdata.co.jp>

Closes #9437 from sarutak/SPARK-11480.
---
 core/src/main/scala/org/apache/spark/rdd/AsyncRDDActions.scala | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/core/src/main/scala/org/apache/spark/rdd/AsyncRDDActions.scala b/core/src/main/scala/org/apache/spark/rdd/AsyncRDDActions.scala
index ca1eb1f4e4a9a..d5e853613b05b 100644
--- a/core/src/main/scala/org/apache/spark/rdd/AsyncRDDActions.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/AsyncRDDActions.scala
@@ -66,6 +66,7 @@ class AsyncRDDActions[T: ClassTag](self: RDD[T]) extends Serializable with Loggi
    */
   def takeAsync(num: Int): FutureAction[Seq[T]] = self.withScope {
     val f = new ComplexFutureAction[Seq[T]]
+    val callSite = self.context.getCallSite
 
     f.run {
       // This is a blocking action so we should use "AsyncRDDActions.futureExecutionContext" which
@@ -73,6 +74,7 @@ class AsyncRDDActions[T: ClassTag](self: RDD[T]) extends Serializable with Loggi
       val results = new ArrayBuffer[T](num)
       val totalParts = self.partitions.length
       var partsScanned = 0
+      self.context.setCallSite(callSite)
       while (results.size < num && partsScanned < totalParts) {
         // The number of partitions to try in this iteration. It is ok for this number to be
         // greater than totalParts because we actually cap it at totalParts in runJob.

From 33a0ec93771ef5c3b388165b07cfab9014918d3b Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Mon, 16 Nov 2015 17:00:18 -0800
Subject: [PATCH 272/518] [SPARK-11710] Document new memory management model

Author: Andrew Or <andrew@databricks.com>

Closes #9676 from andrewor14/memory-management-docs.
---
 docs/configuration.md | 13 +++++++----
 docs/tuning.md        | 54 ++++++++++++++++++++++++++++---------------
 2 files changed, 44 insertions(+), 23 deletions(-)

diff --git a/docs/configuration.md b/docs/configuration.md
index d961f43acf4ab..c496146e3ed63 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -722,17 +722,20 @@ Apart from these, the following properties are also available, and may be useful
     Fraction of the heap space used for execution and storage. The lower this is, the more
     frequently spills and cached data eviction occur. The purpose of this config is to set
     aside memory for internal metadata, user data structures, and imprecise size estimation
-    in the case of sparse, unusually large records.
+    in the case of sparse, unusually large records. Leaving this at the default value is
+    recommended. For more detail, see <a href="tuning.html#memory-management-overview">
+    this description</a>.
   </td>
 </tr>
 <tr>
   <td><code>spark.memory.storageFraction</code></td>
   <td>0.5</td>
   <td>
-    T​he size of the storage region within the space set aside by
-    <code>s​park.memory.fraction</code>. This region is not statically reserved, but dynamically
-    allocated as cache requests come in. ​Cached data may be evicted only if total storage exceeds
-    this region.
+    Amount of storage memory immune to eviction, expressed as a fraction of the size of the
+    region set aside by <code>s​park.memory.fraction</code>. The higher this is, the less
+    working memory may be available to execution and tasks may spill to disk more often.
+    Leaving this at the default value is recommended. For more detail, see
+    <a href="tuning.html#memory-management-overview">this description</a>.
   </td>
 </tr>
 <tr>
diff --git a/docs/tuning.md b/docs/tuning.md
index 879340a01544f..a8fe7a4532798 100644
--- a/docs/tuning.md
+++ b/docs/tuning.md
@@ -88,9 +88,39 @@ than the "raw" data inside their fields. This is due to several reasons:
   but also pointers (typically 8 bytes each) to the next object in the list.
 * Collections of primitive types often store them as "boxed" objects such as `java.lang.Integer`.
 
-This section will discuss how to determine the memory usage of your objects, and how to improve
-it -- either by changing your data structures, or by storing data in a serialized format.
-We will then cover tuning Spark's cache size and the Java garbage collector.
+This section will start with an overview of memory management in Spark, then discuss specific
+strategies the user can take to make more efficient use of memory in his/her application. In
+particular, we will describe how to determine the memory usage of your objects, and how to
+improve it -- either by changing your data structures, or by storing data in a serialized
+format. We will then cover tuning Spark's cache size and the Java garbage collector.
+
+## Memory Management Overview
+
+Memory usage in Spark largely falls under one of two categories: execution and storage.
+Execution memory refers to that used for computation in shuffles, joins, sorts and aggregations,
+while storage memory refers to that used for caching and propagating internal data across the
+cluster. In Spark, execution and storage share a unified region (M). When no execution memory is
+used, storage can acquire all the available memory and vice versa. Execution may evict storage
+if necessary, but only until total storage memory usage falls under a certain threshold (R).
+In other words, `R` describes a subregion within `M` where cached blocks are never evicted.
+Storage may not evict execution due to complexities in implementation.
+
+This design ensures several desirable properties. First, applications that do not use caching
+can use the entire space for execution, obviating unnecessary disk spills. Second, applications
+that do use caching can reserve a minimum storage space (R) where their data blocks are immune
+to being evicted. Lastly, this approach provides reasonable out-of-the-box performance for a
+variety of workloads without requiring user expertise of how memory is divided internally.
+
+Although there are two relevant configurations, the typical user should not need to adjust them
+as the default values are applicable to most workloads:
+
+* `spark.memory.fraction` expresses the size of `M` as a fraction of the total JVM heap space
+(default 0.75). The rest of the space (25%) is reserved for user data structures, internal
+metadata in Spark, and safeguarding against OOM errors in the case of sparse and unusually
+large records.
+* `spark.memory.storageFraction` expresses the size of `R` as a fraction of `M` (default 0.5).
+`R` is the storage space within `M` where cached blocks immune to being evicted by execution.
+
 
 ## Determining Memory Consumption
 
@@ -151,18 +181,6 @@ time spent GC. This can be done by adding `-verbose:gc -XX:+PrintGCDetails -XX:+
 each time a garbage collection occurs. Note these logs will be on your cluster's worker nodes (in the `stdout` files in
 their work directories), *not* on your driver program.
 
-**Cache Size Tuning**
-
-One important configuration parameter for GC is the amount of memory that should be used for caching RDDs.
-By default, Spark uses 60% of the configured executor memory (`spark.executor.memory`) to
-cache RDDs. This means that 40% of memory is available for any objects created during task execution.
-
-In case your tasks slow down and you find that your JVM is garbage-collecting frequently or running out of
-memory, lowering this value will help reduce the memory consumption. To change this to, say, 50%, you can call
-`conf.set("spark.storage.memoryFraction", "0.5")` on your SparkConf. Combined with the use of serialized caching,
-using a smaller cache should be sufficient to mitigate most of the garbage collection problems.
-In case you are interested in further tuning the Java GC, continue reading below.
-
 **Advanced GC Tuning**
 
 To further tune garbage collection, we first need to understand some basic information about memory management in the JVM:
@@ -183,9 +201,9 @@ temporary objects created during task execution. Some steps which may be useful
 * Check if there are too many garbage collections by collecting GC stats. If a full GC is invoked multiple times for
   before a task completes, it means that there isn't enough memory available for executing tasks.
 
-* In the GC stats that are printed, if the OldGen is close to being full, reduce the amount of memory used for caching.
-  This can be done using the `spark.storage.memoryFraction` property. It is better to cache fewer objects than to slow
-  down task execution!
+* In the GC stats that are printed, if the OldGen is close to being full, reduce the amount of
+  memory used for caching by lowering `spark.memory.storageFraction`; it is better to cache fewer
+  objects than to slow down task execution!
 
 * If there are too many minor collections but not many major GCs, allocating more memory for Eden would help. You
   can set the size of the Eden to be an over-estimate of how much memory each task will need. If the size of Eden

From bd10eb81c98e5e9df453f721943a3e82d9f74ae4 Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Mon, 16 Nov 2015 17:02:21 -0800
Subject: [PATCH 273/518] [EXAMPLE][MINOR] Add missing awaitTermination in
 click stream example

Author: jerryshao <sshao@hortonworks.com>

Closes #9730 from jerryshao/clickstream-fix.
---
 .../spark/examples/streaming/clickstream/PageViewStream.scala | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/src/main/scala/org/apache/spark/examples/streaming/clickstream/PageViewStream.scala b/examples/src/main/scala/org/apache/spark/examples/streaming/clickstream/PageViewStream.scala
index ec7d39da8b2e9..4ef238606f82e 100644
--- a/examples/src/main/scala/org/apache/spark/examples/streaming/clickstream/PageViewStream.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/streaming/clickstream/PageViewStream.scala
@@ -18,7 +18,6 @@
 // scalastyle:off println
 package org.apache.spark.examples.streaming.clickstream
 
-import org.apache.spark.SparkContext._
 import org.apache.spark.streaming.{Seconds, StreamingContext}
 import org.apache.spark.examples.streaming.StreamingExamples
 // scalastyle:off
@@ -88,7 +87,7 @@ object PageViewStream {
 
     // An external dataset we want to join to this stream
     val userList = ssc.sparkContext.parallelize(
-       Map(1 -> "Patrick Wendell", 2->"Reynold Xin", 3->"Matei Zaharia").toSeq)
+       Map(1 -> "Patrick Wendell", 2 -> "Reynold Xin", 3 -> "Matei Zaharia").toSeq)
 
     metric match {
       case "pageCounts" => pageCounts.print()
@@ -106,6 +105,7 @@ object PageViewStream {
     }
 
     ssc.start()
+    ssc.awaitTermination()
   }
 }
 // scalastyle:on println

From 1c5475f1401d2233f4c61f213d1e2c2ee9673067 Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Mon, 16 Nov 2015 17:12:39 -0800
Subject: [PATCH 274/518] [SPARK-11612][ML] Pipeline and PipelineModel
 persistence

Pipeline and PipelineModel extend Readable and Writable.  Persistence succeeds only when all stages are Writable.

Note: This PR reinstates tests for other read/write functionality.  It should probably not get merged until [https://issues.apache.org/jira/browse/SPARK-11672] gets fixed.

CC: mengxr

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9674 from jkbradley/pipeline-io.
---
 .../scala/org/apache/spark/ml/Pipeline.scala  | 175 +++++++++++++++++-
 .../org/apache/spark/ml/util/ReadWrite.scala  |   4 +-
 .../org/apache/spark/ml/PipelineSuite.scala   | 120 +++++++++++-
 .../spark/ml/util/DefaultReadWriteTest.scala  |  25 ++-
 4 files changed, 306 insertions(+), 18 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
index a3e59401c5cfb..25f0c696f42be 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
@@ -22,12 +22,19 @@ import java.{util => ju}
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ListBuffer
 
-import org.apache.spark.Logging
+import org.apache.hadoop.fs.Path
+import org.json4s._
+import org.json4s.jackson.JsonMethods._
+
+import org.apache.spark.{SparkContext, Logging}
 import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.ml.param.{Param, ParamMap, Params}
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util.Reader
+import org.apache.spark.ml.util.Writer
+import org.apache.spark.ml.util._
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.Utils
 
 /**
  * :: DeveloperApi ::
@@ -82,7 +89,7 @@ abstract class PipelineStage extends Params with Logging {
  * an identity transformer.
  */
 @Experimental
-class Pipeline(override val uid: String) extends Estimator[PipelineModel] {
+class Pipeline(override val uid: String) extends Estimator[PipelineModel] with Writable {
 
   def this() = this(Identifiable.randomUID("pipeline"))
 
@@ -166,6 +173,131 @@ class Pipeline(override val uid: String) extends Estimator[PipelineModel] {
       "Cannot have duplicate components in a pipeline.")
     theStages.foldLeft(schema)((cur, stage) => stage.transformSchema(cur))
   }
+
+  override def write: Writer = new Pipeline.PipelineWriter(this)
+}
+
+object Pipeline extends Readable[Pipeline] {
+
+  override def read: Reader[Pipeline] = new PipelineReader
+
+  override def load(path: String): Pipeline = read.load(path)
+
+  private[ml] class PipelineWriter(instance: Pipeline) extends Writer {
+
+    SharedReadWrite.validateStages(instance.getStages)
+
+    override protected def saveImpl(path: String): Unit =
+      SharedReadWrite.saveImpl(instance, instance.getStages, sc, path)
+  }
+
+  private[ml] class PipelineReader extends Reader[Pipeline] {
+
+    /** Checked against metadata when loading model */
+    private val className = "org.apache.spark.ml.Pipeline"
+
+    override def load(path: String): Pipeline = {
+      val (uid: String, stages: Array[PipelineStage]) = SharedReadWrite.load(className, sc, path)
+      new Pipeline(uid).setStages(stages)
+    }
+  }
+
+  /** Methods for [[Reader]] and [[Writer]] shared between [[Pipeline]] and [[PipelineModel]] */
+  private[ml] object SharedReadWrite {
+
+    import org.json4s.JsonDSL._
+
+    /** Check that all stages are Writable */
+    def validateStages(stages: Array[PipelineStage]): Unit = {
+      stages.foreach {
+        case stage: Writable => // good
+        case other =>
+          throw new UnsupportedOperationException("Pipeline write will fail on this Pipeline" +
+            s" because it contains a stage which does not implement Writable. Non-Writable stage:" +
+            s" ${other.uid} of type ${other.getClass}")
+      }
+    }
+
+    /**
+     * Save metadata and stages for a [[Pipeline]] or [[PipelineModel]]
+     *  - save metadata to path/metadata
+     *  - save stages to stages/IDX_UID
+     */
+    def saveImpl(
+        instance: Params,
+        stages: Array[PipelineStage],
+        sc: SparkContext,
+        path: String): Unit = {
+      // Copied and edited from DefaultParamsWriter.saveMetadata
+      // TODO: modify DefaultParamsWriter.saveMetadata to avoid duplication
+      val uid = instance.uid
+      val cls = instance.getClass.getName
+      val stageUids = stages.map(_.uid)
+      val jsonParams = List("stageUids" -> parse(compact(render(stageUids.toSeq))))
+      val metadata = ("class" -> cls) ~
+        ("timestamp" -> System.currentTimeMillis()) ~
+        ("sparkVersion" -> sc.version) ~
+        ("uid" -> uid) ~
+        ("paramMap" -> jsonParams)
+      val metadataPath = new Path(path, "metadata").toString
+      val metadataJson = compact(render(metadata))
+      sc.parallelize(Seq(metadataJson), 1).saveAsTextFile(metadataPath)
+
+      // Save stages
+      val stagesDir = new Path(path, "stages").toString
+      stages.zipWithIndex.foreach { case (stage: Writable, idx: Int) =>
+        stage.write.save(getStagePath(stage.uid, idx, stages.length, stagesDir))
+      }
+    }
+
+    /**
+     * Load metadata and stages for a [[Pipeline]] or [[PipelineModel]]
+     * @return  (UID, list of stages)
+     */
+    def load(
+        expectedClassName: String,
+        sc: SparkContext,
+        path: String): (String, Array[PipelineStage]) = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, expectedClassName)
+
+      implicit val format = DefaultFormats
+      val stagesDir = new Path(path, "stages").toString
+      val stageUids: Array[String] = metadata.params match {
+        case JObject(pairs) =>
+          if (pairs.length != 1) {
+            // Should not happen unless file is corrupted or we have a bug.
+            throw new RuntimeException(
+              s"Pipeline read expected 1 Param (stageUids), but found ${pairs.length}.")
+          }
+          pairs.head match {
+            case ("stageUids", jsonValue) =>
+              jsonValue.extract[Seq[String]].toArray
+            case (paramName, jsonValue) =>
+              // Should not happen unless file is corrupted or we have a bug.
+              throw new RuntimeException(s"Pipeline read encountered unexpected Param $paramName" +
+                s" in metadata: ${metadata.metadataStr}")
+          }
+        case _ =>
+          throw new IllegalArgumentException(
+            s"Cannot recognize JSON metadata: ${metadata.metadataStr}.")
+      }
+      val stages: Array[PipelineStage] = stageUids.zipWithIndex.map { case (stageUid, idx) =>
+        val stagePath = SharedReadWrite.getStagePath(stageUid, idx, stageUids.length, stagesDir)
+        val stageMetadata = DefaultParamsReader.loadMetadata(stagePath, sc)
+        val cls = Utils.classForName(stageMetadata.className)
+        cls.getMethod("read").invoke(null).asInstanceOf[Reader[PipelineStage]].load(stagePath)
+      }
+      (metadata.uid, stages)
+    }
+
+    /** Get path for saving the given stage. */
+    def getStagePath(stageUid: String, stageIdx: Int, numStages: Int, stagesDir: String): String = {
+      val stageIdxDigits = numStages.toString.length
+      val idxFormat = s"%0${stageIdxDigits}d"
+      val stageDir = idxFormat.format(stageIdx) + "_" + stageUid
+      new Path(stagesDir, stageDir).toString
+    }
+  }
 }
 
 /**
@@ -176,7 +308,7 @@ class Pipeline(override val uid: String) extends Estimator[PipelineModel] {
 class PipelineModel private[ml] (
     override val uid: String,
     val stages: Array[Transformer])
-  extends Model[PipelineModel] with Logging {
+  extends Model[PipelineModel] with Writable with Logging {
 
   /** A Java/Python-friendly auxiliary constructor. */
   private[ml] def this(uid: String, stages: ju.List[Transformer]) = {
@@ -200,4 +332,39 @@ class PipelineModel private[ml] (
   override def copy(extra: ParamMap): PipelineModel = {
     new PipelineModel(uid, stages.map(_.copy(extra))).setParent(parent)
   }
+
+  override def write: Writer = new PipelineModel.PipelineModelWriter(this)
+}
+
+object PipelineModel extends Readable[PipelineModel] {
+
+  import Pipeline.SharedReadWrite
+
+  override def read: Reader[PipelineModel] = new PipelineModelReader
+
+  override def load(path: String): PipelineModel = read.load(path)
+
+  private[ml] class PipelineModelWriter(instance: PipelineModel) extends Writer {
+
+    SharedReadWrite.validateStages(instance.stages.asInstanceOf[Array[PipelineStage]])
+
+    override protected def saveImpl(path: String): Unit = SharedReadWrite.saveImpl(instance,
+      instance.stages.asInstanceOf[Array[PipelineStage]], sc, path)
+  }
+
+  private[ml] class PipelineModelReader extends Reader[PipelineModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = "org.apache.spark.ml.PipelineModel"
+
+    override def load(path: String): PipelineModel = {
+      val (uid: String, stages: Array[PipelineStage]) = SharedReadWrite.load(className, sc, path)
+      val transformers = stages map {
+        case stage: Transformer => stage
+        case other => throw new RuntimeException(s"PipelineModel.read loaded a stage but found it" +
+          s" was not a Transformer.  Bad stage ${other.uid} of type ${other.getClass}")
+      }
+      new PipelineModel(uid, transformers)
+    }
+  }
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index ca896ed6106c4..3169c9e9af5be 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -164,6 +164,8 @@ trait Readable[T] {
 
   /**
    * Reads an ML instance from the input path, a shortcut of `read.load(path)`.
+   *
+   * Note: Implementing classes should override this to be Java-friendly.
    */
   @Since("1.6.0")
   def load(path: String): T = read.load(path)
@@ -190,7 +192,7 @@ private[ml] object DefaultParamsWriter {
    *  - timestamp
    *  - sparkVersion
    *  - uid
-   *  - paramMap
+   *  - paramMap: These must be encodable using [[org.apache.spark.ml.param.Param.jsonEncode()]].
    */
   def saveMetadata(instance: Params, path: String, sc: SparkContext): Unit = {
     val uid = instance.uid
diff --git a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
index 1f2c9b75b617b..484026b1ba9ad 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
@@ -17,19 +17,25 @@
 
 package org.apache.spark.ml
 
+import java.io.File
+
 import scala.collection.JavaConverters._
 
+import org.apache.hadoop.fs.{FileSystem, Path}
 import org.mockito.Matchers.{any, eq => meq}
 import org.mockito.Mockito.when
 import org.scalatest.mock.MockitoSugar.mock
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.Pipeline.SharedReadWrite
 import org.apache.spark.ml.feature.HashingTF
-import org.apache.spark.ml.param.ParamMap
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.param.{IntParam, ParamMap}
+import org.apache.spark.ml.util._
+import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.types.StructType
 
-class PipelineSuite extends SparkFunSuite {
+class PipelineSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   abstract class MyModel extends Model[MyModel]
 
@@ -111,4 +117,112 @@ class PipelineSuite extends SparkFunSuite {
     assert(pipelineModel1.uid === "pipeline1")
     assert(pipelineModel1.stages === stages)
   }
+
+  test("Pipeline read/write") {
+    val writableStage = new WritableStage("writableStage").setIntParam(56)
+    val pipeline = new Pipeline().setStages(Array(writableStage))
+
+    val pipeline2 = testDefaultReadWrite(pipeline, testParams = false)
+    assert(pipeline2.getStages.length === 1)
+    assert(pipeline2.getStages(0).isInstanceOf[WritableStage])
+    val writableStage2 = pipeline2.getStages(0).asInstanceOf[WritableStage]
+    assert(writableStage.getIntParam === writableStage2.getIntParam)
+  }
+
+  test("Pipeline read/write with non-Writable stage") {
+    val unWritableStage = new UnWritableStage("unwritableStage")
+    val unWritablePipeline = new Pipeline().setStages(Array(unWritableStage))
+    withClue("Pipeline.write should fail when Pipeline contains non-Writable stage") {
+      intercept[UnsupportedOperationException] {
+        unWritablePipeline.write
+      }
+    }
+  }
+
+  test("PipelineModel read/write") {
+    val writableStage = new WritableStage("writableStage").setIntParam(56)
+    val pipeline =
+      new PipelineModel("pipeline_89329327", Array(writableStage.asInstanceOf[Transformer]))
+
+    val pipeline2 = testDefaultReadWrite(pipeline, testParams = false)
+    assert(pipeline2.stages.length === 1)
+    assert(pipeline2.stages(0).isInstanceOf[WritableStage])
+    val writableStage2 = pipeline2.stages(0).asInstanceOf[WritableStage]
+    assert(writableStage.getIntParam === writableStage2.getIntParam)
+
+    val path = new File(tempDir, pipeline.uid).getPath
+    val stagesDir = new Path(path, "stages").toString
+    val expectedStagePath = SharedReadWrite.getStagePath(writableStage.uid, 0, 1, stagesDir)
+    assert(FileSystem.get(sc.hadoopConfiguration).exists(new Path(expectedStagePath)),
+      s"Expected stage 0 of 1 with uid ${writableStage.uid} in Pipeline with uid ${pipeline.uid}" +
+        s" to be saved to path: $expectedStagePath")
+  }
+
+  test("PipelineModel read/write: getStagePath") {
+    val stageUid = "myStage"
+    val stagesDir = new Path("pipeline", "stages").toString
+    def testStage(stageIdx: Int, numStages: Int, expectedPrefix: String): Unit = {
+      val path = SharedReadWrite.getStagePath(stageUid, stageIdx, numStages, stagesDir)
+      val expected = new Path(stagesDir, expectedPrefix + "_" + stageUid).toString
+      assert(path === expected)
+    }
+    testStage(0, 1, "0")
+    testStage(0, 9, "0")
+    testStage(0, 10, "00")
+    testStage(1, 10, "01")
+    testStage(12, 999, "012")
+  }
+
+  test("PipelineModel read/write with non-Writable stage") {
+    val unWritableStage = new UnWritableStage("unwritableStage")
+    val unWritablePipeline =
+      new PipelineModel("pipeline_328957", Array(unWritableStage.asInstanceOf[Transformer]))
+    withClue("PipelineModel.write should fail when PipelineModel contains non-Writable stage") {
+      intercept[UnsupportedOperationException] {
+        unWritablePipeline.write
+      }
+    }
+  }
+}
+
+
+/** Used to test [[Pipeline]] with [[Writable]] stages */
+class WritableStage(override val uid: String) extends Transformer with Writable {
+
+  final val intParam: IntParam = new IntParam(this, "intParam", "doc")
+
+  def getIntParam: Int = $(intParam)
+
+  def setIntParam(value: Int): this.type = set(intParam, value)
+
+  setDefault(intParam -> 0)
+
+  override def copy(extra: ParamMap): WritableStage = defaultCopy(extra)
+
+  override def write: Writer = new DefaultParamsWriter(this)
+
+  override def transform(dataset: DataFrame): DataFrame = dataset
+
+  override def transformSchema(schema: StructType): StructType = schema
+}
+
+object WritableStage extends Readable[WritableStage] {
+
+  override def read: Reader[WritableStage] = new DefaultParamsReader[WritableStage]
+
+  override def load(path: String): WritableStage = read.load(path)
+}
+
+/** Used to test [[Pipeline]] with non-[[Writable]] stages */
+class UnWritableStage(override val uid: String) extends Transformer {
+
+  final val intParam: IntParam = new IntParam(this, "intParam", "doc")
+
+  setDefault(intParam -> 0)
+
+  override def copy(extra: ParamMap): UnWritableStage = defaultCopy(extra)
+
+  override def transform(dataset: DataFrame): DataFrame = dataset
+
+  override def transformSchema(schema: StructType): StructType = schema
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
index cac4bd9aa3ab8..c37f0503f1332 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
@@ -30,10 +30,13 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
   /**
    * Checks "overwrite" option and params.
    * @param instance ML instance to test saving/loading
+   * @param testParams  If true, then test values of Params.  Otherwise, just test overwrite option.
    * @tparam T ML instance type
    * @return  Instance loaded from file
    */
-  def testDefaultReadWrite[T <: Params with Writable](instance: T): T = {
+  def testDefaultReadWrite[T <: Params with Writable](
+      instance: T,
+      testParams: Boolean = true): T = {
     val uid = instance.uid
     val path = new File(tempDir, uid).getPath
 
@@ -46,16 +49,18 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
     val newInstance = loader.load(path)
 
     assert(newInstance.uid === instance.uid)
-    instance.params.foreach { p =>
-      if (instance.isDefined(p)) {
-        (instance.getOrDefault(p), newInstance.getOrDefault(p)) match {
-          case (Array(values), Array(newValues)) =>
-            assert(values === newValues, s"Values do not match on param ${p.name}.")
-          case (value, newValue) =>
-            assert(value === newValue, s"Values do not match on param ${p.name}.")
+    if (testParams) {
+      instance.params.foreach { p =>
+        if (instance.isDefined(p)) {
+          (instance.getOrDefault(p), newInstance.getOrDefault(p)) match {
+            case (Array(values), Array(newValues)) =>
+              assert(values === newValues, s"Values do not match on param ${p.name}.")
+            case (value, newValue) =>
+              assert(value === newValue, s"Values do not match on param ${p.name}.")
+          }
+        } else {
+          assert(!newInstance.isDefined(p), s"Param ${p.name} shouldn't be defined.")
         }
-      } else {
-        assert(!newInstance.isDefined(p), s"Param ${p.name} shouldn't be defined.")
       }
     }
 

From 540bf58f18328c68107d6c616ffd70f3a4640054 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Mon, 16 Nov 2015 17:28:11 -0800
Subject: [PATCH 275/518] [SPARK-11617][NETWORK] Fix leak in
 TransportFrameDecoder.

The code was using the wrong API to add data to the internal composite
buffer, causing buffers to leak in certain situations. Use the right
API and enhance the tests to catch memory leaks.

Also, avoid reusing the composite buffers when downstream handlers keep
references to them; this seems to cause a few different issues even though
the ref counting code seems to be correct, so instead pay the cost of copying
a few bytes when that situation happens.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9619 from vanzin/SPARK-11617.
---
 .../network/util/TransportFrameDecoder.java   |  47 ++++--
 .../util/TransportFrameDecoderSuite.java      | 145 +++++++++++++++---
 2 files changed, 151 insertions(+), 41 deletions(-)

diff --git a/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java b/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
index 272ea84e6180d..5889562dd9705 100644
--- a/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
+++ b/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
@@ -56,32 +56,43 @@ public void channelRead(ChannelHandlerContext ctx, Object data) throws Exception
       buffer = in.alloc().compositeBuffer();
     }
 
-    buffer.writeBytes(in);
+    buffer.addComponent(in).writerIndex(buffer.writerIndex() + in.readableBytes());
 
     while (buffer.isReadable()) {
-      feedInterceptor();
-      if (interceptor != null) {
-        continue;
-      }
+      discardReadBytes();
+      if (!feedInterceptor()) {
+        ByteBuf frame = decodeNext();
+        if (frame == null) {
+          break;
+        }
 
-      ByteBuf frame = decodeNext();
-      if (frame != null) {
         ctx.fireChannelRead(frame);
-      } else {
-        break;
       }
     }
 
-    // We can't discard read sub-buffers if there are other references to the buffer (e.g.
-    // through slices used for framing). This assumes that code that retains references
-    // will call retain() from the thread that called "fireChannelRead()" above, otherwise
-    // ref counting will go awry.
-    if (buffer != null && buffer.refCnt() == 1) {
+    discardReadBytes();
+  }
+
+  private void discardReadBytes() {
+    // If the buffer's been retained by downstream code, then make a copy of the remaining
+    // bytes into a new buffer. Otherwise, just discard stale components.
+    if (buffer.refCnt() > 1) {
+      CompositeByteBuf newBuffer = buffer.alloc().compositeBuffer();
+
+      if (buffer.readableBytes() > 0) {
+        ByteBuf spillBuf = buffer.alloc().buffer(buffer.readableBytes());
+        spillBuf.writeBytes(buffer);
+        newBuffer.addComponent(spillBuf).writerIndex(spillBuf.readableBytes());
+      }
+
+      buffer.release();
+      buffer = newBuffer;
+    } else {
       buffer.discardReadComponents();
     }
   }
 
-  protected ByteBuf decodeNext() throws Exception {
+  private ByteBuf decodeNext() throws Exception {
     if (buffer.readableBytes() < LENGTH_SIZE) {
       return null;
     }
@@ -127,10 +138,14 @@ public void setInterceptor(Interceptor interceptor) {
     this.interceptor = interceptor;
   }
 
-  private void feedInterceptor() throws Exception {
+  /**
+   * @return Whether the interceptor is still active after processing the data.
+   */
+  private boolean feedInterceptor() throws Exception {
     if (interceptor != null && !interceptor.handle(buffer)) {
       interceptor = null;
     }
+    return interceptor != null;
   }
 
   public static interface Interceptor {
diff --git a/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java b/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java
index ca74f0a00cf9d..19475c21ffce9 100644
--- a/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java
@@ -18,41 +18,36 @@
 package org.apache.spark.network.util;
 
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelHandlerContext;
+import org.junit.AfterClass;
 import org.junit.Test;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
 
 public class TransportFrameDecoderSuite {
 
+  private static Random RND = new Random();
+
+  @AfterClass
+  public static void cleanup() {
+    RND = null;
+  }
+
   @Test
   public void testFrameDecoding() throws Exception {
-    Random rnd = new Random();
     TransportFrameDecoder decoder = new TransportFrameDecoder();
-    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
-
-    final int frameCount = 100;
-    ByteBuf data = Unpooled.buffer();
-    try {
-      for (int i = 0; i < frameCount; i++) {
-        byte[] frame = new byte[1024 * (rnd.nextInt(31) + 1)];
-        data.writeLong(frame.length + 8);
-        data.writeBytes(frame);
-      }
-
-      while (data.isReadable()) {
-        int size = rnd.nextInt(16 * 1024) + 256;
-        decoder.channelRead(ctx, data.readSlice(Math.min(data.readableBytes(), size)));
-      }
-
-      verify(ctx, times(frameCount)).fireChannelRead(any(ByteBuf.class));
-    } finally {
-      data.release();
-    }
+    ChannelHandlerContext ctx = mockChannelHandlerContext();
+    ByteBuf data = createAndFeedFrames(100, decoder, ctx);
+    verifyAndCloseDecoder(decoder, ctx, data);
   }
 
   @Test
@@ -60,7 +55,7 @@ public void testInterception() throws Exception {
     final int interceptedReads = 3;
     TransportFrameDecoder decoder = new TransportFrameDecoder();
     TransportFrameDecoder.Interceptor interceptor = spy(new MockInterceptor(interceptedReads));
-    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+    ChannelHandlerContext ctx = mockChannelHandlerContext();
 
     byte[] data = new byte[8];
     ByteBuf len = Unpooled.copyLong(8 + data.length);
@@ -70,16 +65,56 @@ public void testInterception() throws Exception {
       decoder.setInterceptor(interceptor);
       for (int i = 0; i < interceptedReads; i++) {
         decoder.channelRead(ctx, dataBuf);
-        dataBuf.release();
+        assertEquals(0, dataBuf.refCnt());
         dataBuf = Unpooled.wrappedBuffer(data);
       }
       decoder.channelRead(ctx, len);
       decoder.channelRead(ctx, dataBuf);
       verify(interceptor, times(interceptedReads)).handle(any(ByteBuf.class));
       verify(ctx).fireChannelRead(any(ByteBuffer.class));
+      assertEquals(0, len.refCnt());
+      assertEquals(0, dataBuf.refCnt());
     } finally {
-      len.release();
-      dataBuf.release();
+      release(len);
+      release(dataBuf);
+    }
+  }
+
+  @Test
+  public void testRetainedFrames() throws Exception {
+    TransportFrameDecoder decoder = new TransportFrameDecoder();
+
+    final AtomicInteger count = new AtomicInteger();
+    final List<ByteBuf> retained = new ArrayList<>();
+
+    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+    when(ctx.fireChannelRead(any())).thenAnswer(new Answer<Void>() {
+      @Override
+      public Void answer(InvocationOnMock in) {
+        // Retain a few frames but not others.
+        ByteBuf buf = (ByteBuf) in.getArguments()[0];
+        if (count.incrementAndGet() % 2 == 0) {
+          retained.add(buf);
+        } else {
+          buf.release();
+        }
+        return null;
+      }
+    });
+
+    ByteBuf data = createAndFeedFrames(100, decoder, ctx);
+    try {
+      // Verify all retained buffers are readable.
+      for (ByteBuf b : retained) {
+        byte[] tmp = new byte[b.readableBytes()];
+        b.readBytes(tmp);
+        b.release();
+      }
+      verifyAndCloseDecoder(decoder, ctx, data);
+    } finally {
+      for (ByteBuf b : retained) {
+        release(b);
+      }
     }
   }
 
@@ -100,6 +135,47 @@ public void testLargeFrame() throws Exception {
     testInvalidFrame(Integer.MAX_VALUE + 9);
   }
 
+  /**
+   * Creates a number of randomly sized frames and feed them to the given decoder, verifying
+   * that the frames were read.
+   */
+  private ByteBuf createAndFeedFrames(
+      int frameCount,
+      TransportFrameDecoder decoder,
+      ChannelHandlerContext ctx) throws Exception {
+    ByteBuf data = Unpooled.buffer();
+    for (int i = 0; i < frameCount; i++) {
+      byte[] frame = new byte[1024 * (RND.nextInt(31) + 1)];
+      data.writeLong(frame.length + 8);
+      data.writeBytes(frame);
+    }
+
+    try {
+      while (data.isReadable()) {
+        int size = RND.nextInt(4 * 1024) + 256;
+        decoder.channelRead(ctx, data.readSlice(Math.min(data.readableBytes(), size)).retain());
+      }
+
+      verify(ctx, times(frameCount)).fireChannelRead(any(ByteBuf.class));
+    } catch (Exception e) {
+      release(data);
+      throw e;
+    }
+    return data;
+  }
+
+  private void verifyAndCloseDecoder(
+      TransportFrameDecoder decoder,
+      ChannelHandlerContext ctx,
+      ByteBuf data) throws Exception {
+    try {
+      decoder.channelInactive(ctx);
+      assertTrue("There shouldn't be dangling references to the data.", data.release());
+    } finally {
+      release(data);
+    }
+  }
+
   private void testInvalidFrame(long size) throws Exception {
     TransportFrameDecoder decoder = new TransportFrameDecoder();
     ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
@@ -111,6 +187,25 @@ private void testInvalidFrame(long size) throws Exception {
     }
   }
 
+  private ChannelHandlerContext mockChannelHandlerContext() {
+    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+    when(ctx.fireChannelRead(any())).thenAnswer(new Answer<Void>() {
+      @Override
+      public Void answer(InvocationOnMock in) {
+        ByteBuf buf = (ByteBuf) in.getArguments()[0];
+        buf.release();
+        return null;
+      }
+    });
+    return ctx;
+  }
+
+  private void release(ByteBuf buf) {
+    if (buf.refCnt() > 0) {
+      buf.release(buf.refCnt());
+    }
+  }
+
   private static class MockInterceptor implements TransportFrameDecoder.Interceptor {
 
     private int remainingReads;

From fbad920dbfd6f389dea852cdc159cacb0f4f6997 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Mon, 16 Nov 2015 20:47:46 -0800
Subject: [PATCH 276/518] [SPARK-11768][SPARK-9196][SQL] Support now function
 in SQL (alias for current_timestamp).

This patch adds an alias for current_timestamp (now function).

Also fixes SPARK-9196 to re-enable the test case for current_timestamp.

Author: Reynold Xin <rxin@databricks.com>

Closes #9753 from rxin/SPARK-11768.
---
 .../catalyst/analysis/FunctionRegistry.scala   |  1 +
 .../apache/spark/sql/DateFunctionsSuite.scala  | 18 ++++++++++++------
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index a8f4d257acd0a..f9c04d7ec0b0c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -244,6 +244,7 @@ object FunctionRegistry {
     expression[AddMonths]("add_months"),
     expression[CurrentDate]("current_date"),
     expression[CurrentTimestamp]("current_timestamp"),
+    expression[CurrentTimestamp]("now"),
     expression[DateDiff]("datediff"),
     expression[DateAdd]("date_add"),
     expression[DateFormatClass]("date_format"),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
index 1266d534cc5b3..241cbd0115070 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
@@ -38,15 +38,21 @@ class DateFunctionsSuite extends QueryTest with SharedSQLContext {
     assert(d0 <= d1 && d1 <= d2 && d2 <= d3 && d3 - d0 <= 1)
   }
 
-  // This is a bad test. SPARK-9196 will fix it and re-enable it.
-  ignore("function current_timestamp") {
+  test("function current_timestamp and now") {
     val df1 = Seq((1, 2), (3, 1)).toDF("a", "b")
     checkAnswer(df1.select(countDistinct(current_timestamp())), Row(1))
+
     // Execution in one query should return the same value
-    checkAnswer(sql("""SELECT CURRENT_TIMESTAMP() = CURRENT_TIMESTAMP()"""),
-      Row(true))
-    assert(math.abs(sql("""SELECT CURRENT_TIMESTAMP()""").collect().head.getTimestamp(
-      0).getTime - System.currentTimeMillis()) < 5000)
+    checkAnswer(sql("""SELECT CURRENT_TIMESTAMP() = CURRENT_TIMESTAMP()"""), Row(true))
+
+    // Current timestamp should return the current timestamp ...
+    val before = System.currentTimeMillis
+    val got = sql("SELECT CURRENT_TIMESTAMP()").collect().head.getTimestamp(0).getTime
+    val after = System.currentTimeMillis
+    assert(got >= before && got <= after)
+
+    // Now alias
+    checkAnswer(sql("""SELECT CURRENT_TIMESTAMP() = NOW()"""), Row(true))
   }
 
   val sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss")

From 75d202073143d5a7f943890d8682b5b0cf9e3092 Mon Sep 17 00:00:00 2001
From: hyukjinkwon <gurwls223@gmail.com>
Date: Tue, 17 Nov 2015 14:35:00 +0800
Subject: [PATCH 277/518] [SPARK-11694][FOLLOW-UP] Clean up imports, use a
 common function for metadata and add a test for FIXED_LEN_BYTE_ARRAY

As discussed https://github.com/apache/spark/pull/9660 https://github.com/apache/spark/pull/9060, I cleaned up unused imports, added a test for fixed-length byte array and used a common function for writing metadata for Parquet.

For the test for fixed-length byte array, I have tested and checked the encoding types with [parquet-tools](https://github.com/Parquet/parquet-mr/tree/master/parquet-tools).

Author: hyukjinkwon <gurwls223@gmail.com>

Closes #9754 from HyukjinKwon/SPARK-11694-followup.
---
 .../test/resources/dec-in-fixed-len.parquet   | Bin 0 -> 460 bytes
 .../datasources/parquet/ParquetIOSuite.scala  |  42 +++++++-----------
 2 files changed, 15 insertions(+), 27 deletions(-)
 create mode 100644 sql/core/src/test/resources/dec-in-fixed-len.parquet

diff --git a/sql/core/src/test/resources/dec-in-fixed-len.parquet b/sql/core/src/test/resources/dec-in-fixed-len.parquet
new file mode 100644
index 0000000000000000000000000000000000000000..6ad37d5639511cdb430f33fa6165eb70cd9034c0
GIT binary patch
literal 460
zcmZuu%SyvQ6rI#oO3~7Vw<M577-(rBhCY%QsvCDMM2ff&RLXQ_qS2&}JbX~P^8;M^
zF@As_;wSV2ypakj4s*|)nRD+w2WEJ9+Jzzr5XZ0s<i`^E`5i+sb9l*Ng-|%p;SFKK
z$l)0_K?U^^3?LXaoJLHwgp{&cr#fDUp%hEC-r4;t0u(g!Xe*hq_->T&8yPVu5U4^0
z3hK5WJW4SN<zyrzZ&YyJ4yhZOuIbIP4ppsDBif%Tz=HaRg1cEh_(wbE45I03dm3L`
zY{pR2`jWSjd1#@Jc_vdUo<3NSOqpOQQ*ReEeVR!+ph+5MJf$Y}0~v~h4lWjS!lIn|
z<rOQ|9cVm=SlH@z%w8LxCs`P>aflr_N}BXrgbo(V<<*j?`)dokQEarvSr7`N-{ZFH
k+I`P<B=gu_$KicP*phobcO9E!_Pwqv+T=tbWQhOx35yqQV*mgE

literal 0
HcmV?d00001

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index a148facd056a5..177ab42f7767c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql.execution.datasources.parquet
 
-import java.util.Collections
-
 import org.apache.parquet.column.{Encoding, ParquetProperties}
 
 import scala.collection.JavaConverters._
@@ -33,7 +31,7 @@ import org.apache.parquet.example.data.{Group, GroupWriter}
 import org.apache.parquet.hadoop._
 import org.apache.parquet.hadoop.api.WriteSupport
 import org.apache.parquet.hadoop.api.WriteSupport.WriteContext
-import org.apache.parquet.hadoop.metadata.{CompressionCodecName, FileMetaData, ParquetMetadata}
+import org.apache.parquet.hadoop.metadata.CompressionCodecName
 import org.apache.parquet.io.api.RecordConsumer
 import org.apache.parquet.schema.{MessageType, MessageTypeParser}
 
@@ -243,15 +241,9 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
       """.stripMargin)
 
     withTempPath { location =>
-      val extraMetadata = Map.empty[String, String].asJava
-      val fileMetadata = new FileMetaData(parquetSchema, extraMetadata, "Spark")
       val path = new Path(location.getCanonicalPath)
-      val footer = List(
-        new Footer(path, new ParquetMetadata(fileMetadata, Collections.emptyList()))
-      ).asJava
-
-      ParquetFileWriter.writeMetadataFile(sparkContext.hadoopConfiguration, path, footer)
-
+      val conf = sparkContext.hadoopConfiguration
+      writeMetadata(parquetSchema, path, conf)
       val errorMessage = intercept[Throwable] {
         sqlContext.read.parquet(path.toString).printSchema()
       }.toString
@@ -267,20 +259,14 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
         |}
       """.stripMargin)
 
+    val expectedSparkTypes = Seq(StringType, BinaryType)
+
     withTempPath { location =>
-      val extraMetadata = Map.empty[String, String].asJava
-      val fileMetadata = new FileMetaData(parquetSchema, extraMetadata, "Spark")
       val path = new Path(location.getCanonicalPath)
-      val footer = List(
-        new Footer(path, new ParquetMetadata(fileMetadata, Collections.emptyList()))
-      ).asJava
-
-      ParquetFileWriter.writeMetadataFile(sparkContext.hadoopConfiguration, path, footer)
-
-      val jsonDataType = sqlContext.read.parquet(path.toString).schema(0).dataType
-      assert(jsonDataType === StringType)
-      val bsonDataType = sqlContext.read.parquet(path.toString).schema(1).dataType
-      assert(bsonDataType === BinaryType)
+      val conf = sparkContext.hadoopConfiguration
+      writeMetadata(parquetSchema, path, conf)
+      val sparkTypes = sqlContext.read.parquet(path.toString).schema.map(_.dataType)
+      assert(sparkTypes === expectedSparkTypes)
     }
   }
 
@@ -607,10 +593,12 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
       sqlContext.range(1 << 4).select('id % 10 cast DecimalType(10, 2) as 'i64_dec))
   }
 
-  // TODO Adds test case for reading dictionary encoded decimals written as `FIXED_LEN_BYTE_ARRAY`
-  // The Parquet writer version Spark 1.6 and prior versions use is `PARQUET_1_0`, which doesn't
-  // provide dictionary encoding support for `FIXED_LEN_BYTE_ARRAY`.  Should add a test here once
-  // we upgrade to `PARQUET_2_0`.
+  test("read dictionary encoded decimals written as FIXED_LEN_BYTE_ARRAY") {
+    checkAnswer(
+      // Decimal column in this file is encoded using plain dictionary
+      readResourceParquetFile("dec-in-fixed-len.parquet"),
+      sqlContext.range(1 << 4).select('id % 10 cast DecimalType(10, 2) as 'fixed_len_dec))
+  }
 }
 
 class JobCommitFailureParquetOutputCommitter(outputPath: Path, context: TaskAttemptContext)

From e01865af0d5ebe11033de46c388c5c583876c187 Mon Sep 17 00:00:00 2001
From: Kevin Yu <qyu@us.ibm.com>
Date: Mon, 16 Nov 2015 22:54:29 -0800
Subject: [PATCH 278/518] [SPARK-11447][SQL] change NullType to StringType
 during binaryComparison between NullType and StringType

During executing PromoteStrings rule, if one side of binaryComparison is StringType and the other side is not StringType, the current code will promote(cast) the StringType to DoubleType, and if the StringType doesn't contain the numbers, it will get null value. So if it is doing <=> (NULL-safe equal) with Null, it will not filter anything, caused the problem reported by this jira.

I proposal to the changes through this PR, can you review my code changes ?

This problem only happen for <=>, other operators works fine.

scala> val filteredDF = df.filter(df("column") > (new Column(Literal(null))))
filteredDF: org.apache.spark.sql.DataFrame = [column: string]

scala> filteredDF.show
+------+
|column|
+------+
+------+

scala> val filteredDF = df.filter(df("column") === (new Column(Literal(null))))
filteredDF: org.apache.spark.sql.DataFrame = [column: string]

scala> filteredDF.show
+------+
|column|
+------+
+------+

scala> df.registerTempTable("DF")

scala> sqlContext.sql("select * from DF where 'column' = NULL")
res27: org.apache.spark.sql.DataFrame = [column: string]

scala> res27.show
+------+
|column|
+------+
+------+

Author: Kevin Yu <qyu@us.ibm.com>

Closes #9720 from kevinyu98/working_on_spark-11447.
---
 .../sql/catalyst/analysis/HiveTypeCoercion.scala      |  6 ++++++
 .../org/apache/spark/sql/ColumnExpressionSuite.scala  | 11 +++++++++++
 2 files changed, 17 insertions(+)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
index 92188ee54fd28..f90fc3cc12189 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercion.scala
@@ -281,6 +281,12 @@ object HiveTypeCoercion {
       case p @ BinaryComparison(left @ DateType(), right @ TimestampType()) =>
         p.makeCopy(Array(Cast(left, StringType), Cast(right, StringType)))
 
+      // Checking NullType
+      case p @ BinaryComparison(left @ StringType(), right @ NullType()) =>
+        p.makeCopy(Array(left, Literal.create(null, StringType)))
+      case p @ BinaryComparison(left @ NullType(), right @ StringType()) =>
+        p.makeCopy(Array(Literal.create(null, StringType), right))
+
       case p @ BinaryComparison(left @ StringType(), right) if right.dataType != StringType =>
         p.makeCopy(Array(Cast(left, DoubleType), right))
       case p @ BinaryComparison(left, right @ StringType()) if left.dataType != StringType =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
index 3eae3f6d85066..38c0eb589f965 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala
@@ -368,6 +368,17 @@ class ColumnExpressionSuite extends QueryTest with SharedSQLContext {
     checkAnswer(
       nullData.filter($"a" <=> $"b"),
       Row(1, 1) :: Row(null, null) :: Nil)
+
+    val nullData2 = sqlContext.createDataFrame(sparkContext.parallelize(
+        Row("abc") ::
+        Row(null)  ::
+        Row("xyz") :: Nil),
+        StructType(Seq(StructField("a", StringType, true))))
+
+    checkAnswer(
+      nullData2.filter($"a" <=> null),
+      Row(null) :: Nil)
+
   }
 
   test(">") {

From d79d8b08ff69b30b02fe87839e695e29bfea5ace Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Mon, 16 Nov 2015 23:16:17 -0800
Subject: [PATCH 279/518] [MINOR] [SQL] Fix randomly generated ArrayData in
 RowEncoderSuite

The randomly generated ArrayData used for the UDT `ExamplePoint` in `RowEncoderSuite` sometimes doesn't have enough elements. In this case, this test will fail. This patch is to fix it.

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #9757 from viirya/fix-randomgenerated-udt.
---
 .../spark/sql/catalyst/encoders/RowEncoderSuite.scala    | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
index c868ddec1bab2..46c6e0d98d349 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst.encoders
 
+import scala.util.Random
+
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.{RandomDataGenerator, Row}
 import org.apache.spark.sql.catalyst.util.{GenericArrayData, ArrayData}
@@ -59,7 +61,12 @@ class ExamplePointUDT extends UserDefinedType[ExamplePoint] {
   override def deserialize(datum: Any): ExamplePoint = {
     datum match {
       case values: ArrayData =>
-        new ExamplePoint(values.getDouble(0), values.getDouble(1))
+        if (values.numElements() > 1) {
+          new ExamplePoint(values.getDouble(0), values.getDouble(1))
+        } else {
+          val random = new Random()
+          new ExamplePoint(random.nextDouble(), random.nextDouble())
+        }
     }
   }
 

From fa13301ae440c4c9594280f236bcca11b62fdd29 Mon Sep 17 00:00:00 2001
From: Cheng Lian <lian@databricks.com>
Date: Tue, 17 Nov 2015 18:11:08 +0800
Subject: [PATCH 280/518] [SPARK-11191][SQL][FOLLOW-UP] Cleans up unnecessary
 anonymous HiveFunctionRegistry

According to discussion in PR #9664, the anonymous `HiveFunctionRegistry` in `HiveContext` can be removed now.

Author: Cheng Lian <lian@databricks.com>

Closes #9737 from liancheng/spark-11191.follow-up.
---
 .../scala/org/apache/spark/sql/hive/HiveContext.scala  | 10 +---------
 .../scala/org/apache/spark/sql/hive/hiveUDFs.scala     |  7 +++++--
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index 0c473799cc991..2004f24ad26c6 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -454,15 +454,7 @@ class HiveContext private[hive](
   // Note that HiveUDFs will be overridden by functions registered in this context.
   @transient
   override protected[sql] lazy val functionRegistry: FunctionRegistry =
-    new HiveFunctionRegistry(FunctionRegistry.builtin.copy(), this) {
-      override def lookupFunction(name: String, children: Seq[Expression]): Expression = {
-        // Hive Registry need current database to lookup function
-        // TODO: the current database of executionHive should be consistent with metadataHive
-        executionHive.withHiveState {
-          super.lookupFunction(name, children)
-        }
-      }
-    }
+    new HiveFunctionRegistry(FunctionRegistry.builtin.copy(), this.executionHive)
 
   // The Hive UDF current_database() is foldable, will be evaluated by optimizer, but the optimizer
   // can't access the SessionState of metadataHive.
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
index e6fe2ad5f23b6..2e8c026259efe 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFs.scala
@@ -43,16 +43,19 @@ import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.ArrayData
 import org.apache.spark.sql.hive.HiveShim._
+import org.apache.spark.sql.hive.client.ClientWrapper
 import org.apache.spark.sql.types._
 
 
 private[hive] class HiveFunctionRegistry(
     underlying: analysis.FunctionRegistry,
-    hiveContext: HiveContext)
+    executionHive: ClientWrapper)
   extends analysis.FunctionRegistry with HiveInspectors {
 
   def getFunctionInfo(name: String): FunctionInfo = {
-    hiveContext.executionHive.withHiveState {
+    // Hive Registry need current database to lookup function
+    // TODO: the current database of executionHive should be consistent with metadataHive
+    executionHive.withHiveState {
       FunctionRegistry.getFunctionInfo(name)
     }
   }

From 7276fa9aa9d2eccb6aebd5c690ac334699142f1e Mon Sep 17 00:00:00 2001
From: "yangping.wu" <wyphao.2007@163.com>
Date: Tue, 17 Nov 2015 14:11:34 +0000
Subject: [PATCH 281/518] [SPARK-11751] Doc describe error in the "Spark
 Streaming Programming Guide" page

In the **[Task Launching Overheads](http://spark.apache.org/docs/latest/streaming-programming-guide.html#task-launching-overheads)** section,
>Task Serialization: Using Kryo serialization for serializing tasks can reduce the task sizes, and therefore reduce the time taken to send them to the slaves.

as we known **Task Serialization** is configuration by **spark.closure.serializer** parameter, but currently only the Java serializer is supported. If we set **spark.closure.serializer** to **org.apache.spark.serializer.KryoSerializer**, then this will throw a exception.

Author: yangping.wu <wyphao.2007@163.com>

Closes #9734 from 397090770/397090770-patch-1.
---
 docs/streaming-programming-guide.md | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/docs/streaming-programming-guide.md b/docs/streaming-programming-guide.md
index e9a27f446a898..96b36b7a73209 100644
--- a/docs/streaming-programming-guide.md
+++ b/docs/streaming-programming-guide.md
@@ -2001,8 +2001,7 @@ If the number of tasks launched per second is high (say, 50 or more per second),
 of sending out tasks to the slaves may be significant and will make it hard to achieve sub-second
 latencies. The overhead can be reduced by the following changes:
 
-* **Task Serialization**: Using Kryo serialization for serializing tasks can reduce the task
-  sizes, and therefore reduce the time taken to send them to the slaves.
+* **Task Serialization**: Using Kryo serialization for serializing tasks can reduce the task sizes, and therefore reduce the time taken to send them to the slaves. This is controlled by the ```spark.closure.serializer``` property. However, at this time, Kryo serialization cannot be enabled for closure serialization. This may be resolved in a future release.
 
 * **Execution mode**: Running Spark in Standalone mode or coarse-grained Mesos mode leads to
   better task launch times than the fine-grained Mesos mode. Please refer to the

From 15cc36b7786e2d9a460bf565893236edd2ad993e Mon Sep 17 00:00:00 2001
From: Philipp Hoffmann <mail@philipphoffmann.de>
Date: Tue, 17 Nov 2015 14:13:13 +0000
Subject: [PATCH 282/518] [SPARK-11779][DOCS] Fix reference to deprecated
 MESOS_NATIVE_LIBRARY

MESOS_NATIVE_LIBRARY was renamed in favor of MESOS_NATIVE_JAVA_LIBRARY. This commit fixes the reference in the documentation.

Author: Philipp Hoffmann <mail@philipphoffmann.de>

Closes #9768 from philipphoffmann/patch-2.
---
 docs/running-on-mesos.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/running-on-mesos.md b/docs/running-on-mesos.md
index ec5a44d79212b..5be208cf3461e 100644
--- a/docs/running-on-mesos.md
+++ b/docs/running-on-mesos.md
@@ -278,7 +278,7 @@ See the [configuration page](configuration.html) for information on Spark config
     Set the name of the docker image that the Spark executors will run in. The selected
     image must have Spark installed, as well as a compatible version of the Mesos library.
     The installed path of Spark in the image can be specified with <code>spark.mesos.executor.home</code>;
-    the installed path of the Mesos library can be specified with <code>spark.executorEnv.MESOS_NATIVE_LIBRARY</code>.
+    the installed path of the Mesos library can be specified with <code>spark.executorEnv.MESOS_NATIVE_JAVA_LIBRARY</code>.
   </td>
 </tr>
 <tr>

From 6fc2740ebb59aca1aa0ee1e93658a7e4e69de33c Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Tue, 17 Nov 2015 10:01:33 -0800
Subject: [PATCH 283/518] [SPARK-11744][LAUNCHER] Fix print version throw
 exception when using pyspark shell

Exception details can be seen here (https://issues.apache.org/jira/browse/SPARK-11744).

Author: jerryshao <sshao@hortonworks.com>

Closes #9721 from jerryshao/SPARK-11744.
---
 .../launcher/SparkSubmitCommandBuilder.java     | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java b/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java
index 39b46e0db8cc2..312df0b269f32 100644
--- a/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java
+++ b/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java
@@ -77,7 +77,7 @@ class SparkSubmitCommandBuilder extends AbstractCommandBuilder {
   }
 
   final List<String> sparkArgs;
-  private final boolean printHelp;
+  private final boolean printInfo;
 
   /**
    * Controls whether mixing spark-submit arguments with app arguments is allowed. This is needed
@@ -88,7 +88,7 @@ class SparkSubmitCommandBuilder extends AbstractCommandBuilder {
 
   SparkSubmitCommandBuilder() {
     this.sparkArgs = new ArrayList<String>();
-    this.printHelp = false;
+    this.printInfo = false;
   }
 
   SparkSubmitCommandBuilder(List<String> args) {
@@ -108,14 +108,14 @@ class SparkSubmitCommandBuilder extends AbstractCommandBuilder {
 
     OptionParser parser = new OptionParser();
     parser.parse(submitArgs);
-    this.printHelp = parser.helpRequested;
+    this.printInfo = parser.infoRequested;
   }
 
   @Override
   public List<String> buildCommand(Map<String, String> env) throws IOException {
-    if (PYSPARK_SHELL_RESOURCE.equals(appResource) && !printHelp) {
+    if (PYSPARK_SHELL_RESOURCE.equals(appResource) && !printInfo) {
       return buildPySparkShellCommand(env);
-    } else if (SPARKR_SHELL_RESOURCE.equals(appResource) && !printHelp) {
+    } else if (SPARKR_SHELL_RESOURCE.equals(appResource) && !printInfo) {
       return buildSparkRCommand(env);
     } else {
       return buildSparkSubmitCommand(env);
@@ -311,7 +311,7 @@ private boolean isThriftServer(String mainClass) {
 
   private class OptionParser extends SparkSubmitOptionParser {
 
-    boolean helpRequested = false;
+    boolean infoRequested = false;
 
     @Override
     protected boolean handle(String opt, String value) {
@@ -344,7 +344,10 @@ protected boolean handle(String opt, String value) {
           appResource = specialClasses.get(value);
         }
       } else if (opt.equals(HELP) || opt.equals(USAGE_ERROR)) {
-        helpRequested = true;
+        infoRequested = true;
+        sparkArgs.add(opt);
+      } else if (opt.equals(VERSION)) {
+        infoRequested = true;
         sparkArgs.add(opt);
       } else {
         sparkArgs.add(opt);

From cc567b6634c3142125526f4875795c1b1e862838 Mon Sep 17 00:00:00 2001
From: Chris Bannister <chris.bannister@swiftkey.com>
Date: Tue, 17 Nov 2015 10:03:46 -0800
Subject: [PATCH 284/518] [SPARK-11695][CORE] Set s3a credentials

Set s3a credentials when creating a new default hadoop configuration.

Author: Chris Bannister <chris.bannister@swiftkey.com>

Closes #9663 from Zariel/set-s3a-creds.
---
 .../org/apache/spark/deploy/SparkHadoopUtil.scala   | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkHadoopUtil.scala b/core/src/main/scala/org/apache/spark/deploy/SparkHadoopUtil.scala
index d606b80c03c98..59e90564b3516 100644
--- a/core/src/main/scala/org/apache/spark/deploy/SparkHadoopUtil.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/SparkHadoopUtil.scala
@@ -92,10 +92,15 @@ class SparkHadoopUtil extends Logging {
       // Explicitly check for S3 environment variables
       if (System.getenv("AWS_ACCESS_KEY_ID") != null &&
           System.getenv("AWS_SECRET_ACCESS_KEY") != null) {
-        hadoopConf.set("fs.s3.awsAccessKeyId", System.getenv("AWS_ACCESS_KEY_ID"))
-        hadoopConf.set("fs.s3n.awsAccessKeyId", System.getenv("AWS_ACCESS_KEY_ID"))
-        hadoopConf.set("fs.s3.awsSecretAccessKey", System.getenv("AWS_SECRET_ACCESS_KEY"))
-        hadoopConf.set("fs.s3n.awsSecretAccessKey", System.getenv("AWS_SECRET_ACCESS_KEY"))
+        val keyId = System.getenv("AWS_ACCESS_KEY_ID")
+        val accessKey = System.getenv("AWS_SECRET_ACCESS_KEY")
+
+        hadoopConf.set("fs.s3.awsAccessKeyId", keyId)
+        hadoopConf.set("fs.s3n.awsAccessKeyId", keyId)
+        hadoopConf.set("fs.s3a.access.key", keyId)
+        hadoopConf.set("fs.s3.awsSecretAccessKey", accessKey)
+        hadoopConf.set("fs.s3n.awsSecretAccessKey", accessKey)
+        hadoopConf.set("fs.s3a.secret.key", accessKey)
       }
       // Copy any "spark.hadoop.foo=bar" system properties into conf as "foo=bar"
       conf.getAll.foreach { case (key, value) =>

From 21fac5434174389e8b83a2f11341fa7c9e360bfd Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Tue, 17 Nov 2015 10:17:16 -0800
Subject: [PATCH 285/518] [SPARK-11766][MLLIB] add toJson/fromJson to
 Vector/Vectors

This is to support JSON serialization of Param[Vector] in the pipeline API. It could be used for other purposes too. The schema is the same as `VectorUDT`. jkbradley

Author: Xiangrui Meng <meng@databricks.com>

Closes #9751 from mengxr/SPARK-11766.
---
 .../apache/spark/mllib/linalg/Vectors.scala   | 45 +++++++++++++++++++
 .../spark/mllib/linalg/VectorsSuite.scala     | 17 +++++++
 project/MimaExcludes.scala                    |  4 ++
 3 files changed, 66 insertions(+)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala b/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala
index bd9badc03c345..4dcf351df43fa 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala
@@ -24,6 +24,9 @@ import scala.annotation.varargs
 import scala.collection.JavaConverters._
 
 import breeze.linalg.{DenseVector => BDV, SparseVector => BSV, Vector => BV}
+import org.json4s.DefaultFormats
+import org.json4s.JsonDSL._
+import org.json4s.jackson.JsonMethods.{compact, render, parse => parseJson}
 
 import org.apache.spark.SparkException
 import org.apache.spark.annotation.{AlphaComponent, Since}
@@ -171,6 +174,12 @@ sealed trait Vector extends Serializable {
    */
   @Since("1.5.0")
   def argmax: Int
+
+  /**
+   * Converts the vector to a JSON string.
+   */
+  @Since("1.6.0")
+  def toJson: String
 }
 
 /**
@@ -339,6 +348,27 @@ object Vectors {
     parseNumeric(NumericParser.parse(s))
   }
 
+  /**
+   * Parses the JSON representation of a vector into a [[Vector]].
+   */
+  @Since("1.6.0")
+  def fromJson(json: String): Vector = {
+    implicit val formats = DefaultFormats
+    val jValue = parseJson(json)
+    (jValue \ "type").extract[Int] match {
+      case 0 => // sparse
+        val size = (jValue \ "size").extract[Int]
+        val indices = (jValue \ "indices").extract[Seq[Int]].toArray
+        val values = (jValue \ "values").extract[Seq[Double]].toArray
+        sparse(size, indices, values)
+      case 1 => // dense
+        val values = (jValue \ "values").extract[Seq[Double]].toArray
+        dense(values)
+      case _ =>
+        throw new IllegalArgumentException(s"Cannot parse $json into a vector.")
+    }
+  }
+
   private[mllib] def parseNumeric(any: Any): Vector = {
     any match {
       case values: Array[Double] =>
@@ -650,6 +680,12 @@ class DenseVector @Since("1.0.0") (
       maxIdx
     }
   }
+
+  @Since("1.6.0")
+  override def toJson: String = {
+    val jValue = ("type" -> 1) ~ ("values" -> values.toSeq)
+    compact(render(jValue))
+  }
 }
 
 @Since("1.3.0")
@@ -837,6 +873,15 @@ class SparseVector @Since("1.0.0") (
     }.unzip
     new SparseVector(selectedIndices.length, sliceInds.toArray, sliceVals.toArray)
   }
+
+  @Since("1.6.0")
+  override def toJson: String = {
+    val jValue = ("type" -> 0) ~
+      ("size" -> size) ~
+      ("indices" -> indices.toSeq) ~
+      ("values" -> values.toSeq)
+    compact(render(jValue))
+  }
 }
 
 @Since("1.3.0")
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala
index 6508ddeba4206..f895e2a8e4afb 100644
--- a/mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/mllib/linalg/VectorsSuite.scala
@@ -20,6 +20,7 @@ package org.apache.spark.mllib.linalg
 import scala.util.Random
 
 import breeze.linalg.{DenseMatrix => BDM, squaredDistance => breezeSquaredDistance}
+import org.json4s.jackson.JsonMethods.{parse => parseJson}
 
 import org.apache.spark.{Logging, SparkException, SparkFunSuite}
 import org.apache.spark.mllib.util.TestingUtils._
@@ -374,4 +375,20 @@ class VectorsSuite extends SparkFunSuite with Logging {
     assert(v.slice(Array(2, 0)) === new SparseVector(2, Array(0), Array(2.2)))
     assert(v.slice(Array(2, 0, 3, 4)) === new SparseVector(4, Array(0, 3), Array(2.2, 4.4)))
   }
+
+  test("toJson/fromJson") {
+    val sv0 = Vectors.sparse(0, Array.empty, Array.empty)
+    val sv1 = Vectors.sparse(1, Array.empty, Array.empty)
+    val sv2 = Vectors.sparse(2, Array(1), Array(2.0))
+    val dv0 = Vectors.dense(Array.empty[Double])
+    val dv1 = Vectors.dense(1.0)
+    val dv2 = Vectors.dense(0.0, 2.0)
+    for (v <- Seq(sv0, sv1, sv2, dv0, dv1, dv2)) {
+      val json = v.toJson
+      parseJson(json) // `json` should be a valid JSON string
+      val u = Vectors.fromJson(json)
+      assert(u.getClass === v.getClass, "toJson/fromJson should preserve vector types.")
+      assert(u === v, "toJson/fromJson should preserve vector values.")
+    }
+  }
 }
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 50220790d1f84..815951822c1ef 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -137,6 +137,10 @@ object MimaExcludes {
       ) ++ Seq (
         ProblemFilters.exclude[MissingMethodProblem](
           "org.apache.spark.status.api.v1.ApplicationInfo.this")
+      ) ++ Seq(
+        // SPARK-11766 add toJson to Vector
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.mllib.linalg.Vector.toJson")
       )
     case v if v.startsWith("1.5") =>
       Seq(

From e8833dd12c71b23a242727e86684d2d868ff84b3 Mon Sep 17 00:00:00 2001
From: mayuanwen <mayuanwen@qiyi.com>
Date: Tue, 17 Nov 2015 11:15:46 -0800
Subject: [PATCH 286/518] [SPARK-11679][SQL] Invoking method " apply(fields:
 java.util.List[StructField])" in "StructType" gets ClassCastException

In the previous method, fields.toArray will cast java.util.List[StructField] into Array[Object] which can not cast into Array[StructField], thus when invoking this method will throw "java.lang.ClassCastException: [Ljava.lang.Object; cannot be cast to [Lorg.apache.spark.sql.types.StructField;"
I directly cast java.util.List[StructField] into Array[StructField]  in this patch.

Author: mayuanwen <mayuanwen@qiyi.com>

Closes #9649 from jackieMaKing/Spark-11679.
---
 .../org/apache/spark/sql/types/StructType.scala     |  3 ++-
 .../org/apache/spark/sql/JavaDataFrameSuite.java    | 13 +++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/StructType.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/StructType.scala
index 11fce4beaf55f..9778df271ddd5 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/StructType.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/StructType.scala
@@ -328,7 +328,8 @@ object StructType extends AbstractDataType {
   def apply(fields: Seq[StructField]): StructType = StructType(fields.toArray)
 
   def apply(fields: java.util.List[StructField]): StructType = {
-    StructType(fields.toArray.asInstanceOf[Array[StructField]])
+    import scala.collection.JavaConverters._
+    StructType(fields.asScala)
   }
 
   protected[sql] def fromAttributes(attributes: Seq[Attribute]): StructType =
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index d191b50fa802e..567bdddece80e 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -22,6 +22,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.ArrayList;
 
 import scala.collection.JavaConverters;
 import scala.collection.Seq;
@@ -209,6 +210,18 @@ public void testCreateDataFromFromList() {
     Assert.assertEquals(1, result.length);
   }
 
+  @Test
+  public void testCreateStructTypeFromList(){
+    List<StructField> fields1 = new ArrayList<>();
+    fields1.add(new StructField("id", DataTypes.StringType, true, Metadata.empty()));
+    StructType schema1 = StructType$.MODULE$.apply(fields1);
+    Assert.assertEquals(0, schema1.fieldIndex("id"));
+
+    List<StructField> fields2 = Arrays.asList(new StructField("id", DataTypes.StringType, true, Metadata.empty()));
+    StructType schema2 = StructType$.MODULE$.apply(fields2);
+    Assert.assertEquals(0, schema2.fieldIndex("id"));
+  }
+
   private static final Comparator<Row> crosstabRowComparator = new Comparator<Row>() {
     @Override
     public int compare(Row row1, Row row2) {

From 7b1407c7b95c43299a30e891748824c4bc47e43b Mon Sep 17 00:00:00 2001
From: Cheng Lian <lian@databricks.com>
Date: Tue, 17 Nov 2015 11:17:52 -0800
Subject: [PATCH 287/518] [SPARK-11089][SQL] Adds option for disabling
 multi-session in Thrift server

This PR adds a new option `spark.sql.hive.thriftServer.singleSession` for disabling multi-session support in the Thrift server.

Note that this option is added as a Spark configuration (retrieved from `SparkConf`) rather than Spark SQL configuration (retrieved from `SQLConf`). This is because all SQL configurations are session-ized. Since multi-session support is by default on, no JDBC connection can modify global configurations like the newly added one.

Author: Cheng Lian <lian@databricks.com>

Closes #9740 from liancheng/spark-11089.single-session-option.
---
 docs/sql-programming-guide.md                 | 14 +++++
 .../thriftserver/SparkSQLSessionManager.scala |  6 ++-
 .../HiveThriftServer2Suites.scala             | 51 ++++++++++++++++++-
 .../apache/spark/sql/hive/HiveContext.scala   |  3 ++
 4 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/docs/sql-programming-guide.md b/docs/sql-programming-guide.md
index 6e02d6564b002..e347754055e79 100644
--- a/docs/sql-programming-guide.md
+++ b/docs/sql-programming-guide.md
@@ -2051,6 +2051,20 @@ options.
 
 # Migration Guide
 
+## Upgrading From Spark SQL 1.5 to 1.6
+
+ - From Spark 1.6, by default the Thrift server runs in multi-session mode.  Which means each JDBC/ODBC
+   connection owns a copy of their own SQL configuration and temporary function registry.  Cached
+   tables are still shared though.  If you prefer to run the Thrift server in the old single-session
+   mode, please set option `spark.sql.hive.thriftServer.singleSession` to `true`.  You may either add
+   this option to `spark-defaults.conf`, or pass it to `start-thriftserver.sh` via `--conf`:
+
+   {% highlight bash %}
+   ./sbin/start-thriftserver.sh \
+     --conf spark.sql.hive.thriftServer.singleSession=true \
+     ...
+   {% endhighlight %}
+
 ## Upgrading From Spark SQL 1.4 to 1.5
 
  - Optimized execution using manually managed memory (Tungsten) is now enabled by default, along with
diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLSessionManager.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLSessionManager.scala
index 33aaead3fbf96..af4fcdf021bd4 100644
--- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLSessionManager.scala
+++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLSessionManager.scala
@@ -66,7 +66,11 @@ private[hive] class SparkSQLSessionManager(hiveServer: HiveServer2, hiveContext:
     val session = super.getSession(sessionHandle)
     HiveThriftServer2.listener.onSessionCreated(
       session.getIpAddress, sessionHandle.getSessionId.toString, session.getUsername)
-    val ctx = hiveContext.newSession()
+    val ctx = if (hiveContext.hiveThriftServerSingleSession) {
+      hiveContext
+    } else {
+      hiveContext.newSession()
+    }
     ctx.setConf("spark.sql.hive.version", HiveContext.hiveExecutionVersion)
     sparkSqlOperationManager.sessionToContexts += sessionHandle -> ctx
     sessionHandle
diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
index eb1895f263d70..1dd898aa38350 100644
--- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
+++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
@@ -41,7 +41,6 @@ import org.apache.thrift.transport.TSocket
 import org.scalatest.BeforeAndAfterAll
 
 import org.apache.spark.sql.hive.HiveContext
-import org.apache.spark.sql.hive.test.TestHive
 import org.apache.spark.sql.test.ProcessTestUtils.ProcessOutputCapturer
 import org.apache.spark.util.Utils
 import org.apache.spark.{Logging, SparkFunSuite}
@@ -510,6 +509,53 @@ class HiveThriftBinaryServerSuite extends HiveThriftJdbcTest {
   }
 }
 
+class SingleSessionSuite extends HiveThriftJdbcTest {
+  override def mode: ServerMode.Value = ServerMode.binary
+
+  override protected def extraConf: Seq[String] =
+    "--conf spark.sql.hive.thriftServer.singleSession=true" :: Nil
+
+  test("test single session") {
+    withMultipleConnectionJdbcStatement(
+      { statement =>
+        val jarPath = "../hive/src/test/resources/TestUDTF.jar"
+        val jarURL = s"file://${System.getProperty("user.dir")}/$jarPath"
+
+        // Configurations and temporary functions added in this session should be visible to all
+        // the other sessions.
+        Seq(
+          "SET foo=bar",
+          s"ADD JAR $jarURL",
+          s"""CREATE TEMPORARY FUNCTION udtf_count2
+              |AS 'org.apache.spark.sql.hive.execution.GenericUDTFCount2'
+           """.stripMargin
+        ).foreach(statement.execute)
+      },
+
+      { statement =>
+        val rs1 = statement.executeQuery("SET foo")
+
+        assert(rs1.next())
+        assert(rs1.getString(1) === "foo")
+        assert(rs1.getString(2) === "bar")
+
+        val rs2 = statement.executeQuery("DESCRIBE FUNCTION udtf_count2")
+
+        assert(rs2.next())
+        assert(rs2.getString(1) === "Function: udtf_count2")
+
+        assert(rs2.next())
+        assertResult("Class: org.apache.spark.sql.hive.execution.GenericUDTFCount2") {
+          rs2.getString(1)
+        }
+
+        assert(rs2.next())
+        assert(rs2.getString(1) === "Usage: To be added.")
+      }
+    )
+  }
+}
+
 class HiveThriftHttpServerSuite extends HiveThriftJdbcTest {
   override def mode: ServerMode.Value = ServerMode.http
 
@@ -600,6 +646,8 @@ abstract class HiveThriftServer2Test extends SparkFunSuite with BeforeAndAfterAl
   private var logTailingProcess: Process = _
   private var diagnosisBuffer: ArrayBuffer[String] = ArrayBuffer.empty[String]
 
+  protected def extraConf: Seq[String] = Nil
+
   protected def serverStartCommand(port: Int) = {
     val portConf = if (mode == ServerMode.binary) {
       ConfVars.HIVE_SERVER2_THRIFT_PORT
@@ -635,6 +683,7 @@ abstract class HiveThriftServer2Test extends SparkFunSuite with BeforeAndAfterAl
        |  --driver-class-path $driverClassPath
        |  --driver-java-options -Dlog4j.debug
        |  --conf spark.ui.enabled=false
+       |  ${extraConf.mkString("\n")}
      """.stripMargin.split("\\s+").toSeq
   }
 
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index 2004f24ad26c6..c0bb5af7d5c85 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -190,6 +190,9 @@ class HiveContext private[hive](
    */
   protected[hive] def hiveThriftServerAsync: Boolean = getConf(HIVE_THRIFT_SERVER_ASYNC)
 
+  protected[hive] def hiveThriftServerSingleSession: Boolean =
+    sc.conf.get("spark.sql.hive.thriftServer.singleSession", "false").toBoolean
+
   @transient
   protected[sql] lazy val substitutor = new VariableSubstitution()
 

From 0158ff7737d10e68be2e289533241da96b496e89 Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Tue, 17 Nov 2015 11:23:54 -0800
Subject: [PATCH 288/518] [SPARK-8658][SQL][FOLLOW-UP] AttributeReference's
 equals method compares all the members

Based on the comment of cloud-fan in https://github.com/apache/spark/pull/9216, update the AttributeReference's hashCode function by including the hashCode of the other attributes including name, nullable and qualifiers.

Here, I am not 100% sure if we should include name in the hashCode calculation, since the original hashCode calculation does not include it.

marmbrus cloud-fan Please review if the changes are good.

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9761 from gatorsmile/hashCodeNamedExpression.
---
 .../spark/sql/catalyst/expressions/namedExpressions.scala   | 5 ++++-
 .../expressions/SubexpressionEliminationSuite.scala         | 6 +++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
index e3daddace241d..00b7970bd16c6 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
@@ -212,9 +212,12 @@ case class AttributeReference(
   override def hashCode: Int = {
     // See http://stackoverflow.com/questions/113511/hash-code-implementation
     var h = 17
-    h = h * 37 + exprId.hashCode()
+    h = h * 37 + name.hashCode()
     h = h * 37 + dataType.hashCode()
+    h = h * 37 + nullable.hashCode()
     h = h * 37 + metadata.hashCode()
+    h = h * 37 + exprId.hashCode()
+    h = h * 37 + qualifiers.hashCode()
     h
   }
 
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
index 9de066e99d637..a61297b2c0395 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala
@@ -25,14 +25,18 @@ class SubexpressionEliminationSuite extends SparkFunSuite {
     val a: AttributeReference = AttributeReference("name", IntegerType)()
     val b1 = a.withName("name2").withExprId(id)
     val b2 = a.withExprId(id)
+    val b3 = a.withQualifiers("qualifierName" :: Nil)
 
     assert(b1 != b2)
     assert(a != b1)
     assert(b1.semanticEquals(b2))
     assert(!b1.semanticEquals(a))
     assert(a.hashCode != b1.hashCode)
-    assert(b1.hashCode == b2.hashCode)
+    assert(b1.hashCode != b2.hashCode)
     assert(b1.semanticHash() == b2.semanticHash())
+    assert(a != b3)
+    assert(a.hashCode != b3.hashCode)
+    assert(a.semanticEquals(b3))
   }
 
   test("Expression Equivalence - basic") {

From d9251496640a77568a1e9ed5045ce2dfba4b437b Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Tue, 17 Nov 2015 11:29:02 -0800
Subject: [PATCH 289/518] [SPARK-10186][SQL] support postgre array type in
 JDBCRDD

Add ARRAY support to `PostgresDialect`.

Nested ARRAY is not allowed for now because it's hard to get the array dimension info. See http://stackoverflow.com/questions/16619113/how-to-get-array-base-type-in-postgres-via-jdbc

Thanks for the initial work from mariusvniekerk !

Close https://github.com/apache/spark/pull/9137

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9662 from cloud-fan/postgre.
---
 .../sql/jdbc/PostgresIntegrationSuite.scala   | 44 +++++++----
 .../execution/datasources/jdbc/JDBCRDD.scala  | 76 +++++++++++++-----
 .../datasources/jdbc/JdbcUtils.scala          | 77 ++++++++++---------
 .../apache/spark/sql/jdbc/JdbcDialects.scala  |  2 +-
 .../spark/sql/jdbc/PostgresDialect.scala      | 43 ++++++++---
 5 files changed, 157 insertions(+), 85 deletions(-)

diff --git a/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
index 164a7f396280c..2e18d0a2baa1c 100644
--- a/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
+++ b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
@@ -20,6 +20,8 @@ package org.apache.spark.sql.jdbc
 import java.sql.Connection
 import java.util.Properties
 
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.catalyst.expressions.{Literal, If}
 import org.apache.spark.tags.DockerTest
 
 @DockerTest
@@ -37,28 +39,32 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite {
   override def dataPreparation(conn: Connection): Unit = {
     conn.prepareStatement("CREATE DATABASE foo").executeUpdate()
     conn.setCatalog("foo")
-    conn.prepareStatement("CREATE TABLE bar (a text, b integer, c double precision, d bigint, "
-      + "e bit(1), f bit(10), g bytea, h boolean, i inet, j cidr)").executeUpdate()
+    conn.prepareStatement("CREATE TABLE bar (c0 text, c1 integer, c2 double precision, c3 bigint, "
+      + "c4 bit(1), c5 bit(10), c6 bytea, c7 boolean, c8 inet, c9 cidr, "
+      + "c10 integer[], c11 text[])").executeUpdate()
     conn.prepareStatement("INSERT INTO bar VALUES ('hello', 42, 1.25, 123456789012345, B'0', "
-      + "B'1000100101', E'\\\\xDEADBEEF', true, '172.16.0.42', '192.168.0.0/16')").executeUpdate()
+      + "B'1000100101', E'\\\\xDEADBEEF', true, '172.16.0.42', '192.168.0.0/16', "
+      + """'{1, 2}', '{"a", null, "b"}')""").executeUpdate()
   }
 
   test("Type mapping for various types") {
     val df = sqlContext.read.jdbc(jdbcUrl, "bar", new Properties)
     val rows = df.collect()
     assert(rows.length == 1)
-    val types = rows(0).toSeq.map(x => x.getClass.toString)
-    assert(types.length == 10)
-    assert(types(0).equals("class java.lang.String"))
-    assert(types(1).equals("class java.lang.Integer"))
-    assert(types(2).equals("class java.lang.Double"))
-    assert(types(3).equals("class java.lang.Long"))
-    assert(types(4).equals("class java.lang.Boolean"))
-    assert(types(5).equals("class [B"))
-    assert(types(6).equals("class [B"))
-    assert(types(7).equals("class java.lang.Boolean"))
-    assert(types(8).equals("class java.lang.String"))
-    assert(types(9).equals("class java.lang.String"))
+    val types = rows(0).toSeq.map(x => x.getClass)
+    assert(types.length == 12)
+    assert(classOf[String].isAssignableFrom(types(0)))
+    assert(classOf[java.lang.Integer].isAssignableFrom(types(1)))
+    assert(classOf[java.lang.Double].isAssignableFrom(types(2)))
+    assert(classOf[java.lang.Long].isAssignableFrom(types(3)))
+    assert(classOf[java.lang.Boolean].isAssignableFrom(types(4)))
+    assert(classOf[Array[Byte]].isAssignableFrom(types(5)))
+    assert(classOf[Array[Byte]].isAssignableFrom(types(6)))
+    assert(classOf[java.lang.Boolean].isAssignableFrom(types(7)))
+    assert(classOf[String].isAssignableFrom(types(8)))
+    assert(classOf[String].isAssignableFrom(types(9)))
+    assert(classOf[Seq[Int]].isAssignableFrom(types(10)))
+    assert(classOf[Seq[String]].isAssignableFrom(types(11)))
     assert(rows(0).getString(0).equals("hello"))
     assert(rows(0).getInt(1) == 42)
     assert(rows(0).getDouble(2) == 1.25)
@@ -72,11 +78,17 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite {
     assert(rows(0).getBoolean(7) == true)
     assert(rows(0).getString(8) == "172.16.0.42")
     assert(rows(0).getString(9) == "192.168.0.0/16")
+    assert(rows(0).getSeq(10) == Seq(1, 2))
+    assert(rows(0).getSeq(11) == Seq("a", null, "b"))
   }
 
   test("Basic write test") {
     val df = sqlContext.read.jdbc(jdbcUrl, "bar", new Properties)
-    df.write.jdbc(jdbcUrl, "public.barcopy", new Properties)
     // Test only that it doesn't crash.
+    df.write.jdbc(jdbcUrl, "public.barcopy", new Properties)
+    // Test write null values.
+    df.select(df.queryExecution.analyzed.output.map { a =>
+      Column(If(Literal(true), Literal(null), a)).as(a.name)
+    }: _*).write.jdbc(jdbcUrl, "public.barcopy2", new Properties)
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
index 018a009fbda6d..89c850ce238d7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
@@ -25,7 +25,7 @@ import org.apache.commons.lang3.StringUtils
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.SpecificMutableRow
-import org.apache.spark.sql.catalyst.util.DateTimeUtils
+import org.apache.spark.sql.catalyst.util.{GenericArrayData, DateTimeUtils}
 import org.apache.spark.sql.jdbc.JdbcDialects
 import org.apache.spark.sql.sources._
 import org.apache.spark.sql.types._
@@ -324,25 +324,27 @@ private[sql] class JDBCRDD(
   case object StringConversion extends JDBCConversion
   case object TimestampConversion extends JDBCConversion
   case object BinaryConversion extends JDBCConversion
+  case class ArrayConversion(elementConversion: JDBCConversion) extends JDBCConversion
 
   /**
    * Maps a StructType to a type tag list.
    */
-  def getConversions(schema: StructType): Array[JDBCConversion] = {
-    schema.fields.map(sf => sf.dataType match {
-      case BooleanType => BooleanConversion
-      case DateType => DateConversion
-      case DecimalType.Fixed(p, s) => DecimalConversion(p, s)
-      case DoubleType => DoubleConversion
-      case FloatType => FloatConversion
-      case IntegerType => IntegerConversion
-      case LongType =>
-        if (sf.metadata.contains("binarylong")) BinaryLongConversion else LongConversion
-      case StringType => StringConversion
-      case TimestampType => TimestampConversion
-      case BinaryType => BinaryConversion
-      case _ => throw new IllegalArgumentException(s"Unsupported field $sf")
-    }).toArray
+  def getConversions(schema: StructType): Array[JDBCConversion] =
+    schema.fields.map(sf => getConversions(sf.dataType, sf.metadata))
+
+  private def getConversions(dt: DataType, metadata: Metadata): JDBCConversion = dt match {
+    case BooleanType => BooleanConversion
+    case DateType => DateConversion
+    case DecimalType.Fixed(p, s) => DecimalConversion(p, s)
+    case DoubleType => DoubleConversion
+    case FloatType => FloatConversion
+    case IntegerType => IntegerConversion
+    case LongType => if (metadata.contains("binarylong")) BinaryLongConversion else LongConversion
+    case StringType => StringConversion
+    case TimestampType => TimestampConversion
+    case BinaryType => BinaryConversion
+    case ArrayType(et, _) => ArrayConversion(getConversions(et, metadata))
+    case _ => throw new IllegalArgumentException(s"Unsupported type ${dt.simpleString}")
   }
 
   /**
@@ -420,16 +422,44 @@ private[sql] class JDBCRDD(
                 mutableRow.update(i, null)
               }
             case BinaryConversion => mutableRow.update(i, rs.getBytes(pos))
-            case BinaryLongConversion => {
+            case BinaryLongConversion =>
               val bytes = rs.getBytes(pos)
               var ans = 0L
               var j = 0
               while (j < bytes.size) {
                 ans = 256 * ans + (255 & bytes(j))
-                j = j + 1;
+                j = j + 1
               }
               mutableRow.setLong(i, ans)
-            }
+            case ArrayConversion(elementConversion) =>
+              val array = rs.getArray(pos).getArray
+              if (array != null) {
+                val data = elementConversion match {
+                  case TimestampConversion =>
+                    array.asInstanceOf[Array[java.sql.Timestamp]].map { timestamp =>
+                      nullSafeConvert(timestamp, DateTimeUtils.fromJavaTimestamp)
+                    }
+                  case StringConversion =>
+                    array.asInstanceOf[Array[java.lang.String]]
+                      .map(UTF8String.fromString)
+                  case DateConversion =>
+                    array.asInstanceOf[Array[java.sql.Date]].map { date =>
+                      nullSafeConvert(date, DateTimeUtils.fromJavaDate)
+                    }
+                  case DecimalConversion(p, s) =>
+                    array.asInstanceOf[Array[java.math.BigDecimal]].map { decimal =>
+                      nullSafeConvert[java.math.BigDecimal](decimal, d => Decimal(d, p, s))
+                    }
+                  case BinaryLongConversion =>
+                    throw new IllegalArgumentException(s"Unsupported array element conversion $i")
+                  case _: ArrayConversion =>
+                    throw new IllegalArgumentException("Nested arrays unsupported")
+                  case _ => array.asInstanceOf[Array[Any]]
+                }
+                mutableRow.update(i, new GenericArrayData(data))
+              } else {
+                mutableRow.update(i, null)
+              }
           }
           if (rs.wasNull) mutableRow.setNullAt(i)
           i = i + 1
@@ -488,4 +518,12 @@ private[sql] class JDBCRDD(
       nextValue
     }
   }
+
+  private def nullSafeConvert[T](input: T, f: T => Any): Any = {
+    if (input == null) {
+      null
+    } else {
+      f(input)
+    }
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
index f89d55b20e212..32d28e59377a1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
@@ -23,7 +23,7 @@ import java.util.Properties
 import scala.util.Try
 
 import org.apache.spark.Logging
-import org.apache.spark.sql.jdbc.JdbcDialects
+import org.apache.spark.sql.jdbc.{JdbcDialect, JdbcType, JdbcDialects}
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.{DataFrame, Row}
 
@@ -72,6 +72,35 @@ object JdbcUtils extends Logging {
     conn.prepareStatement(sql.toString())
   }
 
+  /**
+   * Retrieve standard jdbc types.
+   * @param dt The datatype (e.g. [[org.apache.spark.sql.types.StringType]])
+   * @return The default JdbcType for this DataType
+   */
+  def getCommonJDBCType(dt: DataType): Option[JdbcType] = {
+    dt match {
+      case IntegerType => Option(JdbcType("INTEGER", java.sql.Types.INTEGER))
+      case LongType => Option(JdbcType("BIGINT", java.sql.Types.BIGINT))
+      case DoubleType => Option(JdbcType("DOUBLE PRECISION", java.sql.Types.DOUBLE))
+      case FloatType => Option(JdbcType("REAL", java.sql.Types.FLOAT))
+      case ShortType => Option(JdbcType("INTEGER", java.sql.Types.SMALLINT))
+      case ByteType => Option(JdbcType("BYTE", java.sql.Types.TINYINT))
+      case BooleanType => Option(JdbcType("BIT(1)", java.sql.Types.BIT))
+      case StringType => Option(JdbcType("TEXT", java.sql.Types.CLOB))
+      case BinaryType => Option(JdbcType("BLOB", java.sql.Types.BLOB))
+      case TimestampType => Option(JdbcType("TIMESTAMP", java.sql.Types.TIMESTAMP))
+      case DateType => Option(JdbcType("DATE", java.sql.Types.DATE))
+      case t: DecimalType => Option(
+        JdbcType(s"DECIMAL(${t.precision},${t.scale})", java.sql.Types.DECIMAL))
+      case _ => None
+    }
+  }
+
+  private def getJdbcType(dt: DataType, dialect: JdbcDialect): JdbcType = {
+    dialect.getJDBCType(dt).orElse(getCommonJDBCType(dt)).getOrElse(
+      throw new IllegalArgumentException(s"Can't get JDBC type for ${dt.simpleString}"))
+  }
+
   /**
    * Saves a partition of a DataFrame to the JDBC database.  This is done in
    * a single database transaction in order to avoid repeatedly inserting
@@ -92,7 +121,8 @@ object JdbcUtils extends Logging {
       iterator: Iterator[Row],
       rddSchema: StructType,
       nullTypes: Array[Int],
-      batchSize: Int): Iterator[Byte] = {
+      batchSize: Int,
+      dialect: JdbcDialect): Iterator[Byte] = {
     val conn = getConnection()
     var committed = false
     try {
@@ -121,6 +151,11 @@ object JdbcUtils extends Logging {
                 case TimestampType => stmt.setTimestamp(i + 1, row.getAs[java.sql.Timestamp](i))
                 case DateType => stmt.setDate(i + 1, row.getAs[java.sql.Date](i))
                 case t: DecimalType => stmt.setBigDecimal(i + 1, row.getDecimal(i))
+                case ArrayType(et, _) =>
+                  val array = conn.createArrayOf(
+                    getJdbcType(et, dialect).databaseTypeDefinition.toLowerCase,
+                    row.getSeq[AnyRef](i).toArray)
+                  stmt.setArray(i + 1, array)
                 case _ => throw new IllegalArgumentException(
                   s"Can't translate non-null value for field $i")
               }
@@ -169,23 +204,7 @@ object JdbcUtils extends Logging {
     val dialect = JdbcDialects.get(url)
     df.schema.fields foreach { field => {
       val name = field.name
-      val typ: String =
-        dialect.getJDBCType(field.dataType).map(_.databaseTypeDefinition).getOrElse(
-          field.dataType match {
-            case IntegerType => "INTEGER"
-            case LongType => "BIGINT"
-            case DoubleType => "DOUBLE PRECISION"
-            case FloatType => "REAL"
-            case ShortType => "INTEGER"
-            case ByteType => "BYTE"
-            case BooleanType => "BIT(1)"
-            case StringType => "TEXT"
-            case BinaryType => "BLOB"
-            case TimestampType => "TIMESTAMP"
-            case DateType => "DATE"
-            case t: DecimalType => s"DECIMAL(${t.precision},${t.scale})"
-            case _ => throw new IllegalArgumentException(s"Don't know how to save $field to JDBC")
-          })
+      val typ: String = getJdbcType(field.dataType, dialect).databaseTypeDefinition
       val nullable = if (field.nullable) "" else "NOT NULL"
       sb.append(s", $name $typ $nullable")
     }}
@@ -202,23 +221,7 @@ object JdbcUtils extends Logging {
       properties: Properties = new Properties()) {
     val dialect = JdbcDialects.get(url)
     val nullTypes: Array[Int] = df.schema.fields.map { field =>
-      dialect.getJDBCType(field.dataType).map(_.jdbcNullType).getOrElse(
-        field.dataType match {
-          case IntegerType => java.sql.Types.INTEGER
-          case LongType => java.sql.Types.BIGINT
-          case DoubleType => java.sql.Types.DOUBLE
-          case FloatType => java.sql.Types.REAL
-          case ShortType => java.sql.Types.INTEGER
-          case ByteType => java.sql.Types.INTEGER
-          case BooleanType => java.sql.Types.BIT
-          case StringType => java.sql.Types.CLOB
-          case BinaryType => java.sql.Types.BLOB
-          case TimestampType => java.sql.Types.TIMESTAMP
-          case DateType => java.sql.Types.DATE
-          case t: DecimalType => java.sql.Types.DECIMAL
-          case _ => throw new IllegalArgumentException(
-            s"Can't translate null value for field $field")
-        })
+      getJdbcType(field.dataType, dialect).jdbcNullType
     }
 
     val rddSchema = df.schema
@@ -226,7 +229,7 @@ object JdbcUtils extends Logging {
     val getConnection: () => Connection = JDBCRDD.getConnector(driver, url, properties)
     val batchSize = properties.getProperty("batchsize", "1000").toInt
     df.foreachPartition { iterator =>
-      savePartition(getConnection, table, iterator, rddSchema, nullTypes, batchSize)
+      savePartition(getConnection, table, iterator, rddSchema, nullTypes, batchSize, dialect)
     }
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
index 14bfea4e3e287..b3b2cb6178c52 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
@@ -51,7 +51,7 @@ case class JdbcType(databaseTypeDefinition : String, jdbcNullType : Int)
  * for the given Catalyst type.
  */
 @DeveloperApi
-abstract class JdbcDialect {
+abstract class JdbcDialect extends Serializable {
   /**
    * Check if this dialect instance can handle a certain jdbc url.
    * @param url the jdbc url.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
index e701a7fcd9e16..ed3faa1268635 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.jdbc
 
 import java.sql.Types
 
+import org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils
 import org.apache.spark.sql.types._
 
 
@@ -29,22 +30,40 @@ private object PostgresDialect extends JdbcDialect {
   override def getCatalystType(
       sqlType: Int, typeName: String, size: Int, md: MetadataBuilder): Option[DataType] = {
     if (sqlType == Types.BIT && typeName.equals("bit") && size != 1) {
-      Option(BinaryType)
-    } else if (sqlType == Types.OTHER && typeName.equals("cidr")) {
-      Option(StringType)
-    } else if (sqlType == Types.OTHER && typeName.equals("inet")) {
-      Option(StringType)
-    } else if (sqlType == Types.OTHER && typeName.equals("json")) {
-      Option(StringType)
-    } else if (sqlType == Types.OTHER && typeName.equals("jsonb")) {
-      Option(StringType)
+      Some(BinaryType)
+    } else if (sqlType == Types.OTHER) {
+      toCatalystType(typeName).filter(_ == StringType)
+    } else if (sqlType == Types.ARRAY && typeName.length > 1 && typeName(0) == '_') {
+      toCatalystType(typeName.drop(1)).map(ArrayType(_))
     } else None
   }
 
+  // TODO: support more type names.
+  private def toCatalystType(typeName: String): Option[DataType] = typeName match {
+    case "bool" => Some(BooleanType)
+    case "bit" => Some(BinaryType)
+    case "int2" => Some(ShortType)
+    case "int4" => Some(IntegerType)
+    case "int8" | "oid" => Some(LongType)
+    case "float4" => Some(FloatType)
+    case "money" | "float8" => Some(DoubleType)
+    case "text" | "varchar" | "char" | "cidr" | "inet" | "json" | "jsonb" | "uuid" =>
+      Some(StringType)
+    case "bytea" => Some(BinaryType)
+    case "timestamp" | "timestamptz" | "time" | "timetz" => Some(TimestampType)
+    case "date" => Some(DateType)
+    case "numeric" => Some(DecimalType.SYSTEM_DEFAULT)
+    case _ => None
+  }
+
   override def getJDBCType(dt: DataType): Option[JdbcType] = dt match {
-    case StringType => Some(JdbcType("TEXT", java.sql.Types.CHAR))
-    case BinaryType => Some(JdbcType("BYTEA", java.sql.Types.BINARY))
-    case BooleanType => Some(JdbcType("BOOLEAN", java.sql.Types.BOOLEAN))
+    case StringType => Some(JdbcType("TEXT", Types.CHAR))
+    case BinaryType => Some(JdbcType("BYTEA", Types.BINARY))
+    case BooleanType => Some(JdbcType("BOOLEAN", Types.BOOLEAN))
+    case ArrayType(et, _) if et.isInstanceOf[AtomicType] =>
+      getJDBCType(et).map(_.databaseTypeDefinition)
+        .orElse(JdbcUtils.getCommonJDBCType(et).map(_.databaseTypeDefinition))
+        .map(typeName => JdbcType(s"$typeName[]", java.sql.Types.ARRAY))
     case _ => None
   }
 

From d98d1cb000c8c4e391d73ae86efd09f15e5d165c Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Tue, 17 Nov 2015 12:43:56 -0800
Subject: [PATCH 290/518] [SPARK-11769][ML] Add save, load to all basic
 Transformers

This excludes Estimators and ones which include Vector and other non-basic types for Params or data.  This adds:
* Bucketizer
* DCT
* HashingTF
* Interaction
* NGram
* Normalizer
* OneHotEncoder
* PolynomialExpansion
* QuantileDiscretizer
* RFormula
* SQLTransformer
* StopWordsRemover
* StringIndexer
* Tokenizer
* VectorAssembler
* VectorSlicer

CC: mengxr

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9755 from jkbradley/transformer-io.
---
 .../apache/spark/ml/feature/Binarizer.scala   |  8 ++++-
 .../apache/spark/ml/feature/Bucketizer.scala  | 22 ++++++++----
 .../org/apache/spark/ml/feature/DCT.scala     | 19 ++++++++--
 .../apache/spark/ml/feature/HashingTF.scala   | 20 +++++++++--
 .../apache/spark/ml/feature/Interaction.scala | 29 ++++++++++++---
 .../org/apache/spark/ml/feature/NGram.scala   | 19 ++++++++--
 .../apache/spark/ml/feature/Normalizer.scala  | 20 +++++++++--
 .../spark/ml/feature/OneHotEncoder.scala      | 19 ++++++++--
 .../ml/feature/PolynomialExpansion.scala      | 20 ++++++++---
 .../ml/feature/QuantileDiscretizer.scala      | 22 ++++++++----
 .../spark/ml/feature/SQLTransformer.scala     | 27 ++++++++++++--
 .../spark/ml/feature/StopWordsRemover.scala   | 19 ++++++++--
 .../spark/ml/feature/StringIndexer.scala      | 22 +++++++++---
 .../apache/spark/ml/feature/Tokenizer.scala   | 35 ++++++++++++++++---
 .../spark/ml/feature/VectorAssembler.scala    | 18 +++++++---
 .../spark/ml/feature/VectorSlicer.scala       | 22 ++++++++----
 .../spark/ml/feature/BinarizerSuite.scala     |  8 ++---
 .../spark/ml/feature/BucketizerSuite.scala    | 12 +++++--
 .../apache/spark/ml/feature/DCTSuite.scala    | 11 +++++-
 .../spark/ml/feature/HashingTFSuite.scala     | 11 +++++-
 .../spark/ml/feature/InteractionSuite.scala   | 10 +++++-
 .../apache/spark/ml/feature/NGramSuite.scala  | 11 +++++-
 .../spark/ml/feature/NormalizerSuite.scala    | 11 +++++-
 .../spark/ml/feature/OneHotEncoderSuite.scala | 12 ++++++-
 .../ml/feature/PolynomialExpansionSuite.scala | 12 ++++++-
 .../ml/feature/QuantileDiscretizerSuite.scala | 13 ++++++-
 .../ml/feature/SQLTransformerSuite.scala      | 10 +++++-
 .../ml/feature/StopWordsRemoverSuite.scala    | 14 +++++++-
 .../spark/ml/feature/StringIndexerSuite.scala | 13 +++++--
 .../spark/ml/feature/TokenizerSuite.scala     | 25 +++++++++++--
 .../ml/feature/VectorAssemblerSuite.scala     | 11 +++++-
 .../spark/ml/feature/VectorSlicerSuite.scala  | 12 ++++++-
 32 files changed, 453 insertions(+), 84 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
index e5c25574d4b11..e2be6547d8f00 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.Transformer
 import org.apache.spark.ml.attribute.BinaryAttribute
 import org.apache.spark.ml.param._
@@ -87,10 +87,16 @@ final class Binarizer(override val uid: String)
 
   override def copy(extra: ParamMap): Binarizer = defaultCopy(extra)
 
+  @Since("1.6.0")
   override def write: Writer = new DefaultParamsWriter(this)
 }
 
+@Since("1.6.0")
 object Binarizer extends Readable[Binarizer] {
 
+  @Since("1.6.0")
   override def read: Reader[Binarizer] = new DefaultParamsReader[Binarizer]
+
+  @Since("1.6.0")
+  override def load(path: String): Binarizer = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala
index 6fdf25b015b0b..7095fbd70aa07 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala
@@ -20,12 +20,12 @@ package org.apache.spark.ml.feature
 import java.{util => ju}
 
 import org.apache.spark.SparkException
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.Model
 import org.apache.spark.ml.attribute.NominalAttribute
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
@@ -36,7 +36,7 @@ import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
  */
 @Experimental
 final class Bucketizer(override val uid: String)
-  extends Model[Bucketizer] with HasInputCol with HasOutputCol {
+  extends Model[Bucketizer] with HasInputCol with HasOutputCol with Writable {
 
   def this() = this(Identifiable.randomUID("bucketizer"))
 
@@ -93,11 +93,15 @@ final class Bucketizer(override val uid: String)
   override def copy(extra: ParamMap): Bucketizer = {
     defaultCopy[Bucketizer](extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
 }
 
-private[feature] object Bucketizer {
+object Bucketizer extends Readable[Bucketizer] {
+
   /** We require splits to be of length >= 3 and to be in strictly increasing order. */
-  def checkSplits(splits: Array[Double]): Boolean = {
+  private[feature] def checkSplits(splits: Array[Double]): Boolean = {
     if (splits.length < 3) {
       false
     } else {
@@ -115,7 +119,7 @@ private[feature] object Bucketizer {
    * Binary searching in several buckets to place each data point.
    * @throws SparkException if a feature is < splits.head or > splits.last
    */
-  def binarySearchForBuckets(splits: Array[Double], feature: Double): Double = {
+  private[feature] def binarySearchForBuckets(splits: Array[Double], feature: Double): Double = {
     if (feature == splits.last) {
       splits.length - 2
     } else {
@@ -134,4 +138,10 @@ private[feature] object Bucketizer {
       }
     }
   }
+
+  @Since("1.6.0")
+  override def read: Reader[Bucketizer] = new DefaultParamsReader[Bucketizer]
+
+  @Since("1.6.0")
+  override def load(path: String): Bucketizer = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala
index 228347635c92b..6ea5a616173ee 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala
@@ -19,10 +19,10 @@ package org.apache.spark.ml.feature
 
 import edu.emory.mathcs.jtransforms.dct._
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.UnaryTransformer
 import org.apache.spark.ml.param.BooleanParam
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT, Vectors}
 import org.apache.spark.sql.types.DataType
 
@@ -37,7 +37,7 @@ import org.apache.spark.sql.types.DataType
  */
 @Experimental
 class DCT(override val uid: String)
-  extends UnaryTransformer[Vector, Vector, DCT] {
+  extends UnaryTransformer[Vector, Vector, DCT] with Writable {
 
   def this() = this(Identifiable.randomUID("dct"))
 
@@ -69,4 +69,17 @@ class DCT(override val uid: String)
   }
 
   override protected def outputDataType: DataType = new VectorUDT
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object DCT extends Readable[DCT] {
+
+  @Since("1.6.0")
+  override def read: Reader[DCT] = new DefaultParamsReader[DCT]
+
+  @Since("1.6.0")
+  override def load(path: String): DCT = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala
index 319d23e46cef4..6d2ea675f5617 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala
@@ -17,12 +17,12 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.Transformer
 import org.apache.spark.ml.attribute.AttributeGroup
 import org.apache.spark.ml.param.{IntParam, ParamMap, ParamValidators}
 import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.feature
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions.{col, udf}
@@ -33,7 +33,8 @@ import org.apache.spark.sql.types.{ArrayType, StructType}
  * Maps a sequence of terms to their term frequencies using the hashing trick.
  */
 @Experimental
-class HashingTF(override val uid: String) extends Transformer with HasInputCol with HasOutputCol {
+class HashingTF(override val uid: String)
+  extends Transformer with HasInputCol with HasOutputCol with Writable {
 
   def this() = this(Identifiable.randomUID("hashingTF"))
 
@@ -76,4 +77,17 @@ class HashingTF(override val uid: String) extends Transformer with HasInputCol w
   }
 
   override def copy(extra: ParamMap): HashingTF = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object HashingTF extends Readable[HashingTF] {
+
+  @Since("1.6.0")
+  override def read: Reader[HashingTF] = new DefaultParamsReader[HashingTF]
+
+  @Since("1.6.0")
+  override def load(path: String): HashingTF = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala
index 37f7862476cfe..9df6b311cc9da 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala
@@ -20,11 +20,11 @@ package org.apache.spark.ml.feature
 import scala.collection.mutable.ArrayBuilder
 
 import org.apache.spark.SparkException
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.attribute._
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.ml.Transformer
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT, Vectors}
 import org.apache.spark.sql.{DataFrame, Row}
@@ -42,24 +42,30 @@ import org.apache.spark.sql.types._
  * `Vector(6, 8)` if all input features were numeric. If the first feature was instead nominal
  * with four categories, the output would then be `Vector(0, 0, 0, 0, 3, 4, 0, 0)`.
  */
+@Since("1.6.0")
 @Experimental
-class Interaction(override val uid: String) extends Transformer
-  with HasInputCols with HasOutputCol {
+class Interaction @Since("1.6.0") (override val uid: String) extends Transformer
+  with HasInputCols with HasOutputCol with Writable {
 
+  @Since("1.6.0")
   def this() = this(Identifiable.randomUID("interaction"))
 
   /** @group setParam */
+  @Since("1.6.0")
   def setInputCols(values: Array[String]): this.type = set(inputCols, values)
 
   /** @group setParam */
+  @Since("1.6.0")
   def setOutputCol(value: String): this.type = set(outputCol, value)
 
   // optimistic schema; does not contain any ML attributes
+  @Since("1.6.0")
   override def transformSchema(schema: StructType): StructType = {
     validateParams()
     StructType(schema.fields :+ StructField($(outputCol), new VectorUDT, false))
   }
 
+  @Since("1.6.0")
   override def transform(dataset: DataFrame): DataFrame = {
     validateParams()
     val inputFeatures = $(inputCols).map(c => dataset.schema(c))
@@ -208,14 +214,29 @@ class Interaction(override val uid: String) extends Transformer
     }
   }
 
+  @Since("1.6.0")
   override def copy(extra: ParamMap): Interaction = defaultCopy(extra)
 
+  @Since("1.6.0")
   override def validateParams(): Unit = {
     require(get(inputCols).isDefined, "Input cols must be defined first.")
     require(get(outputCol).isDefined, "Output col must be defined first.")
     require($(inputCols).length > 0, "Input cols must have non-zero length.")
     require($(inputCols).distinct.length == $(inputCols).length, "Input cols must be distinct.")
   }
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object Interaction extends Readable[Interaction] {
+
+  @Since("1.6.0")
+  override def read: Reader[Interaction] = new DefaultParamsReader[Interaction]
+
+  @Since("1.6.0")
+  override def load(path: String): Interaction = read.load(path)
 }
 
 /**
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala
index 8de10eb51f923..4a17acd95199f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala
@@ -17,10 +17,10 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.UnaryTransformer
 import org.apache.spark.ml.param._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.sql.types.{ArrayType, DataType, StringType}
 
 /**
@@ -36,7 +36,7 @@ import org.apache.spark.sql.types.{ArrayType, DataType, StringType}
  */
 @Experimental
 class NGram(override val uid: String)
-  extends UnaryTransformer[Seq[String], Seq[String], NGram] {
+  extends UnaryTransformer[Seq[String], Seq[String], NGram] with Writable {
 
   def this() = this(Identifiable.randomUID("ngram"))
 
@@ -66,4 +66,17 @@ class NGram(override val uid: String)
   }
 
   override protected def outputDataType: DataType = new ArrayType(StringType, false)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object NGram extends Readable[NGram] {
+
+  @Since("1.6.0")
+  override def read: Reader[NGram] = new DefaultParamsReader[NGram]
+
+  @Since("1.6.0")
+  override def load(path: String): NGram = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala
index 8282e5ffa17f7..9df6a091d5058 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala
@@ -17,10 +17,10 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.UnaryTransformer
 import org.apache.spark.ml.param.{DoubleParam, ParamValidators}
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.feature
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT}
 import org.apache.spark.sql.types.DataType
@@ -30,7 +30,8 @@ import org.apache.spark.sql.types.DataType
  * Normalize a vector to have unit norm using the given p-norm.
  */
 @Experimental
-class Normalizer(override val uid: String) extends UnaryTransformer[Vector, Vector, Normalizer] {
+class Normalizer(override val uid: String)
+  extends UnaryTransformer[Vector, Vector, Normalizer] with Writable {
 
   def this() = this(Identifiable.randomUID("normalizer"))
 
@@ -55,4 +56,17 @@ class Normalizer(override val uid: String) extends UnaryTransformer[Vector, Vect
   }
 
   override protected def outputDataType: DataType = new VectorUDT()
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object Normalizer extends Readable[Normalizer] {
+
+  @Since("1.6.0")
+  override def read: Reader[Normalizer] = new DefaultParamsReader[Normalizer]
+
+  @Since("1.6.0")
+  override def load(path: String): Normalizer = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
index 9c60d4084ec46..4e2adfaafa21e 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
@@ -17,12 +17,12 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.Transformer
 import org.apache.spark.ml.attribute._
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg.Vectors
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions.{col, udf}
@@ -44,7 +44,7 @@ import org.apache.spark.sql.types.{DoubleType, StructType}
  */
 @Experimental
 class OneHotEncoder(override val uid: String) extends Transformer
-  with HasInputCol with HasOutputCol {
+  with HasInputCol with HasOutputCol with Writable {
 
   def this() = this(Identifiable.randomUID("oneHot"))
 
@@ -165,4 +165,17 @@ class OneHotEncoder(override val uid: String) extends Transformer
   }
 
   override def copy(extra: ParamMap): OneHotEncoder = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object OneHotEncoder extends Readable[OneHotEncoder] {
+
+  @Since("1.6.0")
+  override def read: Reader[OneHotEncoder] = new DefaultParamsReader[OneHotEncoder]
+
+  @Since("1.6.0")
+  override def load(path: String): OneHotEncoder = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala
index d85e468562d4a..49415398325fd 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala
@@ -19,10 +19,10 @@ package org.apache.spark.ml.feature
 
 import scala.collection.mutable
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.UnaryTransformer
 import org.apache.spark.ml.param.{ParamMap, IntParam, ParamValidators}
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg._
 import org.apache.spark.sql.types.DataType
 
@@ -36,7 +36,7 @@ import org.apache.spark.sql.types.DataType
  */
 @Experimental
 class PolynomialExpansion(override val uid: String)
-  extends UnaryTransformer[Vector, Vector, PolynomialExpansion] {
+  extends UnaryTransformer[Vector, Vector, PolynomialExpansion] with Writable {
 
   def this() = this(Identifiable.randomUID("poly"))
 
@@ -63,6 +63,9 @@ class PolynomialExpansion(override val uid: String)
   override protected def outputDataType: DataType = new VectorUDT()
 
   override def copy(extra: ParamMap): PolynomialExpansion = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 /**
@@ -77,7 +80,8 @@ class PolynomialExpansion(override val uid: String)
  * To handle sparsity, if c is zero, we can skip all monomials that contain it. We remember the
  * current index and increment it properly for sparse input.
  */
-private[feature] object PolynomialExpansion {
+@Since("1.6.0")
+object PolynomialExpansion extends Readable[PolynomialExpansion] {
 
   private def choose(n: Int, k: Int): Int = {
     Range(n, n - k, -1).product / Range(k, 1, -1).product
@@ -169,11 +173,17 @@ private[feature] object PolynomialExpansion {
     new SparseVector(polySize - 1, polyIndices.result(), polyValues.result())
   }
 
-  def expand(v: Vector, degree: Int): Vector = {
+  private[feature] def expand(v: Vector, degree: Int): Vector = {
     v match {
       case dv: DenseVector => expand(dv, degree)
       case sv: SparseVector => expand(sv, degree)
       case _ => throw new IllegalArgumentException
     }
   }
+
+  @Since("1.6.0")
+  override def read: Reader[PolynomialExpansion] = new DefaultParamsReader[PolynomialExpansion]
+
+  @Since("1.6.0")
+  override def load(path: String): PolynomialExpansion = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/QuantileDiscretizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/QuantileDiscretizer.scala
index 46b836da9cfde..2da5c966d2967 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/QuantileDiscretizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/QuantileDiscretizer.scala
@@ -20,7 +20,7 @@ package org.apache.spark.ml.feature
 import scala.collection.mutable
 
 import org.apache.spark.Logging
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml._
 import org.apache.spark.ml.attribute.NominalAttribute
 import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
@@ -60,7 +60,7 @@ private[feature] trait QuantileDiscretizerBase extends Params with HasInputCol w
  */
 @Experimental
 final class QuantileDiscretizer(override val uid: String)
-  extends Estimator[Bucketizer] with QuantileDiscretizerBase {
+  extends Estimator[Bucketizer] with QuantileDiscretizerBase with Writable {
 
   def this() = this(Identifiable.randomUID("quantileDiscretizer"))
 
@@ -93,13 +93,17 @@ final class QuantileDiscretizer(override val uid: String)
   }
 
   override def copy(extra: ParamMap): QuantileDiscretizer = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
 }
 
-private[feature] object QuantileDiscretizer extends Logging {
+@Since("1.6.0")
+object QuantileDiscretizer extends Readable[QuantileDiscretizer] with Logging {
   /**
    * Sampling from the given dataset to collect quantile statistics.
    */
-  def getSampledInput(dataset: DataFrame, numBins: Int): Array[Row] = {
+  private[feature] def getSampledInput(dataset: DataFrame, numBins: Int): Array[Row] = {
     val totalSamples = dataset.count()
     require(totalSamples > 0,
       "QuantileDiscretizer requires non-empty input dataset but was given an empty input.")
@@ -111,6 +115,7 @@ private[feature] object QuantileDiscretizer extends Logging {
   /**
    * Compute split points with respect to the sample distribution.
    */
+  private[feature]
   def findSplitCandidates(samples: Array[Double], numSplits: Int): Array[Double] = {
     val valueCountMap = samples.foldLeft(Map.empty[Double, Int]) { (m, x) =>
       m + ((x, m.getOrElse(x, 0) + 1))
@@ -150,7 +155,7 @@ private[feature] object QuantileDiscretizer extends Logging {
    * Adjust split candidates to proper splits by: adding positive/negative infinity to both sides as
    * needed, and adding a default split value of 0 if no good candidates are found.
    */
-  def getSplits(candidates: Array[Double]): Array[Double] = {
+  private[feature] def getSplits(candidates: Array[Double]): Array[Double] = {
     val effectiveValues = if (candidates.size != 0) {
       if (candidates.head == Double.NegativeInfinity
         && candidates.last == Double.PositiveInfinity) {
@@ -172,5 +177,10 @@ private[feature] object QuantileDiscretizer extends Logging {
       Array(Double.NegativeInfinity) ++ effectiveValues ++ Array(Double.PositiveInfinity)
     }
   }
-}
 
+  @Since("1.6.0")
+  override def read: Reader[QuantileDiscretizer] = new DefaultParamsReader[QuantileDiscretizer]
+
+  @Since("1.6.0")
+  override def load(path: String): QuantileDiscretizer = read.load(path)
+}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala
index 95e4305638730..c115064ff301a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala
@@ -18,10 +18,10 @@
 package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkContext
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.param.{ParamMap, Param}
 import org.apache.spark.ml.Transformer
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.sql.{SQLContext, DataFrame, Row}
 import org.apache.spark.sql.types.StructType
 
@@ -32,24 +32,30 @@ import org.apache.spark.sql.types.StructType
  * where '__THIS__' represents the underlying table of the input dataset.
  */
 @Experimental
-class SQLTransformer (override val uid: String) extends Transformer {
+@Since("1.6.0")
+class SQLTransformer @Since("1.6.0") (override val uid: String) extends Transformer with Writable {
 
+  @Since("1.6.0")
   def this() = this(Identifiable.randomUID("sql"))
 
   /**
    * SQL statement parameter. The statement is provided in string form.
    * @group param
    */
+  @Since("1.6.0")
   final val statement: Param[String] = new Param[String](this, "statement", "SQL statement")
 
   /** @group setParam */
+  @Since("1.6.0")
   def setStatement(value: String): this.type = set(statement, value)
 
   /** @group getParam */
+  @Since("1.6.0")
   def getStatement: String = $(statement)
 
   private val tableIdentifier: String = "__THIS__"
 
+  @Since("1.6.0")
   override def transform(dataset: DataFrame): DataFrame = {
     val tableName = Identifiable.randomUID(uid)
     dataset.registerTempTable(tableName)
@@ -58,6 +64,7 @@ class SQLTransformer (override val uid: String) extends Transformer {
     outputDF
   }
 
+  @Since("1.6.0")
   override def transformSchema(schema: StructType): StructType = {
     val sc = SparkContext.getOrCreate()
     val sqlContext = SQLContext.getOrCreate(sc)
@@ -68,5 +75,19 @@ class SQLTransformer (override val uid: String) extends Transformer {
     outputSchema
   }
 
+  @Since("1.6.0")
   override def copy(extra: ParamMap): SQLTransformer = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object SQLTransformer extends Readable[SQLTransformer] {
+
+  @Since("1.6.0")
+  override def read: Reader[SQLTransformer] = new DefaultParamsReader[SQLTransformer]
+
+  @Since("1.6.0")
+  override def load(path: String): SQLTransformer = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala
index 2a79582625e9a..f1146988dcc7c 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala
@@ -17,11 +17,11 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.Transformer
 import org.apache.spark.ml.param.{BooleanParam, ParamMap, StringArrayParam}
 import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions.{col, udf}
 import org.apache.spark.sql.types.{ArrayType, StringType, StructField, StructType}
@@ -86,7 +86,7 @@ private[spark] object StopWords {
  */
 @Experimental
 class StopWordsRemover(override val uid: String)
-  extends Transformer with HasInputCol with HasOutputCol {
+  extends Transformer with HasInputCol with HasOutputCol with Writable {
 
   def this() = this(Identifiable.randomUID("stopWords"))
 
@@ -154,4 +154,17 @@ class StopWordsRemover(override val uid: String)
   }
 
   override def copy(extra: ParamMap): StopWordsRemover = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object StopWordsRemover extends Readable[StopWordsRemover] {
+
+  @Since("1.6.0")
+  override def read: Reader[StopWordsRemover] = new DefaultParamsReader[StopWordsRemover]
+
+  @Since("1.6.0")
+  override def load(path: String): StopWordsRemover = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
index 486274cd75a14..f782a272d11db 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
@@ -18,13 +18,13 @@
 package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkException
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.attribute.{Attribute, NominalAttribute}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
 import org.apache.spark.ml.Transformer
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
@@ -188,9 +188,8 @@ class StringIndexerModel (
  * @see [[StringIndexer]] for converting strings into indices
  */
 @Experimental
-class IndexToString private[ml] (
-  override val uid: String) extends Transformer
-    with HasInputCol with HasOutputCol {
+class IndexToString private[ml] (override val uid: String)
+  extends Transformer with HasInputCol with HasOutputCol with Writable {
 
   def this() =
     this(Identifiable.randomUID("idxToStr"))
@@ -257,4 +256,17 @@ class IndexToString private[ml] (
   override def copy(extra: ParamMap): IndexToString = {
     defaultCopy(extra)
   }
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object IndexToString extends Readable[IndexToString] {
+
+  @Since("1.6.0")
+  override def read: Reader[IndexToString] = new DefaultParamsReader[IndexToString]
+
+  @Since("1.6.0")
+  override def load(path: String): IndexToString = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
index 1b82b40caac18..0e4445d1e2fa7 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
@@ -17,10 +17,10 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.UnaryTransformer
 import org.apache.spark.ml.param._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.sql.types.{ArrayType, DataType, StringType}
 
 /**
@@ -30,7 +30,8 @@ import org.apache.spark.sql.types.{ArrayType, DataType, StringType}
  * @see [[RegexTokenizer]]
  */
 @Experimental
-class Tokenizer(override val uid: String) extends UnaryTransformer[String, Seq[String], Tokenizer] {
+class Tokenizer(override val uid: String)
+  extends UnaryTransformer[String, Seq[String], Tokenizer] with Writable {
 
   def this() = this(Identifiable.randomUID("tok"))
 
@@ -45,6 +46,19 @@ class Tokenizer(override val uid: String) extends UnaryTransformer[String, Seq[S
   override protected def outputDataType: DataType = new ArrayType(StringType, true)
 
   override def copy(extra: ParamMap): Tokenizer = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object Tokenizer extends Readable[Tokenizer] {
+
+  @Since("1.6.0")
+  override def read: Reader[Tokenizer] = new DefaultParamsReader[Tokenizer]
+
+  @Since("1.6.0")
+  override def load(path: String): Tokenizer = read.load(path)
 }
 
 /**
@@ -56,7 +70,7 @@ class Tokenizer(override val uid: String) extends UnaryTransformer[String, Seq[S
  */
 @Experimental
 class RegexTokenizer(override val uid: String)
-  extends UnaryTransformer[String, Seq[String], RegexTokenizer] {
+  extends UnaryTransformer[String, Seq[String], RegexTokenizer] with Writable {
 
   def this() = this(Identifiable.randomUID("regexTok"))
 
@@ -131,4 +145,17 @@ class RegexTokenizer(override val uid: String)
   override protected def outputDataType: DataType = new ArrayType(StringType, true)
 
   override def copy(extra: ParamMap): RegexTokenizer = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object RegexTokenizer extends Readable[RegexTokenizer] {
+
+  @Since("1.6.0")
+  override def read: Reader[RegexTokenizer] = new DefaultParamsReader[RegexTokenizer]
+
+  @Since("1.6.0")
+  override def load(path: String): RegexTokenizer = read.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
index 086917fa680f8..7e54205292ca2 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
@@ -20,12 +20,12 @@ package org.apache.spark.ml.feature
 import scala.collection.mutable.ArrayBuilder
 
 import org.apache.spark.SparkException
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.Transformer
 import org.apache.spark.ml.attribute.{Attribute, AttributeGroup, NumericAttribute, UnresolvedAttribute}
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT, Vectors}
 import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.sql.functions._
@@ -37,7 +37,7 @@ import org.apache.spark.sql.types._
  */
 @Experimental
 class VectorAssembler(override val uid: String)
-  extends Transformer with HasInputCols with HasOutputCol {
+  extends Transformer with HasInputCols with HasOutputCol with Writable {
 
   def this() = this(Identifiable.randomUID("vecAssembler"))
 
@@ -120,9 +120,19 @@ class VectorAssembler(override val uid: String)
   }
 
   override def copy(extra: ParamMap): VectorAssembler = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
 }
 
-private object VectorAssembler {
+@Since("1.6.0")
+object VectorAssembler extends Readable[VectorAssembler] {
+
+  @Since("1.6.0")
+  override def read: Reader[VectorAssembler] = new DefaultParamsReader[VectorAssembler]
+
+  @Since("1.6.0")
+  override def load(path: String): VectorAssembler = read.load(path)
 
   private[feature] def assemble(vv: Any*): Vector = {
     val indices = ArrayBuilder.make[Int]
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala
index fb3387d4aa9be..911582b55b574 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala
@@ -17,12 +17,12 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.Transformer
 import org.apache.spark.ml.attribute.{Attribute, AttributeGroup}
 import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
 import org.apache.spark.ml.param.{IntArrayParam, ParamMap, StringArrayParam}
-import org.apache.spark.ml.util.{Identifiable, MetadataUtils, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg._
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions._
@@ -42,7 +42,7 @@ import org.apache.spark.sql.types.StructType
  */
 @Experimental
 final class VectorSlicer(override val uid: String)
-  extends Transformer with HasInputCol with HasOutputCol {
+  extends Transformer with HasInputCol with HasOutputCol with Writable {
 
   def this() = this(Identifiable.randomUID("vectorSlicer"))
 
@@ -151,12 +151,16 @@ final class VectorSlicer(override val uid: String)
   }
 
   override def copy(extra: ParamMap): VectorSlicer = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
 }
 
-private[feature] object VectorSlicer {
+@Since("1.6.0")
+object VectorSlicer extends Readable[VectorSlicer] {
 
   /** Return true if given feature indices are valid */
-  def validIndices(indices: Array[Int]): Boolean = {
+  private[feature] def validIndices(indices: Array[Int]): Boolean = {
     if (indices.isEmpty) {
       true
     } else {
@@ -165,7 +169,13 @@ private[feature] object VectorSlicer {
   }
 
   /** Return true if given feature names are valid */
-  def validNames(names: Array[String]): Boolean = {
+  private[feature] def validNames(names: Array[String]): Boolean = {
     names.forall(_.nonEmpty) && names.length == names.distinct.length
   }
+
+  @Since("1.6.0")
+  override def read: Reader[VectorSlicer] = new DefaultParamsReader[VectorSlicer]
+
+  @Since("1.6.0")
+  override def load(path: String): VectorSlicer = read.load(path)
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
index 9dfa1439cc303..6d2d8fe714444 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/BinarizerSuite.scala
@@ -69,10 +69,10 @@ class BinarizerSuite extends SparkFunSuite with MLlibTestSparkContext with Defau
   }
 
   test("read/write") {
-    val binarizer = new Binarizer()
-      .setInputCol("feature")
-      .setOutputCol("binarized_feature")
+    val t = new Binarizer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
       .setThreshold(0.1)
-    testDefaultReadWrite(binarizer)
+    testDefaultReadWrite(t)
   }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/BucketizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/BucketizerSuite.scala
index 0eba34fda6228..9ea7d431763a1 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/BucketizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/BucketizerSuite.scala
@@ -21,13 +21,13 @@ import scala.util.Random
 
 import org.apache.spark.{SparkException, SparkFunSuite}
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.Vectors
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.{DataFrame, Row}
 
-class BucketizerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class BucketizerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new Bucketizer)
@@ -112,6 +112,14 @@ class BucketizerSuite extends SparkFunSuite with MLlibTestSparkContext {
     val lsResult = Vectors.dense(data.map(x => BucketizerSuite.linearSearchForBuckets(splits, x)))
     assert(bsResult ~== lsResult absTol 1e-5)
   }
+
+  test("read/write") {
+    val t = new Bucketizer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setSplits(Array(0.1, 0.8, 0.9))
+    testDefaultReadWrite(t)
+  }
 }
 
 private object BucketizerSuite extends SparkFunSuite {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/DCTSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/DCTSuite.scala
index 37ed2367c33f7..0f2aafebafe67 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/DCTSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/DCTSuite.scala
@@ -22,6 +22,7 @@ import scala.beans.BeanInfo
 import edu.emory.mathcs.jtransforms.dct.DoubleDCT_1D
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{DataFrame, Row}
@@ -29,7 +30,7 @@ import org.apache.spark.sql.{DataFrame, Row}
 @BeanInfo
 case class DCTTestData(vec: Vector, wantedVec: Vector)
 
-class DCTSuite extends SparkFunSuite with MLlibTestSparkContext {
+class DCTSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("forward transform of discrete cosine matches jTransforms result") {
     val data = Vectors.dense((0 until 128).map(_ => 2D * math.random - 1D).toArray)
@@ -45,6 +46,14 @@ class DCTSuite extends SparkFunSuite with MLlibTestSparkContext {
     testDCT(data, inverse)
   }
 
+  test("read/write") {
+    val t = new DCT()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setInverse(true)
+    testDefaultReadWrite(t)
+  }
+
   private def testDCT(data: Vector, inverse: Boolean): Unit = {
     val expectedResultBuffer = data.toArray.clone()
     if (inverse) {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/HashingTFSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/HashingTFSuite.scala
index 4157b84b29d01..0dcd0f49465ed 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/HashingTFSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/HashingTFSuite.scala
@@ -20,12 +20,13 @@ package org.apache.spark.ml.feature
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.attribute.AttributeGroup
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.util.Utils
 
-class HashingTFSuite extends SparkFunSuite with MLlibTestSparkContext {
+class HashingTFSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new HashingTF)
@@ -50,4 +51,12 @@ class HashingTFSuite extends SparkFunSuite with MLlibTestSparkContext {
       Seq((idx("a"), 2.0), (idx("b"), 2.0), (idx("c"), 1.0), (idx("d"), 1.0)))
     assert(features ~== expected absTol 1e-14)
   }
+
+  test("read/write") {
+    val t = new HashingTF()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setNumFeatures(10)
+    testDefaultReadWrite(t)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/InteractionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/InteractionSuite.scala
index 2beb62ca08233..932d331b472b9 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/InteractionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/InteractionSuite.scala
@@ -19,6 +19,7 @@ package org.apache.spark.ml.feature
 
 import scala.collection.mutable.ArrayBuilder
 
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.{SparkException, SparkFunSuite}
 import org.apache.spark.ml.attribute._
 import org.apache.spark.ml.param.ParamsSuite
@@ -26,7 +27,7 @@ import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.functions.col
 
-class InteractionSuite extends SparkFunSuite with MLlibTestSparkContext {
+class InteractionSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
   test("params") {
     ParamsSuite.checkParams(new Interaction())
   }
@@ -162,4 +163,11 @@ class InteractionSuite extends SparkFunSuite with MLlibTestSparkContext {
         new NumericAttribute(Some("a_2:b_1:c"), Some(9))))
     assert(attrs === expectedAttrs)
   }
+
+  test("read/write") {
+    val t = new Interaction()
+      .setInputCols(Array("myInputCol", "myInputCol2"))
+      .setOutputCol("myOutputCol")
+    testDefaultReadWrite(t)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/NGramSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/NGramSuite.scala
index ab97e3dbc6ee0..58fda29aa1e69 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/NGramSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/NGramSuite.scala
@@ -20,13 +20,14 @@ package org.apache.spark.ml.feature
 import scala.beans.BeanInfo
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{DataFrame, Row}
 
 @BeanInfo
 case class NGramTestData(inputTokens: Array[String], wantedNGrams: Array[String])
 
-class NGramSuite extends SparkFunSuite with MLlibTestSparkContext {
+class NGramSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
   import org.apache.spark.ml.feature.NGramSuite._
 
   test("default behavior yields bigram features") {
@@ -79,6 +80,14 @@ class NGramSuite extends SparkFunSuite with MLlibTestSparkContext {
       )))
     testNGram(nGram, dataset)
   }
+
+  test("read/write") {
+    val t = new NGram()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setN(3)
+    testDefaultReadWrite(t)
+  }
 }
 
 object NGramSuite extends SparkFunSuite {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/NormalizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/NormalizerSuite.scala
index 9f03470b7f328..de3d438ce83be 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/NormalizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/NormalizerSuite.scala
@@ -18,13 +18,14 @@
 package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.{DataFrame, Row, SQLContext}
 
 
-class NormalizerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class NormalizerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   @transient var data: Array[Vector] = _
   @transient var dataFrame: DataFrame = _
@@ -104,6 +105,14 @@ class NormalizerSuite extends SparkFunSuite with MLlibTestSparkContext {
 
     assertValues(result, l1Normalized)
   }
+
+  test("read/write") {
+    val t = new Normalizer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setP(3.0)
+    testDefaultReadWrite(t)
+  }
 }
 
 private object NormalizerSuite {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderSuite.scala
index 321eeb843941c..76d12050f9677 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderSuite.scala
@@ -20,12 +20,14 @@ package org.apache.spark.ml.feature
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.attribute.{AttributeGroup, BinaryAttribute, NominalAttribute}
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.linalg.Vector
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions.col
 
-class OneHotEncoderSuite extends SparkFunSuite with MLlibTestSparkContext {
+class OneHotEncoderSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   def stringIndexed(): DataFrame = {
     val data = sc.parallelize(Seq((0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c")), 2)
@@ -101,4 +103,12 @@ class OneHotEncoderSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(group.getAttr(0) === BinaryAttribute.defaultAttr.withName("0").withIndex(0))
     assert(group.getAttr(1) === BinaryAttribute.defaultAttr.withName("1").withIndex(1))
   }
+
+  test("read/write") {
+    val t = new OneHotEncoder()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setDropLast(false)
+    testDefaultReadWrite(t)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/PolynomialExpansionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/PolynomialExpansionSuite.scala
index 29eebd8960ebc..70892dc57170a 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/PolynomialExpansionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/PolynomialExpansionSuite.scala
@@ -21,12 +21,14 @@ import org.apache.spark.ml.param.ParamsSuite
 import org.scalatest.exceptions.TestFailedException
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.Row
 
-class PolynomialExpansionSuite extends SparkFunSuite with MLlibTestSparkContext {
+class PolynomialExpansionSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new PolynomialExpansion)
@@ -98,5 +100,13 @@ class PolynomialExpansionSuite extends SparkFunSuite with MLlibTestSparkContext
         throw new TestFailedException("Unmatched data types after polynomial expansion", 0)
     }
   }
+
+  test("read/write") {
+    val t = new PolynomialExpansion()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setDegree(3)
+    testDefaultReadWrite(t)
+  }
 }
 
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/QuantileDiscretizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/QuantileDiscretizerSuite.scala
index b2bdd8935f903..3a4f6d235aa6c 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/QuantileDiscretizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/QuantileDiscretizerSuite.scala
@@ -18,11 +18,14 @@
 package org.apache.spark.ml.feature
 
 import org.apache.spark.ml.attribute.{Attribute, NominalAttribute}
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{Row, SQLContext}
 import org.apache.spark.{SparkContext, SparkFunSuite}
 
-class QuantileDiscretizerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class QuantileDiscretizerSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
+
   import org.apache.spark.ml.feature.QuantileDiscretizerSuite._
 
   test("Test quantile discretizer") {
@@ -67,6 +70,14 @@ class QuantileDiscretizerSuite extends SparkFunSuite with MLlibTestSparkContext
       assert(QuantileDiscretizer.getSplits(ori) === res, "Returned splits are invalid.")
     }
   }
+
+  test("read/write") {
+    val t = new QuantileDiscretizer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setNumBuckets(6)
+    testDefaultReadWrite(t)
+  }
 }
 
 private object QuantileDiscretizerSuite extends SparkFunSuite {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/SQLTransformerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/SQLTransformerSuite.scala
index d19052881ae45..553e0b870216c 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/SQLTransformerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/SQLTransformerSuite.scala
@@ -19,9 +19,11 @@ package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 
-class SQLTransformerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class SQLTransformerSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new SQLTransformer())
@@ -41,4 +43,10 @@ class SQLTransformerSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(resultSchema == expected.schema)
     assert(result.collect().toSeq == expected.collect().toSeq)
   }
+
+  test("read/write") {
+    val t = new SQLTransformer()
+      .setStatement("select * from __THIS__")
+    testDefaultReadWrite(t)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StopWordsRemoverSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StopWordsRemoverSuite.scala
index e0d433f566c25..fb217e0c1de93 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/StopWordsRemoverSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/StopWordsRemoverSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{DataFrame, Row}
 
@@ -32,7 +33,9 @@ object StopWordsRemoverSuite extends SparkFunSuite {
   }
 }
 
-class StopWordsRemoverSuite extends SparkFunSuite with MLlibTestSparkContext {
+class StopWordsRemoverSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
+
   import StopWordsRemoverSuite._
 
   test("StopWordsRemover default") {
@@ -77,4 +80,13 @@ class StopWordsRemoverSuite extends SparkFunSuite with MLlibTestSparkContext {
 
     testStopWordsRemover(remover, dataSet)
   }
+
+  test("read/write") {
+    val t = new StopWordsRemover()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setStopWords(Array("the", "a"))
+      .setCaseSensitive(true)
+    testDefaultReadWrite(t)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
index ddcdb5f4212be..be37bfb438833 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
@@ -21,12 +21,13 @@ import org.apache.spark.sql.types.{StringType, StructType, StructField, DoubleTy
 import org.apache.spark.{SparkException, SparkFunSuite}
 import org.apache.spark.ml.attribute.{Attribute, NominalAttribute}
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.functions.col
 
-class StringIndexerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class StringIndexerSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new StringIndexer)
@@ -173,4 +174,12 @@ class StringIndexerSuite extends SparkFunSuite with MLlibTestSparkContext {
     val outSchema = idxToStr.transformSchema(inSchema)
     assert(outSchema("output").dataType === StringType)
   }
+
+  test("read/write") {
+    val t = new IndexToString()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setLabels(Array("a", "b", "c"))
+    testDefaultReadWrite(t)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/TokenizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/TokenizerSuite.scala
index a02992a2407b3..36e8e5d868389 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/TokenizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/TokenizerSuite.scala
@@ -21,20 +21,30 @@ import scala.beans.BeanInfo
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{DataFrame, Row}
 
 @BeanInfo
 case class TokenizerTestData(rawText: String, wantedTokens: Array[String])
 
-class TokenizerSuite extends SparkFunSuite {
+class TokenizerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new Tokenizer)
   }
+
+  test("read/write") {
+    val t = new Tokenizer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+    testDefaultReadWrite(t)
+  }
 }
 
-class RegexTokenizerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class RegexTokenizerSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
+
   import org.apache.spark.ml.feature.RegexTokenizerSuite._
 
   test("params") {
@@ -81,6 +91,17 @@ class RegexTokenizerSuite extends SparkFunSuite with MLlibTestSparkContext {
     ))
     testRegexTokenizer(tokenizer, dataset)
   }
+
+  test("read/write") {
+    val t = new RegexTokenizer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setMinTokenLength(2)
+      .setGaps(false)
+      .setPattern("hi")
+      .setToLowercase(false)
+    testDefaultReadWrite(t)
+  }
 }
 
 object RegexTokenizerSuite extends SparkFunSuite {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorAssemblerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorAssemblerSuite.scala
index bb4d5b983e0d4..fb21ab6b9bf2c 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorAssemblerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorAssemblerSuite.scala
@@ -17,6 +17,7 @@
 
 package org.apache.spark.ml.feature
 
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.{SparkException, SparkFunSuite}
 import org.apache.spark.ml.attribute.{AttributeGroup, NominalAttribute, NumericAttribute}
 import org.apache.spark.ml.param.ParamsSuite
@@ -25,7 +26,8 @@ import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.functions.col
 
-class VectorAssemblerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class VectorAssemblerSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new VectorAssembler)
@@ -101,4 +103,11 @@ class VectorAssemblerSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(features.getAttr(5) === NumericAttribute.defaultAttr.withIndex(5))
     assert(features.getAttr(6) === NumericAttribute.defaultAttr.withIndex(6))
   }
+
+  test("read/write") {
+    val t = new VectorAssembler()
+      .setInputCols(Array("myInputCol", "myInputCol2"))
+      .setOutputCol("myOutputCol")
+    testDefaultReadWrite(t)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorSlicerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorSlicerSuite.scala
index a6c2fba8360dd..74706a23e0936 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorSlicerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorSlicerSuite.scala
@@ -20,12 +20,13 @@ package org.apache.spark.ml.feature
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.attribute.{Attribute, AttributeGroup, NumericAttribute}
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.{DataFrame, Row, SQLContext}
 
-class VectorSlicerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class VectorSlicerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     val slicer = new VectorSlicer
@@ -106,4 +107,13 @@ class VectorSlicerSuite extends SparkFunSuite with MLlibTestSparkContext {
     vectorSlicer.setIndices(Array.empty).setNames(Array("f1", "f4"))
     validateResults(vectorSlicer.transform(df))
   }
+
+  test("read/write") {
+    val t = new VectorSlicer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setIndices(Array(1, 3))
+      .setNames(Array("a", "d"))
+    testDefaultReadWrite(t)
+  }
 }

From 5aca6ad00c9d7fa43c725b8da4a10114a3a77421 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Tue, 17 Nov 2015 12:50:01 -0800
Subject: [PATCH 291/518] [SPARK-11767] [SQL] limit the size of caced batch

Currently the size of cached batch in only controlled by `batchSize` (default value is 10000), which does not work well with the size of serialized columns (for example, complex types). The memory used to build the batch is not accounted, it's easy to OOM (especially after unified memory management).

This PR introduce a hard limit as 4M for total columns (up to 50 columns of uncompressed primitive columns).

This also change the way to grow buffer, double it each time, then trim it once finished.

cc liancheng

Author: Davies Liu <davies@databricks.com>

Closes #9760 from davies/cache_limit.
---
 .../apache/spark/sql/columnar/ColumnBuilder.scala    | 12 ++++++++++--
 .../org/apache/spark/sql/columnar/ColumnStats.scala  |  2 +-
 .../sql/columnar/InMemoryColumnarTableScan.scala     |  6 +++++-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnBuilder.scala
index 7a7345a7e004b..599f30f2d73b4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnBuilder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnBuilder.scala
@@ -73,6 +73,13 @@ private[sql] class BasicColumnBuilder[JvmType](
   }
 
   override def build(): ByteBuffer = {
+    if (buffer.capacity() > buffer.position() * 1.1) {
+      // trim the buffer
+      buffer = ByteBuffer
+        .allocate(buffer.position())
+        .order(ByteOrder.nativeOrder())
+        .put(buffer.array(), 0, buffer.position())
+    }
     buffer.flip().asInstanceOf[ByteBuffer]
   }
 }
@@ -129,7 +136,8 @@ private[sql] class MapColumnBuilder(dataType: MapType)
   extends ComplexColumnBuilder(new ObjectColumnStats(dataType), MAP(dataType))
 
 private[sql] object ColumnBuilder {
-  val DEFAULT_INITIAL_BUFFER_SIZE = 1024 * 1024
+  val DEFAULT_INITIAL_BUFFER_SIZE = 128 * 1024
+  val MAX_BATCH_SIZE_IN_BYTE = 4 * 1024 * 1024L
 
   private[columnar] def ensureFreeSpace(orig: ByteBuffer, size: Int) = {
     if (orig.remaining >= size) {
@@ -137,7 +145,7 @@ private[sql] object ColumnBuilder {
     } else {
       // grow in steps of initial size
       val capacity = orig.capacity()
-      val newSize = capacity + size.max(capacity / 8 + 1)
+      val newSize = capacity + size.max(capacity)
       val pos = orig.position()
 
       ByteBuffer
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnStats.scala b/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnStats.scala
index ba61003ba41c6..91a05650585cf 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnStats.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnStats.scala
@@ -48,7 +48,7 @@ private[sql] class PartitionStatistics(tableSchema: Seq[Attribute]) extends Seri
 private[sql] sealed trait ColumnStats extends Serializable {
   protected var count = 0
   protected var nullCount = 0
-  protected var sizeInBytes = 0L
+  private[sql] var sizeInBytes = 0L
 
   /**
    * Gathers statistics information from `row(ordinal)`.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
index 2cface61e59c4..ae77298e6da2f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
@@ -133,7 +133,9 @@ private[sql] case class InMemoryRelation(
           }.toArray
 
           var rowCount = 0
-          while (rowIterator.hasNext && rowCount < batchSize) {
+          var totalSize = 0L
+          while (rowIterator.hasNext && rowCount < batchSize
+            && totalSize < ColumnBuilder.MAX_BATCH_SIZE_IN_BYTE) {
             val row = rowIterator.next()
 
             // Added for SPARK-6082. This assertion can be useful for scenarios when something
@@ -147,8 +149,10 @@ private[sql] case class InMemoryRelation(
                 s"\nRow content: $row")
 
             var i = 0
+            totalSize = 0
             while (i < row.numFields) {
               columnBuilders(i).appendFrom(row, i)
+              totalSize += columnBuilders(i).columnStats.sizeInBytes
               i += 1
             }
             rowCount += 1

From fa603e08de641df16d066302be5d5f92a60a923e Mon Sep 17 00:00:00 2001
From: Timothy Hunter <timhunter@databricks.com>
Date: Tue, 17 Nov 2015 20:51:20 +0000
Subject: [PATCH 292/518] [SPARK-11732] Removes some MiMa false positives

This adds an extra filter for private or protected classes. We only filter for package private right now.

Author: Timothy Hunter <timhunter@databricks.com>

Closes #9697 from thunterdb/spark-11732.
---
 project/MimaExcludes.scala                                 | 7 +------
 .../scala/org/apache/spark/tools/GenerateMIMAIgnore.scala  | 4 +++-
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 815951822c1ef..8b3bc96801e20 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -54,12 +54,7 @@ object MimaExcludes {
       MimaBuild.excludeSparkClass("streaming.flume.FlumeTestUtils") ++
       MimaBuild.excludeSparkClass("streaming.flume.PollingFlumeTestUtils") ++
       Seq(
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.ml.classification.LogisticCostFun.this"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.ml.classification.LogisticAggregator.add"),
-        ProblemFilters.exclude[MissingMethodProblem](
-          "org.apache.spark.ml.classification.LogisticAggregator.count"),
+        // MiMa does not deal properly with sealed traits
         ProblemFilters.exclude[MissingMethodProblem](
           "org.apache.spark.ml.classification.LogisticRegressionSummary.featuresCol")
       ) ++ Seq(
diff --git a/tools/src/main/scala/org/apache/spark/tools/GenerateMIMAIgnore.scala b/tools/src/main/scala/org/apache/spark/tools/GenerateMIMAIgnore.scala
index a0524cabff2d4..5155daa6d17bf 100644
--- a/tools/src/main/scala/org/apache/spark/tools/GenerateMIMAIgnore.scala
+++ b/tools/src/main/scala/org/apache/spark/tools/GenerateMIMAIgnore.scala
@@ -72,7 +72,9 @@ object GenerateMIMAIgnore {
         val classSymbol = mirror.classSymbol(Class.forName(className, false, classLoader))
         val moduleSymbol = mirror.staticModule(className)
         val directlyPrivateSpark =
-          isPackagePrivate(classSymbol) || isPackagePrivateModule(moduleSymbol)
+          isPackagePrivate(classSymbol) ||
+          isPackagePrivateModule(moduleSymbol) ||
+          classSymbol.isPrivate
         val developerApi = isDeveloperApi(classSymbol) || isDeveloperApi(moduleSymbol)
         val experimental = isExperimental(classSymbol) || isExperimental(moduleSymbol)
         /* Inner classes defined within a private[spark] class or object are effectively

From 328eb49e6222271337e09188853b29c8f32fb157 Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Tue, 17 Nov 2015 13:59:59 -0800
Subject: [PATCH 293/518] [SPARK-11729] Replace example code in
 ml-linear-methods.md using include_example

JIRA link: https://issues.apache.org/jira/browse/SPARK-11729

Author: Xusen Yin <yinxusen@gmail.com>

Closes #9713 from yinxusen/SPARK-11729.
---
 docs/ml-linear-methods.md                     | 218 +-----------------
 ...LinearRegressionWithElasticNetExample.java |  65 ++++++
 .../JavaLogisticRegressionSummaryExample.java |  84 +++++++
 ...gisticRegressionWithElasticNetExample.java |  55 +++++
 .../ml/linear_regression_with_elastic_net.py  |  44 ++++
 .../logistic_regression_with_elastic_net.py   |  44 ++++
 ...inearRegressionWithElasticNetExample.scala |  61 +++++
 .../ml/LogisticRegressionSummaryExample.scala |  77 +++++++
 ...isticRegressionWithElasticNetExample.scala |  53 +++++
 9 files changed, 491 insertions(+), 210 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionSummaryExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionWithElasticNetExample.java
 create mode 100644 examples/src/main/python/ml/linear_regression_with_elastic_net.py
 create mode 100644 examples/src/main/python/ml/logistic_regression_with_elastic_net.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionSummaryExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionWithElasticNetExample.scala

diff --git a/docs/ml-linear-methods.md b/docs/ml-linear-methods.md
index 85edfd373465f..0c13d7d0c82b3 100644
--- a/docs/ml-linear-methods.md
+++ b/docs/ml-linear-methods.md
@@ -57,77 +57,15 @@ $\alpha$ and `regParam` corresponds to $\lambda$.
 <div class="codetabs">
 
 <div data-lang="scala" markdown="1">
-{% highlight scala %}
-import org.apache.spark.ml.classification.LogisticRegression
-
-// Load training data
-val training = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-val lr = new LogisticRegression()
-  .setMaxIter(10)
-  .setRegParam(0.3)
-  .setElasticNetParam(0.8)
-
-// Fit the model
-val lrModel = lr.fit(training)
-
-// Print the coefficients and intercept for logistic regression
-println(s"Coefficients: ${lrModel.coefficients} Intercept: ${lrModel.intercept}")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/LogisticRegressionWithElasticNetExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
-{% highlight java %}
-import org.apache.spark.ml.classification.LogisticRegression;
-import org.apache.spark.ml.classification.LogisticRegressionModel;
-import org.apache.spark.SparkConf;
-import org.apache.spark.SparkContext;
-import org.apache.spark.sql.DataFrame;
-import org.apache.spark.sql.SQLContext;
-
-public class LogisticRegressionWithElasticNetExample {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf()
-      .setAppName("Logistic Regression with Elastic Net Example");
-
-    SparkContext sc = new SparkContext(conf);
-    SQLContext sql = new SQLContext(sc);
-    String path = "data/mllib/sample_libsvm_data.txt";
-
-    // Load training data
-    DataFrame training = sqlContext.read().format("libsvm").load(path);
-
-    LogisticRegression lr = new LogisticRegression()
-      .setMaxIter(10)
-      .setRegParam(0.3)
-      .setElasticNetParam(0.8);
-
-    // Fit the model
-    LogisticRegressionModel lrModel = lr.fit(training);
-
-    // Print the coefficients and intercept for logistic regression
-    System.out.println("Coefficients: " + lrModel.coefficients() + " Intercept: " + lrModel.intercept());
-  }
-}
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaLogisticRegressionWithElasticNetExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
-{% highlight python %}
-from pyspark.ml.classification import LogisticRegression
-
-# Load training data
-training = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-lr = LogisticRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)
-
-# Fit the model
-lrModel = lr.fit(training)
-
-# Print the coefficients and intercept for logistic regression
-print("Coefficients: " + str(lrModel.coefficients))
-print("Intercept: " + str(lrModel.intercept))
-{% endhighlight %}
+{% include_example python/ml/logistic_regression_with_elastic_net.py %}
 </div>
 
 </div>
@@ -152,33 +90,7 @@ This will likely change when multiclass classification is supported.
 
 Continuing the earlier example:
 
-{% highlight scala %}
-import org.apache.spark.ml.classification.BinaryLogisticRegressionSummary
-
-// Extract the summary from the returned LogisticRegressionModel instance trained in the earlier example
-val trainingSummary = lrModel.summary
-
-// Obtain the objective per iteration.
-val objectiveHistory = trainingSummary.objectiveHistory
-objectiveHistory.foreach(loss => println(loss))
-
-// Obtain the metrics useful to judge performance on test data.
-// We cast the summary to a BinaryLogisticRegressionSummary since the problem is a
-// binary classification problem.
-val binarySummary = trainingSummary.asInstanceOf[BinaryLogisticRegressionSummary]
-
-// Obtain the receiver-operating characteristic as a dataframe and areaUnderROC.
-val roc = binarySummary.roc
-roc.show()
-println(binarySummary.areaUnderROC)
-
-// Set the model threshold to maximize F-Measure
-val fMeasure = binarySummary.fMeasureByThreshold
-val maxFMeasure = fMeasure.select(max("F-Measure")).head().getDouble(0)
-val bestThreshold = fMeasure.where($"F-Measure" === maxFMeasure).
-  select("threshold").head().getDouble(0)
-lrModel.setThreshold(bestThreshold)
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/LogisticRegressionSummaryExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
@@ -192,39 +104,7 @@ This will likely change when multiclass classification is supported.
 
 Continuing the earlier example:
 
-{% highlight java %}
-import org.apache.spark.ml.classification.LogisticRegressionTrainingSummary;
-import org.apache.spark.ml.classification.BinaryLogisticRegressionSummary;
-import org.apache.spark.sql.functions;
-
-// Extract the summary from the returned LogisticRegressionModel instance trained in the earlier example
-LogisticRegressionTrainingSummary trainingSummary = lrModel.summary();
-
-// Obtain the loss per iteration.
-double[] objectiveHistory = trainingSummary.objectiveHistory();
-for (double lossPerIteration : objectiveHistory) {
-  System.out.println(lossPerIteration);
-}
-
-// Obtain the metrics useful to judge performance on test data.
-// We cast the summary to a BinaryLogisticRegressionSummary since the problem is a
-// binary classification problem.
-BinaryLogisticRegressionSummary binarySummary = (BinaryLogisticRegressionSummary) trainingSummary;
-
-// Obtain the receiver-operating characteristic as a dataframe and areaUnderROC.
-DataFrame roc = binarySummary.roc();
-roc.show();
-roc.select("FPR").show();
-System.out.println(binarySummary.areaUnderROC());
-
-// Get the threshold corresponding to the maximum F-Measure and rerun LogisticRegression with
-// this selected threshold.
-DataFrame fMeasure = binarySummary.fMeasureByThreshold();
-double maxFMeasure = fMeasure.select(functions.max("F-Measure")).head().getDouble(0);
-double bestThreshold = fMeasure.where(fMeasure.col("F-Measure").equalTo(maxFMeasure)).
-  select("threshold").head().getDouble(0);
-lrModel.setThreshold(bestThreshold);
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaLogisticRegressionSummaryExample.java %}
 </div>
 
 <!--- TODO: Add python model summaries once implemented -->
@@ -244,98 +124,16 @@ regression model and extracting model summary statistics.
 <div class="codetabs">
 
 <div data-lang="scala" markdown="1">
-{% highlight scala %}
-import org.apache.spark.ml.regression.LinearRegression
-
-// Load training data
-val training = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-val lr = new LinearRegression()
-  .setMaxIter(10)
-  .setRegParam(0.3)
-  .setElasticNetParam(0.8)
-
-// Fit the model
-val lrModel = lr.fit(training)
-
-// Print the coefficients and intercept for linear regression
-println(s"Coefficients: ${lrModel.coefficients} Intercept: ${lrModel.intercept}")
-
-// Summarize the model over the training set and print out some metrics
-val trainingSummary = lrModel.summary
-println(s"numIterations: ${trainingSummary.totalIterations}")
-println(s"objectiveHistory: ${trainingSummary.objectiveHistory.toList}")
-trainingSummary.residuals.show()
-println(s"RMSE: ${trainingSummary.rootMeanSquaredError}")
-println(s"r2: ${trainingSummary.r2}")
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
-{% highlight java %}
-import org.apache.spark.ml.regression.LinearRegression;
-import org.apache.spark.ml.regression.LinearRegressionModel;
-import org.apache.spark.ml.regression.LinearRegressionTrainingSummary;
-import org.apache.spark.mllib.linalg.Vectors;
-import org.apache.spark.SparkConf;
-import org.apache.spark.SparkContext;
-import org.apache.spark.sql.DataFrame;
-import org.apache.spark.sql.SQLContext;
-
-public class LinearRegressionWithElasticNetExample {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf()
-      .setAppName("Linear Regression with Elastic Net Example");
-
-    SparkContext sc = new SparkContext(conf);
-    SQLContext sql = new SQLContext(sc);
-    String path = "data/mllib/sample_libsvm_data.txt";
-
-    // Load training data
-    DataFrame training = sqlContext.read().format("libsvm").load(path);
-
-    LinearRegression lr = new LinearRegression()
-      .setMaxIter(10)
-      .setRegParam(0.3)
-      .setElasticNetParam(0.8);
-
-    // Fit the model
-    LinearRegressionModel lrModel = lr.fit(training);
-
-    // Print the coefficients and intercept for linear regression
-    System.out.println("Coefficients: " + lrModel.coefficients() + " Intercept: " + lrModel.intercept());
-
-    // Summarize the model over the training set and print out some metrics
-    LinearRegressionTrainingSummary trainingSummary = lrModel.summary();
-    System.out.println("numIterations: " + trainingSummary.totalIterations());
-    System.out.println("objectiveHistory: " + Vectors.dense(trainingSummary.objectiveHistory()));
-    trainingSummary.residuals().show();
-    System.out.println("RMSE: " + trainingSummary.rootMeanSquaredError());
-    System.out.println("r2: " + trainingSummary.r2());
-  }
-}
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 <!--- TODO: Add python model summaries once implemented -->
-{% highlight python %}
-from pyspark.ml.regression import LinearRegression
-
-# Load training data
-training = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-lr = LinearRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)
-
-# Fit the model
-lrModel = lr.fit(training)
-
-# Print the coefficients and intercept for linear regression
-print("Coefficients: " + str(lrModel.coefficients))
-print("Intercept: " + str(lrModel.intercept))
-
-# Linear regression model summary is not yet supported in Python.
-{% endhighlight %}
+{% include_example python/ml/linear_regression_with_elastic_net.py %}
 </div>
 
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java
new file mode 100644
index 0000000000000..593f8fb3e9fe9
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java
@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
+import org.apache.spark.ml.regression.LinearRegression;
+import org.apache.spark.ml.regression.LinearRegressionModel;
+import org.apache.spark.ml.regression.LinearRegressionTrainingSummary;
+import org.apache.spark.mllib.linalg.Vectors;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+// $example off$
+
+public class JavaLinearRegressionWithElasticNetExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaLinearRegressionWithElasticNetExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // $example on$
+    // Load training data
+    DataFrame training = sqlContext.read().format("libsvm")
+      .load("data/mllib/sample_libsvm_data.txt");
+
+    LinearRegression lr = new LinearRegression()
+      .setMaxIter(10)
+      .setRegParam(0.3)
+      .setElasticNetParam(0.8);
+
+    // Fit the model
+    LinearRegressionModel lrModel = lr.fit(training);
+
+    // Print the coefficients and intercept for linear regression
+    System.out.println("Coefficients: "
+      + lrModel.coefficients() + " Intercept: " + lrModel.intercept());
+
+    // Summarize the model over the training set and print out some metrics
+    LinearRegressionTrainingSummary trainingSummary = lrModel.summary();
+    System.out.println("numIterations: " + trainingSummary.totalIterations());
+    System.out.println("objectiveHistory: " + Vectors.dense(trainingSummary.objectiveHistory()));
+    trainingSummary.residuals().show();
+    System.out.println("RMSE: " + trainingSummary.rootMeanSquaredError());
+    System.out.println("r2: " + trainingSummary.r2());
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionSummaryExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionSummaryExample.java
new file mode 100644
index 0000000000000..986f3b3b28d77
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionSummaryExample.java
@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
+import org.apache.spark.ml.classification.BinaryLogisticRegressionSummary;
+import org.apache.spark.ml.classification.LogisticRegression;
+import org.apache.spark.ml.classification.LogisticRegressionModel;
+import org.apache.spark.ml.classification.LogisticRegressionTrainingSummary;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.sql.functions;
+// $example off$
+
+public class JavaLogisticRegressionSummaryExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaLogisticRegressionSummaryExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // Load training data
+    DataFrame training = sqlContext.read().format("libsvm")
+      .load("data/mllib/sample_libsvm_data.txt");
+
+    LogisticRegression lr = new LogisticRegression()
+      .setMaxIter(10)
+      .setRegParam(0.3)
+      .setElasticNetParam(0.8);
+
+    // Fit the model
+    LogisticRegressionModel lrModel = lr.fit(training);
+
+    // $example on$
+    // Extract the summary from the returned LogisticRegressionModel instance trained in the earlier
+    // example
+    LogisticRegressionTrainingSummary trainingSummary = lrModel.summary();
+
+    // Obtain the loss per iteration.
+    double[] objectiveHistory = trainingSummary.objectiveHistory();
+    for (double lossPerIteration : objectiveHistory) {
+      System.out.println(lossPerIteration);
+    }
+
+    // Obtain the metrics useful to judge performance on test data.
+    // We cast the summary to a BinaryLogisticRegressionSummary since the problem is a binary
+    // classification problem.
+    BinaryLogisticRegressionSummary binarySummary =
+      (BinaryLogisticRegressionSummary) trainingSummary;
+
+    // Obtain the receiver-operating characteristic as a dataframe and areaUnderROC.
+    DataFrame roc = binarySummary.roc();
+    roc.show();
+    roc.select("FPR").show();
+    System.out.println(binarySummary.areaUnderROC());
+
+    // Get the threshold corresponding to the maximum F-Measure and rerun LogisticRegression with
+    // this selected threshold.
+    DataFrame fMeasure = binarySummary.fMeasureByThreshold();
+    double maxFMeasure = fMeasure.select(functions.max("F-Measure")).head().getDouble(0);
+    double bestThreshold = fMeasure.where(fMeasure.col("F-Measure").equalTo(maxFMeasure))
+      .select("threshold").head().getDouble(0);
+    lrModel.setThreshold(bestThreshold);
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionWithElasticNetExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionWithElasticNetExample.java
new file mode 100644
index 0000000000000..1d28279d72a0a
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaLogisticRegressionWithElasticNetExample.java
@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
+import org.apache.spark.ml.classification.LogisticRegression;
+import org.apache.spark.ml.classification.LogisticRegressionModel;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+// $example off$
+
+public class JavaLogisticRegressionWithElasticNetExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaLogisticRegressionWithElasticNetExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // $example on$
+    // Load training data
+    DataFrame training = sqlContext.read().format("libsvm")
+      .load("data/mllib/sample_libsvm_data.txt");
+
+    LogisticRegression lr = new LogisticRegression()
+      .setMaxIter(10)
+      .setRegParam(0.3)
+      .setElasticNetParam(0.8);
+
+    // Fit the model
+    LogisticRegressionModel lrModel = lr.fit(training);
+
+    // Print the coefficients and intercept for logistic regression
+    System.out.println("Coefficients: "
+      + lrModel.coefficients() + " Intercept: " + lrModel.intercept());
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/python/ml/linear_regression_with_elastic_net.py b/examples/src/main/python/ml/linear_regression_with_elastic_net.py
new file mode 100644
index 0000000000000..b0278276330c3
--- /dev/null
+++ b/examples/src/main/python/ml/linear_regression_with_elastic_net.py
@@ -0,0 +1,44 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import print_function
+
+from pyspark import SparkContext
+from pyspark.sql import SQLContext
+# $example on$
+from pyspark.ml.regression import LinearRegression
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="LinearRegressionWithElasticNet")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load training data
+    training = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    lr = LinearRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)
+
+    # Fit the model
+    lrModel = lr.fit(training)
+
+    # Print the coefficients and intercept for linear regression
+    print("Coefficients: " + str(lrModel.coefficients))
+    print("Intercept: " + str(lrModel.intercept))
+    # $example off$
+
+    sc.stop()
diff --git a/examples/src/main/python/ml/logistic_regression_with_elastic_net.py b/examples/src/main/python/ml/logistic_regression_with_elastic_net.py
new file mode 100644
index 0000000000000..b0b1d27e13bb0
--- /dev/null
+++ b/examples/src/main/python/ml/logistic_regression_with_elastic_net.py
@@ -0,0 +1,44 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from __future__ import print_function
+
+from pyspark import SparkContext
+from pyspark.sql import SQLContext
+# $example on$
+from pyspark.ml.classification import LogisticRegression
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="LogisticRegressionWithElasticNet")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load training data
+    training = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    lr = LogisticRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)
+
+    # Fit the model
+    lrModel = lr.fit(training)
+
+    # Print the coefficients and intercept for logistic regression
+    print("Coefficients: " + str(lrModel.coefficients))
+    print("Intercept: " + str(lrModel.intercept))
+    # $example off$
+
+    sc.stop()
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala
new file mode 100644
index 0000000000000..5a51ece6f9ba7
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+// $example on$
+import org.apache.spark.ml.regression.LinearRegression
+// $example off$
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkConf, SparkContext}
+
+object LinearRegressionWithElasticNetExample {
+
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("LinearRegressionWithElasticNetExample")
+    val sc = new SparkContext(conf)
+    val sqlCtx = new SQLContext(sc)
+
+    // $example on$
+    // Load training data
+    val training = sqlCtx.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    val lr = new LinearRegression()
+      .setMaxIter(10)
+      .setRegParam(0.3)
+      .setElasticNetParam(0.8)
+
+    // Fit the model
+    val lrModel = lr.fit(training)
+
+    // Print the coefficients and intercept for linear regression
+    println(s"Coefficients: ${lrModel.coefficients} Intercept: ${lrModel.intercept}")
+
+    // Summarize the model over the training set and print out some metrics
+    val trainingSummary = lrModel.summary
+    println(s"numIterations: ${trainingSummary.totalIterations}")
+    println(s"objectiveHistory: ${trainingSummary.objectiveHistory.toList}")
+    trainingSummary.residuals.show()
+    println(s"RMSE: ${trainingSummary.rootMeanSquaredError}")
+    println(s"r2: ${trainingSummary.r2}")
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionSummaryExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionSummaryExample.scala
new file mode 100644
index 0000000000000..4c420421b670e
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionSummaryExample.scala
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+// $example on$
+import org.apache.spark.ml.classification.{BinaryLogisticRegressionSummary, LogisticRegression}
+// $example off$
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.sql.functions.max
+import org.apache.spark.{SparkConf, SparkContext}
+
+object LogisticRegressionSummaryExample {
+
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("LogisticRegressionSummaryExample")
+    val sc = new SparkContext(conf)
+    val sqlCtx = new SQLContext(sc)
+    import sqlCtx.implicits._
+
+    // Load training data
+    val training = sqlCtx.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    val lr = new LogisticRegression()
+      .setMaxIter(10)
+      .setRegParam(0.3)
+      .setElasticNetParam(0.8)
+
+    // Fit the model
+    val lrModel = lr.fit(training)
+
+    // $example on$
+    // Extract the summary from the returned LogisticRegressionModel instance trained in the earlier
+    // example
+    val trainingSummary = lrModel.summary
+
+    // Obtain the objective per iteration.
+    val objectiveHistory = trainingSummary.objectiveHistory
+    objectiveHistory.foreach(loss => println(loss))
+
+    // Obtain the metrics useful to judge performance on test data.
+    // We cast the summary to a BinaryLogisticRegressionSummary since the problem is a
+    // binary classification problem.
+    val binarySummary = trainingSummary.asInstanceOf[BinaryLogisticRegressionSummary]
+
+    // Obtain the receiver-operating characteristic as a dataframe and areaUnderROC.
+    val roc = binarySummary.roc
+    roc.show()
+    println(binarySummary.areaUnderROC)
+
+    // Set the model threshold to maximize F-Measure
+    val fMeasure = binarySummary.fMeasureByThreshold
+    val maxFMeasure = fMeasure.select(max("F-Measure")).head().getDouble(0)
+    val bestThreshold = fMeasure.where($"F-Measure" === maxFMeasure)
+      .select("threshold").head().getDouble(0)
+    lrModel.setThreshold(bestThreshold)
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionWithElasticNetExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionWithElasticNetExample.scala
new file mode 100644
index 0000000000000..9ee995b52c90b
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/LogisticRegressionWithElasticNetExample.scala
@@ -0,0 +1,53 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+// $example on$
+import org.apache.spark.ml.classification.LogisticRegression
+// $example off$
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkConf, SparkContext}
+
+object LogisticRegressionWithElasticNetExample {
+
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("LogisticRegressionWithElasticNetExample")
+    val sc = new SparkContext(conf)
+    val sqlCtx = new SQLContext(sc)
+
+    // $example on$
+    // Load training data
+    val training = sqlCtx.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    val lr = new LogisticRegression()
+      .setMaxIter(10)
+      .setRegParam(0.3)
+      .setElasticNetParam(0.8)
+
+    // Fit the model
+    val lrModel = lr.fit(training)
+
+    // Print the coefficients and intercept for logistic regression
+    println(s"Coefficients: ${lrModel.coefficients} Intercept: ${lrModel.intercept}")
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:on println

From 6eb7008b7f33a36b06d0615b68cc21ed90ad1d8a Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Tue, 17 Nov 2015 14:03:49 -0800
Subject: [PATCH 294/518] [SPARK-11763][ML] Add save,load to LogisticRegression
 Estimator

Add save/load to LogisticRegression Estimator, and refactor tests a little to make it easier to add similar support to other Estimator, Model pairs.

Moved LogisticRegressionReader/Writer to within LogisticRegressionModel

CC: mengxr

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9749 from jkbradley/lr-io-2.
---
 .../classification/LogisticRegression.scala   | 91 ++++++++++---------
 .../org/apache/spark/ml/util/ReadWrite.scala  |  1 +
 .../org/apache/spark/ml/PipelineSuite.scala   |  7 --
 .../ml/classification/ClassifierSuite.scala   | 32 +++++++
 .../LogisticRegressionSuite.scala             | 37 ++++++--
 .../ProbabilisticClassifierSuite.scala        | 14 +++
 .../spark/ml/util/DefaultReadWriteTest.scala  | 50 +++++++++-
 7 files changed, 173 insertions(+), 59 deletions(-)
 create mode 100644 mllib/src/test/scala/org/apache/spark/ml/classification/ClassifierSuite.scala

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
index a88f52674102c..71c2533bcbf47 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
@@ -157,7 +157,7 @@ private[classification] trait LogisticRegressionParams extends ProbabilisticClas
 @Experimental
 class LogisticRegression(override val uid: String)
   extends ProbabilisticClassifier[Vector, LogisticRegression, LogisticRegressionModel]
-  with LogisticRegressionParams with Logging {
+  with LogisticRegressionParams with Writable with Logging {
 
   def this() = this(Identifiable.randomUID("logreg"))
 
@@ -385,6 +385,12 @@ class LogisticRegression(override val uid: String)
   }
 
   override def copy(extra: ParamMap): LogisticRegression = defaultCopy(extra)
+
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+object LogisticRegression extends Readable[LogisticRegression] {
+  override def read: Reader[LogisticRegression] = new DefaultParamsReader[LogisticRegression]
 }
 
 /**
@@ -517,61 +523,62 @@ class LogisticRegressionModel private[ml] (
    *
    * For [[LogisticRegressionModel]], this does NOT currently save the training [[summary]].
    * An option to save [[summary]] may be added in the future.
+   *
+   * This also does not save the [[parent]] currently.
    */
-  override def write: Writer = new LogisticRegressionWriter(this)
-}
-
-
-/** [[Writer]] instance for [[LogisticRegressionModel]] */
-private[classification] class LogisticRegressionWriter(instance: LogisticRegressionModel)
-  extends Writer with Logging {
-
-  private case class Data(
-      numClasses: Int,
-      numFeatures: Int,
-      intercept: Double,
-      coefficients: Vector)
-
-  override protected def saveImpl(path: String): Unit = {
-    // Save metadata and Params
-    DefaultParamsWriter.saveMetadata(instance, path, sc)
-    // Save model data: numClasses, numFeatures, intercept, coefficients
-    val data = Data(instance.numClasses, instance.numFeatures, instance.intercept,
-      instance.coefficients)
-    val dataPath = new Path(path, "data").toString
-    sqlContext.createDataFrame(Seq(data)).write.format("parquet").save(dataPath)
-  }
+  override def write: Writer = new LogisticRegressionModel.LogisticRegressionModelWriter(this)
 }
 
 
 object LogisticRegressionModel extends Readable[LogisticRegressionModel] {
 
-  override def read: Reader[LogisticRegressionModel] = new LogisticRegressionReader
+  override def read: Reader[LogisticRegressionModel] = new LogisticRegressionModelReader
 
   override def load(path: String): LogisticRegressionModel = read.load(path)
-}
 
+  /** [[Writer]] instance for [[LogisticRegressionModel]] */
+  private[classification] class LogisticRegressionModelWriter(instance: LogisticRegressionModel)
+    extends Writer with Logging {
+
+    private case class Data(
+        numClasses: Int,
+        numFeatures: Int,
+        intercept: Double,
+        coefficients: Vector)
+
+    override protected def saveImpl(path: String): Unit = {
+      // Save metadata and Params
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      // Save model data: numClasses, numFeatures, intercept, coefficients
+      val data = Data(instance.numClasses, instance.numFeatures, instance.intercept,
+        instance.coefficients)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).write.format("parquet").save(dataPath)
+    }
+  }
 
-private[classification] class LogisticRegressionReader extends Reader[LogisticRegressionModel] {
+  private[classification] class LogisticRegressionModelReader
+    extends Reader[LogisticRegressionModel] {
 
-  /** Checked against metadata when loading model */
-  private val className = "org.apache.spark.ml.classification.LogisticRegressionModel"
+    /** Checked against metadata when loading model */
+    private val className = "org.apache.spark.ml.classification.LogisticRegressionModel"
 
-  override def load(path: String): LogisticRegressionModel = {
-    val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+    override def load(path: String): LogisticRegressionModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
 
-    val dataPath = new Path(path, "data").toString
-    val data = sqlContext.read.format("parquet").load(dataPath)
-      .select("numClasses", "numFeatures", "intercept", "coefficients").head()
-    // We will need numClasses, numFeatures in the future for multinomial logreg support.
-    // val numClasses = data.getInt(0)
-    // val numFeatures = data.getInt(1)
-    val intercept = data.getDouble(2)
-    val coefficients = data.getAs[Vector](3)
-    val model = new LogisticRegressionModel(metadata.uid, coefficients, intercept)
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.format("parquet").load(dataPath)
+        .select("numClasses", "numFeatures", "intercept", "coefficients").head()
+      // We will need numClasses, numFeatures in the future for multinomial logreg support.
+      // val numClasses = data.getInt(0)
+      // val numFeatures = data.getInt(1)
+      val intercept = data.getDouble(2)
+      val coefficients = data.getAs[Vector](3)
+      val model = new LogisticRegressionModel(metadata.uid, coefficients, intercept)
 
-    DefaultParamsReader.getAndSetParams(model, metadata)
-    model
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
   }
 }
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index 3169c9e9af5be..dddb72af5ba78 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -217,6 +217,7 @@ private[ml] object DefaultParamsWriter {
  * (json4s-serializable) params and no data. This will not handle more complex params or types with
  * data (e.g., models with coefficients).
  * @tparam T ML instance type
+ * TODO: Consider adding check for correct class name.
  */
 private[ml] class DefaultParamsReader[T] extends Reader[T] {
 
diff --git a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
index 484026b1ba9ad..7f5c3895acb0c 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
@@ -149,13 +149,6 @@ class PipelineSuite extends SparkFunSuite with MLlibTestSparkContext with Defaul
     assert(pipeline2.stages(0).isInstanceOf[WritableStage])
     val writableStage2 = pipeline2.stages(0).asInstanceOf[WritableStage]
     assert(writableStage.getIntParam === writableStage2.getIntParam)
-
-    val path = new File(tempDir, pipeline.uid).getPath
-    val stagesDir = new Path(path, "stages").toString
-    val expectedStagePath = SharedReadWrite.getStagePath(writableStage.uid, 0, 1, stagesDir)
-    assert(FileSystem.get(sc.hadoopConfiguration).exists(new Path(expectedStagePath)),
-      s"Expected stage 0 of 1 with uid ${writableStage.uid} in Pipeline with uid ${pipeline.uid}" +
-        s" to be saved to path: $expectedStagePath")
   }
 
   test("PipelineModel read/write: getStagePath") {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/ClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/ClassifierSuite.scala
new file mode 100644
index 0000000000000..d0e3fe7ad14b6
--- /dev/null
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/ClassifierSuite.scala
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.classification
+
+object ClassifierSuite {
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = Map(
+    "predictionCol" -> "myPrediction",
+    "rawPredictionCol" -> "myRawPrediction"
+  )
+
+}
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
index 51b06b7eb6d53..48ce1bb630685 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
@@ -873,15 +873,34 @@ class LogisticRegressionSuite
   }
 
   test("read/write") {
-    // Set some Params to make sure set Params are serialized.
+    def checkModelData(model: LogisticRegressionModel, model2: LogisticRegressionModel): Unit = {
+      assert(model.intercept === model2.intercept)
+      assert(model.coefficients.toArray === model2.coefficients.toArray)
+      assert(model.numClasses === model2.numClasses)
+      assert(model.numFeatures === model2.numFeatures)
+    }
     val lr = new LogisticRegression()
-      .setElasticNetParam(0.1)
-      .setMaxIter(2)
-      .fit(dataset)
-    val lr2 = testDefaultReadWrite(lr)
-    assert(lr.intercept === lr2.intercept)
-    assert(lr.coefficients.toArray === lr2.coefficients.toArray)
-    assert(lr.numClasses === lr2.numClasses)
-    assert(lr.numFeatures === lr2.numFeatures)
+    testEstimatorAndModelReadWrite(lr, dataset, LogisticRegressionSuite.allParamSettings,
+      checkModelData)
   }
 }
+
+object LogisticRegressionSuite {
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = ProbabilisticClassifierSuite.allParamSettings ++ Map(
+    "probabilityCol" -> "myProbability",
+    "thresholds" -> Array(0.4, 0.6),
+    "regParam" -> 0.01,
+    "elasticNetParam" -> 0.1,
+    "maxIter" -> 2,  // intentionally small
+    "fitIntercept" -> false,
+    "tol" -> 0.8,
+    "standardization" -> false,
+    "threshold" -> 0.6
+  )
+}
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/ProbabilisticClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/ProbabilisticClassifierSuite.scala
index fb5f00e0646c6..cfa75ecf387cd 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/ProbabilisticClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/ProbabilisticClassifierSuite.scala
@@ -57,3 +57,17 @@ class ProbabilisticClassifierSuite extends SparkFunSuite {
     assert(testModel.friendlyPredict(Vectors.dense(Array(1.0, 2.0))) === 1.0)
   }
 }
+
+object ProbabilisticClassifierSuite {
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = ClassifierSuite.allParamSettings ++ Map(
+    "probabilityCol" -> "myProbability",
+    "thresholds" -> Array(0.4, 0.6)
+  )
+
+}
diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
index c37f0503f1332..dd1e8acce9418 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
@@ -22,13 +22,17 @@ import java.io.{File, IOException}
 import org.scalatest.Suite
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.{Model, Estimator}
 import org.apache.spark.ml.param._
 import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.sql.DataFrame
 
 trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
 
   /**
    * Checks "overwrite" option and params.
+   * This saves to and loads from [[tempDir]], but creates a subdirectory with a random name
+   * in order to avoid conflicts from multiple calls to this method.
    * @param instance ML instance to test saving/loading
    * @param testParams  If true, then test values of Params.  Otherwise, just test overwrite option.
    * @tparam T ML instance type
@@ -38,7 +42,10 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
       instance: T,
       testParams: Boolean = true): T = {
     val uid = instance.uid
-    val path = new File(tempDir, uid).getPath
+    val subdirName = Identifiable.randomUID("test")
+
+    val subdir = new File(tempDir, subdirName)
+    val path = new File(subdir, uid).getPath
 
     instance.save(path)
     intercept[IOException] {
@@ -69,6 +76,47 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
     assert(another.uid === instance.uid)
     another
   }
+
+  /**
+   * Default test for Estimator, Model pairs:
+   *  - Explicitly set Params, and train model
+   *  - Test save/load using [[testDefaultReadWrite()]] on Estimator and Model
+   *  - Check Params on Estimator and Model
+   *
+   * This requires that the [[Estimator]] and [[Model]] share the same set of [[Param]]s.
+   * @param estimator  Estimator to test
+   * @param dataset  Dataset to pass to [[Estimator.fit()]]
+   * @param testParams  Set of [[Param]] values to set in estimator
+   * @param checkModelData  Method which takes the original and loaded [[Model]] and compares their
+   *                        data.  This method does not need to check [[Param]] values.
+   * @tparam E  Type of [[Estimator]]
+   * @tparam M  Type of [[Model]] produced by estimator
+   */
+  def testEstimatorAndModelReadWrite[E <: Estimator[M] with Writable, M <: Model[M] with Writable](
+      estimator: E,
+      dataset: DataFrame,
+      testParams: Map[String, Any],
+      checkModelData: (M, M) => Unit): Unit = {
+    // Set some Params to make sure set Params are serialized.
+    testParams.foreach { case (p, v) =>
+      estimator.set(estimator.getParam(p), v)
+    }
+    val model = estimator.fit(dataset)
+
+    // Test Estimator save/load
+    val estimator2 = testDefaultReadWrite(estimator)
+    testParams.foreach { case (p, v) =>
+      val param = estimator.getParam(p)
+      assert(estimator.get(param).get === estimator2.get(param).get)
+    }
+
+    // Test Model save/load
+    val model2 = testDefaultReadWrite(model)
+    testParams.foreach { case (p, v) =>
+      val param = model.getParam(p)
+      assert(model.get(param).get === model2.get(param).get)
+    }
+  }
 }
 
 class MyParams(override val uid: String) extends Params with Writable {

From 3e9e6380236985ec5b51b459f8c61f964a76ff8b Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Tue, 17 Nov 2015 14:04:49 -0800
Subject: [PATCH 295/518] [SPARK-11764][ML] make Param.jsonEncode/jsonDecode
 support Vector

This PR makes the default read/write work with simple transformers/estimators that have params of type `Param[Vector]`. jkbradley

Author: Xiangrui Meng <meng@databricks.com>

Closes #9776 from mengxr/SPARK-11764.
---
 .../org/apache/spark/ml/param/params.scala    | 12 ++++++++--
 .../apache/spark/ml/param/ParamsSuite.scala   | 22 +++++++++++++++----
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/param/params.scala b/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
index c9325709187c5..d182b0a98896c 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/param/params.scala
@@ -29,6 +29,7 @@ import org.json4s.jackson.JsonMethods._
 
 import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.mllib.linalg.{Vector, Vectors}
 
 /**
  * :: DeveloperApi ::
@@ -88,9 +89,11 @@ class Param[T](val parent: String, val name: String, val doc: String, val isVali
     value match {
       case x: String =>
         compact(render(JString(x)))
+      case v: Vector =>
+        v.toJson
       case _ =>
         throw new NotImplementedError(
-          "The default jsonEncode only supports string. " +
+          "The default jsonEncode only supports string and vector. " +
             s"${this.getClass.getName} must override jsonEncode for ${value.getClass.getName}.")
     }
   }
@@ -100,9 +103,14 @@ class Param[T](val parent: String, val name: String, val doc: String, val isVali
     parse(json) match {
       case JString(x) =>
         x.asInstanceOf[T]
+      case JObject(v) =>
+        val keys = v.map(_._1)
+        assert(keys.contains("type") && keys.contains("values"),
+          s"Expect a JSON serialized vector but cannot find fields 'type' and 'values' in $json.")
+        Vectors.fromJson(json).asInstanceOf[T]
       case _ =>
         throw new NotImplementedError(
-          "The default jsonDecode only supports string. " +
+          "The default jsonDecode only supports string and vector. " +
             s"${this.getClass.getName} must override jsonDecode to support its value type.")
     }
   }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/param/ParamsSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/param/ParamsSuite.scala
index eeb03dba2f825..a1878be747ceb 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/param/ParamsSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/param/ParamsSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.ml.param
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.mllib.linalg.{Vector, Vectors}
 
 class ParamsSuite extends SparkFunSuite {
 
@@ -80,7 +81,7 @@ class ParamsSuite extends SparkFunSuite {
       }
     }
 
-    { // StringParam
+    { // Param[String]
       val param = new Param[String](dummy, "name", "doc")
       // Currently we do not support null.
       for (value <- Seq("", "1", "abc", "quote\"", "newline\n")) {
@@ -89,6 +90,19 @@ class ParamsSuite extends SparkFunSuite {
       }
     }
 
+    { // Param[Vector]
+      val param = new Param[Vector](dummy, "name", "doc")
+      val values = Seq(
+        Vectors.dense(Array.empty[Double]),
+        Vectors.dense(0.0, 2.0),
+        Vectors.sparse(0, Array.empty, Array.empty),
+        Vectors.sparse(2, Array(1), Array(2.0)))
+      for (value <- values) {
+        val json = param.jsonEncode(value)
+        assert(param.jsonDecode(json) === value)
+      }
+    }
+
     { // IntArrayParam
       val param = new IntArrayParam(dummy, "name", "doc")
       val values: Seq[Array[Int]] = Seq(
@@ -138,7 +152,7 @@ class ParamsSuite extends SparkFunSuite {
   test("param") {
     val solver = new TestParams()
     val uid = solver.uid
-    import solver.{maxIter, inputCol}
+    import solver.{inputCol, maxIter}
 
     assert(maxIter.name === "maxIter")
     assert(maxIter.doc === "maximum number of iterations (>= 0)")
@@ -181,7 +195,7 @@ class ParamsSuite extends SparkFunSuite {
 
   test("param map") {
     val solver = new TestParams()
-    import solver.{maxIter, inputCol}
+    import solver.{inputCol, maxIter}
 
     val map0 = ParamMap.empty
 
@@ -220,7 +234,7 @@ class ParamsSuite extends SparkFunSuite {
 
   test("params") {
     val solver = new TestParams()
-    import solver.{handleInvalid, maxIter, inputCol}
+    import solver.{handleInvalid, inputCol, maxIter}
 
     val params = solver.params
     assert(params.length === 3)

From 936bc0bcbf957fa1d7cb5cfe88d628c830df5981 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Tue, 17 Nov 2015 14:23:28 -0800
Subject: [PATCH 296/518] [SPARK-11786][CORE] Tone down messages from akka
 error monitor.

There events happen normally during the app's lifecycle, so printing
out ERROR logs all the time is misleading, and can actually affect usability
of interactive shells.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9772 from vanzin/SPARK-11786.
---
 core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala
index 3fad595a0d0b0..059a7e10ec12f 100644
--- a/core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala
@@ -263,7 +263,7 @@ private[akka] class ErrorMonitor extends Actor with ActorLogReceive with Logging
   }
 
   override def receiveWithLogging: Actor.Receive = {
-    case Error(cause: Throwable, _, _, message: String) => logError(message, cause)
+    case Error(cause: Throwable, _, _, message: String) => logDebug(message, cause)
   }
 }
 

From 928d631625297857fb6998fbeb0696917fbfd60f Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Tue, 17 Nov 2015 14:48:29 -0800
Subject: [PATCH 297/518] [SPARK-11740][STREAMING] Fix the race condition of
 two checkpoints in a batch

We will do checkpoint when generating a batch and completing a batch. When the processing time of a batch is greater than the batch interval, checkpointing for completing an old batch may run after checkpointing for generating a new batch. If this happens, checkpoint of an old batch actually has the latest information, so we want to recovery from it. This PR will use the latest checkpoint time as the file name, so that we can always recovery from the latest checkpoint file.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9707 from zsxwing/fix-checkpoint.
---
 python/pyspark/streaming/tests.py             |  9 +++----
 .../apache/spark/streaming/Checkpoint.scala   | 18 +++++++++++--
 .../spark/streaming/CheckpointSuite.scala     | 27 +++++++++++++++++--
 3 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index 2983028413bb8..ff95639146e59 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -753,7 +753,6 @@ def tearDown(self):
         if self.cpd is not None:
             shutil.rmtree(self.cpd)
 
-    @unittest.skip("Enable it when we fix the checkpoint bug")
     def test_get_or_create_and_get_active_or_create(self):
         inputd = tempfile.mkdtemp()
         outputd = tempfile.mkdtemp() + "/"
@@ -822,11 +821,11 @@ def check_output(n):
         # Verify that getOrCreate() uses existing SparkContext
         self.ssc.stop(True, True)
         time.sleep(1)
-        sc = SparkContext(SparkConf())
+        self.sc = SparkContext(conf=SparkConf())
         self.setupCalled = False
         self.ssc = StreamingContext.getOrCreate(self.cpd, setup)
         self.assertFalse(self.setupCalled)
-        self.assertTrue(self.ssc.sparkContext == sc)
+        self.assertTrue(self.ssc.sparkContext == self.sc)
 
         # Verify the getActiveOrCreate() recovers from checkpoint files
         self.ssc.stop(True, True)
@@ -845,11 +844,11 @@ def check_output(n):
         # Verify that getActiveOrCreate() uses existing SparkContext
         self.ssc.stop(True, True)
         time.sleep(1)
-        self.sc = SparkContext(SparkConf())
+        self.sc = SparkContext(conf=SparkConf())
         self.setupCalled = False
         self.ssc = StreamingContext.getActiveOrCreate(self.cpd, setup)
         self.assertFalse(self.setupCalled)
-        self.assertTrue(self.ssc.sparkContext == sc)
+        self.assertTrue(self.ssc.sparkContext == self.sc)
 
         # Verify that getActiveOrCreate() calls setup() in absence of checkpoint files
         self.ssc.stop(True, True)
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala b/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala
index 0cd55d9aec2cd..fd0e8d5d690b6 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/Checkpoint.scala
@@ -187,16 +187,30 @@ class CheckpointWriter(
   private var stopped = false
   private var fs_ : FileSystem = _
 
+  @volatile private var latestCheckpointTime: Time = null
+
   class CheckpointWriteHandler(
       checkpointTime: Time,
       bytes: Array[Byte],
       clearCheckpointDataLater: Boolean) extends Runnable {
     def run() {
+      if (latestCheckpointTime == null || latestCheckpointTime < checkpointTime) {
+        latestCheckpointTime = checkpointTime
+      }
       var attempts = 0
       val startTime = System.currentTimeMillis()
       val tempFile = new Path(checkpointDir, "temp")
-      val checkpointFile = Checkpoint.checkpointFile(checkpointDir, checkpointTime)
-      val backupFile = Checkpoint.checkpointBackupFile(checkpointDir, checkpointTime)
+      // We will do checkpoint when generating a batch and completing a batch. When the processing
+      // time of a batch is greater than the batch interval, checkpointing for completing an old
+      // batch may run after checkpointing of a new batch. If this happens, checkpoint of an old
+      // batch actually has the latest information, so we want to recovery from it. Therefore, we
+      // also use the latest checkpoint time as the file name, so that we can recovery from the
+      // latest checkpoint file.
+      //
+      // Note: there is only one thread writting the checkpoint files, so we don't need to worry
+      // about thread-safety.
+      val checkpointFile = Checkpoint.checkpointFile(checkpointDir, latestCheckpointTime)
+      val backupFile = Checkpoint.checkpointBackupFile(checkpointDir, latestCheckpointTime)
 
       while (attempts < MAX_ATTEMPTS && !stopped) {
         attempts += 1
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala
index 84f5294aa39cc..b1cbc7163bee3 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala
@@ -18,7 +18,6 @@
 package org.apache.spark.streaming
 
 import java.io.{ObjectOutputStream, ByteArrayOutputStream, ByteArrayInputStream, File}
-import org.apache.spark.TestUtils
 
 import scala.collection.mutable.{ArrayBuffer, SynchronizedBuffer}
 import scala.reflect.ClassTag
@@ -30,11 +29,13 @@ import org.apache.hadoop.fs.{FileSystem, Path}
 import org.apache.hadoop.io.{IntWritable, Text}
 import org.apache.hadoop.mapred.TextOutputFormat
 import org.apache.hadoop.mapreduce.lib.output.{TextOutputFormat => NewTextOutputFormat}
+import org.mockito.Mockito.mock
 import org.scalatest.concurrent.Eventually._
 import org.scalatest.time.SpanSugar._
 
+import org.apache.spark.TestUtils
 import org.apache.spark.streaming.dstream.{DStream, FileInputDStream}
-import org.apache.spark.streaming.scheduler.{ConstantEstimator, RateTestInputDStream, RateTestReceiver}
+import org.apache.spark.streaming.scheduler._
 import org.apache.spark.util.{MutableURLClassLoader, Clock, ManualClock, Utils}
 
 /**
@@ -611,6 +612,28 @@ class CheckpointSuite extends TestSuiteBase {
     assert(ois.readObject().asInstanceOf[Class[_]].getName == "[LtestClz;")
   }
 
+  test("SPARK-11267: the race condition of two checkpoints in a batch") {
+    val jobGenerator = mock(classOf[JobGenerator])
+    val checkpointDir = Utils.createTempDir().toString
+    val checkpointWriter =
+      new CheckpointWriter(jobGenerator, conf, checkpointDir, new Configuration())
+    val bytes1 = Array.fill[Byte](10)(1)
+    new checkpointWriter.CheckpointWriteHandler(
+      Time(2000), bytes1, clearCheckpointDataLater = false).run()
+    val bytes2 = Array.fill[Byte](10)(2)
+    new checkpointWriter.CheckpointWriteHandler(
+      Time(1000), bytes2, clearCheckpointDataLater = true).run()
+    val checkpointFiles = Checkpoint.getCheckpointFiles(checkpointDir).reverse.map { path =>
+      new File(path.toUri)
+    }
+    assert(checkpointFiles.size === 2)
+    // Although bytes2 was written with an old time, it contains the latest status, so we should
+    // try to read from it at first.
+    assert(Files.toByteArray(checkpointFiles(0)) === bytes2)
+    assert(Files.toByteArray(checkpointFiles(1)) === bytes1)
+    checkpointWriter.stop()
+  }
+
   /**
    * Tests a streaming operation under checkpointing, by restarting the operation
    * from checkpoint file and verifying whether the final output is correct.

From 965245d087c18edc6c3d5baddeaf83163e32e330 Mon Sep 17 00:00:00 2001
From: Grace <jie.huang@intel.com>
Date: Tue, 17 Nov 2015 15:43:35 -0800
Subject: [PATCH 298/518] [SPARK-9552] Add force control for killExecutors to
 avoid false killing for those busy executors

By using the dynamic allocation, sometimes it occurs false killing for those busy executors. Some executors with assignments will be killed because of being idle for enough time (say 60 seconds). The root cause is that the Task-Launch listener event is asynchronized.

For example, some executors are under assigning tasks, but not sending out the listener notification yet. Meanwhile, the dynamic allocation's executor idle time is up (e.g., 60 seconds). It will trigger killExecutor event at the same time.
 1. the timer expiration starts before the listener event arrives.
 2. Then, the task is going to run on top of that killed/killing executor. It will lead to task failure finally.

Here is the proposal to fix it. We can add the force control for killExecutor. If the force control is not set (i.e., false), we'd better to check if the executor under killing is idle or busy. If the current executor has some assignment, we should not kill that executor and return back false (to indicate killing failure). In dynamic allocation, we'd better to turn off force killing (i.e., force = false), we will meet killing failure if tries to kill a busy executor. And then, the executor timer won't be invalid. Later on, the task assignment event arrives, we can remove the idle timer accordingly. So that we can avoid false killing for those busy executors in dynamic allocation.

For the rest of usages, the end users can decide if to use force killing or not by themselves.  If to turn on that option, the killExecutor will do the action without any status checking.

Author: Grace <jie.huang@intel.com>
Author: Andrew Or <andrew@databricks.com>
Author: Jie Huang <jie.huang@intel.com>

Closes #7888 from GraceH/forcekill.
---
 .../spark/ExecutorAllocationManager.scala     |  1 +
 .../scala/org/apache/spark/SparkContext.scala |  4 +-
 .../spark/scheduler/TaskSchedulerImpl.scala   | 27 +++++++---
 .../CoarseGrainedSchedulerBackend.scala       | 13 +++--
 .../StandaloneDynamicAllocationSuite.scala    | 52 ++++++++++++++++++-
 5 files changed, 82 insertions(+), 15 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala b/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala
index b93536e6536e2..6419218f47c85 100644
--- a/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala
+++ b/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala
@@ -509,6 +509,7 @@ private[spark] class ExecutorAllocationManager(
   private def onExecutorBusy(executorId: String): Unit = synchronized {
     logDebug(s"Clearing idle timer for $executorId because it is now running a task")
     removeTimes.remove(executorId)
+    executorsPendingToRemove.remove(executorId)
   }
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 4bbd0b038c00f..b5645b08f92d4 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -1461,7 +1461,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
   override def killExecutors(executorIds: Seq[String]): Boolean = {
     schedulerBackend match {
       case b: CoarseGrainedSchedulerBackend =>
-        b.killExecutors(executorIds)
+        b.killExecutors(executorIds, replace = false, force = true)
       case _ =>
         logWarning("Killing executors is only supported in coarse-grained mode")
         false
@@ -1499,7 +1499,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
   private[spark] def killAndReplaceExecutor(executorId: String): Boolean = {
     schedulerBackend match {
       case b: CoarseGrainedSchedulerBackend =>
-        b.killExecutors(Seq(executorId), replace = true)
+        b.killExecutors(Seq(executorId), replace = true, force = true)
       case _ =>
         logWarning("Killing executors is only supported in coarse-grained mode")
         false
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index 5f136690f456c..bf0419db1f75e 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -87,8 +87,8 @@ private[spark] class TaskSchedulerImpl(
   // Incrementing task IDs
   val nextTaskId = new AtomicLong(0)
 
-  // Which executor IDs we have executors on
-  val activeExecutorIds = new HashSet[String]
+  // Number of tasks running on each executor
+  private val executorIdToTaskCount = new HashMap[String, Int]
 
   // The set of executors we have on each host; this is used to compute hostsAlive, which
   // in turn is used to decide when we can attain data locality on a given host
@@ -254,6 +254,7 @@ private[spark] class TaskSchedulerImpl(
             val tid = task.taskId
             taskIdToTaskSetManager(tid) = taskSet
             taskIdToExecutorId(tid) = execId
+            executorIdToTaskCount(execId) += 1
             executorsByHost(host) += execId
             availableCpus(i) -= CPUS_PER_TASK
             assert(availableCpus(i) >= 0)
@@ -282,7 +283,7 @@ private[spark] class TaskSchedulerImpl(
     var newExecAvail = false
     for (o <- offers) {
       executorIdToHost(o.executorId) = o.host
-      activeExecutorIds += o.executorId
+      executorIdToTaskCount.getOrElseUpdate(o.executorId, 0)
       if (!executorsByHost.contains(o.host)) {
         executorsByHost(o.host) = new HashSet[String]()
         executorAdded(o.executorId, o.host)
@@ -331,7 +332,8 @@ private[spark] class TaskSchedulerImpl(
         if (state == TaskState.LOST && taskIdToExecutorId.contains(tid)) {
           // We lost this entire executor, so remember that it's gone
           val execId = taskIdToExecutorId(tid)
-          if (activeExecutorIds.contains(execId)) {
+
+          if (executorIdToTaskCount.contains(execId)) {
             removeExecutor(execId,
               SlaveLost(s"Task $tid was lost, so marking the executor as lost as well."))
             failedExecutor = Some(execId)
@@ -341,7 +343,11 @@ private[spark] class TaskSchedulerImpl(
           case Some(taskSet) =>
             if (TaskState.isFinished(state)) {
               taskIdToTaskSetManager.remove(tid)
-              taskIdToExecutorId.remove(tid)
+              taskIdToExecutorId.remove(tid).foreach { execId =>
+                if (executorIdToTaskCount.contains(execId)) {
+                  executorIdToTaskCount(execId) -= 1
+                }
+              }
             }
             if (state == TaskState.FINISHED) {
               taskSet.removeRunningTask(tid)
@@ -462,7 +468,7 @@ private[spark] class TaskSchedulerImpl(
     var failedExecutor: Option[String] = None
 
     synchronized {
-      if (activeExecutorIds.contains(executorId)) {
+      if (executorIdToTaskCount.contains(executorId)) {
         val hostPort = executorIdToHost(executorId)
         logError("Lost executor %s on %s: %s".format(executorId, hostPort, reason))
         removeExecutor(executorId, reason)
@@ -498,7 +504,8 @@ private[spark] class TaskSchedulerImpl(
    * of any running tasks, since the loss reason defines whether we'll fail those tasks.
    */
   private def removeExecutor(executorId: String, reason: ExecutorLossReason) {
-    activeExecutorIds -= executorId
+    executorIdToTaskCount -= executorId
+
     val host = executorIdToHost(executorId)
     val execs = executorsByHost.getOrElse(host, new HashSet)
     execs -= executorId
@@ -535,7 +542,11 @@ private[spark] class TaskSchedulerImpl(
   }
 
   def isExecutorAlive(execId: String): Boolean = synchronized {
-    activeExecutorIds.contains(execId)
+    executorIdToTaskCount.contains(execId)
+  }
+
+  def isExecutorBusy(execId: String): Boolean = synchronized {
+    executorIdToTaskCount.getOrElse(execId, -1) > 0
   }
 
   // By default, rack is unknown
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
index 3373caf0d15eb..6f0c910c009a5 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
@@ -453,7 +453,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
    * @return whether the kill request is acknowledged.
    */
   final override def killExecutors(executorIds: Seq[String]): Boolean = synchronized {
-    killExecutors(executorIds, replace = false)
+    killExecutors(executorIds, replace = false, force = false)
   }
 
   /**
@@ -461,9 +461,13 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
    *
    * @param executorIds identifiers of executors to kill
    * @param replace whether to replace the killed executors with new ones
+   * @param force whether to force kill busy executors
    * @return whether the kill request is acknowledged.
    */
-  final def killExecutors(executorIds: Seq[String], replace: Boolean): Boolean = synchronized {
+  final def killExecutors(
+      executorIds: Seq[String],
+      replace: Boolean,
+      force: Boolean): Boolean = synchronized {
     logInfo(s"Requesting to kill executor(s) ${executorIds.mkString(", ")}")
     val (knownExecutors, unknownExecutors) = executorIds.partition(executorDataMap.contains)
     unknownExecutors.foreach { id =>
@@ -471,7 +475,10 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     }
 
     // If an executor is already pending to be removed, do not kill it again (SPARK-9795)
-    val executorsToKill = knownExecutors.filter { id => !executorsPendingToRemove.contains(id) }
+    // If this executor is busy, do not kill it unless we are told to force kill it (SPARK-9552)
+    val executorsToKill = knownExecutors
+      .filter { id => !executorsPendingToRemove.contains(id) }
+      .filter { id => force || !scheduler.isExecutorBusy(id) }
     executorsPendingToRemove ++= executorsToKill
 
     // If we do not wish to replace the executors we kill, sync the target number of executors
diff --git a/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala b/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala
index d145e78834b1b..2fa795f846667 100644
--- a/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala
@@ -17,10 +17,11 @@
 
 package org.apache.spark.deploy
 
+import scala.collection.mutable
 import scala.concurrent.duration._
 
 import org.mockito.Mockito.{mock, when}
-import org.scalatest.BeforeAndAfterAll
+import org.scalatest.{BeforeAndAfterAll, PrivateMethodTester}
 import org.scalatest.concurrent.Eventually._
 
 import org.apache.spark._
@@ -29,6 +30,7 @@ import org.apache.spark.deploy.master.ApplicationInfo
 import org.apache.spark.deploy.master.Master
 import org.apache.spark.deploy.worker.Worker
 import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef, RpcEnv}
+import org.apache.spark.scheduler.TaskSchedulerImpl
 import org.apache.spark.scheduler.cluster._
 import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.RegisterExecutor
 
@@ -38,7 +40,8 @@ import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.RegisterE
 class StandaloneDynamicAllocationSuite
   extends SparkFunSuite
   with LocalSparkContext
-  with BeforeAndAfterAll {
+  with BeforeAndAfterAll
+  with PrivateMethodTester {
 
   private val numWorkers = 2
   private val conf = new SparkConf()
@@ -404,6 +407,41 @@ class StandaloneDynamicAllocationSuite
     assert(apps.head.getExecutorLimit === 1)
   }
 
+  test("disable force kill for busy executors (SPARK-9552)") {
+    sc = new SparkContext(appConf)
+    val appId = sc.applicationId
+    eventually(timeout(10.seconds), interval(10.millis)) {
+      val apps = getApplications()
+      assert(apps.size === 1)
+      assert(apps.head.id === appId)
+      assert(apps.head.executors.size === 2)
+      assert(apps.head.getExecutorLimit === Int.MaxValue)
+    }
+    var apps = getApplications()
+    // sync executors between the Master and the driver, needed because
+    // the driver refuses to kill executors it does not know about
+    syncExecutors(sc)
+    val executors = getExecutorIds(sc)
+    assert(executors.size === 2)
+
+    // simulate running a task on the executor
+    val getMap = PrivateMethod[mutable.HashMap[String, Int]]('executorIdToTaskCount)
+    val taskScheduler = sc.taskScheduler.asInstanceOf[TaskSchedulerImpl]
+    val executorIdToTaskCount = taskScheduler invokePrivate getMap()
+    executorIdToTaskCount(executors.head) = 1
+    // kill the busy executor without force; this should fail
+    assert(killExecutor(sc, executors.head, force = false))
+    apps = getApplications()
+    assert(apps.head.executors.size === 2)
+
+    // force kill busy executor
+    assert(killExecutor(sc, executors.head, force = true))
+    apps = getApplications()
+    // kill executor successfully
+    assert(apps.head.executors.size === 1)
+
+  }
+
   // ===============================
   // | Utility methods for testing |
   // ===============================
@@ -455,6 +493,16 @@ class StandaloneDynamicAllocationSuite
     sc.killExecutors(getExecutorIds(sc).take(n))
   }
 
+  /** Kill the given executor, specifying whether to force kill it. */
+  private def killExecutor(sc: SparkContext, executorId: String, force: Boolean): Boolean = {
+    syncExecutors(sc)
+    sc.schedulerBackend match {
+      case b: CoarseGrainedSchedulerBackend =>
+        b.killExecutors(Seq(executorId), replace = false, force)
+      case _ => fail("expected coarse grained scheduler")
+    }
+  }
+
   /**
    * Return a list of executor IDs belonging to this application.
    *

From e29656f8e7fa19686b448292e20d8bbf07ab9f11 Mon Sep 17 00:00:00 2001
From: Rohan Bhanderi <rohan.bhanderi@sjsu.edu>
Date: Tue, 17 Nov 2015 15:45:39 -0800
Subject: [PATCH 299/518] [MINOR] Correct comments in JavaDirectKafkaWordCount

Author: Rohan Bhanderi <rohan.bhanderi@sjsu.edu>

Closes #9781 from RohanBhanderi/patch-3.
---
 .../examples/streaming/JavaDirectKafkaWordCount.java      | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKafkaWordCount.java b/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKafkaWordCount.java
index bab9f2478e779..f9a5e7f69ffe1 100644
--- a/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKafkaWordCount.java
+++ b/examples/src/main/java/org/apache/spark/examples/streaming/JavaDirectKafkaWordCount.java
@@ -35,12 +35,12 @@
 
 /**
  * Consumes messages from one or more topics in Kafka and does wordcount.
- * Usage: DirectKafkaWordCount <brokers> <topics>
+ * Usage: JavaDirectKafkaWordCount <brokers> <topics>
  *   <brokers> is a list of one or more Kafka brokers
  *   <topics> is a list of one or more kafka topics to consume from
  *
  * Example:
- *    $ bin/run-example streaming.KafkaWordCount broker1-host:port,broker2-host:port topic1,topic2
+ *    $ bin/run-example streaming.JavaDirectKafkaWordCount broker1-host:port,broker2-host:port topic1,topic2
  */
 
 public final class JavaDirectKafkaWordCount {
@@ -48,7 +48,7 @@ public final class JavaDirectKafkaWordCount {
 
   public static void main(String[] args) {
     if (args.length < 2) {
-      System.err.println("Usage: DirectKafkaWordCount <brokers> <topics>\n" +
+      System.err.println("Usage: JavaDirectKafkaWordCount <brokers> <topics>\n" +
           "  <brokers> is a list of one or more Kafka brokers\n" +
           "  <topics> is a list of one or more kafka topics to consume from\n\n");
       System.exit(1);
@@ -59,7 +59,7 @@ public static void main(String[] args) {
     String brokers = args[0];
     String topics = args[1];
 
-    // Create context with 2 second batch interval
+    // Create context with a 2 seconds batch interval
     SparkConf sparkConf = new SparkConf().setAppName("JavaDirectKafkaWordCount");
     JavaStreamingContext jssc = new JavaStreamingContext(sparkConf, Durations.seconds(2));
 

From 3720b1480c7d050ca20f98d65762224ae5639607 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Tue, 17 Nov 2015 15:47:39 -0800
Subject: [PATCH 300/518] [SPARK-11790][STREAMING][TESTS] Increase the
 connection timeout

Sometimes, EmbeddedZookeeper may need more than 6 seconds to setup up in a slow Jenkins worker. So just increase the timeout, it won't increase the test time if the test passes.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9778 from zsxwing/SPARK-11790.
---
 .../scala/org/apache/spark/streaming/kafka/KafkaTestUtils.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaTestUtils.scala b/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaTestUtils.scala
index c9fd715d3d554..86394ea8a685e 100644
--- a/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaTestUtils.scala
+++ b/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaTestUtils.scala
@@ -52,7 +52,7 @@ private[kafka] class KafkaTestUtils extends Logging {
   // Zookeeper related configurations
   private val zkHost = "localhost"
   private var zkPort: Int = 0
-  private val zkConnectionTimeout = 6000
+  private val zkConnectionTimeout = 60000
   private val zkSessionTimeout = 6000
 
   private var zookeeper: EmbeddedZookeeper = _

From 52c734b589277267be07e245c959199db92aa189 Mon Sep 17 00:00:00 2001
From: Holden Karau <holden@us.ibm.com>
Date: Tue, 17 Nov 2015 15:51:03 -0800
Subject: [PATCH 301/518] [SPARK-11771][YARN][TRIVIAL] maximum memory in yarn
 is controlled by two params have both in error msg

When we exceed the max memory tell users to increase both params instead of just the one.

Author: Holden Karau <holden@us.ibm.com>

Closes #9758 from holdenk/SPARK-11771-maximum-memory-in-yarn-is-controlled-by-two-params-have-both-in-error-msg.
---
 yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
index a3f33d80184a3..ba799884f5689 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
@@ -258,7 +258,8 @@ private[spark] class Client(
     if (executorMem > maxMem) {
       throw new IllegalArgumentException(s"Required executor memory (${args.executorMemory}" +
         s"+$executorMemoryOverhead MB) is above the max threshold ($maxMem MB) of this cluster! " +
-        "Please increase the value of 'yarn.scheduler.maximum-allocation-mb'.")
+        "Please check the values of 'yarn.scheduler.maximum-allocation-mb' and/or " +
+        "'yarn.nodemanager.resource.memory-mb'.")
     }
     val amMem = args.amMemory + amMemoryOverhead
     if (amMem > maxMem) {

From b362d50fca30693f97bd859984157bb8a76d48a1 Mon Sep 17 00:00:00 2001
From: Jacek Lewandowski <lewandowski.jacek@gmail.com>
Date: Tue, 17 Nov 2015 15:57:43 -0800
Subject: [PATCH 302/518] [SPARK-11726] Throw exception on timeout when waiting
 for REST server response

Author: Jacek Lewandowski <lewandowski.jacek@gmail.com>

Closes #9692 from jacek-lewandowski/SPARK-11726.
---
 .../spark/deploy/rest/RestSubmissionClient.scala   | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/rest/RestSubmissionClient.scala b/core/src/main/scala/org/apache/spark/deploy/rest/RestSubmissionClient.scala
index 957a928bc402b..f0dd667ea1b26 100644
--- a/core/src/main/scala/org/apache/spark/deploy/rest/RestSubmissionClient.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/rest/RestSubmissionClient.scala
@@ -19,16 +19,19 @@ package org.apache.spark.deploy.rest
 
 import java.io.{DataOutputStream, FileNotFoundException}
 import java.net.{ConnectException, HttpURLConnection, SocketException, URL}
+import java.util.concurrent.TimeoutException
 import javax.servlet.http.HttpServletResponse
 
 import scala.collection.mutable
+import scala.concurrent.duration._
+import scala.concurrent.{Await, Future}
 import scala.io.Source
 
 import com.fasterxml.jackson.core.JsonProcessingException
 import com.google.common.base.Charsets
 
-import org.apache.spark.{Logging, SparkConf, SPARK_VERSION => sparkVersion}
 import org.apache.spark.util.Utils
+import org.apache.spark.{Logging, SPARK_VERSION => sparkVersion, SparkConf}
 
 /**
  * A client that submits applications to a [[RestSubmissionServer]].
@@ -225,7 +228,8 @@ private[spark] class RestSubmissionClient(master: String) extends Logging {
    * Exposed for testing.
    */
   private[rest] def readResponse(connection: HttpURLConnection): SubmitRestProtocolResponse = {
-    try {
+    import scala.concurrent.ExecutionContext.Implicits.global
+    val responseFuture = Future {
       val dataStream =
         if (connection.getResponseCode == HttpServletResponse.SC_OK) {
           connection.getInputStream
@@ -251,11 +255,15 @@ private[spark] class RestSubmissionClient(master: String) extends Logging {
           throw new SubmitRestProtocolException(
             s"Message received from server was not a response:\n${unexpected.toJson}")
       }
-    } catch {
+    }
+
+    try { Await.result(responseFuture, 10.seconds) } catch {
       case unreachable @ (_: FileNotFoundException | _: SocketException) =>
         throw new SubmitRestConnectionException("Unable to connect to server", unreachable)
       case malformed @ (_: JsonProcessingException | _: SubmitRestProtocolException) =>
         throw new SubmitRestProtocolException("Malformed response received from server", malformed)
+      case timeout: TimeoutException =>
+        throw new SubmitRestConnectionException("No response from server", timeout)
     }
   }
 

From 75a292291062783129d02607302f91c85655975e Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Tue, 17 Nov 2015 16:57:52 -0800
Subject: [PATCH 303/518] [SPARK-9065][STREAMING][PYSPARK] Add MessageHandler
 for Kafka Python API

Fixed the merge conflicts in #7410

Closes #7410

Author: Shixiong Zhu <shixiong@databricks.com>
Author: jerryshao <saisai.shao@intel.com>
Author: jerryshao <sshao@hortonworks.com>

Closes #9742 from zsxwing/pr7410.
---
 .../spark/streaming/kafka/KafkaUtils.scala    | 245 ++++++++++++------
 project/MimaExcludes.scala                    |   6 +
 python/pyspark/streaming/kafka.py             | 111 +++++++-
 python/pyspark/streaming/tests.py             |  35 +++
 4 files changed, 299 insertions(+), 98 deletions(-)

diff --git a/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaUtils.scala b/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaUtils.scala
index 3128222077537..ad2fb8aa5f24c 100644
--- a/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaUtils.scala
+++ b/external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaUtils.scala
@@ -17,25 +17,29 @@
 
 package org.apache.spark.streaming.kafka
 
+import java.io.OutputStream
 import java.lang.{Integer => JInt, Long => JLong}
 import java.util.{List => JList, Map => JMap, Set => JSet}
 
 import scala.collection.JavaConverters._
 import scala.reflect.ClassTag
 
+import com.google.common.base.Charsets.UTF_8
 import kafka.common.TopicAndPartition
 import kafka.message.MessageAndMetadata
-import kafka.serializer.{Decoder, DefaultDecoder, StringDecoder}
+import kafka.serializer.{DefaultDecoder, Decoder, StringDecoder}
+import net.razorvine.pickle.{Opcodes, Pickler, IObjectPickler}
 
 import org.apache.spark.api.java.function.{Function => JFunction}
-import org.apache.spark.api.java.{JavaPairRDD, JavaRDD, JavaSparkContext}
+import org.apache.spark.streaming.util.WriteAheadLogUtils
+import org.apache.spark.{SparkContext, SparkException}
+import org.apache.spark.api.java.{JavaSparkContext, JavaPairRDD, JavaRDD}
+import org.apache.spark.api.python.SerDeUtil
 import org.apache.spark.rdd.RDD
 import org.apache.spark.storage.StorageLevel
 import org.apache.spark.streaming.StreamingContext
-import org.apache.spark.streaming.api.java.{JavaInputDStream, JavaPairInputDStream, JavaPairReceiverInputDStream, JavaStreamingContext}
-import org.apache.spark.streaming.dstream.{InputDStream, ReceiverInputDStream}
-import org.apache.spark.streaming.util.WriteAheadLogUtils
-import org.apache.spark.{SparkContext, SparkException}
+import org.apache.spark.streaming.api.java._
+import org.apache.spark.streaming.dstream.{DStream, InputDStream, ReceiverInputDStream}
 
 object KafkaUtils {
   /**
@@ -184,6 +188,27 @@ object KafkaUtils {
     }
   }
 
+  private[kafka] def getFromOffsets(
+      kc: KafkaCluster,
+      kafkaParams: Map[String, String],
+      topics: Set[String]
+    ): Map[TopicAndPartition, Long] = {
+    val reset = kafkaParams.get("auto.offset.reset").map(_.toLowerCase)
+    val result = for {
+      topicPartitions <- kc.getPartitions(topics).right
+      leaderOffsets <- (if (reset == Some("smallest")) {
+        kc.getEarliestLeaderOffsets(topicPartitions)
+      } else {
+        kc.getLatestLeaderOffsets(topicPartitions)
+      }).right
+    } yield {
+      leaderOffsets.map { case (tp, lo) =>
+          (tp, lo.offset)
+      }
+    }
+    KafkaCluster.checkErrors(result)
+  }
+
   /**
    * Create a RDD from Kafka using offset ranges for each topic and partition.
    *
@@ -246,7 +271,7 @@ object KafkaUtils {
       // This could be avoided by refactoring KafkaRDD.leaders and KafkaCluster to use Broker
       leaders.map {
         case (tp: TopicAndPartition, Broker(host, port)) => (tp, (host, port))
-      }.toMap
+      }
     }
     val cleanedHandler = sc.clean(messageHandler)
     checkOffsets(kc, offsetRanges)
@@ -406,23 +431,9 @@ object KafkaUtils {
   ): InputDStream[(K, V)] = {
     val messageHandler = (mmd: MessageAndMetadata[K, V]) => (mmd.key, mmd.message)
     val kc = new KafkaCluster(kafkaParams)
-    val reset = kafkaParams.get("auto.offset.reset").map(_.toLowerCase)
-
-    val result = for {
-      topicPartitions <- kc.getPartitions(topics).right
-      leaderOffsets <- (if (reset == Some("smallest")) {
-        kc.getEarliestLeaderOffsets(topicPartitions)
-      } else {
-        kc.getLatestLeaderOffsets(topicPartitions)
-      }).right
-    } yield {
-      val fromOffsets = leaderOffsets.map { case (tp, lo) =>
-          (tp, lo.offset)
-      }
-      new DirectKafkaInputDStream[K, V, KD, VD, (K, V)](
-        ssc, kafkaParams, fromOffsets, messageHandler)
-    }
-    KafkaCluster.checkErrors(result)
+    val fromOffsets = getFromOffsets(kc, kafkaParams, topics)
+    new DirectKafkaInputDStream[K, V, KD, VD, (K, V)](
+      ssc, kafkaParams, fromOffsets, messageHandler)
   }
 
   /**
@@ -550,6 +561,8 @@ object KafkaUtils {
  * takes care of known parameters instead of passing them from Python
  */
 private[kafka] class KafkaUtilsPythonHelper {
+  import KafkaUtilsPythonHelper._
+
   def createStream(
       jssc: JavaStreamingContext,
       kafkaParams: JMap[String, String],
@@ -566,86 +579,92 @@ private[kafka] class KafkaUtilsPythonHelper {
       storageLevel)
   }
 
-  def createRDD(
+  def createRDDWithoutMessageHandler(
       jsc: JavaSparkContext,
       kafkaParams: JMap[String, String],
       offsetRanges: JList[OffsetRange],
-      leaders: JMap[TopicAndPartition, Broker]): JavaPairRDD[Array[Byte], Array[Byte]] = {
-    val messageHandler = new JFunction[MessageAndMetadata[Array[Byte], Array[Byte]],
-      (Array[Byte], Array[Byte])] {
-      def call(t1: MessageAndMetadata[Array[Byte], Array[Byte]]): (Array[Byte], Array[Byte]) =
-        (t1.key(), t1.message())
-    }
+      leaders: JMap[TopicAndPartition, Broker]): JavaRDD[(Array[Byte], Array[Byte])] = {
+    val messageHandler =
+      (mmd: MessageAndMetadata[Array[Byte], Array[Byte]]) => (mmd.key, mmd.message)
+    new JavaRDD(createRDD(jsc, kafkaParams, offsetRanges, leaders, messageHandler))
+  }
 
-    val jrdd = KafkaUtils.createRDD[
-      Array[Byte],
-      Array[Byte],
-      DefaultDecoder,
-      DefaultDecoder,
-      (Array[Byte], Array[Byte])](
-        jsc,
-        classOf[Array[Byte]],
-        classOf[Array[Byte]],
-        classOf[DefaultDecoder],
-        classOf[DefaultDecoder],
-        classOf[(Array[Byte], Array[Byte])],
-        kafkaParams,
-        offsetRanges.toArray(new Array[OffsetRange](offsetRanges.size())),
-        leaders,
-        messageHandler
-      )
-    new JavaPairRDD(jrdd.rdd)
+  def createRDDWithMessageHandler(
+      jsc: JavaSparkContext,
+      kafkaParams: JMap[String, String],
+      offsetRanges: JList[OffsetRange],
+      leaders: JMap[TopicAndPartition, Broker]): JavaRDD[Array[Byte]] = {
+    val messageHandler = (mmd: MessageAndMetadata[Array[Byte], Array[Byte]]) =>
+      new PythonMessageAndMetadata(
+        mmd.topic, mmd.partition, mmd.offset, mmd.key(), mmd.message())
+    val rdd = createRDD(jsc, kafkaParams, offsetRanges, leaders, messageHandler).
+        mapPartitions(picklerIterator)
+    new JavaRDD(rdd)
   }
 
-  def createDirectStream(
+  private def createRDD[V: ClassTag](
+      jsc: JavaSparkContext,
+      kafkaParams: JMap[String, String],
+      offsetRanges: JList[OffsetRange],
+      leaders: JMap[TopicAndPartition, Broker],
+      messageHandler: MessageAndMetadata[Array[Byte], Array[Byte]] => V): RDD[V] = {
+    KafkaUtils.createRDD[Array[Byte], Array[Byte], DefaultDecoder, DefaultDecoder, V](
+      jsc.sc,
+      kafkaParams.asScala.toMap,
+      offsetRanges.toArray(new Array[OffsetRange](offsetRanges.size())),
+      leaders.asScala.toMap,
+      messageHandler
+    )
+  }
+
+  def createDirectStreamWithoutMessageHandler(
+      jssc: JavaStreamingContext,
+      kafkaParams: JMap[String, String],
+      topics: JSet[String],
+      fromOffsets: JMap[TopicAndPartition, JLong]): JavaDStream[(Array[Byte], Array[Byte])] = {
+    val messageHandler =
+      (mmd: MessageAndMetadata[Array[Byte], Array[Byte]]) => (mmd.key, mmd.message)
+    new JavaDStream(createDirectStream(jssc, kafkaParams, topics, fromOffsets, messageHandler))
+  }
+
+  def createDirectStreamWithMessageHandler(
       jssc: JavaStreamingContext,
       kafkaParams: JMap[String, String],
       topics: JSet[String],
-      fromOffsets: JMap[TopicAndPartition, JLong]
-    ): JavaPairInputDStream[Array[Byte], Array[Byte]] = {
+      fromOffsets: JMap[TopicAndPartition, JLong]): JavaDStream[Array[Byte]] = {
+    val messageHandler = (mmd: MessageAndMetadata[Array[Byte], Array[Byte]]) =>
+      new PythonMessageAndMetadata(mmd.topic, mmd.partition, mmd.offset, mmd.key(), mmd.message())
+    val stream = createDirectStream(jssc, kafkaParams, topics, fromOffsets, messageHandler).
+      mapPartitions(picklerIterator)
+    new JavaDStream(stream)
+  }
 
-    if (!fromOffsets.isEmpty) {
+  private def createDirectStream[V: ClassTag](
+      jssc: JavaStreamingContext,
+      kafkaParams: JMap[String, String],
+      topics: JSet[String],
+      fromOffsets: JMap[TopicAndPartition, JLong],
+      messageHandler: MessageAndMetadata[Array[Byte], Array[Byte]] => V): DStream[V] = {
+
+    val currentFromOffsets = if (!fromOffsets.isEmpty) {
       val topicsFromOffsets = fromOffsets.keySet().asScala.map(_.topic)
       if (topicsFromOffsets != topics.asScala.toSet) {
         throw new IllegalStateException(
           s"The specified topics: ${topics.asScala.toSet.mkString(" ")} " +
           s"do not equal to the topic from offsets: ${topicsFromOffsets.mkString(" ")}")
       }
-    }
-
-    if (fromOffsets.isEmpty) {
-      KafkaUtils.createDirectStream[Array[Byte], Array[Byte], DefaultDecoder, DefaultDecoder](
-        jssc,
-        classOf[Array[Byte]],
-        classOf[Array[Byte]],
-        classOf[DefaultDecoder],
-        classOf[DefaultDecoder],
-        kafkaParams,
-        topics)
+      Map(fromOffsets.asScala.mapValues { _.longValue() }.toSeq: _*)
     } else {
-      val messageHandler = new JFunction[MessageAndMetadata[Array[Byte], Array[Byte]],
-        (Array[Byte], Array[Byte])] {
-        def call(t1: MessageAndMetadata[Array[Byte], Array[Byte]]): (Array[Byte], Array[Byte]) =
-          (t1.key(), t1.message())
-      }
-
-      val jstream = KafkaUtils.createDirectStream[
-        Array[Byte],
-        Array[Byte],
-        DefaultDecoder,
-        DefaultDecoder,
-        (Array[Byte], Array[Byte])](
-          jssc,
-          classOf[Array[Byte]],
-          classOf[Array[Byte]],
-          classOf[DefaultDecoder],
-          classOf[DefaultDecoder],
-          classOf[(Array[Byte], Array[Byte])],
-          kafkaParams,
-          fromOffsets,
-          messageHandler)
-      new JavaPairInputDStream(jstream.inputDStream)
+      val kc = new KafkaCluster(Map(kafkaParams.asScala.toSeq: _*))
+      KafkaUtils.getFromOffsets(
+        kc, Map(kafkaParams.asScala.toSeq: _*), Set(topics.asScala.toSeq: _*))
     }
+
+    KafkaUtils.createDirectStream[Array[Byte], Array[Byte], DefaultDecoder, DefaultDecoder, V](
+      jssc.ssc,
+      Map(kafkaParams.asScala.toSeq: _*),
+      Map(currentFromOffsets.toSeq: _*),
+      messageHandler)
   }
 
   def createOffsetRange(topic: String, partition: JInt, fromOffset: JLong, untilOffset: JLong
@@ -669,3 +688,57 @@ private[kafka] class KafkaUtilsPythonHelper {
     kafkaRDD.offsetRanges.toSeq.asJava
   }
 }
+
+private object KafkaUtilsPythonHelper {
+  private var initialized = false
+
+  def initialize(): Unit = {
+    SerDeUtil.initialize()
+    synchronized {
+      if (!initialized) {
+        new PythonMessageAndMetadataPickler().register()
+        initialized = true
+      }
+    }
+  }
+
+  initialize()
+
+  def picklerIterator(iter: Iterator[Any]): Iterator[Array[Byte]] = {
+    new SerDeUtil.AutoBatchedPickler(iter)
+  }
+
+  case class PythonMessageAndMetadata(
+      topic: String,
+      partition: JInt,
+      offset: JLong,
+      key: Array[Byte],
+      message: Array[Byte])
+
+  class PythonMessageAndMetadataPickler extends IObjectPickler {
+    private val module = "pyspark.streaming.kafka"
+
+    def register(): Unit = {
+      Pickler.registerCustomPickler(classOf[PythonMessageAndMetadata], this)
+      Pickler.registerCustomPickler(this.getClass, this)
+    }
+
+    def pickle(obj: Object, out: OutputStream, pickler: Pickler) {
+      if (obj == this) {
+        out.write(Opcodes.GLOBAL)
+        out.write(s"$module\nKafkaMessageAndMetadata\n".getBytes(UTF_8))
+      } else {
+        pickler.save(this)
+        val msgAndMetaData = obj.asInstanceOf[PythonMessageAndMetadata]
+        out.write(Opcodes.MARK)
+        pickler.save(msgAndMetaData.topic)
+        pickler.save(msgAndMetaData.partition)
+        pickler.save(msgAndMetaData.offset)
+        pickler.save(msgAndMetaData.key)
+        pickler.save(msgAndMetaData.message)
+        out.write(Opcodes.TUPLE)
+        out.write(Opcodes.REDUCE)
+      }
+    }
+  }
+}
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 8b3bc96801e20..eb70d27c34c20 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -136,6 +136,12 @@ object MimaExcludes {
         // SPARK-11766 add toJson to Vector
         ProblemFilters.exclude[MissingMethodProblem](
           "org.apache.spark.mllib.linalg.Vector.toJson")
+      ) ++ Seq(
+        // SPARK-9065 Support message handler in Kafka Python API
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.streaming.kafka.KafkaUtilsPythonHelper.createDirectStream"),
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.streaming.kafka.KafkaUtilsPythonHelper.createRDD")
       )
     case v if v.startsWith("1.5") =>
       Seq(
diff --git a/python/pyspark/streaming/kafka.py b/python/pyspark/streaming/kafka.py
index 06e159172ab51..cdf97ec73aaf9 100644
--- a/python/pyspark/streaming/kafka.py
+++ b/python/pyspark/streaming/kafka.py
@@ -19,12 +19,14 @@
 
 from pyspark.rdd import RDD
 from pyspark.storagelevel import StorageLevel
-from pyspark.serializers import PairDeserializer, NoOpSerializer
+from pyspark.serializers import AutoBatchedSerializer, PickleSerializer, PairDeserializer, \
+    NoOpSerializer
 from pyspark.streaming import DStream
 from pyspark.streaming.dstream import TransformedDStream
 from pyspark.streaming.util import TransformFunction
 
-__all__ = ['Broker', 'KafkaUtils', 'OffsetRange', 'TopicAndPartition', 'utf8_decoder']
+__all__ = ['Broker', 'KafkaMessageAndMetadata', 'KafkaUtils', 'OffsetRange',
+           'TopicAndPartition', 'utf8_decoder']
 
 
 def utf8_decoder(s):
@@ -82,7 +84,8 @@ def createStream(ssc, zkQuorum, groupId, topics, kafkaParams=None,
 
     @staticmethod
     def createDirectStream(ssc, topics, kafkaParams, fromOffsets=None,
-                           keyDecoder=utf8_decoder, valueDecoder=utf8_decoder):
+                           keyDecoder=utf8_decoder, valueDecoder=utf8_decoder,
+                           messageHandler=None):
         """
         .. note:: Experimental
 
@@ -107,6 +110,8 @@ def createDirectStream(ssc, topics, kafkaParams, fromOffsets=None,
                             point of the stream.
         :param keyDecoder:  A function used to decode key (default is utf8_decoder).
         :param valueDecoder:  A function used to decode value (default is utf8_decoder).
+        :param messageHandler: A function used to convert KafkaMessageAndMetadata. You can assess
+                               meta using messageHandler (default is None).
         :return: A DStream object
         """
         if fromOffsets is None:
@@ -116,6 +121,14 @@ def createDirectStream(ssc, topics, kafkaParams, fromOffsets=None,
         if not isinstance(kafkaParams, dict):
             raise TypeError("kafkaParams should be dict")
 
+        def funcWithoutMessageHandler(k_v):
+            return (keyDecoder(k_v[0]), valueDecoder(k_v[1]))
+
+        def funcWithMessageHandler(m):
+            m._set_key_decoder(keyDecoder)
+            m._set_value_decoder(valueDecoder)
+            return messageHandler(m)
+
         try:
             helperClass = ssc._jvm.java.lang.Thread.currentThread().getContextClassLoader() \
                 .loadClass("org.apache.spark.streaming.kafka.KafkaUtilsPythonHelper")
@@ -123,20 +136,28 @@ def createDirectStream(ssc, topics, kafkaParams, fromOffsets=None,
 
             jfromOffsets = dict([(k._jTopicAndPartition(helper),
                                   v) for (k, v) in fromOffsets.items()])
-            jstream = helper.createDirectStream(ssc._jssc, kafkaParams, set(topics), jfromOffsets)
+            if messageHandler is None:
+                ser = PairDeserializer(NoOpSerializer(), NoOpSerializer())
+                func = funcWithoutMessageHandler
+                jstream = helper.createDirectStreamWithoutMessageHandler(
+                    ssc._jssc, kafkaParams, set(topics), jfromOffsets)
+            else:
+                ser = AutoBatchedSerializer(PickleSerializer())
+                func = funcWithMessageHandler
+                jstream = helper.createDirectStreamWithMessageHandler(
+                    ssc._jssc, kafkaParams, set(topics), jfromOffsets)
         except Py4JJavaError as e:
             if 'ClassNotFoundException' in str(e.java_exception):
                 KafkaUtils._printErrorMsg(ssc.sparkContext)
             raise e
 
-        ser = PairDeserializer(NoOpSerializer(), NoOpSerializer())
-        stream = DStream(jstream, ssc, ser) \
-            .map(lambda k_v: (keyDecoder(k_v[0]), valueDecoder(k_v[1])))
+        stream = DStream(jstream, ssc, ser).map(func)
         return KafkaDStream(stream._jdstream, ssc, stream._jrdd_deserializer)
 
     @staticmethod
     def createRDD(sc, kafkaParams, offsetRanges, leaders=None,
-                  keyDecoder=utf8_decoder, valueDecoder=utf8_decoder):
+                  keyDecoder=utf8_decoder, valueDecoder=utf8_decoder,
+                  messageHandler=None):
         """
         .. note:: Experimental
 
@@ -149,6 +170,8 @@ def createRDD(sc, kafkaParams, offsetRanges, leaders=None,
             map, in which case leaders will be looked up on the driver.
         :param keyDecoder:  A function used to decode key (default is utf8_decoder)
         :param valueDecoder:  A function used to decode value (default is utf8_decoder)
+        :param messageHandler: A function used to convert KafkaMessageAndMetadata. You can assess
+                               meta using messageHandler (default is None).
         :return: A RDD object
         """
         if leaders is None:
@@ -158,6 +181,14 @@ def createRDD(sc, kafkaParams, offsetRanges, leaders=None,
         if not isinstance(offsetRanges, list):
             raise TypeError("offsetRanges should be list")
 
+        def funcWithoutMessageHandler(k_v):
+            return (keyDecoder(k_v[0]), valueDecoder(k_v[1]))
+
+        def funcWithMessageHandler(m):
+            m._set_key_decoder(keyDecoder)
+            m._set_value_decoder(valueDecoder)
+            return messageHandler(m)
+
         try:
             helperClass = sc._jvm.java.lang.Thread.currentThread().getContextClassLoader() \
                 .loadClass("org.apache.spark.streaming.kafka.KafkaUtilsPythonHelper")
@@ -165,15 +196,21 @@ def createRDD(sc, kafkaParams, offsetRanges, leaders=None,
             joffsetRanges = [o._jOffsetRange(helper) for o in offsetRanges]
             jleaders = dict([(k._jTopicAndPartition(helper),
                               v._jBroker(helper)) for (k, v) in leaders.items()])
-            jrdd = helper.createRDD(sc._jsc, kafkaParams, joffsetRanges, jleaders)
+            if messageHandler is None:
+                jrdd = helper.createRDDWithoutMessageHandler(
+                    sc._jsc, kafkaParams, joffsetRanges, jleaders)
+                ser = PairDeserializer(NoOpSerializer(), NoOpSerializer())
+                rdd = RDD(jrdd, sc, ser).map(funcWithoutMessageHandler)
+            else:
+                jrdd = helper.createRDDWithMessageHandler(
+                    sc._jsc, kafkaParams, joffsetRanges, jleaders)
+                rdd = RDD(jrdd, sc).map(funcWithMessageHandler)
         except Py4JJavaError as e:
             if 'ClassNotFoundException' in str(e.java_exception):
                 KafkaUtils._printErrorMsg(sc)
             raise e
 
-        ser = PairDeserializer(NoOpSerializer(), NoOpSerializer())
-        rdd = RDD(jrdd, sc, ser).map(lambda k_v: (keyDecoder(k_v[0]), valueDecoder(k_v[1])))
-        return KafkaRDD(rdd._jrdd, rdd.ctx, rdd._jrdd_deserializer)
+        return KafkaRDD(rdd._jrdd, sc, rdd._jrdd_deserializer)
 
     @staticmethod
     def _printErrorMsg(sc):
@@ -365,3 +402,53 @@ def _jdstream(self):
         dstream = self._sc._jvm.PythonTransformedDStream(self.prev._jdstream.dstream(), jfunc)
         self._jdstream_val = dstream.asJavaDStream()
         return self._jdstream_val
+
+
+class KafkaMessageAndMetadata(object):
+    """
+    Kafka message and metadata information. Including topic, partition, offset and message
+    """
+
+    def __init__(self, topic, partition, offset, key, message):
+        """
+        Python wrapper of Kafka MessageAndMetadata
+        :param topic: topic name of this Kafka message
+        :param partition: partition id of this Kafka message
+        :param offset: Offset of this Kafka message in the specific partition
+        :param key: key payload of this Kafka message, can be null if this Kafka message has no key
+                    specified, the return data is undecoded bytearry.
+        :param message: actual message payload of this Kafka message, the return data is
+                        undecoded bytearray.
+        """
+        self.topic = topic
+        self.partition = partition
+        self.offset = offset
+        self._rawKey = key
+        self._rawMessage = message
+        self._keyDecoder = utf8_decoder
+        self._valueDecoder = utf8_decoder
+
+    def __str__(self):
+        return "KafkaMessageAndMetadata(topic: %s, partition: %d, offset: %d, key and message...)" \
+               % (self.topic, self.partition, self.offset)
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __reduce__(self):
+        return (KafkaMessageAndMetadata,
+                (self.topic, self.partition, self.offset, self._rawKey, self._rawMessage))
+
+    def _set_key_decoder(self, decoder):
+        self._keyDecoder = decoder
+
+    def _set_value_decoder(self, decoder):
+        self._valueDecoder = decoder
+
+    @property
+    def key(self):
+        return self._keyDecoder(self._rawKey)
+
+    @property
+    def message(self):
+        return self._valueDecoder(self._rawMessage)
diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index ff95639146e59..0bcd1f15532b5 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -1042,6 +1042,41 @@ def test_topic_and_partition_equality(self):
         self.assertNotEqual(topic_and_partition_a, topic_and_partition_c)
         self.assertNotEqual(topic_and_partition_a, topic_and_partition_d)
 
+    @unittest.skipIf(sys.version >= "3", "long type not support")
+    def test_kafka_rdd_message_handler(self):
+        """Test Python direct Kafka RDD MessageHandler."""
+        topic = self._randomTopic()
+        sendData = {"a": 1, "b": 1, "c": 2}
+        offsetRanges = [OffsetRange(topic, 0, long(0), long(sum(sendData.values())))]
+        kafkaParams = {"metadata.broker.list": self._kafkaTestUtils.brokerAddress()}
+
+        def getKeyAndDoubleMessage(m):
+            return m and (m.key, m.message * 2)
+
+        self._kafkaTestUtils.createTopic(topic)
+        self._kafkaTestUtils.sendMessages(topic, sendData)
+        rdd = KafkaUtils.createRDD(self.sc, kafkaParams, offsetRanges,
+                                   messageHandler=getKeyAndDoubleMessage)
+        self._validateRddResult({"aa": 1, "bb": 1, "cc": 2}, rdd)
+
+    @unittest.skipIf(sys.version >= "3", "long type not support")
+    def test_kafka_direct_stream_message_handler(self):
+        """Test the Python direct Kafka stream MessageHandler."""
+        topic = self._randomTopic()
+        sendData = {"a": 1, "b": 2, "c": 3}
+        kafkaParams = {"metadata.broker.list": self._kafkaTestUtils.brokerAddress(),
+                       "auto.offset.reset": "smallest"}
+
+        self._kafkaTestUtils.createTopic(topic)
+        self._kafkaTestUtils.sendMessages(topic, sendData)
+
+        def getKeyAndDoubleMessage(m):
+            return m and (m.key, m.message * 2)
+
+        stream = KafkaUtils.createDirectStream(self.ssc, [topic], kafkaParams,
+                                               messageHandler=getKeyAndDoubleMessage)
+        self._validateStreamResult({"aa": 1, "bb": 2, "cc": 3}, stream)
+
 
 class FlumeStreamTests(PySparkStreamingTestCase):
     timeout = 20  # seconds

From ed8d1531f93f697c54bbaecefe08c37c32b0d391 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 17 Nov 2015 19:02:44 -0800
Subject: [PATCH 304/518] [SPARK-11793][SQL] Dataset should set the resolved
 encoders internally for maps.

I also wrote a test case -- but unfortunately the test case is not working due to SPARK-11795.

Author: Reynold Xin <rxin@databricks.com>

Closes #9784 from rxin/SPARK-11503.
---
 .../src/main/scala/org/apache/spark/sql/Dataset.scala |  3 ++-
 .../scala/org/apache/spark/sql/DatasetSuite.scala     | 11 +++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 4cc3aa2465f2e..bd01dd4dc5799 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -199,11 +199,12 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def mapPartitions[U : Encoder](func: Iterator[T] => Iterator[U]): Dataset[U] = {
+    encoderFor[T].assertUnresolved()
     new Dataset[U](
       sqlContext,
       MapPartitions[T, U](
         func,
-        encoderFor[T],
+        resolvedTEncoder,
         encoderFor[U],
         encoderFor[U].schema.toAttributes,
         logicalPlan))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index c23dd46d3767b..a3922340ccc9a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -73,6 +73,17 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       ("a", 2), ("b", 3), ("c", 4))
   }
 
+  ignore("Dataset should set the resolved encoders internally for maps") {
+    // TODO: Enable this once we fix SPARK-11793.
+    val ds: Dataset[(ClassData, Long)] = Seq(ClassData("one", 1), ClassData("two", 2)).toDS()
+        .map(c => ClassData(c.a, c.b + 1))
+        .groupBy(p => p).count()
+
+    checkAnswer(
+      ds,
+      (ClassData("one", 1), 1L), (ClassData("two", 2), 1L))
+  }
+
   test("select") {
     val ds = Seq(("a", 1) , ("b", 2), ("c", 3)).toDS()
     checkAnswer(

From bf25f9bdfc7bd8533890c7df1b35afa912dc6d3d Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Tue, 17 Nov 2015 19:39:39 -0800
Subject: [PATCH 305/518] [SPARK-11016] Move RoaringBitmap to explicit Kryo
 serializer

Fix the serialization of RoaringBitmap with Kyro serializer

This PR came from https://github.com/metamx/spark/pull/1, thanks to drcrallen

Author: Davies Liu <davies@databricks.com>
Author: Charles Allen <charles@allen-net.com>

Closes #9748 from davies/SPARK-11016.
---
 .../spark/serializer/KryoSerializer.scala     | 64 ++++++++++++++++---
 1 file changed, 55 insertions(+), 9 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
index c5195c1143a8f..1bcb3175a3016 100644
--- a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
+++ b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.serializer
 
-import java.io.{EOFException, IOException, InputStream, OutputStream}
+import java.io.{EOFException, IOException, InputStream, OutputStream, DataInput, DataOutput}
 import java.nio.ByteBuffer
 import javax.annotation.Nullable
 
@@ -25,12 +25,12 @@ import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
 
-import com.esotericsoftware.kryo.{Kryo, KryoException}
+import com.esotericsoftware.kryo.{Kryo, KryoException, Serializer => KryoClassSerializer}
 import com.esotericsoftware.kryo.io.{Input => KryoInput, Output => KryoOutput}
 import com.esotericsoftware.kryo.serializers.{JavaSerializer => KryoJavaSerializer}
 import com.twitter.chill.{AllScalaRegistrar, EmptyScalaKryoInstantiator}
 import org.apache.avro.generic.{GenericData, GenericRecord}
-import org.roaringbitmap.{ArrayContainer, BitmapContainer, RoaringArray, RoaringBitmap}
+import org.roaringbitmap.RoaringBitmap
 
 import org.apache.spark._
 import org.apache.spark.api.python.PythonBroadcast
@@ -94,6 +94,9 @@ class KryoSerializer(conf: SparkConf)
     for (cls <- KryoSerializer.toRegister) {
       kryo.register(cls)
     }
+    for ((cls, ser) <- KryoSerializer.toRegisterSerializer) {
+      kryo.register(cls, ser)
+    }
 
     // For results returned by asJavaIterable. See JavaIterableWrapperSerializer.
     kryo.register(JavaIterableWrapperSerializer.wrapperClass, new JavaIterableWrapperSerializer)
@@ -363,12 +366,6 @@ private[serializer] object KryoSerializer {
     classOf[StorageLevel],
     classOf[CompressedMapStatus],
     classOf[HighlyCompressedMapStatus],
-    classOf[RoaringBitmap],
-    classOf[RoaringArray],
-    classOf[RoaringArray.Element],
-    classOf[Array[RoaringArray.Element]],
-    classOf[ArrayContainer],
-    classOf[BitmapContainer],
     classOf[CompactBuffer[_]],
     classOf[BlockManagerId],
     classOf[Array[Byte]],
@@ -377,6 +374,55 @@ private[serializer] object KryoSerializer {
     classOf[BoundedPriorityQueue[_]],
     classOf[SparkConf]
   )
+
+  private val toRegisterSerializer = Map[Class[_], KryoClassSerializer[_]](
+    classOf[RoaringBitmap] -> new KryoClassSerializer[RoaringBitmap]() {
+      override def write(kryo: Kryo, output: KryoOutput, bitmap: RoaringBitmap): Unit = {
+        bitmap.serialize(new KryoOutputDataOutputBridge(output))
+      }
+      override def read(kryo: Kryo, input: KryoInput, cls: Class[RoaringBitmap]): RoaringBitmap = {
+        val ret = new RoaringBitmap
+        ret.deserialize(new KryoInputDataInputBridge(input))
+        ret
+      }
+    }
+  )
+}
+
+private[serializer] class KryoInputDataInputBridge(input: KryoInput) extends DataInput {
+  override def readLong(): Long = input.readLong()
+  override def readChar(): Char = input.readChar()
+  override def readFloat(): Float = input.readFloat()
+  override def readByte(): Byte = input.readByte()
+  override def readShort(): Short = input.readShort()
+  override def readUTF(): String = input.readString() // readString in kryo does utf8
+  override def readInt(): Int = input.readInt()
+  override def readUnsignedShort(): Int = input.readShortUnsigned()
+  override def skipBytes(n: Int): Int = input.skip(n.toLong).toInt
+  override def readFully(b: Array[Byte]): Unit = input.read(b)
+  override def readFully(b: Array[Byte], off: Int, len: Int): Unit = input.read(b, off, len)
+  override def readLine(): String = throw new UnsupportedOperationException("readLine")
+  override def readBoolean(): Boolean = input.readBoolean()
+  override def readUnsignedByte(): Int = input.readByteUnsigned()
+  override def readDouble(): Double = input.readDouble()
+}
+
+private[serializer] class KryoOutputDataOutputBridge(output: KryoOutput) extends DataOutput {
+  override def writeFloat(v: Float): Unit = output.writeFloat(v)
+  // There is no "readChars" counterpart, except maybe "readLine", which is not supported
+  override def writeChars(s: String): Unit = throw new UnsupportedOperationException("writeChars")
+  override def writeDouble(v: Double): Unit = output.writeDouble(v)
+  override def writeUTF(s: String): Unit = output.writeString(s) // writeString in kryo does UTF8
+  override def writeShort(v: Int): Unit = output.writeShort(v)
+  override def writeInt(v: Int): Unit = output.writeInt(v)
+  override def writeBoolean(v: Boolean): Unit = output.writeBoolean(v)
+  override def write(b: Int): Unit = output.write(b)
+  override def write(b: Array[Byte]): Unit = output.write(b)
+  override def write(b: Array[Byte], off: Int, len: Int): Unit = output.write(b, off, len)
+  override def writeBytes(s: String): Unit = output.writeString(s)
+  override def writeChar(v: Int): Unit = output.writeChar(v.toChar)
+  override def writeLong(v: Long): Unit = output.writeLong(v)
+  override def writeByte(v: Int): Unit = output.writeByte(v)
 }
 
 /**

From e33053ee0015025bbcfddb20cc9216c225bbe624 Mon Sep 17 00:00:00 2001
From: Kent Yao <yaooqinn@hotmail.com>
Date: Tue, 17 Nov 2015 19:44:29 -0800
Subject: [PATCH 306/518] [SPARK-11583] [CORE] MapStatus Using RoaringBitmap
 More Properly

This PR upgrade the version of RoaringBitmap to 0.5.10, to optimize the memory layout, will be much smaller when most of blocks are empty.

This PR is based on #9661 (fix conflicts), see all of the comments at https://github.com/apache/spark/pull/9661 .

Author: Kent Yao <yaooqinn@hotmail.com>
Author: Davies Liu <davies@databricks.com>
Author: Charles Allen <charles@allen-net.com>

Closes #9746 from davies/roaring_mapstatus.
---
 .../apache/spark/scheduler/MapStatus.scala    |  5 +--
 .../spark/serializer/KryoSerializer.scala     |  6 ++--
 .../spark/scheduler/MapStatusSuite.scala      | 31 +++++++++++++++++++
 pom.xml                                       |  2 +-
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala b/core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala
index 1efce124c0a6b..b2e9a97129f08 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala
@@ -122,8 +122,7 @@ private[spark] class CompressedMapStatus(
 
 /**
  * A [[MapStatus]] implementation that only stores the average size of non-empty blocks,
- * plus a bitmap for tracking which blocks are empty.  During serialization, this bitmap
- * is compressed.
+ * plus a bitmap for tracking which blocks are empty.
  *
  * @param loc location where the task is being executed
  * @param numNonEmptyBlocks the number of non-empty blocks
@@ -194,6 +193,8 @@ private[spark] object HighlyCompressedMapStatus {
     } else {
       0
     }
+    emptyBlocks.trim()
+    emptyBlocks.runOptimize()
     new HighlyCompressedMapStatus(loc, numNonEmptyBlocks, emptyBlocks, avgSize)
   }
 }
diff --git a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
index 1bcb3175a3016..d5ba690ed04be 100644
--- a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
+++ b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.serializer
 
-import java.io.{EOFException, IOException, InputStream, OutputStream, DataInput, DataOutput}
+import java.io.{DataInput, DataOutput, EOFException, IOException, InputStream, OutputStream}
 import java.nio.ByteBuffer
 import javax.annotation.Nullable
 
@@ -25,9 +25,9 @@ import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
 
-import com.esotericsoftware.kryo.{Kryo, KryoException, Serializer => KryoClassSerializer}
 import com.esotericsoftware.kryo.io.{Input => KryoInput, Output => KryoOutput}
 import com.esotericsoftware.kryo.serializers.{JavaSerializer => KryoJavaSerializer}
+import com.esotericsoftware.kryo.{Kryo, KryoException, Serializer => KryoClassSerializer}
 import com.twitter.chill.{AllScalaRegistrar, EmptyScalaKryoInstantiator}
 import org.apache.avro.generic.{GenericData, GenericRecord}
 import org.roaringbitmap.RoaringBitmap
@@ -38,8 +38,8 @@ import org.apache.spark.broadcast.HttpBroadcast
 import org.apache.spark.network.util.ByteUnit
 import org.apache.spark.scheduler.{CompressedMapStatus, HighlyCompressedMapStatus}
 import org.apache.spark.storage._
-import org.apache.spark.util.{Utils, BoundedPriorityQueue, SerializableConfiguration, SerializableJobConf}
 import org.apache.spark.util.collection.CompactBuffer
+import org.apache.spark.util.{BoundedPriorityQueue, SerializableConfiguration, SerializableJobConf, Utils}
 
 /**
  * A Spark serializer that uses the [[https://code.google.com/p/kryo/ Kryo serialization library]].
diff --git a/core/src/test/scala/org/apache/spark/scheduler/MapStatusSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/MapStatusSuite.scala
index b8e466fab4506..15c8de61b8240 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/MapStatusSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/MapStatusSuite.scala
@@ -21,6 +21,7 @@ import org.apache.spark.storage.BlockManagerId
 
 import org.apache.spark.{SparkConf, SparkFunSuite}
 import org.apache.spark.serializer.JavaSerializer
+import org.roaringbitmap.RoaringBitmap
 
 import scala.util.Random
 
@@ -97,4 +98,34 @@ class MapStatusSuite extends SparkFunSuite {
     val buf = ser.newInstance().serialize(status)
     ser.newInstance().deserialize[MapStatus](buf)
   }
+
+  test("RoaringBitmap: runOptimize succeeded") {
+    val r = new RoaringBitmap
+    (1 to 200000).foreach(i =>
+      if (i % 200 != 0) {
+        r.add(i)
+      }
+    )
+    val size1 = r.getSizeInBytes
+    val success = r.runOptimize()
+    r.trim()
+    val size2 = r.getSizeInBytes
+    assert(size1 > size2)
+    assert(success)
+  }
+
+  test("RoaringBitmap: runOptimize failed") {
+    val r = new RoaringBitmap
+    (1 to 200000).foreach(i =>
+      if (i % 200 == 0) {
+        r.add(i)
+      }
+    )
+    val size1 = r.getSizeInBytes
+    val success = r.runOptimize()
+    r.trim()
+    val size2 = r.getSizeInBytes
+    assert(size1 === size2)
+    assert(!success)
+  }
 }
diff --git a/pom.xml b/pom.xml
index 2a8a445057174..940e2d8740bf1 100644
--- a/pom.xml
+++ b/pom.xml
@@ -637,7 +637,7 @@
       <dependency>
         <groupId>org.roaringbitmap</groupId>
         <artifactId>RoaringBitmap</artifactId>
-        <version>0.4.5</version>
+        <version>0.5.11</version>
       </dependency>
       <dependency>
         <groupId>commons-net</groupId>

From 98be8169f07eb0f1b8f01776c71d0e1ed3d5e4d5 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Tue, 17 Nov 2015 19:50:02 -0800
Subject: [PATCH 307/518] [SPARK-11737] [SQL] Fix serialization of UTF8String
 with Kyro

The default implementation of serialization UTF8String with Kyro may be not correct (BYTE_ARRAY_OFFSET could be different across JVM)

Author: Davies Liu <davies@databricks.com>

Closes #9704 from davies/kyro_string.
---
 unsafe/pom.xml                                |  4 ++++
 .../apache/spark/unsafe/types/UTF8String.java | 24 +++++++++++++++++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/unsafe/pom.xml b/unsafe/pom.xml
index caf1f77890b58..a1c1111364ee8 100644
--- a/unsafe/pom.xml
+++ b/unsafe/pom.xml
@@ -36,6 +36,10 @@
   </properties>
 
   <dependencies>
+    <dependency>
+      <groupId>com.twitter</groupId>
+      <artifactId>chill_${scala.binary.version}</artifactId>
+    </dependency>
 
     <!-- Core dependencies -->
     <dependency>
diff --git a/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java b/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java
index b7aecb5102ba6..4bd3fd7772079 100644
--- a/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java
+++ b/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java
@@ -24,6 +24,11 @@
 import java.util.Arrays;
 import java.util.Map;
 
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.KryoSerializable;
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.Output;
+
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.array.ByteArrayMethods;
 
@@ -38,9 +43,9 @@
  * <p>
  * Note: This is not designed for general use cases, should not be used outside SQL.
  */
-public final class UTF8String implements Comparable<UTF8String>, Externalizable {
+public final class UTF8String implements Comparable<UTF8String>, Externalizable, KryoSerializable {
 
-  // These are only updated by readExternal()
+  // These are only updated by readExternal() or read()
   @Nonnull
   private Object base;
   private long offset;
@@ -1003,4 +1008,19 @@ public void readExternal(ObjectInput in) throws IOException, ClassNotFoundExcept
     in.readFully((byte[]) base);
   }
 
+  @Override
+  public void write(Kryo kryo, Output out) {
+    byte[] bytes = getBytes();
+    out.writeInt(bytes.length);
+    out.write(bytes);
+  }
+
+  @Override
+  public void read(Kryo kryo, Input in) {
+    this.offset = BYTE_ARRAY_OFFSET;
+    this.numBytes = in.readInt();
+    this.base = new byte[numBytes];
+    in.read((byte[]) base);
+  }
+
 }

From 91f4b6f2db12650dfc33a576803ba8aeccf935dd Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 17 Nov 2015 21:40:58 -0800
Subject: [PATCH 308/518] [SPARK-11797][SQL] collect, first, and take should
 use encoders for serialization

They were previously using Spark's default serializer for serialization.

Author: Reynold Xin <rxin@databricks.com>

Closes #9787 from rxin/SPARK-11797.
---
 .../scala/org/apache/spark/sql/Dataset.scala  | 17 +++++++----
 .../org/apache/spark/sql/DatasetSuite.scala   | 30 ++++++++++++++++++-
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index bd01dd4dc5799..718ed812dd64c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -22,6 +22,7 @@ import scala.collection.JavaConverters._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
 import org.apache.spark.api.java.function._
+import org.apache.spark.sql.catalyst.InternalRow
 
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
@@ -199,7 +200,6 @@ class Dataset[T] private[sql](
    * @since 1.6.0
    */
   def mapPartitions[U : Encoder](func: Iterator[T] => Iterator[U]): Dataset[U] = {
-    encoderFor[T].assertUnresolved()
     new Dataset[U](
       sqlContext,
       MapPartitions[T, U](
@@ -519,7 +519,7 @@ class Dataset[T] private[sql](
    * Returns the first element in this [[Dataset]].
    * @since 1.6.0
    */
-  def first(): T = rdd.first()
+  def first(): T = take(1).head
 
   /**
    * Returns an array that contains all the elements in this [[Dataset]].
@@ -530,7 +530,14 @@ class Dataset[T] private[sql](
    * For Java API, use [[collectAsList]].
    * @since 1.6.0
    */
-  def collect(): Array[T] = rdd.collect()
+  def collect(): Array[T] = {
+    // This is different from Dataset.rdd in that it collects Rows, and then runs the encoders
+    // to convert the rows into objects of type T.
+    val tEnc = resolvedTEncoder
+    val input = queryExecution.analyzed.output
+    val bound = tEnc.bind(input)
+    queryExecution.toRdd.map(_.copy()).collect().map(bound.fromRow)
+  }
 
   /**
    * Returns an array that contains all the elements in this [[Dataset]].
@@ -541,7 +548,7 @@ class Dataset[T] private[sql](
    * For Java API, use [[collectAsList]].
    * @since 1.6.0
    */
-  def collectAsList(): java.util.List[T] = rdd.collect().toSeq.asJava
+  def collectAsList(): java.util.List[T] = collect().toSeq.asJava
 
   /**
    * Returns the first `num` elements of this [[Dataset]] as an array.
@@ -551,7 +558,7 @@ class Dataset[T] private[sql](
    *
    * @since 1.6.0
    */
-  def take(num: Int): Array[T] = rdd.take(num)
+  def take(num: Int): Array[T] = withPlan(Limit(Literal(num), _)).collect()
 
   /**
    * Returns the first `num` elements of this [[Dataset]] as an array.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index a3922340ccc9a..ea29428c55088 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql
 
+import java.io.{ObjectInput, ObjectOutput, Externalizable}
+
 import scala.language.postfixOps
 
 import org.apache.spark.sql.functions._
@@ -24,6 +26,20 @@ import org.apache.spark.sql.test.SharedSQLContext
 
 case class ClassData(a: String, b: Int)
 
+/**
+ * A class used to test serialization using encoders. This class throws exceptions when using
+ * Java serialization -- so the only way it can be "serialized" is through our encoders.
+ */
+case class NonSerializableCaseClass(value: String) extends Externalizable {
+  override def readExternal(in: ObjectInput): Unit = {
+    throw new UnsupportedOperationException
+  }
+
+  override def writeExternal(out: ObjectOutput): Unit = {
+    throw new UnsupportedOperationException
+  }
+}
+
 class DatasetSuite extends QueryTest with SharedSQLContext {
   import testImplicits._
 
@@ -41,6 +57,16 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       1, 1, 1)
   }
 
+  test("collect, first, and take should use encoders for serialization") {
+    val item = NonSerializableCaseClass("abcd")
+    val ds = Seq(item).toDS()
+    assert(ds.collect().head == item)
+    assert(ds.collectAsList().get(0) == item)
+    assert(ds.first() == item)
+    assert(ds.take(1).head == item)
+    assert(ds.takeAsList(1).get(0) == item)
+  }
+
   test("as tuple") {
     val data = Seq(("a", 1), ("b", 2)).toDF("a", "b")
     checkAnswer(
@@ -75,6 +101,8 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
 
   ignore("Dataset should set the resolved encoders internally for maps") {
     // TODO: Enable this once we fix SPARK-11793.
+    // We inject a group by here to make sure this test case is future proof
+    // when we implement better pipelining and local execution mode.
     val ds: Dataset[(ClassData, Long)] = Seq(ClassData("one", 1), ClassData("two", 2)).toDS()
         .map(c => ClassData(c.a, c.b + 1))
         .groupBy(p => p).count()
@@ -219,7 +247,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       ("a", 30), ("b", 3), ("c", 1))
   }
 
-  test("groupBy function, fatMap") {
+  test("groupBy function, flatMap") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy(v => (v._1, "word"))
     val agged = grouped.flatMap { case (g, iter) => Iterator(g._1, iter.map(_._2).sum.toString) }

From 8fb775ba874dd0488667bf299a7b49760062dc00 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Tue, 17 Nov 2015 22:13:15 -0800
Subject: [PATCH 309/518] [SPARK-11755][R] SparkR should export "predict"
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The bug described at [SPARK-11755](https://issues.apache.org/jira/browse/SPARK-11755), after exporting ```predict``` we can both get the help information from the SparkR and base R package like the following:
```Java
> help(predict)
Help on topic ‘predict’ was found in the following packages:

  Package               Library
  SparkR                /Users/yanboliang/data/trunk2/spark/R/lib
  stats                 /Library/Frameworks/R.framework/Versions/3.2/Resources/library

Choose one

1: Make predictions from a model {SparkR}
2: Model Predictions {stats}
```

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9732 from yanboliang/spark-11755.
---
 R/pkg/R/generics.R | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 612e639f8ad99..afdeffc2abd83 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -1054,6 +1054,10 @@ setGeneric("year", function(x) { standardGeneric("year") })
 #' @export
 setGeneric("glm")
 
+#' @rdname predict
+#' @export
+setGeneric("predict", function(object, ...) { standardGeneric("predict") })
+
 #' @rdname rbind
 #' @export
 setGeneric("rbind", signature = "...")

From 446738e51fcda50cf2dc44123ff6bf12a1611dc0 Mon Sep 17 00:00:00 2001
From: tedyu <yuzhihong@gmail.com>
Date: Tue, 17 Nov 2015 22:47:53 -0800
Subject: [PATCH 310/518] [SPARK-11761] Prevent the call to
 StreamingContext#stop() in the listener bus's thread

See discussion toward the tail of https://github.com/apache/spark/pull/9723
From zsxwing :
```
The user should not call stop or other long-time work in a listener since it will block the listener thread, and prevent from stopping SparkContext/StreamingContext.

I cannot see an approach since we need to stop the listener bus's thread before stopping SparkContext/StreamingContext totally.
```
Proposed solution is to prevent the call to StreamingContext#stop() in the listener bus's thread.

Author: tedyu <yuzhihong@gmail.com>

Closes #9741 from tedyu/master.
---
 .../spark/util/AsynchronousListenerBus.scala  | 46 +++++++++++--------
 .../spark/streaming/StreamingContext.scala    |  6 ++-
 .../streaming/StreamingListenerSuite.scala    | 34 ++++++++++++++
 3 files changed, 67 insertions(+), 19 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala b/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
index c20627b056bef..6c1fca71f2281 100644
--- a/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
+++ b/core/src/main/scala/org/apache/spark/util/AsynchronousListenerBus.scala
@@ -19,6 +19,7 @@ package org.apache.spark.util
 
 import java.util.concurrent._
 import java.util.concurrent.atomic.AtomicBoolean
+import scala.util.DynamicVariable
 
 import org.apache.spark.SparkContext
 
@@ -60,25 +61,27 @@ private[spark] abstract class AsynchronousListenerBus[L <: AnyRef, E](name: Stri
   private val listenerThread = new Thread(name) {
     setDaemon(true)
     override def run(): Unit = Utils.tryOrStopSparkContext(sparkContext) {
-      while (true) {
-        eventLock.acquire()
-        self.synchronized {
-          processingEvent = true
-        }
-        try {
-          val event = eventQueue.poll
-          if (event == null) {
-            // Get out of the while loop and shutdown the daemon thread
-            if (!stopped.get) {
-              throw new IllegalStateException("Polling `null` from eventQueue means" +
-                " the listener bus has been stopped. So `stopped` must be true")
-            }
-            return
-          }
-          postToAll(event)
-        } finally {
+      AsynchronousListenerBus.withinListenerThread.withValue(true) {
+        while (true) {
+          eventLock.acquire()
           self.synchronized {
-            processingEvent = false
+            processingEvent = true
+          }
+          try {
+            val event = eventQueue.poll
+            if (event == null) {
+              // Get out of the while loop and shutdown the daemon thread
+              if (!stopped.get) {
+                throw new IllegalStateException("Polling `null` from eventQueue means" +
+                  " the listener bus has been stopped. So `stopped` must be true")
+              }
+              return
+            }
+            postToAll(event)
+          } finally {
+            self.synchronized {
+              processingEvent = false
+            }
           }
         }
       }
@@ -177,3 +180,10 @@ private[spark] abstract class AsynchronousListenerBus[L <: AnyRef, E](name: Stri
    */
   def onDropEvent(event: E): Unit
 }
+
+private[spark] object AsynchronousListenerBus {
+  /* Allows for Context to check whether stop() call is made within listener thread
+  */
+  val withinListenerThread: DynamicVariable[Boolean] = new DynamicVariable[Boolean](false)
+}
+
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala b/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala
index 97113835f3bd0..aee172a4f549a 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/StreamingContext.scala
@@ -44,7 +44,7 @@ import org.apache.spark.streaming.dstream._
 import org.apache.spark.streaming.receiver.{ActorReceiver, ActorSupervisorStrategy, Receiver}
 import org.apache.spark.streaming.scheduler.{JobScheduler, StreamingListener}
 import org.apache.spark.streaming.ui.{StreamingJobProgressListener, StreamingTab}
-import org.apache.spark.util.{CallSite, ShutdownHookManager, ThreadUtils, Utils}
+import org.apache.spark.util.{AsynchronousListenerBus, CallSite, ShutdownHookManager, ThreadUtils, Utils}
 
 /**
  * Main entry point for Spark Streaming functionality. It provides methods used to create
@@ -693,6 +693,10 @@ class StreamingContext private[streaming] (
    */
   def stop(stopSparkContext: Boolean, stopGracefully: Boolean): Unit = {
     var shutdownHookRefToRemove: AnyRef = null
+    if (AsynchronousListenerBus.withinListenerThread.value) {
+      throw new SparkException("Cannot stop StreamingContext within listener thread of" +
+        " AsynchronousListenerBus")
+    }
     synchronized {
       try {
         state match {
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/StreamingListenerSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/StreamingListenerSuite.scala
index 5dc0472c7770c..df4575ab25aad 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/StreamingListenerSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/StreamingListenerSuite.scala
@@ -21,6 +21,7 @@ import scala.collection.mutable.{ArrayBuffer, HashMap, SynchronizedBuffer, Synch
 import scala.concurrent.Future
 import scala.concurrent.ExecutionContext.Implicits.global
 
+import org.apache.spark.SparkException
 import org.apache.spark.storage.StorageLevel
 import org.apache.spark.streaming.dstream.DStream
 import org.apache.spark.streaming.receiver.Receiver
@@ -161,6 +162,14 @@ class StreamingListenerSuite extends TestSuiteBase with Matchers {
     }
   }
 
+  test("don't call ssc.stop in listener") {
+    ssc = new StreamingContext("local[2]", "ssc", Milliseconds(1000))
+    val inputStream = ssc.receiverStream(new StreamingListenerSuiteReceiver)
+    inputStream.foreachRDD(_.count)
+
+    startStreamingContextAndCallStop(ssc)
+  }
+
   test("onBatchCompleted with successful batch") {
     ssc = new StreamingContext("local[2]", "test", Milliseconds(1000))
     val inputStream = ssc.receiverStream(new StreamingListenerSuiteReceiver)
@@ -207,6 +216,17 @@ class StreamingListenerSuite extends TestSuiteBase with Matchers {
     assert(failureReasons(1).contains("This is another failed job"))
   }
 
+  private def startStreamingContextAndCallStop(_ssc: StreamingContext): Unit = {
+    val contextStoppingCollector = new StreamingContextStoppingCollector(_ssc)
+    _ssc.addStreamingListener(contextStoppingCollector)
+    val batchCounter = new BatchCounter(_ssc)
+    _ssc.start()
+    // Make sure running at least one batch
+    batchCounter.waitUntilBatchesCompleted(expectedNumCompletedBatches = 1, timeout = 10000)
+    _ssc.stop()
+    assert(contextStoppingCollector.sparkExSeen)
+  }
+
   private def startStreamingContextAndCollectFailureReasons(
       _ssc: StreamingContext, isFailed: Boolean = false): Map[Int, String] = {
     val failureReasonsCollector = new FailureReasonsCollector()
@@ -320,3 +340,17 @@ class FailureReasonsCollector extends StreamingListener {
     }
   }
 }
+/**
+ * A StreamingListener that calls StreamingContext.stop().
+ */
+class StreamingContextStoppingCollector(val ssc: StreamingContext) extends StreamingListener {
+  @volatile var sparkExSeen = false
+  override def onBatchCompleted(batchCompleted: StreamingListenerBatchCompleted) {
+    try {
+      ssc.stop()
+    } catch {
+      case se: SparkException =>
+        sparkExSeen = true
+    }
+  }
+}

From 67a5132c21bc8338adbae80b33b85e8fa0ddda34 Mon Sep 17 00:00:00 2001
From: RoyGaoVLIS <roygao@zju.edu.cn>
Date: Tue, 17 Nov 2015 23:00:49 -0800
Subject: [PATCH 311/518] [SPARK-7013][ML][TEST] Add unit test for spark.ml
 StandardScaler

I have added unit test for ML's StandardScaler By comparing with R's output, please review  for me.
Thx.

Author: RoyGaoVLIS <roygao@zju.edu.cn>

Closes #6665 from RoyGao/7013.
---
 .../ml/feature/StandardScalerSuite.scala      | 108 ++++++++++++++++++
 1 file changed, 108 insertions(+)
 create mode 100644 mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala

diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala
new file mode 100644
index 0000000000000..879a3ae875004
--- /dev/null
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.feature
+
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector, Vectors}
+import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.mllib.util.TestingUtils._
+import org.apache.spark.sql.{DataFrame, Row, SQLContext}
+
+class StandardScalerSuite extends SparkFunSuite with MLlibTestSparkContext{
+
+  @transient var data: Array[Vector] = _
+  @transient var resWithStd: Array[Vector] = _
+  @transient var resWithMean: Array[Vector] = _
+  @transient var resWithBoth: Array[Vector] = _
+
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+
+    data = Array(
+      Vectors.dense(-2.0, 2.3, 0.0),
+      Vectors.dense(0.0, -5.1, 1.0),
+      Vectors.dense(1.7, -0.6, 3.3)
+    )
+    resWithMean = Array(
+      Vectors.dense(-1.9, 3.433333333333, -1.433333333333),
+      Vectors.dense(0.1, -3.966666666667, -0.433333333333),
+      Vectors.dense(1.8, 0.533333333333, 1.866666666667)
+    )
+    resWithStd = Array(
+      Vectors.dense(-1.079898494312, 0.616834091415, 0.0),
+      Vectors.dense(0.0, -1.367762550529, 0.590968109266),
+      Vectors.dense(0.917913720165, -0.160913241239, 1.950194760579)
+    )
+    resWithBoth = Array(
+      Vectors.dense(-1.0259035695965, 0.920781324866, -0.8470542899497),
+      Vectors.dense(0.0539949247156, -1.063815317078, -0.256086180682),
+      Vectors.dense(0.9719086448809, 0.143033992212, 1.103140470631)
+    )
+  }
+
+  def assertResult(dataframe: DataFrame): Unit = {
+    dataframe.select("standarded_features", "expected").collect().foreach {
+      case Row(vector1: Vector, vector2: Vector) =>
+        assert(vector1 ~== vector2 absTol 1E-5,
+          "The vector value is not correct after standardization.")
+    }
+  }
+
+  test("Standardization with default parameter") {
+    val df0 = sqlContext.createDataFrame(data.zip(resWithStd)).toDF("features", "expected")
+
+    val standardscaler0 = new StandardScaler()
+      .setInputCol("features")
+      .setOutputCol("standarded_features")
+      .fit(df0)
+
+    assertResult(standardscaler0.transform(df0))
+  }
+
+  test("Standardization with setter") {
+    val df1 = sqlContext.createDataFrame(data.zip(resWithBoth)).toDF("features", "expected")
+    val df2 = sqlContext.createDataFrame(data.zip(resWithMean)).toDF("features", "expected")
+    val df3 = sqlContext.createDataFrame(data.zip(data)).toDF("features", "expected")
+
+    val standardscaler1 = new StandardScaler()
+      .setInputCol("features")
+      .setOutputCol("standarded_features")
+      .setWithMean(true)
+      .setWithStd(true)
+      .fit(df1)
+
+    val standardscaler2 = new StandardScaler()
+      .setInputCol("features")
+      .setOutputCol("standarded_features")
+      .setWithMean(true)
+      .setWithStd(false)
+      .fit(df2)
+
+    val standardscaler3 = new StandardScaler()
+      .setInputCol("features")
+      .setOutputCol("standarded_features")
+      .setWithMean(false)
+      .setWithStd(false)
+      .fit(df3)
+
+    assertResult(standardscaler1.transform(df1))
+    assertResult(standardscaler2.transform(df2))
+    assertResult(standardscaler3.transform(df3))
+  }
+}

From 2f191c66b668fc97f82f44fd8336b6a4488c2f5d Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Tue, 17 Nov 2015 23:14:05 -0800
Subject: [PATCH 312/518] [SPARK-11643] [SQL] parse year with leading zero

Support the years between 0 <= year < 1000

Author: Davies Liu <davies@databricks.com>

Closes #9701 from davies/leading_zero.
---
 .../sql/catalyst/util/DateTimeUtils.scala     | 20 +++++++++++++++++--
 .../catalyst/util/DateTimeUtilsSuite.scala    | 17 +++++++++++++---
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
index 8fb3f41f1bd6a..17a5527f3fb29 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
@@ -241,6 +241,10 @@ object DateTimeUtils {
           i += 3
         } else if (i < 2) {
           if (b == '-') {
+            if (i == 0 && j != 4) {
+              // year should have exact four digits
+              return None
+            }
             segments(i) = currentSegmentValue
             currentSegmentValue = 0
             i += 1
@@ -308,13 +312,17 @@ object DateTimeUtils {
     }
 
     segments(i) = currentSegmentValue
+    if (!justTime && i == 0 && j != 4) {
+      // year should have exact four digits
+      return None
+    }
 
     while (digitsMilli < 6) {
       segments(6) *= 10
       digitsMilli += 1
     }
 
-    if (!justTime && (segments(0) < 1000 || segments(0) > 9999 || segments(1) < 1 ||
+    if (!justTime && (segments(0) < 0 || segments(0) > 9999 || segments(1) < 1 ||
         segments(1) > 12 || segments(2) < 1 || segments(2) > 31)) {
       return None
     }
@@ -368,6 +376,10 @@ object DateTimeUtils {
     while (j < bytes.length && (i < 3 && !(bytes(j) == ' ' || bytes(j) == 'T'))) {
       val b = bytes(j)
       if (i < 2 && b == '-') {
+        if (i == 0 && j != 4) {
+          // year should have exact four digits
+          return None
+        }
         segments(i) = currentSegmentValue
         currentSegmentValue = 0
         i += 1
@@ -381,8 +393,12 @@ object DateTimeUtils {
       }
       j += 1
     }
+    if (i == 0 && j != 4) {
+      // year should have exact four digits
+      return None
+    }
     segments(i) = currentSegmentValue
-    if (segments(0) < 1000 || segments(0) > 9999 || segments(1) < 1 || segments(1) > 12 ||
+    if (segments(0) < 0 || segments(0) > 9999 || segments(1) < 1 || segments(1) > 12 ||
         segments(2) < 1 || segments(2) > 31) {
       return None
     }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
index 60d45422bc9b8..faca128badfd6 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
@@ -110,6 +110,10 @@ class DateTimeUtilsSuite extends SparkFunSuite {
     c.set(Calendar.MILLISECOND, 0)
     assert(stringToDate(UTF8String.fromString("2015")).get ===
       millisToDays(c.getTimeInMillis))
+    c.set(1, 0, 1, 0, 0, 0)
+    c.set(Calendar.MILLISECOND, 0)
+    assert(stringToDate(UTF8String.fromString("0001")).get ===
+      millisToDays(c.getTimeInMillis))
     c = Calendar.getInstance()
     c.set(2015, 2, 1, 0, 0, 0)
     c.set(Calendar.MILLISECOND, 0)
@@ -134,11 +138,15 @@ class DateTimeUtilsSuite extends SparkFunSuite {
     assert(stringToDate(UTF8String.fromString("2015.03.18")).isEmpty)
     assert(stringToDate(UTF8String.fromString("20150318")).isEmpty)
     assert(stringToDate(UTF8String.fromString("2015-031-8")).isEmpty)
+    assert(stringToDate(UTF8String.fromString("02015-03-18")).isEmpty)
+    assert(stringToDate(UTF8String.fromString("015-03-18")).isEmpty)
+    assert(stringToDate(UTF8String.fromString("015")).isEmpty)
+    assert(stringToDate(UTF8String.fromString("02015")).isEmpty)
   }
 
   test("string to time") {
     // Tests with UTC.
-    var c = Calendar.getInstance(TimeZone.getTimeZone("UTC"))
+    val c = Calendar.getInstance(TimeZone.getTimeZone("UTC"))
     c.set(Calendar.MILLISECOND, 0)
 
     c.set(1900, 0, 1, 0, 0, 0)
@@ -174,9 +182,9 @@ class DateTimeUtilsSuite extends SparkFunSuite {
     c.set(Calendar.MILLISECOND, 0)
     assert(stringToTimestamp(UTF8String.fromString("1969-12-31 16:00:00")).get ===
       c.getTimeInMillis * 1000)
-    c.set(2015, 0, 1, 0, 0, 0)
+    c.set(1, 0, 1, 0, 0, 0)
     c.set(Calendar.MILLISECOND, 0)
-    assert(stringToTimestamp(UTF8String.fromString("2015")).get ===
+    assert(stringToTimestamp(UTF8String.fromString("0001")).get ===
       c.getTimeInMillis * 1000)
     c = Calendar.getInstance()
     c.set(2015, 2, 1, 0, 0, 0)
@@ -319,6 +327,7 @@ class DateTimeUtilsSuite extends SparkFunSuite {
       UTF8String.fromString("2011-05-06 07:08:09.1000")).get === c.getTimeInMillis * 1000)
 
     assert(stringToTimestamp(UTF8String.fromString("238")).isEmpty)
+    assert(stringToTimestamp(UTF8String.fromString("00238")).isEmpty)
     assert(stringToTimestamp(UTF8String.fromString("2015-03-18 123142")).isEmpty)
     assert(stringToTimestamp(UTF8String.fromString("2015-03-18T123123")).isEmpty)
     assert(stringToTimestamp(UTF8String.fromString("2015-03-18X")).isEmpty)
@@ -326,6 +335,8 @@ class DateTimeUtilsSuite extends SparkFunSuite {
     assert(stringToTimestamp(UTF8String.fromString("2015.03.18")).isEmpty)
     assert(stringToTimestamp(UTF8String.fromString("20150318")).isEmpty)
     assert(stringToTimestamp(UTF8String.fromString("2015-031-8")).isEmpty)
+    assert(stringToTimestamp(UTF8String.fromString("02015-01-18")).isEmpty)
+    assert(stringToTimestamp(UTF8String.fromString("015-01-18")).isEmpty)
     assert(stringToTimestamp(
       UTF8String.fromString("2015-03-18T12:03.17-20:0")).isEmpty)
     assert(stringToTimestamp(

From 9154f89befb7a33d4853cea95efd7dc6b25d033b Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Tue, 17 Nov 2015 23:44:06 -0800
Subject: [PATCH 313/518] [SPARK-11728] Replace example code in ml-ensembles.md
 using include_example

JIRA issue https://issues.apache.org/jira/browse/SPARK-11728.

The ml-ensembles.md file contains `OneVsRestExample`. Instead of writing new code files of two `OneVsRestExample`s, I use two existing files in the examples directory, they are `OneVsRestExample.scala` and `JavaOneVsRestExample.scala`.

Author: Xusen Yin <yinxusen@gmail.com>

Closes #9716 from yinxusen/SPARK-11728.
---
 docs/ml-ensembles.md                          | 754 +-----------------
 ...aGradientBoostedTreeClassifierExample.java | 102 +++
 ...vaGradientBoostedTreeRegressorExample.java |  90 +++
 .../examples/ml/JavaOneVsRestExample.java     |   4 +
 .../ml/JavaRandomForestClassifierExample.java | 101 +++
 .../ml/JavaRandomForestRegressorExample.java  |  90 +++
 ...radient_boosted_tree_classifier_example.py |  77 ++
 ...gradient_boosted_tree_regressor_example.py |  74 ++
 .../ml/random_forest_classifier_example.py    |  77 ++
 .../ml/random_forest_regressor_example.py     |  74 ++
 ...GradientBoostedTreeClassifierExample.scala |  97 +++
 .../GradientBoostedTreeRegressorExample.scala |  85 ++
 .../spark/examples/ml/OneVsRestExample.scala  |   4 +
 .../ml/RandomForestClassifierExample.scala    |  97 +++
 .../ml/RandomForestRegressorExample.scala     |  84 ++
 15 files changed, 1070 insertions(+), 740 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeClassifierExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeRegressorExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestClassifierExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestRegressorExample.java
 create mode 100644 examples/src/main/python/ml/gradient_boosted_tree_classifier_example.py
 create mode 100644 examples/src/main/python/ml/gradient_boosted_tree_regressor_example.py
 create mode 100644 examples/src/main/python/ml/random_forest_classifier_example.py
 create mode 100644 examples/src/main/python/ml/random_forest_regressor_example.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/GradientBoostedTreeClassifierExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/GradientBoostedTreeRegressorExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/RandomForestClassifierExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/RandomForestRegressorExample.scala

diff --git a/docs/ml-ensembles.md b/docs/ml-ensembles.md
index ce15f5e6466ec..f6c3c30d5334f 100644
--- a/docs/ml-ensembles.md
+++ b/docs/ml-ensembles.md
@@ -115,194 +115,21 @@ We use two feature transformers to prepare the data; these help index categories
 
 Refer to the [Scala API docs](api/scala/index.html#org.apache.spark.ml.classification.RandomForestClassifier) for more details.
 
-{% highlight scala %}
-import org.apache.spark.ml.Pipeline
-import org.apache.spark.ml.classification.RandomForestClassifier
-import org.apache.spark.ml.classification.RandomForestClassificationModel
-import org.apache.spark.ml.feature.{StringIndexer, IndexToString, VectorIndexer}
-import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
-
-// Load and parse the data file, converting it to a DataFrame.
-val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-// Index labels, adding metadata to the label column.
-// Fit on whole dataset to include all labels in index.
-val labelIndexer = new StringIndexer()
-  .setInputCol("label")
-  .setOutputCol("indexedLabel")
-  .fit(data)
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-val featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data)
-
-// Split the data into training and test sets (30% held out for testing)
-val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
-
-// Train a RandomForest model.
-val rf = new RandomForestClassifier()
-  .setLabelCol("indexedLabel")
-  .setFeaturesCol("indexedFeatures")
-  .setNumTrees(10)
-
-// Convert indexed labels back to original labels.
-val labelConverter = new IndexToString()
-  .setInputCol("prediction")
-  .setOutputCol("predictedLabel")
-  .setLabels(labelIndexer.labels)
-
-// Chain indexers and forest in a Pipeline
-val pipeline = new Pipeline()
-  .setStages(Array(labelIndexer, featureIndexer, rf, labelConverter))
-
-// Train model.  This also runs the indexers.
-val model = pipeline.fit(trainingData)
-
-// Make predictions.
-val predictions = model.transform(testData)
-
-// Select example rows to display.
-predictions.select("predictedLabel", "label", "features").show(5)
-
-// Select (prediction, true label) and compute test error
-val evaluator = new MulticlassClassificationEvaluator()
-  .setLabelCol("indexedLabel")
-  .setPredictionCol("prediction")
-  .setMetricName("precision")
-val accuracy = evaluator.evaluate(predictions)
-println("Test Error = " + (1.0 - accuracy))
-
-val rfModel = model.stages(2).asInstanceOf[RandomForestClassificationModel]
-println("Learned classification forest model:\n" + rfModel.toDebugString)
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/RandomForestClassifierExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 
 Refer to the [Java API docs](api/java/org/apache/spark/ml/classification/RandomForestClassifier.html) for more details.
 
-{% highlight java %}
-import org.apache.spark.ml.Pipeline;
-import org.apache.spark.ml.PipelineModel;
-import org.apache.spark.ml.PipelineStage;
-import org.apache.spark.ml.classification.RandomForestClassifier;
-import org.apache.spark.ml.classification.RandomForestClassificationModel;
-import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
-import org.apache.spark.ml.feature.*;
-import org.apache.spark.sql.DataFrame;
-
-// Load and parse the data file, converting it to a DataFrame.
-DataFrame data = sqlContext.read().format("libsvm")
-  .load("data/mllib/sample_libsvm_data.txt");
-
-// Index labels, adding metadata to the label column.
-// Fit on whole dataset to include all labels in index.
-StringIndexerModel labelIndexer = new StringIndexer()
-  .setInputCol("label")
-  .setOutputCol("indexedLabel")
-  .fit(data);
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-VectorIndexerModel featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data);
-
-// Split the data into training and test sets (30% held out for testing)
-DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
-DataFrame trainingData = splits[0];
-DataFrame testData = splits[1];
-
-// Train a RandomForest model.
-RandomForestClassifier rf = new RandomForestClassifier()
-  .setLabelCol("indexedLabel")
-  .setFeaturesCol("indexedFeatures");
-
-// Convert indexed labels back to original labels.
-IndexToString labelConverter = new IndexToString()
-  .setInputCol("prediction")
-  .setOutputCol("predictedLabel")
-  .setLabels(labelIndexer.labels());
-
-// Chain indexers and forest in a Pipeline
-Pipeline pipeline = new Pipeline()
-  .setStages(new PipelineStage[] {labelIndexer, featureIndexer, rf, labelConverter});
-
-// Train model.  This also runs the indexers.
-PipelineModel model = pipeline.fit(trainingData);
-
-// Make predictions.
-DataFrame predictions = model.transform(testData);
-
-// Select example rows to display.
-predictions.select("predictedLabel", "label", "features").show(5);
-
-// Select (prediction, true label) and compute test error
-MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
-  .setLabelCol("indexedLabel")
-  .setPredictionCol("prediction")
-  .setMetricName("precision");
-double accuracy = evaluator.evaluate(predictions);
-System.out.println("Test Error = " + (1.0 - accuracy));
-
-RandomForestClassificationModel rfModel =
-  (RandomForestClassificationModel)(model.stages()[2]);
-System.out.println("Learned classification forest model:\n" + rfModel.toDebugString());
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaRandomForestClassifierExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 
 Refer to the [Python API docs](api/python/pyspark.ml.html#pyspark.ml.classification.RandomForestClassifier) for more details.
 
-{% highlight python %}
-from pyspark.ml import Pipeline
-from pyspark.ml.classification import RandomForestClassifier
-from pyspark.ml.feature import StringIndexer, VectorIndexer
-from pyspark.ml.evaluation import MulticlassClassificationEvaluator
-
-# Load and parse the data file, converting it to a DataFrame.
-data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-# Index labels, adding metadata to the label column.
-# Fit on whole dataset to include all labels in index.
-labelIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel").fit(data)
-# Automatically identify categorical features, and index them.
-# Set maxCategories so features with > 4 distinct values are treated as continuous.
-featureIndexer =\
-    VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
-
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a RandomForest model.
-rf = RandomForestClassifier(labelCol="indexedLabel", featuresCol="indexedFeatures")
-
-# Chain indexers and forest in a Pipeline
-pipeline = Pipeline(stages=[labelIndexer, featureIndexer, rf])
-
-# Train model.  This also runs the indexers.
-model = pipeline.fit(trainingData)
-
-# Make predictions.
-predictions = model.transform(testData)
-
-# Select example rows to display.
-predictions.select("prediction", "indexedLabel", "features").show(5)
-
-# Select (prediction, true label) and compute test error
-evaluator = MulticlassClassificationEvaluator(
-    labelCol="indexedLabel", predictionCol="prediction", metricName="precision")
-accuracy = evaluator.evaluate(predictions)
-print "Test Error = %g" % (1.0 - accuracy)
-
-rfModel = model.stages[2]
-print rfModel # summary only
-{% endhighlight %}
+{% include_example python/ml/random_forest_classifier_example.py %}
 </div>
 </div>
 
@@ -316,167 +143,21 @@ We use a feature transformer to index categorical features, adding metadata to t
 
 Refer to the [Scala API docs](api/scala/index.html#org.apache.spark.ml.regression.RandomForestRegressor) for more details.
 
-{% highlight scala %}
-import org.apache.spark.ml.Pipeline
-import org.apache.spark.ml.regression.RandomForestRegressor
-import org.apache.spark.ml.regression.RandomForestRegressionModel
-import org.apache.spark.ml.feature.VectorIndexer
-import org.apache.spark.ml.evaluation.RegressionEvaluator
-
-// Load and parse the data file, converting it to a DataFrame.
-val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-val featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data)
-
-// Split the data into training and test sets (30% held out for testing)
-val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
-
-// Train a RandomForest model.
-val rf = new RandomForestRegressor()
-  .setLabelCol("label")
-  .setFeaturesCol("indexedFeatures")
-
-// Chain indexer and forest in a Pipeline
-val pipeline = new Pipeline()
-  .setStages(Array(featureIndexer, rf))
-
-// Train model.  This also runs the indexer.
-val model = pipeline.fit(trainingData)
-
-// Make predictions.
-val predictions = model.transform(testData)
-
-// Select example rows to display.
-predictions.select("prediction", "label", "features").show(5)
-
-// Select (prediction, true label) and compute test error
-val evaluator = new RegressionEvaluator()
-  .setLabelCol("label")
-  .setPredictionCol("prediction")
-  .setMetricName("rmse")
-val rmse = evaluator.evaluate(predictions)
-println("Root Mean Squared Error (RMSE) on test data = " + rmse)
-
-val rfModel = model.stages(1).asInstanceOf[RandomForestRegressionModel]
-println("Learned regression forest model:\n" + rfModel.toDebugString)
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/RandomForestRegressorExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 
 Refer to the [Java API docs](api/java/org/apache/spark/ml/regression/RandomForestRegressor.html) for more details.
 
-{% highlight java %}
-import org.apache.spark.ml.Pipeline;
-import org.apache.spark.ml.PipelineModel;
-import org.apache.spark.ml.PipelineStage;
-import org.apache.spark.ml.evaluation.RegressionEvaluator;
-import org.apache.spark.ml.feature.VectorIndexer;
-import org.apache.spark.ml.feature.VectorIndexerModel;
-import org.apache.spark.ml.regression.RandomForestRegressionModel;
-import org.apache.spark.ml.regression.RandomForestRegressor;
-import org.apache.spark.sql.DataFrame;
-
-// Load and parse the data file, converting it to a DataFrame.
-DataFrame data = sqlContext.read().format("libsvm")
-  .load("data/mllib/sample_libsvm_data.txt");
-
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-VectorIndexerModel featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data);
-
-// Split the data into training and test sets (30% held out for testing)
-DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
-DataFrame trainingData = splits[0];
-DataFrame testData = splits[1];
-
-// Train a RandomForest model.
-RandomForestRegressor rf = new RandomForestRegressor()
-  .setLabelCol("label")
-  .setFeaturesCol("indexedFeatures");
-
-// Chain indexer and forest in a Pipeline
-Pipeline pipeline = new Pipeline()
-  .setStages(new PipelineStage[] {featureIndexer, rf});
-
-// Train model.  This also runs the indexer.
-PipelineModel model = pipeline.fit(trainingData);
-
-// Make predictions.
-DataFrame predictions = model.transform(testData);
-
-// Select example rows to display.
-predictions.select("prediction", "label", "features").show(5);
-
-// Select (prediction, true label) and compute test error
-RegressionEvaluator evaluator = new RegressionEvaluator()
-  .setLabelCol("label")
-  .setPredictionCol("prediction")
-  .setMetricName("rmse");
-double rmse = evaluator.evaluate(predictions);
-System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse);
-
-RandomForestRegressionModel rfModel =
-  (RandomForestRegressionModel)(model.stages()[1]);
-System.out.println("Learned regression forest model:\n" + rfModel.toDebugString());
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaRandomForestRegressorExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 
 Refer to the [Python API docs](api/python/pyspark.ml.html#pyspark.ml.regression.RandomForestRegressor) for more details.
 
-{% highlight python %}
-from pyspark.ml import Pipeline
-from pyspark.ml.regression import RandomForestRegressor
-from pyspark.ml.feature import VectorIndexer
-from pyspark.ml.evaluation import RegressionEvaluator
-
-# Load and parse the data file, converting it to a DataFrame.
-data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-# Automatically identify categorical features, and index them.
-# Set maxCategories so features with > 4 distinct values are treated as continuous.
-featureIndexer =\
-    VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
-
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a RandomForest model.
-rf = RandomForestRegressor(featuresCol="indexedFeatures")
-
-# Chain indexer and forest in a Pipeline
-pipeline = Pipeline(stages=[featureIndexer, rf])
-
-# Train model.  This also runs the indexer.
-model = pipeline.fit(trainingData)
-
-# Make predictions.
-predictions = model.transform(testData)
-
-# Select example rows to display.
-predictions.select("prediction", "label", "features").show(5)
-
-# Select (prediction, true label) and compute test error
-evaluator = RegressionEvaluator(
-    labelCol="label", predictionCol="prediction", metricName="rmse")
-rmse = evaluator.evaluate(predictions)
-print "Root Mean Squared Error (RMSE) on test data = %g" % rmse
-
-rfModel = model.stages[1]
-print rfModel # summary only
-{% endhighlight %}
+{% include_example python/ml/random_forest_regressor_example.py %}
 </div>
 </div>
 
@@ -560,194 +241,21 @@ We use two feature transformers to prepare the data; these help index categories
 
 Refer to the [Scala API docs](api/scala/index.html#org.apache.spark.ml.classification.GBTClassifier) for more details.
 
-{% highlight scala %}
-import org.apache.spark.ml.Pipeline
-import org.apache.spark.ml.classification.GBTClassifier
-import org.apache.spark.ml.classification.GBTClassificationModel
-import org.apache.spark.ml.feature.{StringIndexer, IndexToString, VectorIndexer}
-import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
-
-// Load and parse the data file, converting it to a DataFrame.
-val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-// Index labels, adding metadata to the label column.
-// Fit on whole dataset to include all labels in index.
-val labelIndexer = new StringIndexer()
-  .setInputCol("label")
-  .setOutputCol("indexedLabel")
-  .fit(data)
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-val featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data)
-
-// Split the data into training and test sets (30% held out for testing)
-val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
-
-// Train a GBT model.
-val gbt = new GBTClassifier()
-  .setLabelCol("indexedLabel")
-  .setFeaturesCol("indexedFeatures")
-  .setMaxIter(10)
-
-// Convert indexed labels back to original labels.
-val labelConverter = new IndexToString()
-  .setInputCol("prediction")
-  .setOutputCol("predictedLabel")
-  .setLabels(labelIndexer.labels)
-
-// Chain indexers and GBT in a Pipeline
-val pipeline = new Pipeline()
-  .setStages(Array(labelIndexer, featureIndexer, gbt, labelConverter))
-
-// Train model.  This also runs the indexers.
-val model = pipeline.fit(trainingData)
-
-// Make predictions.
-val predictions = model.transform(testData)
-
-// Select example rows to display.
-predictions.select("predictedLabel", "label", "features").show(5)
-
-// Select (prediction, true label) and compute test error
-val evaluator = new MulticlassClassificationEvaluator()
-  .setLabelCol("indexedLabel")
-  .setPredictionCol("prediction")
-  .setMetricName("precision")
-val accuracy = evaluator.evaluate(predictions)
-println("Test Error = " + (1.0 - accuracy))
-
-val gbtModel = model.stages(2).asInstanceOf[GBTClassificationModel]
-println("Learned classification GBT model:\n" + gbtModel.toDebugString)
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/GradientBoostedTreeClassifierExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 
 Refer to the [Java API docs](api/java/org/apache/spark/ml/classification/GBTClassifier.html) for more details.
 
-{% highlight java %}
-import org.apache.spark.ml.Pipeline;
-import org.apache.spark.ml.PipelineModel;
-import org.apache.spark.ml.PipelineStage;
-import org.apache.spark.ml.classification.GBTClassifier;
-import org.apache.spark.ml.classification.GBTClassificationModel;
-import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
-import org.apache.spark.ml.feature.*;
-import org.apache.spark.sql.DataFrame;
-
-// Load and parse the data file, converting it to a DataFrame.
-DataFrame data sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
-
-// Index labels, adding metadata to the label column.
-// Fit on whole dataset to include all labels in index.
-StringIndexerModel labelIndexer = new StringIndexer()
-  .setInputCol("label")
-  .setOutputCol("indexedLabel")
-  .fit(data);
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-VectorIndexerModel featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data);
-
-// Split the data into training and test sets (30% held out for testing)
-DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
-DataFrame trainingData = splits[0];
-DataFrame testData = splits[1];
-
-// Train a GBT model.
-GBTClassifier gbt = new GBTClassifier()
-  .setLabelCol("indexedLabel")
-  .setFeaturesCol("indexedFeatures")
-  .setMaxIter(10);
-
-// Convert indexed labels back to original labels.
-IndexToString labelConverter = new IndexToString()
-  .setInputCol("prediction")
-  .setOutputCol("predictedLabel")
-  .setLabels(labelIndexer.labels());
-
-// Chain indexers and GBT in a Pipeline
-Pipeline pipeline = new Pipeline()
-  .setStages(new PipelineStage[] {labelIndexer, featureIndexer, gbt, labelConverter});
-
-// Train model.  This also runs the indexers.
-PipelineModel model = pipeline.fit(trainingData);
-
-// Make predictions.
-DataFrame predictions = model.transform(testData);
-
-// Select example rows to display.
-predictions.select("predictedLabel", "label", "features").show(5);
-
-// Select (prediction, true label) and compute test error
-MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
-  .setLabelCol("indexedLabel")
-  .setPredictionCol("prediction")
-  .setMetricName("precision");
-double accuracy = evaluator.evaluate(predictions);
-System.out.println("Test Error = " + (1.0 - accuracy));
-
-GBTClassificationModel gbtModel =
-  (GBTClassificationModel)(model.stages()[2]);
-System.out.println("Learned classification GBT model:\n" + gbtModel.toDebugString());
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaGradientBoostedTreeClassifierExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 
 Refer to the [Python API docs](api/python/pyspark.ml.html#pyspark.ml.classification.GBTClassifier) for more details.
 
-{% highlight python %}
-from pyspark.ml import Pipeline
-from pyspark.ml.classification import GBTClassifier
-from pyspark.ml.feature import StringIndexer, VectorIndexer
-from pyspark.ml.evaluation import MulticlassClassificationEvaluator
-
-# Load and parse the data file, converting it to a DataFrame.
-data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-# Index labels, adding metadata to the label column.
-# Fit on whole dataset to include all labels in index.
-labelIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel").fit(data)
-# Automatically identify categorical features, and index them.
-# Set maxCategories so features with > 4 distinct values are treated as continuous.
-featureIndexer =\
-    VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
-
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a GBT model.
-gbt = GBTClassifier(labelCol="indexedLabel", featuresCol="indexedFeatures", maxIter=10)
-
-# Chain indexers and GBT in a Pipeline
-pipeline = Pipeline(stages=[labelIndexer, featureIndexer, gbt])
-
-# Train model.  This also runs the indexers.
-model = pipeline.fit(trainingData)
-
-# Make predictions.
-predictions = model.transform(testData)
-
-# Select example rows to display.
-predictions.select("prediction", "indexedLabel", "features").show(5)
-
-# Select (prediction, true label) and compute test error
-evaluator = MulticlassClassificationEvaluator(
-    labelCol="indexedLabel", predictionCol="prediction", metricName="precision")
-accuracy = evaluator.evaluate(predictions)
-print "Test Error = %g" % (1.0 - accuracy)
-
-gbtModel = model.stages[2]
-print gbtModel # summary only
-{% endhighlight %}
+{% include_example python/ml/gradient_boosted_tree_classifier_example.py %}
 </div>
 </div>
 
@@ -761,168 +269,21 @@ be true in general.
 
 Refer to the [Scala API docs](api/scala/index.html#org.apache.spark.ml.regression.GBTRegressor) for more details.
 
-{% highlight scala %}
-import org.apache.spark.ml.Pipeline
-import org.apache.spark.ml.regression.GBTRegressor
-import org.apache.spark.ml.regression.GBTRegressionModel
-import org.apache.spark.ml.feature.VectorIndexer
-import org.apache.spark.ml.evaluation.RegressionEvaluator
-
-// Load and parse the data file, converting it to a DataFrame.
-val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-val featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data)
-
-// Split the data into training and test sets (30% held out for testing)
-val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
-
-// Train a GBT model.
-val gbt = new GBTRegressor()
-  .setLabelCol("label")
-  .setFeaturesCol("indexedFeatures")
-  .setMaxIter(10)
-
-// Chain indexer and GBT in a Pipeline
-val pipeline = new Pipeline()
-  .setStages(Array(featureIndexer, gbt))
-
-// Train model.  This also runs the indexer.
-val model = pipeline.fit(trainingData)
-
-// Make predictions.
-val predictions = model.transform(testData)
-
-// Select example rows to display.
-predictions.select("prediction", "label", "features").show(5)
-
-// Select (prediction, true label) and compute test error
-val evaluator = new RegressionEvaluator()
-  .setLabelCol("label")
-  .setPredictionCol("prediction")
-  .setMetricName("rmse")
-val rmse = evaluator.evaluate(predictions)
-println("Root Mean Squared Error (RMSE) on test data = " + rmse)
-
-val gbtModel = model.stages(1).asInstanceOf[GBTRegressionModel]
-println("Learned regression GBT model:\n" + gbtModel.toDebugString)
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/GradientBoostedTreeRegressorExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 
 Refer to the [Java API docs](api/java/org/apache/spark/ml/regression/GBTRegressor.html) for more details.
 
-{% highlight java %}
-import org.apache.spark.ml.Pipeline;
-import org.apache.spark.ml.PipelineModel;
-import org.apache.spark.ml.PipelineStage;
-import org.apache.spark.ml.evaluation.RegressionEvaluator;
-import org.apache.spark.ml.feature.VectorIndexer;
-import org.apache.spark.ml.feature.VectorIndexerModel;
-import org.apache.spark.ml.regression.GBTRegressionModel;
-import org.apache.spark.ml.regression.GBTRegressor;
-import org.apache.spark.sql.DataFrame;
-
-// Load and parse the data file, converting it to a DataFrame.
-DataFrame data = sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
-
-// Automatically identify categorical features, and index them.
-// Set maxCategories so features with > 4 distinct values are treated as continuous.
-VectorIndexerModel featureIndexer = new VectorIndexer()
-  .setInputCol("features")
-  .setOutputCol("indexedFeatures")
-  .setMaxCategories(4)
-  .fit(data);
-
-// Split the data into training and test sets (30% held out for testing)
-DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
-DataFrame trainingData = splits[0];
-DataFrame testData = splits[1];
-
-// Train a GBT model.
-GBTRegressor gbt = new GBTRegressor()
-  .setLabelCol("label")
-  .setFeaturesCol("indexedFeatures")
-  .setMaxIter(10);
-
-// Chain indexer and GBT in a Pipeline
-Pipeline pipeline = new Pipeline()
-  .setStages(new PipelineStage[] {featureIndexer, gbt});
-
-// Train model.  This also runs the indexer.
-PipelineModel model = pipeline.fit(trainingData);
-
-// Make predictions.
-DataFrame predictions = model.transform(testData);
-
-// Select example rows to display.
-predictions.select("prediction", "label", "features").show(5);
-
-// Select (prediction, true label) and compute test error
-RegressionEvaluator evaluator = new RegressionEvaluator()
-  .setLabelCol("label")
-  .setPredictionCol("prediction")
-  .setMetricName("rmse");
-double rmse = evaluator.evaluate(predictions);
-System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse);
-
-GBTRegressionModel gbtModel =
-  (GBTRegressionModel)(model.stages()[1]);
-System.out.println("Learned regression GBT model:\n" + gbtModel.toDebugString());
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaGradientBoostedTreeRegressorExample.java %}
 </div>
 
 <div data-lang="python" markdown="1">
 
 Refer to the [Python API docs](api/python/pyspark.ml.html#pyspark.ml.regression.GBTRegressor) for more details.
 
-{% highlight python %}
-from pyspark.ml import Pipeline
-from pyspark.ml.regression import GBTRegressor
-from pyspark.ml.feature import VectorIndexer
-from pyspark.ml.evaluation import RegressionEvaluator
-
-# Load and parse the data file, converting it to a DataFrame.
-data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-# Automatically identify categorical features, and index them.
-# Set maxCategories so features with > 4 distinct values are treated as continuous.
-featureIndexer =\
-    VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
-
-# Split the data into training and test sets (30% held out for testing)
-(trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-# Train a GBT model.
-gbt = GBTRegressor(featuresCol="indexedFeatures", maxIter=10)
-
-# Chain indexer and GBT in a Pipeline
-pipeline = Pipeline(stages=[featureIndexer, gbt])
-
-# Train model.  This also runs the indexer.
-model = pipeline.fit(trainingData)
-
-# Make predictions.
-predictions = model.transform(testData)
-
-# Select example rows to display.
-predictions.select("prediction", "label", "features").show(5)
-
-# Select (prediction, true label) and compute test error
-evaluator = RegressionEvaluator(
-    labelCol="label", predictionCol="prediction", metricName="rmse")
-rmse = evaluator.evaluate(predictions)
-print "Root Mean Squared Error (RMSE) on test data = %g" % rmse
-
-gbtModel = model.stages[1]
-print gbtModel # summary only
-{% endhighlight %}
+{% include_example python/ml/gradient_boosted_tree_regressor_example.py %}
 </div>
 </div>
 
@@ -945,100 +306,13 @@ The example below demonstrates how to load the
 
 Refer to the [Scala API docs](api/scala/index.html#org.apache.spark.ml.classifier.OneVsRest) for more details.
 
-{% highlight scala %}
-import org.apache.spark.ml.classification.{LogisticRegression, OneVsRest}
-import org.apache.spark.mllib.evaluation.MulticlassMetrics
-import org.apache.spark.sql.{Row, SQLContext}
-
-val sqlContext = new SQLContext(sc)
-
-// parse data into dataframe
-val data = sqlContext.read.format("libsvm")
-  .load("data/mllib/sample_multiclass_classification_data.txt")
-val Array(train, test) = data.randomSplit(Array(0.7, 0.3))
-
-// instantiate multiclass learner and train
-val ovr = new OneVsRest().setClassifier(new LogisticRegression)
-
-val ovrModel = ovr.fit(train)
-
-// score model on test data
-val predictions = ovrModel.transform(test).select("prediction", "label")
-val predictionsAndLabels = predictions.map {case Row(p: Double, l: Double) => (p, l)}
-
-// compute confusion matrix
-val metrics = new MulticlassMetrics(predictionsAndLabels)
-println(metrics.confusionMatrix)
-
-// the Iris DataSet has three classes
-val numClasses = 3
-
-println("label\tfpr\n")
-(0 until numClasses).foreach { index =>
-  val label = index.toDouble
-  println(label + "\t" + metrics.falsePositiveRate(label))
-}
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/ml/OneVsRestExample.scala %}
 </div>
 
 <div data-lang="java" markdown="1">
 
 Refer to the [Java API docs](api/java/org/apache/spark/ml/classification/OneVsRest.html) for more details.
 
-{% highlight java %}
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.ml.classification.LogisticRegression;
-import org.apache.spark.ml.classification.OneVsRest;
-import org.apache.spark.ml.classification.OneVsRestModel;
-import org.apache.spark.mllib.evaluation.MulticlassMetrics;
-import org.apache.spark.mllib.linalg.Matrix;
-import org.apache.spark.sql.DataFrame;
-import org.apache.spark.sql.SQLContext;
-
-SparkConf conf = new SparkConf().setAppName("JavaOneVsRestExample");
-JavaSparkContext jsc = new JavaSparkContext(conf);
-SQLContext jsql = new SQLContext(jsc);
-
-DataFrame dataFrame = sqlContext.read().format("libsvm")
-  .load("data/mllib/sample_multiclass_classification_data.txt");
-
-DataFrame[] splits = dataFrame.randomSplit(new double[] {0.7, 0.3}, 12345);
-DataFrame train = splits[0];
-DataFrame test = splits[1];
-
-// instantiate the One Vs Rest Classifier
-OneVsRest ovr = new OneVsRest().setClassifier(new LogisticRegression());
-
-// train the multiclass model
-OneVsRestModel ovrModel = ovr.fit(train.cache());
-
-// score the model on test data
-DataFrame predictions = ovrModel
-  .transform(test)
-  .select("prediction", "label");
-
-// obtain metrics
-MulticlassMetrics metrics = new MulticlassMetrics(predictions);
-Matrix confusionMatrix = metrics.confusionMatrix();
-
-// output the Confusion Matrix
-System.out.println("Confusion Matrix");
-System.out.println(confusionMatrix);
-
-// compute the false positive rate per label
-System.out.println();
-System.out.println("label\tfpr\n");
-
-// the Iris DataSet has three classes
-int numClasses = 3;
-for (int index = 0; index < numClasses; index++) {
-  double label = (double) index;
-  System.out.print(label);
-  System.out.print("\t");
-  System.out.print(metrics.falsePositiveRate(label));
-  System.out.println();
-}
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/ml/JavaOneVsRestExample.java %}
 </div>
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeClassifierExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeClassifierExample.java
new file mode 100644
index 0000000000000..848fe6566c1ec
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeClassifierExample.java
@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
+import org.apache.spark.ml.Pipeline;
+import org.apache.spark.ml.PipelineModel;
+import org.apache.spark.ml.PipelineStage;
+import org.apache.spark.ml.classification.GBTClassificationModel;
+import org.apache.spark.ml.classification.GBTClassifier;
+import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
+import org.apache.spark.ml.feature.*;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+// $example off$
+
+public class JavaGradientBoostedTreeClassifierExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaGradientBoostedTreeClassifierExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    DataFrame data = sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
+
+    // Index labels, adding metadata to the label column.
+    // Fit on whole dataset to include all labels in index.
+    StringIndexerModel labelIndexer = new StringIndexer()
+      .setInputCol("label")
+      .setOutputCol("indexedLabel")
+      .fit(data);
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    VectorIndexerModel featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data);
+
+    // Split the data into training and test sets (30% held out for testing)
+    DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
+    DataFrame trainingData = splits[0];
+    DataFrame testData = splits[1];
+
+    // Train a GBT model.
+    GBTClassifier gbt = new GBTClassifier()
+      .setLabelCol("indexedLabel")
+      .setFeaturesCol("indexedFeatures")
+      .setMaxIter(10);
+
+    // Convert indexed labels back to original labels.
+    IndexToString labelConverter = new IndexToString()
+      .setInputCol("prediction")
+      .setOutputCol("predictedLabel")
+      .setLabels(labelIndexer.labels());
+
+    // Chain indexers and GBT in a Pipeline
+    Pipeline pipeline = new Pipeline()
+      .setStages(new PipelineStage[] {labelIndexer, featureIndexer, gbt, labelConverter});
+
+    // Train model.  This also runs the indexers.
+    PipelineModel model = pipeline.fit(trainingData);
+
+    // Make predictions.
+    DataFrame predictions = model.transform(testData);
+
+    // Select example rows to display.
+    predictions.select("predictedLabel", "label", "features").show(5);
+
+    // Select (prediction, true label) and compute test error
+    MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
+      .setLabelCol("indexedLabel")
+      .setPredictionCol("prediction")
+      .setMetricName("precision");
+    double accuracy = evaluator.evaluate(predictions);
+    System.out.println("Test Error = " + (1.0 - accuracy));
+
+    GBTClassificationModel gbtModel = (GBTClassificationModel)(model.stages()[2]);
+    System.out.println("Learned classification GBT model:\n" + gbtModel.toDebugString());
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeRegressorExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeRegressorExample.java
new file mode 100644
index 0000000000000..1f67b0842db0d
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaGradientBoostedTreeRegressorExample.java
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
+import org.apache.spark.ml.Pipeline;
+import org.apache.spark.ml.PipelineModel;
+import org.apache.spark.ml.PipelineStage;
+import org.apache.spark.ml.evaluation.RegressionEvaluator;
+import org.apache.spark.ml.feature.VectorIndexer;
+import org.apache.spark.ml.feature.VectorIndexerModel;
+import org.apache.spark.ml.regression.GBTRegressionModel;
+import org.apache.spark.ml.regression.GBTRegressor;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+// $example off$
+
+public class JavaGradientBoostedTreeRegressorExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaGradientBoostedTreeRegressorExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    DataFrame data = sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
+
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    VectorIndexerModel featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data);
+
+    // Split the data into training and test sets (30% held out for testing)
+    DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
+    DataFrame trainingData = splits[0];
+    DataFrame testData = splits[1];
+
+    // Train a GBT model.
+    GBTRegressor gbt = new GBTRegressor()
+      .setLabelCol("label")
+      .setFeaturesCol("indexedFeatures")
+      .setMaxIter(10);
+
+    // Chain indexer and GBT in a Pipeline
+    Pipeline pipeline = new Pipeline().setStages(new PipelineStage[] {featureIndexer, gbt});
+
+    // Train model.  This also runs the indexer.
+    PipelineModel model = pipeline.fit(trainingData);
+
+    // Make predictions.
+    DataFrame predictions = model.transform(testData);
+
+    // Select example rows to display.
+    predictions.select("prediction", "label", "features").show(5);
+
+    // Select (prediction, true label) and compute test error
+    RegressionEvaluator evaluator = new RegressionEvaluator()
+      .setLabelCol("label")
+      .setPredictionCol("prediction")
+      .setMetricName("rmse");
+    double rmse = evaluator.evaluate(predictions);
+    System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse);
+
+    GBTRegressionModel gbtModel = (GBTRegressionModel)(model.stages()[1]);
+    System.out.println("Learned regression GBT model:\n" + gbtModel.toDebugString());
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java
index f0d92a56bee73..42374e77acf01 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaOneVsRestExample.java
@@ -21,6 +21,7 @@
 
 import org.apache.spark.SparkConf;
 import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
 import org.apache.spark.ml.classification.LogisticRegression;
 import org.apache.spark.ml.classification.OneVsRest;
 import org.apache.spark.ml.classification.OneVsRestModel;
@@ -31,6 +32,7 @@
 import org.apache.spark.sql.DataFrame;
 import org.apache.spark.sql.SQLContext;
 import org.apache.spark.sql.types.StructField;
+// $example off$
 
 /**
  * An example runner for Multiclass to Binary Reduction with One Vs Rest.
@@ -61,6 +63,7 @@ public static void main(String[] args) {
     JavaSparkContext jsc = new JavaSparkContext(conf);
     SQLContext jsql = new SQLContext(jsc);
 
+    // $example on$
     // configure the base classifier
     LogisticRegression classifier = new LogisticRegression()
       .setMaxIter(params.maxIter)
@@ -125,6 +128,7 @@ public static void main(String[] args) {
     System.out.println(confusionMatrix);
     System.out.println();
     System.out.println(results);
+    // $example off$
 
     jsc.stop();
   }
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestClassifierExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestClassifierExample.java
new file mode 100644
index 0000000000000..5a62496660290
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestClassifierExample.java
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
+import org.apache.spark.ml.Pipeline;
+import org.apache.spark.ml.PipelineModel;
+import org.apache.spark.ml.PipelineStage;
+import org.apache.spark.ml.classification.RandomForestClassificationModel;
+import org.apache.spark.ml.classification.RandomForestClassifier;
+import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator;
+import org.apache.spark.ml.feature.*;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+// $example off$
+
+public class JavaRandomForestClassifierExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaRandomForestClassifierExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    DataFrame data = sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
+
+    // Index labels, adding metadata to the label column.
+    // Fit on whole dataset to include all labels in index.
+    StringIndexerModel labelIndexer = new StringIndexer()
+      .setInputCol("label")
+      .setOutputCol("indexedLabel")
+      .fit(data);
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    VectorIndexerModel featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data);
+
+    // Split the data into training and test sets (30% held out for testing)
+    DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
+    DataFrame trainingData = splits[0];
+    DataFrame testData = splits[1];
+
+    // Train a RandomForest model.
+    RandomForestClassifier rf = new RandomForestClassifier()
+      .setLabelCol("indexedLabel")
+      .setFeaturesCol("indexedFeatures");
+
+    // Convert indexed labels back to original labels.
+    IndexToString labelConverter = new IndexToString()
+      .setInputCol("prediction")
+      .setOutputCol("predictedLabel")
+      .setLabels(labelIndexer.labels());
+
+    // Chain indexers and forest in a Pipeline
+    Pipeline pipeline = new Pipeline()
+      .setStages(new PipelineStage[] {labelIndexer, featureIndexer, rf, labelConverter});
+
+    // Train model. This also runs the indexers.
+    PipelineModel model = pipeline.fit(trainingData);
+
+    // Make predictions.
+    DataFrame predictions = model.transform(testData);
+
+    // Select example rows to display.
+    predictions.select("predictedLabel", "label", "features").show(5);
+
+    // Select (prediction, true label) and compute test error
+    MulticlassClassificationEvaluator evaluator = new MulticlassClassificationEvaluator()
+      .setLabelCol("indexedLabel")
+      .setPredictionCol("prediction")
+      .setMetricName("precision");
+    double accuracy = evaluator.evaluate(predictions);
+    System.out.println("Test Error = " + (1.0 - accuracy));
+
+    RandomForestClassificationModel rfModel = (RandomForestClassificationModel)(model.stages()[2]);
+    System.out.println("Learned classification forest model:\n" + rfModel.toDebugString());
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestRegressorExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestRegressorExample.java
new file mode 100644
index 0000000000000..05782a0724a77
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaRandomForestRegressorExample.java
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaSparkContext;
+// $example on$
+import org.apache.spark.ml.Pipeline;
+import org.apache.spark.ml.PipelineModel;
+import org.apache.spark.ml.PipelineStage;
+import org.apache.spark.ml.evaluation.RegressionEvaluator;
+import org.apache.spark.ml.feature.VectorIndexer;
+import org.apache.spark.ml.feature.VectorIndexerModel;
+import org.apache.spark.ml.regression.RandomForestRegressionModel;
+import org.apache.spark.ml.regression.RandomForestRegressor;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.SQLContext;
+// $example off$
+
+public class JavaRandomForestRegressorExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("JavaRandomForestRegressorExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    DataFrame data = sqlContext.read().format("libsvm").load("data/mllib/sample_libsvm_data.txt");
+
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    VectorIndexerModel featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data);
+
+    // Split the data into training and test sets (30% held out for testing)
+    DataFrame[] splits = data.randomSplit(new double[] {0.7, 0.3});
+    DataFrame trainingData = splits[0];
+    DataFrame testData = splits[1];
+
+    // Train a RandomForest model.
+    RandomForestRegressor rf = new RandomForestRegressor()
+      .setLabelCol("label")
+      .setFeaturesCol("indexedFeatures");
+
+    // Chain indexer and forest in a Pipeline
+    Pipeline pipeline = new Pipeline()
+      .setStages(new PipelineStage[] {featureIndexer, rf});
+
+    // Train model.  This also runs the indexer.
+    PipelineModel model = pipeline.fit(trainingData);
+
+    // Make predictions.
+    DataFrame predictions = model.transform(testData);
+
+    // Select example rows to display.
+    predictions.select("prediction", "label", "features").show(5);
+
+    // Select (prediction, true label) and compute test error
+    RegressionEvaluator evaluator = new RegressionEvaluator()
+      .setLabelCol("label")
+      .setPredictionCol("prediction")
+      .setMetricName("rmse");
+    double rmse = evaluator.evaluate(predictions);
+    System.out.println("Root Mean Squared Error (RMSE) on test data = " + rmse);
+
+    RandomForestRegressionModel rfModel = (RandomForestRegressionModel)(model.stages()[1]);
+    System.out.println("Learned regression forest model:\n" + rfModel.toDebugString());
+    // $example off$
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/python/ml/gradient_boosted_tree_classifier_example.py b/examples/src/main/python/ml/gradient_boosted_tree_classifier_example.py
new file mode 100644
index 0000000000000..028497651fbf9
--- /dev/null
+++ b/examples/src/main/python/ml/gradient_boosted_tree_classifier_example.py
@@ -0,0 +1,77 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Gradient Boosted Tree Classifier Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext, SQLContext
+# $example on$
+from pyspark.ml import Pipeline
+from pyspark.ml.classification import GBTClassifier
+from pyspark.ml.feature import StringIndexer, VectorIndexer
+from pyspark.ml.evaluation import MulticlassClassificationEvaluator
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="gradient_boosted_tree_classifier_example")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load and parse the data file, converting it to a DataFrame.
+    data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    # Index labels, adding metadata to the label column.
+    # Fit on whole dataset to include all labels in index.
+    labelIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel").fit(data)
+    # Automatically identify categorical features, and index them.
+    # Set maxCategories so features with > 4 distinct values are treated as continuous.
+    featureIndexer =\
+        VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
+
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a GBT model.
+    gbt = GBTClassifier(labelCol="indexedLabel", featuresCol="indexedFeatures", maxIter=10)
+
+    # Chain indexers and GBT in a Pipeline
+    pipeline = Pipeline(stages=[labelIndexer, featureIndexer, gbt])
+
+    # Train model.  This also runs the indexers.
+    model = pipeline.fit(trainingData)
+
+    # Make predictions.
+    predictions = model.transform(testData)
+
+    # Select example rows to display.
+    predictions.select("prediction", "indexedLabel", "features").show(5)
+
+    # Select (prediction, true label) and compute test error
+    evaluator = MulticlassClassificationEvaluator(
+        labelCol="indexedLabel", predictionCol="prediction", metricName="precision")
+    accuracy = evaluator.evaluate(predictions)
+    print("Test Error = %g" % (1.0 - accuracy))
+
+    gbtModel = model.stages[2]
+    print(gbtModel)  # summary only
+    # $example off$
+
+    sc.stop()
diff --git a/examples/src/main/python/ml/gradient_boosted_tree_regressor_example.py b/examples/src/main/python/ml/gradient_boosted_tree_regressor_example.py
new file mode 100644
index 0000000000000..4246e133a9030
--- /dev/null
+++ b/examples/src/main/python/ml/gradient_boosted_tree_regressor_example.py
@@ -0,0 +1,74 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Gradient Boosted Tree Regressor Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext, SQLContext
+# $example on$
+from pyspark.ml import Pipeline
+from pyspark.ml.regression import GBTRegressor
+from pyspark.ml.feature import VectorIndexer
+from pyspark.ml.evaluation import RegressionEvaluator
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="gradient_boosted_tree_regressor_example")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load and parse the data file, converting it to a DataFrame.
+    data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    # Automatically identify categorical features, and index them.
+    # Set maxCategories so features with > 4 distinct values are treated as continuous.
+    featureIndexer =\
+        VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
+
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a GBT model.
+    gbt = GBTRegressor(featuresCol="indexedFeatures", maxIter=10)
+
+    # Chain indexer and GBT in a Pipeline
+    pipeline = Pipeline(stages=[featureIndexer, gbt])
+
+    # Train model.  This also runs the indexer.
+    model = pipeline.fit(trainingData)
+
+    # Make predictions.
+    predictions = model.transform(testData)
+
+    # Select example rows to display.
+    predictions.select("prediction", "label", "features").show(5)
+
+    # Select (prediction, true label) and compute test error
+    evaluator = RegressionEvaluator(
+        labelCol="label", predictionCol="prediction", metricName="rmse")
+    rmse = evaluator.evaluate(predictions)
+    print("Root Mean Squared Error (RMSE) on test data = %g" % rmse)
+
+    gbtModel = model.stages[1]
+    print(gbtModel)  # summary only
+    # $example off$
+
+    sc.stop()
diff --git a/examples/src/main/python/ml/random_forest_classifier_example.py b/examples/src/main/python/ml/random_forest_classifier_example.py
new file mode 100644
index 0000000000000..b3530d4f41c8e
--- /dev/null
+++ b/examples/src/main/python/ml/random_forest_classifier_example.py
@@ -0,0 +1,77 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Random Forest Classifier Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext, SQLContext
+# $example on$
+from pyspark.ml import Pipeline
+from pyspark.ml.classification import RandomForestClassifier
+from pyspark.ml.feature import StringIndexer, VectorIndexer
+from pyspark.ml.evaluation import MulticlassClassificationEvaluator
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="random_forest_classifier_example")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load and parse the data file, converting it to a DataFrame.
+    data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    # Index labels, adding metadata to the label column.
+    # Fit on whole dataset to include all labels in index.
+    labelIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel").fit(data)
+    # Automatically identify categorical features, and index them.
+    # Set maxCategories so features with > 4 distinct values are treated as continuous.
+    featureIndexer =\
+        VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
+
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a RandomForest model.
+    rf = RandomForestClassifier(labelCol="indexedLabel", featuresCol="indexedFeatures")
+
+    # Chain indexers and forest in a Pipeline
+    pipeline = Pipeline(stages=[labelIndexer, featureIndexer, rf])
+
+    # Train model.  This also runs the indexers.
+    model = pipeline.fit(trainingData)
+
+    # Make predictions.
+    predictions = model.transform(testData)
+
+    # Select example rows to display.
+    predictions.select("prediction", "indexedLabel", "features").show(5)
+
+    # Select (prediction, true label) and compute test error
+    evaluator = MulticlassClassificationEvaluator(
+        labelCol="indexedLabel", predictionCol="prediction", metricName="precision")
+    accuracy = evaluator.evaluate(predictions)
+    print("Test Error = %g" % (1.0 - accuracy))
+
+    rfModel = model.stages[2]
+    print(rfModel)  # summary only
+    # $example off$
+
+    sc.stop()
diff --git a/examples/src/main/python/ml/random_forest_regressor_example.py b/examples/src/main/python/ml/random_forest_regressor_example.py
new file mode 100644
index 0000000000000..b59c7c9414841
--- /dev/null
+++ b/examples/src/main/python/ml/random_forest_regressor_example.py
@@ -0,0 +1,74 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Random Forest Regressor Example.
+"""
+from __future__ import print_function
+
+import sys
+
+from pyspark import SparkContext, SQLContext
+# $example on$
+from pyspark.ml import Pipeline
+from pyspark.ml.regression import RandomForestRegressor
+from pyspark.ml.feature import VectorIndexer
+from pyspark.ml.evaluation import RegressionEvaluator
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="random_forest_regressor_example")
+    sqlContext = SQLContext(sc)
+
+    # $example on$
+    # Load and parse the data file, converting it to a DataFrame.
+    data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    # Automatically identify categorical features, and index them.
+    # Set maxCategories so features with > 4 distinct values are treated as continuous.
+    featureIndexer =\
+        VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4).fit(data)
+
+    # Split the data into training and test sets (30% held out for testing)
+    (trainingData, testData) = data.randomSplit([0.7, 0.3])
+
+    # Train a RandomForest model.
+    rf = RandomForestRegressor(featuresCol="indexedFeatures")
+
+    # Chain indexer and forest in a Pipeline
+    pipeline = Pipeline(stages=[featureIndexer, rf])
+
+    # Train model.  This also runs the indexer.
+    model = pipeline.fit(trainingData)
+
+    # Make predictions.
+    predictions = model.transform(testData)
+
+    # Select example rows to display.
+    predictions.select("prediction", "label", "features").show(5)
+
+    # Select (prediction, true label) and compute test error
+    evaluator = RegressionEvaluator(
+        labelCol="label", predictionCol="prediction", metricName="rmse")
+    rmse = evaluator.evaluate(predictions)
+    print("Root Mean Squared Error (RMSE) on test data = %g" % rmse)
+
+    rfModel = model.stages[1]
+    print(rfModel)  # summary only
+    # $example off$
+
+    sc.stop()
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/GradientBoostedTreeClassifierExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/GradientBoostedTreeClassifierExample.scala
new file mode 100644
index 0000000000000..474af7db4b49b
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/GradientBoostedTreeClassifierExample.scala
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkConf, SparkContext}
+// $example on$
+import org.apache.spark.ml.Pipeline
+import org.apache.spark.ml.classification.{GBTClassificationModel, GBTClassifier}
+import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
+import org.apache.spark.ml.feature.{IndexToString, StringIndexer, VectorIndexer}
+// $example off$
+
+object GradientBoostedTreeClassifierExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("GradientBoostedTreeClassifierExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    // Index labels, adding metadata to the label column.
+    // Fit on whole dataset to include all labels in index.
+    val labelIndexer = new StringIndexer()
+      .setInputCol("label")
+      .setOutputCol("indexedLabel")
+      .fit(data)
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    val featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data)
+
+    // Split the data into training and test sets (30% held out for testing)
+    val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
+
+    // Train a GBT model.
+    val gbt = new GBTClassifier()
+      .setLabelCol("indexedLabel")
+      .setFeaturesCol("indexedFeatures")
+      .setMaxIter(10)
+
+    // Convert indexed labels back to original labels.
+    val labelConverter = new IndexToString()
+      .setInputCol("prediction")
+      .setOutputCol("predictedLabel")
+      .setLabels(labelIndexer.labels)
+
+    // Chain indexers and GBT in a Pipeline
+    val pipeline = new Pipeline()
+      .setStages(Array(labelIndexer, featureIndexer, gbt, labelConverter))
+
+    // Train model.  This also runs the indexers.
+    val model = pipeline.fit(trainingData)
+
+    // Make predictions.
+    val predictions = model.transform(testData)
+
+    // Select example rows to display.
+    predictions.select("predictedLabel", "label", "features").show(5)
+
+    // Select (prediction, true label) and compute test error
+    val evaluator = new MulticlassClassificationEvaluator()
+      .setLabelCol("indexedLabel")
+      .setPredictionCol("prediction")
+      .setMetricName("precision")
+    val accuracy = evaluator.evaluate(predictions)
+    println("Test Error = " + (1.0 - accuracy))
+
+    val gbtModel = model.stages(2).asInstanceOf[GBTClassificationModel]
+    println("Learned classification GBT model:\n" + gbtModel.toDebugString)
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/GradientBoostedTreeRegressorExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/GradientBoostedTreeRegressorExample.scala
new file mode 100644
index 0000000000000..da1cd9c2ce525
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/GradientBoostedTreeRegressorExample.scala
@@ -0,0 +1,85 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkConf, SparkContext}
+// $example on$
+import org.apache.spark.ml.Pipeline
+import org.apache.spark.ml.evaluation.RegressionEvaluator
+import org.apache.spark.ml.feature.VectorIndexer
+import org.apache.spark.ml.regression.{GBTRegressionModel, GBTRegressor}
+// $example off$
+
+object GradientBoostedTreeRegressorExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("GradientBoostedTreeRegressorExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    val featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data)
+
+    // Split the data into training and test sets (30% held out for testing)
+    val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
+
+    // Train a GBT model.
+    val gbt = new GBTRegressor()
+      .setLabelCol("label")
+      .setFeaturesCol("indexedFeatures")
+      .setMaxIter(10)
+
+    // Chain indexer and GBT in a Pipeline
+    val pipeline = new Pipeline()
+      .setStages(Array(featureIndexer, gbt))
+
+    // Train model.  This also runs the indexer.
+    val model = pipeline.fit(trainingData)
+
+    // Make predictions.
+    val predictions = model.transform(testData)
+
+    // Select example rows to display.
+    predictions.select("prediction", "label", "features").show(5)
+
+    // Select (prediction, true label) and compute test error
+    val evaluator = new RegressionEvaluator()
+      .setLabelCol("label")
+      .setPredictionCol("prediction")
+      .setMetricName("rmse")
+    val rmse = evaluator.evaluate(predictions)
+    println("Root Mean Squared Error (RMSE) on test data = " + rmse)
+
+    val gbtModel = model.stages(1).asInstanceOf[GBTRegressionModel]
+    println("Learned regression GBT model:\n" + gbtModel.toDebugString)
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/OneVsRestExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/OneVsRestExample.scala
index 8e4f1b09a24b5..b46faea5713fb 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/OneVsRestExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/OneVsRestExample.scala
@@ -23,12 +23,14 @@ import java.util.concurrent.TimeUnit.{NANOSECONDS => NANO}
 import scopt.OptionParser
 
 import org.apache.spark.{SparkContext, SparkConf}
+// $example on$
 import org.apache.spark.examples.mllib.AbstractParams
 import org.apache.spark.ml.classification.{OneVsRest, LogisticRegression}
 import org.apache.spark.ml.util.MetadataUtils
 import org.apache.spark.mllib.evaluation.MulticlassMetrics
 import org.apache.spark.mllib.linalg.Vector
 import org.apache.spark.sql.DataFrame
+// $example off$
 import org.apache.spark.sql.SQLContext
 
 /**
@@ -112,6 +114,7 @@ object OneVsRestExample {
     val sc = new SparkContext(conf)
     val sqlContext = new SQLContext(sc)
 
+    // $example on$
     val inputData = sqlContext.read.format("libsvm").load(params.input)
     // compute the train/test split: if testInput is not provided use part of input.
     val data = params.testInput match {
@@ -172,6 +175,7 @@ object OneVsRestExample {
     println("label\tfpr")
 
     println(fprs.map {case (label, fpr) => label + "\t" + fpr}.mkString("\n"))
+    // $example off$
 
     sc.stop()
   }
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestClassifierExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestClassifierExample.scala
new file mode 100644
index 0000000000000..e79176ca6ca1c
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestClassifierExample.scala
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkConf, SparkContext}
+// $example on$
+import org.apache.spark.ml.Pipeline
+import org.apache.spark.ml.classification.{RandomForestClassificationModel, RandomForestClassifier}
+import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
+import org.apache.spark.ml.feature.{IndexToString, StringIndexer, VectorIndexer}
+// $example off$
+
+object RandomForestClassifierExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("RandomForestClassifierExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    // Index labels, adding metadata to the label column.
+    // Fit on whole dataset to include all labels in index.
+    val labelIndexer = new StringIndexer()
+      .setInputCol("label")
+      .setOutputCol("indexedLabel")
+      .fit(data)
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    val featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data)
+
+    // Split the data into training and test sets (30% held out for testing)
+    val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
+
+    // Train a RandomForest model.
+    val rf = new RandomForestClassifier()
+      .setLabelCol("indexedLabel")
+      .setFeaturesCol("indexedFeatures")
+      .setNumTrees(10)
+
+    // Convert indexed labels back to original labels.
+    val labelConverter = new IndexToString()
+      .setInputCol("prediction")
+      .setOutputCol("predictedLabel")
+      .setLabels(labelIndexer.labels)
+
+    // Chain indexers and forest in a Pipeline
+    val pipeline = new Pipeline()
+      .setStages(Array(labelIndexer, featureIndexer, rf, labelConverter))
+
+    // Train model.  This also runs the indexers.
+    val model = pipeline.fit(trainingData)
+
+    // Make predictions.
+    val predictions = model.transform(testData)
+
+    // Select example rows to display.
+    predictions.select("predictedLabel", "label", "features").show(5)
+
+    // Select (prediction, true label) and compute test error
+    val evaluator = new MulticlassClassificationEvaluator()
+      .setLabelCol("indexedLabel")
+      .setPredictionCol("prediction")
+      .setMetricName("precision")
+    val accuracy = evaluator.evaluate(predictions)
+    println("Test Error = " + (1.0 - accuracy))
+
+    val rfModel = model.stages(2).asInstanceOf[RandomForestClassificationModel]
+    println("Learned classification forest model:\n" + rfModel.toDebugString)
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestRegressorExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestRegressorExample.scala
new file mode 100644
index 0000000000000..acec1437a1af5
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/RandomForestRegressorExample.scala
@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.ml
+
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkConf, SparkContext}
+// $example on$
+import org.apache.spark.ml.Pipeline
+import org.apache.spark.ml.evaluation.RegressionEvaluator
+import org.apache.spark.ml.feature.VectorIndexer
+import org.apache.spark.ml.regression.{RandomForestRegressionModel, RandomForestRegressor}
+// $example off$
+
+object RandomForestRegressorExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("RandomForestRegressorExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+
+    // $example on$
+    // Load and parse the data file, converting it to a DataFrame.
+    val data = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+
+    // Automatically identify categorical features, and index them.
+    // Set maxCategories so features with > 4 distinct values are treated as continuous.
+    val featureIndexer = new VectorIndexer()
+      .setInputCol("features")
+      .setOutputCol("indexedFeatures")
+      .setMaxCategories(4)
+      .fit(data)
+
+    // Split the data into training and test sets (30% held out for testing)
+    val Array(trainingData, testData) = data.randomSplit(Array(0.7, 0.3))
+
+    // Train a RandomForest model.
+    val rf = new RandomForestRegressor()
+      .setLabelCol("label")
+      .setFeaturesCol("indexedFeatures")
+
+    // Chain indexer and forest in a Pipeline
+    val pipeline = new Pipeline()
+      .setStages(Array(featureIndexer, rf))
+
+    // Train model.  This also runs the indexer.
+    val model = pipeline.fit(trainingData)
+
+    // Make predictions.
+    val predictions = model.transform(testData)
+
+    // Select example rows to display.
+    predictions.select("prediction", "label", "features").show(5)
+
+    // Select (prediction, true label) and compute test error
+    val evaluator = new RegressionEvaluator()
+      .setLabelCol("label")
+      .setPredictionCol("prediction")
+      .setMetricName("rmse")
+    val rmse = evaluator.evaluate(predictions)
+    println("Root Mean Squared Error (RMSE) on test data = " + rmse)
+
+    val rfModel = model.stages(1).asInstanceOf[RandomForestRegressionModel]
+    println("Learned regression forest model:\n" + rfModel.toDebugString)
+    // $example off$
+
+    sc.stop()
+  }
+}
+// scalastyle:on println

From 8019f66df5c65e21d6e4e7e8fbfb7d0471ba3e37 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Tue, 17 Nov 2015 23:51:05 -0800
Subject: [PATCH 314/518] [SPARK-10186][SQL][FOLLOW-UP] simplify test

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9783 from cloud-fan/postgre.
---
 .../org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
index 2e18d0a2baa1c..6eb6b3391a4a4 100644
--- a/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
+++ b/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/PostgresIntegrationSuite.scala
@@ -88,7 +88,7 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationSuite {
     df.write.jdbc(jdbcUrl, "public.barcopy", new Properties)
     // Test write null values.
     df.select(df.queryExecution.analyzed.output.map { a =>
-      Column(If(Literal(true), Literal(null), a)).as(a.name)
+      Column(Literal.create(null, a.dataType)).as(a.name)
     }: _*).write.jdbc(jdbcUrl, "public.barcopy2", new Properties)
   }
 }

From 5e2b44474c2b838bebeffe5ba5cd72961b0cd31e Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 18 Nov 2015 00:09:29 -0800
Subject: [PATCH 315/518] [SPARK-11802][SQL] Kryo-based encoder for opaque
 types in Datasets

I also found a bug with self-joins returning incorrect results in the Dataset API. Two test cases attached and filed SPARK-11803.

Author: Reynold Xin <rxin@databricks.com>

Closes #9789 from rxin/SPARK-11802.
---
 .../scala/org/apache/spark/sql/Encoder.scala  | 31 +++++++-
 .../catalyst/encoders/ExpressionEncoder.scala |  4 +-
 .../catalyst/encoders/ProductEncoder.scala    |  2 +-
 .../sql/catalyst/expressions/objects.scala    | 69 +++++++++++++++++-
 .../catalyst/encoders/FlatEncoderSuite.scala  | 18 +++++
 .../scala/org/apache/spark/sql/Dataset.scala  |  6 ++
 .../org/apache/spark/sql/GroupedDataset.scala |  1 -
 .../org/apache/spark/sql/DatasetSuite.scala   | 70 +++++++++++++++----
 8 files changed, 178 insertions(+), 23 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index c8b017e251637..79c2255641c06 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -17,10 +17,11 @@
 
 package org.apache.spark.sql
 
-import scala.reflect.ClassTag
+import scala.reflect.{ClassTag, classTag}
 
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor}
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.catalyst.expressions.{DeserializeWithKryo, BoundReference, SerializeWithKryo}
+import org.apache.spark.sql.types._
 
 /**
  * Used to convert a JVM object of type `T` to and from the internal Spark SQL representation.
@@ -37,7 +38,33 @@ trait Encoder[T] extends Serializable {
   def clsTag: ClassTag[T]
 }
 
+/**
+ * Methods for creating encoders.
+ */
 object Encoders {
+
+  /**
+   * (Scala-specific) Creates an encoder that serializes objects of type T using Kryo.
+   * This encoder maps T into a single byte array (binary) field.
+   */
+  def kryo[T: ClassTag]: Encoder[T] = {
+    val ser = SerializeWithKryo(BoundReference(0, ObjectType(classOf[AnyRef]), nullable = true))
+    val deser = DeserializeWithKryo[T](BoundReference(0, BinaryType, nullable = true), classTag[T])
+    ExpressionEncoder[T](
+      schema = new StructType().add("value", BinaryType),
+      flat = true,
+      toRowExpressions = Seq(ser),
+      fromRowExpression = deser,
+      clsTag = classTag[T]
+    )
+  }
+
+  /**
+   * Creates an encoder that serializes objects of type T using Kryo.
+   * This encoder maps T into a single byte array (binary) field.
+   */
+  def kryo[T](clazz: Class[T]): Encoder[T] = kryo(ClassTag[T](clazz))
+
   def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder(flat = true)
   def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder(flat = true)
   def SHORT: Encoder[java.lang.Short] = ExpressionEncoder(flat = true)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index 9a1a8f5cbbdc3..b977f278c5b5c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -161,7 +161,9 @@ case class ExpressionEncoder[T](
 
   @transient
   private lazy val extractProjection = GenerateUnsafeProjection.generate(toRowExpressions)
-  private val inputRow = new GenericMutableRow(1)
+
+  @transient
+  private lazy val inputRow = new GenericMutableRow(1)
 
   @transient
   private lazy val constructProjection = GenerateSafeProjection.generate(fromRowExpression :: Nil)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
index 414adb21168ed..55c4ee11b20f4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
@@ -230,7 +230,7 @@ object ProductEncoder {
           Invoke(inputObject, "booleanValue", BooleanType)
 
         case other =>
-          throw new UnsupportedOperationException(s"Extractor for type $other is not supported")
+          throw new UnsupportedOperationException(s"Encoder for type $other is not supported")
       }
     }
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index 5cd19de68391c..489c6126f8cd3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -17,13 +17,15 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
+import scala.language.existentials
+import scala.reflect.ClassTag
+
+import org.apache.spark.SparkConf
+import org.apache.spark.serializer.{KryoSerializerInstance, KryoSerializer}
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.analysis.SimpleAnalyzer
 import org.apache.spark.sql.catalyst.plans.logical.{Project, LocalRelation}
 import org.apache.spark.sql.catalyst.util.GenericArrayData
-
-import scala.language.existentials
-
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.codegen.{GeneratedExpressionCode, CodeGenContext}
 import org.apache.spark.sql.types._
@@ -514,3 +516,64 @@ case class GetInternalRowField(child: Expression, ordinal: Int, dataType: DataTy
     """
   }
 }
+
+/** Serializes an input object using Kryo serializer. */
+case class SerializeWithKryo(child: Expression) extends UnaryExpression {
+
+  override def eval(input: InternalRow): Any =
+    throw new UnsupportedOperationException("Only code-generated evaluation is supported")
+
+  override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    val input = child.gen(ctx)
+    val kryo = ctx.freshName("kryoSerializer")
+    val kryoClass = classOf[KryoSerializer].getName
+    val kryoInstanceClass = classOf[KryoSerializerInstance].getName
+    val sparkConfClass = classOf[SparkConf].getName
+    ctx.addMutableState(
+      kryoInstanceClass,
+      kryo,
+      s"$kryo = ($kryoInstanceClass) new $kryoClass(new $sparkConfClass()).newInstance();")
+
+    s"""
+      ${input.code}
+      final boolean ${ev.isNull} = ${input.isNull};
+      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+      if (!${ev.isNull}) {
+        ${ev.value} = $kryo.serialize(${input.value}, null).array();
+      }
+     """
+  }
+
+  override def dataType: DataType = BinaryType
+}
+
+/**
+ * Deserializes an input object using Kryo serializer. Note that the ClassTag is not an implicit
+ * parameter because TreeNode cannot copy implicit parameters.
+ */
+case class DeserializeWithKryo[T](child: Expression, tag: ClassTag[T]) extends UnaryExpression {
+
+  override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    val input = child.gen(ctx)
+    val kryo = ctx.freshName("kryoSerializer")
+    val kryoClass = classOf[KryoSerializer].getName
+    val kryoInstanceClass = classOf[KryoSerializerInstance].getName
+    val sparkConfClass = classOf[SparkConf].getName
+    ctx.addMutableState(
+      kryoInstanceClass,
+      kryo,
+      s"$kryo = ($kryoInstanceClass) new $kryoClass(new $sparkConfClass()).newInstance();")
+
+    s"""
+      ${input.code}
+      final boolean ${ev.isNull} = ${input.isNull};
+      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+      if (!${ev.isNull}) {
+        ${ev.value} = (${ctx.javaType(dataType)})
+          $kryo.deserialize(java.nio.ByteBuffer.wrap(${input.value}), null);
+      }
+     """
+  }
+
+  override def dataType: DataType = ObjectType(tag.runtimeClass)
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
index 55821c4370684..2729db84897a2 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.catalyst.encoders
 
 import java.sql.{Date, Timestamp}
+import org.apache.spark.sql.Encoders
 
 class FlatEncoderSuite extends ExpressionEncoderSuite {
   encodeDecodeTest(false, FlatEncoder[Boolean], "primitive boolean")
@@ -71,4 +72,21 @@ class FlatEncoderSuite extends ExpressionEncoderSuite {
   encodeDecodeTest(Map(1 -> "a", 2 -> null), FlatEncoder[Map[Int, String]], "map with null")
   encodeDecodeTest(Map(1 -> Map("a" -> 1), 2 -> Map("b" -> 2)),
     FlatEncoder[Map[Int, Map[String, Int]]], "map of map")
+
+  // Kryo encoders
+  encodeDecodeTest(
+    "hello",
+    Encoders.kryo[String].asInstanceOf[ExpressionEncoder[String]],
+    "kryo string")
+  encodeDecodeTest(
+    new NotJavaSerializable(15),
+    Encoders.kryo[NotJavaSerializable].asInstanceOf[ExpressionEncoder[NotJavaSerializable]],
+    "kryo object serialization")
+}
+
+
+class NotJavaSerializable(val value: Int) {
+  override def equals(other: Any): Boolean = {
+    this.value == other.asInstanceOf[NotJavaSerializable].value
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 718ed812dd64c..817c20fdbb9f3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -147,6 +147,12 @@ class Dataset[T] private[sql](
     }
   }
 
+  /**
+   * Returns the number of elements in the [[Dataset]].
+   * @since 1.6.0
+   */
+  def count(): Long = toDF().count()
+
   /* *********************** *
    *  Functional Operations  *
    * *********************** */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 467cd42b9b8dc..c66162ee2148a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -17,7 +17,6 @@
 
 package org.apache.spark.sql
 
-
 import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index ea29428c55088..a522894c374f9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -24,21 +24,6 @@ import scala.language.postfixOps
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
 
-case class ClassData(a: String, b: Int)
-
-/**
- * A class used to test serialization using encoders. This class throws exceptions when using
- * Java serialization -- so the only way it can be "serialized" is through our encoders.
- */
-case class NonSerializableCaseClass(value: String) extends Externalizable {
-  override def readExternal(in: ObjectInput): Unit = {
-    throw new UnsupportedOperationException
-  }
-
-  override def writeExternal(out: ObjectOutput): Unit = {
-    throw new UnsupportedOperationException
-  }
-}
 
 class DatasetSuite extends QueryTest with SharedSQLContext {
   import testImplicits._
@@ -362,8 +347,63 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     checkAnswer(joined, ("2", 2))
   }
 
+  ignore("self join") {
+    val ds = Seq("1", "2").toDS().as("a")
+    val joined = ds.joinWith(ds, lit(true))
+    checkAnswer(joined, ("1", "1"), ("1", "2"), ("2", "1"), ("2", "2"))
+  }
+
   test("toString") {
     val ds = Seq((1, 2)).toDS()
     assert(ds.toString == "[_1: int, _2: int]")
   }
+
+  test("kryo encoder") {
+    implicit val kryoEncoder = Encoders.kryo[KryoData]
+    val ds = sqlContext.createDataset(Seq(KryoData(1), KryoData(2)))
+
+    assert(ds.groupBy(p => p).count().collect().toSeq ==
+      Seq((KryoData(1), 1L), (KryoData(2), 1L)))
+  }
+
+  ignore("kryo encoder self join") {
+    implicit val kryoEncoder = Encoders.kryo[KryoData]
+    val ds = sqlContext.createDataset(Seq(KryoData(1), KryoData(2)))
+    assert(ds.joinWith(ds, lit(true)).collect().toSet ==
+      Set(
+        (KryoData(1), KryoData(1)),
+        (KryoData(1), KryoData(2)),
+        (KryoData(2), KryoData(1)),
+        (KryoData(2), KryoData(2))))
+  }
+}
+
+
+case class ClassData(a: String, b: Int)
+
+/**
+ * A class used to test serialization using encoders. This class throws exceptions when using
+ * Java serialization -- so the only way it can be "serialized" is through our encoders.
+ */
+case class NonSerializableCaseClass(value: String) extends Externalizable {
+  override def readExternal(in: ObjectInput): Unit = {
+    throw new UnsupportedOperationException
+  }
+
+  override def writeExternal(out: ObjectOutput): Unit = {
+    throw new UnsupportedOperationException
+  }
+}
+
+/** Used to test Kryo encoder. */
+class KryoData(val a: Int) {
+  override def equals(other: Any): Boolean = {
+    a == other.asInstanceOf[KryoData].a
+  }
+  override def hashCode: Int = a
+  override def toString: String = s"KryoData($a)"
+}
+
+object KryoData {
+  def apply(a: Int): KryoData = new KryoData(a)
 }

From 1714350bddd78cd1398e1a816f675ab729001081 Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Wed, 18 Nov 2015 00:42:52 -0800
Subject: [PATCH 316/518] [SPARK-11792][SQL] SizeEstimator cannot provide a
 good size estimation of UnsafeHashedRelations

https://issues.apache.org/jira/browse/SPARK-11792

Right now, SizeEstimator will "think" a small UnsafeHashedRelation is several GBs.

Author: Yin Huai <yhuai@databricks.com>

Closes #9788 from yhuai/SPARK-11792.
---
 .../spark/memory/TaskMemoryManager.java       |  3 +++
 .../org/apache/spark/util/SizeEstimator.scala | 26 ++++++++++++++++---
 .../spark/util/SizeEstimatorSuite.scala       | 22 ++++++++++++++++
 .../sql/execution/joins/HashedRelation.scala  | 10 +++++--
 4 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
index 5f743b28857b4..d31eb449eb82e 100644
--- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
+++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java
@@ -215,6 +215,9 @@ public void showMemoryUsage() {
       logger.info(
         "{} bytes of memory were used by task {} but are not associated with specific consumers",
         memoryNotAccountedFor, taskAttemptId);
+      logger.info(
+        "{} bytes of memory are used for execution and {} bytes of memory are used for storage",
+        memoryManager.executionMemoryUsed(), memoryManager.storageMemoryUsed());
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/util/SizeEstimator.scala b/core/src/main/scala/org/apache/spark/util/SizeEstimator.scala
index 23ee4eff0881b..c3a2675ee5f45 100644
--- a/core/src/main/scala/org/apache/spark/util/SizeEstimator.scala
+++ b/core/src/main/scala/org/apache/spark/util/SizeEstimator.scala
@@ -31,6 +31,16 @@ import org.apache.spark.Logging
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.util.collection.OpenHashSet
 
+/**
+ * A trait that allows a class to give [[SizeEstimator]] more accurate size estimation.
+ * When a class extends it, [[SizeEstimator]] will query the `estimatedSize` first.
+ * If `estimatedSize` does not return [[None]], [[SizeEstimator]] will use the returned size
+ * as the size of the object. Otherwise, [[SizeEstimator]] will do the estimation work.
+ */
+private[spark] trait SizeEstimation {
+  def estimatedSize: Option[Long]
+}
+
 /**
  * :: DeveloperApi ::
  * Estimates the sizes of Java objects (number of bytes of memory they occupy), for use in
@@ -199,10 +209,18 @@ object SizeEstimator extends Logging {
       // the size estimator since it references the whole REPL. Do nothing in this case. In
       // general all ClassLoaders and Classes will be shared between objects anyway.
     } else {
-      val classInfo = getClassInfo(cls)
-      state.size += alignSize(classInfo.shellSize)
-      for (field <- classInfo.pointerFields) {
-        state.enqueue(field.get(obj))
+      val estimatedSize = obj match {
+        case s: SizeEstimation => s.estimatedSize
+        case _ => None
+      }
+      if (estimatedSize.isDefined) {
+        state.size += estimatedSize.get
+      } else {
+        val classInfo = getClassInfo(cls)
+        state.size += alignSize(classInfo.shellSize)
+        for (field <- classInfo.pointerFields) {
+          state.enqueue(field.get(obj))
+        }
       }
     }
   }
diff --git a/core/src/test/scala/org/apache/spark/util/SizeEstimatorSuite.scala b/core/src/test/scala/org/apache/spark/util/SizeEstimatorSuite.scala
index 20550178fb1bd..9b6261af123e6 100644
--- a/core/src/test/scala/org/apache/spark/util/SizeEstimatorSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/SizeEstimatorSuite.scala
@@ -60,6 +60,18 @@ class DummyString(val arr: Array[Char]) {
   @transient val hash32: Int = 0
 }
 
+class DummyClass8 extends SizeEstimation {
+  val x: Int = 0
+
+  override def estimatedSize: Option[Long] = Some(2015)
+}
+
+class DummyClass9 extends SizeEstimation {
+  val x: Int = 0
+
+  override def estimatedSize: Option[Long] = None
+}
+
 class SizeEstimatorSuite
   extends SparkFunSuite
   with BeforeAndAfterEach
@@ -214,4 +226,14 @@ class SizeEstimatorSuite
     // Class should be 32 bytes on s390x if recognised as 64 bit platform
     assertResult(32)(SizeEstimator.estimate(new DummyClass7))
   }
+
+  test("SizeEstimation can provide the estimated size") {
+    // DummyClass8 provides its size estimation.
+    assertResult(2015)(SizeEstimator.estimate(new DummyClass8))
+    assertResult(20206)(SizeEstimator.estimate(Array.fill(10)(new DummyClass8)))
+
+    // DummyClass9 does not provide its size estimation.
+    assertResult(16)(SizeEstimator.estimate(new DummyClass9))
+    assertResult(216)(SizeEstimator.estimate(Array.fill(10)(new DummyClass9)))
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala
index cc8abb1ba463c..49ae09bf53782 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala
@@ -30,7 +30,7 @@ import org.apache.spark.sql.execution.metric.{LongSQLMetric, SQLMetrics}
 import org.apache.spark.unsafe.Platform
 import org.apache.spark.unsafe.map.BytesToBytesMap
 import org.apache.spark.unsafe.memory.MemoryLocation
-import org.apache.spark.util.Utils
+import org.apache.spark.util.{SizeEstimation, Utils}
 import org.apache.spark.util.collection.CompactBuffer
 import org.apache.spark.{SparkConf, SparkEnv}
 
@@ -189,7 +189,9 @@ private[execution] object HashedRelation {
  */
 private[joins] final class UnsafeHashedRelation(
     private var hashTable: JavaHashMap[UnsafeRow, CompactBuffer[UnsafeRow]])
-  extends HashedRelation with Externalizable {
+  extends HashedRelation
+  with SizeEstimation
+  with Externalizable {
 
   private[joins] def this() = this(null)  // Needed for serialization
 
@@ -215,6 +217,10 @@ private[joins] final class UnsafeHashedRelation(
     }
   }
 
+  override def estimatedSize: Option[Long] = {
+    Option(binaryMap).map(_.getTotalMemoryConsumption)
+  }
+
   override def get(key: InternalRow): Seq[InternalRow] = {
     val unsafeKey = key.asInstanceOf[UnsafeRow]
 

From b8f4379ba1c5c1a8f3b4c88bd97031dc8ad2dfea Mon Sep 17 00:00:00 2001
From: somideshmukh <somilde@us.ibm.com>
Date: Wed, 18 Nov 2015 08:51:01 +0000
Subject: [PATCH 317/518] [SPARK-10946][SQL] JDBC - Use Statement.executeUpdate
 instead of PreparedStatement.executeUpdate for DDLs

New changes with JDBCRDD

Author: somideshmukh <somilde@us.ibm.com>

Closes #9733 from somideshmukh/SomilBranch-1.1.
---
 .../src/main/scala/org/apache/spark/sql/DataFrameWriter.scala   | 2 +-
 .../apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala
index e63a4d5e8b10b..03867beb78224 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala
@@ -297,7 +297,7 @@ final class DataFrameWriter private[sql](df: DataFrame) {
       if (!tableExists) {
         val schema = JdbcUtils.schemaString(df, url)
         val sql = s"CREATE TABLE $table ($schema)"
-        conn.prepareStatement(sql).executeUpdate()
+        conn.createStatement.executeUpdate(sql)
       }
     } finally {
       conn.close()
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
index 32d28e59377a1..7375a5c09123f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
@@ -55,7 +55,7 @@ object JdbcUtils extends Logging {
    * Drops a table from the JDBC database.
    */
   def dropTable(conn: Connection, table: String): Unit = {
-    conn.prepareStatement(s"DROP TABLE $table").executeUpdate()
+    conn.createStatement.executeUpdate(s"DROP TABLE $table")
   }
 
   /**

From e62820c85fe02c70f9ed51b2e68d41ff8cfecd40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-Baptiste=20Onofr=C3=A9?= <jbonofre@apache.org>
Date: Wed, 18 Nov 2015 08:57:58 +0000
Subject: [PATCH 318/518] [SPARK-6541] Sort executors by ID (numeric)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

"Force" the executor ID sort with Int.

Author: Jean-Baptiste Onofré <jbonofre@apache.org>

Closes #9165 from jbonofre/SPARK-6541.
---
 .../org/apache/spark/ui/static/sorttable.js         |  2 +-
 .../org/apache/spark/ui/jobs/ExecutorTable.scala    | 13 +++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/core/src/main/resources/org/apache/spark/ui/static/sorttable.js b/core/src/main/resources/org/apache/spark/ui/static/sorttable.js
index dde6069000bc4..a73d9a5cbc215 100644
--- a/core/src/main/resources/org/apache/spark/ui/static/sorttable.js
+++ b/core/src/main/resources/org/apache/spark/ui/static/sorttable.js
@@ -89,7 +89,7 @@ sorttable = {
         // make it clickable to sort
         headrow[i].sorttable_columnindex = i;
         headrow[i].sorttable_tbody = table.tBodies[0];
-        dean_addEvent(headrow[i],"click", function(e) {
+        dean_addEvent(headrow[i],"click", sorttable.innerSortFunction = function(e) {
 
           if (this.className.search(/\bsorttable_sorted\b/) != -1) {
             // if we're already sorted by this column, just 
diff --git a/core/src/main/scala/org/apache/spark/ui/jobs/ExecutorTable.scala b/core/src/main/scala/org/apache/spark/ui/jobs/ExecutorTable.scala
index be144f6065baa..1268f44596f8a 100644
--- a/core/src/main/scala/org/apache/spark/ui/jobs/ExecutorTable.scala
+++ b/core/src/main/scala/org/apache/spark/ui/jobs/ExecutorTable.scala
@@ -18,7 +18,7 @@
 package org.apache.spark.ui.jobs
 
 import scala.collection.mutable
-import scala.xml.Node
+import scala.xml.{Unparsed, Node}
 
 import org.apache.spark.ui.{ToolTips, UIUtils}
 import org.apache.spark.ui.jobs.UIData.StageUIData
@@ -52,7 +52,7 @@ private[ui] class ExecutorTable(stageId: Int, stageAttemptId: Int, parent: Stage
 
     <table class={UIUtils.TABLE_CLASS_STRIPED_SORTABLE}>
       <thead>
-        <th>Executor ID</th>
+        <th id="executorid">Executor ID</th>
         <th>Address</th>
         <th>Task Time</th>
         <th>Total Tasks</th>
@@ -89,6 +89,15 @@ private[ui] class ExecutorTable(stageId: Int, stageAttemptId: Int, parent: Stage
         {createExecutorTable()}
       </tbody>
     </table>
+    <script>
+      {Unparsed {
+        """
+          |      window.onload = function() {
+          |        sorttable.innerSortFunction.apply(document.getElementById('executorid'), [])
+          |      };
+        """.stripMargin
+      }}
+    </script>
   }
 
   private def createExecutorTable() : Seq[Node] = {

From 9631ca35275b0ce8a5219f975907ac36ed11f528 Mon Sep 17 00:00:00 2001
From: Sean Owen <sowen@cloudera.com>
Date: Wed, 18 Nov 2015 08:59:20 +0000
Subject: [PATCH 319/518] [SPARK-11652][CORE] Remote code execution with
 InvokerTransformer

Update to Commons Collections 3.2.2 to avoid any potential remote code execution vulnerability

Author: Sean Owen <sowen@cloudera.com>

Closes #9731 from srowen/SPARK-11652.
---
 pom.xml | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pom.xml b/pom.xml
index 940e2d8740bf1..ad849112ce76c 100644
--- a/pom.xml
+++ b/pom.xml
@@ -162,6 +162,8 @@
     <!--  commons-httpclient/commons-httpclient-->
     <httpclient.classic.version>3.1</httpclient.classic.version>
     <commons.math3.version>3.4.1</commons.math3.version>
+    <!-- managed up from 3.2.1 for SPARK-11652 -->
+    <commons.collections.version>3.2.2</commons.collections.version>
     <scala.version>2.10.5</scala.version>
     <scala.binary.version>2.10</scala.binary.version>
     <jline.version>${scala.version}</jline.version>
@@ -475,6 +477,11 @@
         <artifactId>commons-math3</artifactId>
         <version>${commons.math3.version}</version>
       </dependency>
+      <dependency>
+        <groupId>org.apache.commons</groupId>
+        <artifactId>commons-collections</artifactId>
+        <version>${commons.collections.version}</version>
+      </dependency>
       <dependency>
         <groupId>org.apache.ivy</groupId>
         <artifactId>ivy</artifactId>

From 1429e0a2b562469146b6fa06051c85a00092e5b8 Mon Sep 17 00:00:00 2001
From: Viveka Kulharia <vivkul@iitk.ac.in>
Date: Wed, 18 Nov 2015 09:10:15 +0000
Subject: [PATCH 320/518] rmse was wrongly calculated

It was multiplying with U instaed of dividing by U

Author: Viveka Kulharia <vivkul@iitk.ac.in>

Closes #9771 from vivkul/patch-1.
---
 examples/src/main/python/als.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/src/main/python/als.py b/examples/src/main/python/als.py
index 1c3a787bd0e94..205ca02962bee 100755
--- a/examples/src/main/python/als.py
+++ b/examples/src/main/python/als.py
@@ -36,7 +36,7 @@
 
 def rmse(R, ms, us):
     diff = R - ms * us.T
-    return np.sqrt(np.sum(np.power(diff, 2)) / M * U)
+    return np.sqrt(np.sum(np.power(diff, 2)) / (M * U))
 
 
 def update(i, vec, mat, ratings):

From 3a6807fdf07b0e73d76502a6bd91ad979fde8b61 Mon Sep 17 00:00:00 2001
From: Jeff Zhang <zjffdu@apache.org>
Date: Wed, 18 Nov 2015 08:18:54 -0800
Subject: [PATCH 321/518] =?UTF-8?q?[SPARK-11804]=20[PYSPARK]=20Exception?=
 =?UTF-8?q?=20raise=20when=20using=20Jdbc=20predicates=20opt=E2=80=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…ion in PySpark

Author: Jeff Zhang <zjffdu@apache.org>

Closes #9791 from zjffdu/SPARK-11804.
---
 python/pyspark/sql/readwriter.py | 10 +++++-----
 python/pyspark/sql/utils.py      | 13 +++++++++++++
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py
index 7b8ddb9feba34..e8f0d7ec77035 100644
--- a/python/pyspark/sql/readwriter.py
+++ b/python/pyspark/sql/readwriter.py
@@ -26,6 +26,7 @@
 from pyspark.rdd import ignore_unicode_prefix
 from pyspark.sql.column import _to_seq
 from pyspark.sql.types import *
+from pyspark.sql import utils
 
 __all__ = ["DataFrameReader", "DataFrameWriter"]
 
@@ -131,9 +132,7 @@ def load(self, path=None, format=None, schema=None, **options):
             if type(path) == list:
                 paths = path
                 gateway = self._sqlContext._sc._gateway
-                jpaths = gateway.new_array(gateway.jvm.java.lang.String, len(paths))
-                for i in range(0, len(paths)):
-                    jpaths[i] = paths[i]
+                jpaths = utils.toJArray(gateway, gateway.jvm.java.lang.String, paths)
                 return self._df(self._jreader.load(jpaths))
             else:
                 return self._df(self._jreader.load(path))
@@ -269,8 +268,9 @@ def jdbc(self, url, table, column=None, lowerBound=None, upperBound=None, numPar
             return self._df(self._jreader.jdbc(url, table, column, int(lowerBound), int(upperBound),
                                                int(numPartitions), jprop))
         if predicates is not None:
-            arr = self._sqlContext._sc._jvm.PythonUtils.toArray(predicates)
-            return self._df(self._jreader.jdbc(url, table, arr, jprop))
+            gateway = self._sqlContext._sc._gateway
+            jpredicates = utils.toJArray(gateway, gateway.jvm.java.lang.String, predicates)
+            return self._df(self._jreader.jdbc(url, table, jpredicates, jprop))
         return self._df(self._jreader.jdbc(url, table, jprop))
 
 
diff --git a/python/pyspark/sql/utils.py b/python/pyspark/sql/utils.py
index c4fda8bd3b891..b0a0373372d20 100644
--- a/python/pyspark/sql/utils.py
+++ b/python/pyspark/sql/utils.py
@@ -71,3 +71,16 @@ def install_exception_handler():
     patched = capture_sql_exception(original)
     # only patch the one used in in py4j.java_gateway (call Java API)
     py4j.java_gateway.get_return_value = patched
+
+
+def toJArray(gateway, jtype, arr):
+    """
+    Convert python list to java type array
+    :param gateway: Py4j Gateway
+    :param jtype: java type of element in array
+    :param arr: python type list
+    """
+    jarr = gateway.new_array(jtype, len(arr))
+    for i in range(0, len(arr)):
+        jarr[i] = arr[i]
+    return jarr

From a97d6f3a5861e9f2bbe36957e3b39f835f3e214c Mon Sep 17 00:00:00 2001
From: zero323 <matthew.szymkiewicz@gmail.com>
Date: Wed, 18 Nov 2015 08:32:03 -0800
Subject: [PATCH 322/518] [SPARK-11281][SPARKR] Add tests covering the issue.

The goal of this PR is to add tests covering the issue to ensure that is was resolved by [SPARK-11086](https://issues.apache.org/jira/browse/SPARK-11086).

Author: zero323 <matthew.szymkiewicz@gmail.com>

Closes #9743 from zero323/SPARK-11281-tests.
---
 R/pkg/inst/tests/test_sparkSQL.R | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 8ff06276599e2..87ab33f6384b1 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -229,7 +229,7 @@ test_that("create DataFrame from list or data.frame", {
   df <- createDataFrame(sqlContext, l, c("a", "b"))
   expect_equal(columns(df), c("a", "b"))
 
-  l <- list(list(a=1, b=2), list(a=3, b=4))
+  l <- list(list(a = 1, b = 2), list(a = 3, b = 4))
   df <- createDataFrame(sqlContext, l)
   expect_equal(columns(df), c("a", "b"))
 
@@ -292,11 +292,15 @@ test_that("create DataFrame with complex types", {
 })
 
 test_that("create DataFrame from a data.frame with complex types", {
-  ldf <- data.frame(row.names=1:2)
+  ldf <- data.frame(row.names = 1:2)
   ldf$a_list <- list(list(1, 2), list(3, 4))
+  ldf$an_envir <- c(as.environment(list(a = 1, b = 2)), as.environment(list(c = 3)))
+
   sdf <- createDataFrame(sqlContext, ldf)
+  collected <- collect(sdf)
 
-  expect_equivalent(ldf, collect(sdf))
+  expect_identical(ldf[, 1, FALSE], collected[, 1, FALSE])
+  expect_equal(ldf$an_envir, collected$an_envir)
 })
 
 # For test map type and struct type in DataFrame

From 224723e6a8b198ef45d6c5ca5d2f9c61188ada8f Mon Sep 17 00:00:00 2001
From: Sun Rui <rui.sun@intel.com>
Date: Wed, 18 Nov 2015 08:41:45 -0800
Subject: [PATCH 323/518] [SPARK-11773][SPARKR] Implement collection functions
 in SparkR.

Author: Sun Rui <rui.sun@intel.com>

Closes #9764 from sun-rui/SPARK-11773.
---
 R/pkg/NAMESPACE                  |   2 +
 R/pkg/R/DataFrame.R              |   2 +-
 R/pkg/R/functions.R              | 109 ++++++++++++++++++++++---------
 R/pkg/R/generics.R               |  10 ++-
 R/pkg/R/utils.R                  |   2 +-
 R/pkg/inst/tests/test_sparkSQL.R |  10 +++
 6 files changed, 100 insertions(+), 35 deletions(-)

diff --git a/R/pkg/NAMESPACE b/R/pkg/NAMESPACE
index 2ee7d6f94f1bc..260c9edce62e0 100644
--- a/R/pkg/NAMESPACE
+++ b/R/pkg/NAMESPACE
@@ -98,6 +98,7 @@ exportMethods("%in%",
               "add_months",
               "alias",
               "approxCountDistinct",
+              "array_contains",
               "asc",
               "ascii",
               "asin",
@@ -215,6 +216,7 @@ exportMethods("%in%",
               "sinh",
               "size",
               "skewness",
+              "sort_array",
               "soundex",
               "stddev",
               "stddev_pop",
diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index fd105ba5bc9bb..34177e3cdd94f 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -2198,4 +2198,4 @@ setMethod("coltypes",
             rTypes[naIndices] <- types[naIndices]
 
             rTypes
-          })
\ No newline at end of file
+          })
diff --git a/R/pkg/R/functions.R b/R/pkg/R/functions.R
index 3d0255a62f155..ff0f438045c14 100644
--- a/R/pkg/R/functions.R
+++ b/R/pkg/R/functions.R
@@ -373,22 +373,6 @@ setMethod("exp",
             column(jc)
           })
 
-#' explode
-#'
-#' Creates a new row for each element in the given array or map column.
-#'
-#' @rdname explode
-#' @name explode
-#' @family collection_funcs
-#' @export
-#' @examples \dontrun{explode(df$c)}
-setMethod("explode",
-          signature(x = "Column"),
-          function(x) {
-            jc <- callJStatic("org.apache.spark.sql.functions", "explode", x@jc)
-            column(jc)
-          })
-
 #' expm1
 #'
 #' Computes the exponential of the given value minus one.
@@ -980,22 +964,6 @@ setMethod("sinh",
             column(jc)
           })
 
-#' size
-#'
-#' Returns length of array or map.
-#'
-#' @rdname size
-#' @name size
-#' @family collection_funcs
-#' @export
-#' @examples \dontrun{size(df$c)}
-setMethod("size",
-          signature(x = "Column"),
-          function(x) {
-            jc <- callJStatic("org.apache.spark.sql.functions", "size", x@jc)
-            column(jc)
-          })
-
 #' skewness
 #'
 #' Aggregate function: returns the skewness of the values in a group.
@@ -2365,3 +2333,80 @@ setMethod("rowNumber",
             jc <- callJStatic("org.apache.spark.sql.functions", "rowNumber")
             column(jc)
           })
+
+###################### Collection functions######################
+
+#' array_contains
+#'
+#' Returns true if the array contain the value.
+#'
+#' @param x A Column
+#' @param value A value to be checked if contained in the column
+#' @rdname array_contains
+#' @name array_contains
+#' @family collection_funcs
+#' @export
+#' @examples \dontrun{array_contains(df$c, 1)}
+setMethod("array_contains",
+          signature(x = "Column", value = "ANY"),
+          function(x, value) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "array_contains", x@jc, value)
+            column(jc)
+          })
+
+#' explode
+#'
+#' Creates a new row for each element in the given array or map column.
+#'
+#' @rdname explode
+#' @name explode
+#' @family collection_funcs
+#' @export
+#' @examples \dontrun{explode(df$c)}
+setMethod("explode",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "explode", x@jc)
+            column(jc)
+          })
+
+#' size
+#'
+#' Returns length of array or map.
+#'
+#' @rdname size
+#' @name size
+#' @family collection_funcs
+#' @export
+#' @examples \dontrun{size(df$c)}
+setMethod("size",
+          signature(x = "Column"),
+          function(x) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "size", x@jc)
+            column(jc)
+          })
+
+#' sort_array
+#'
+#' Sorts the input array for the given column in ascending order,
+#' according to the natural ordering of the array elements.
+#'
+#' @param x A Column to sort
+#' @param asc A logical flag indicating the sorting order.
+#'            TRUE, sorting is in ascending order.
+#'            FALSE, sorting is in descending order.
+#' @rdname sort_array
+#' @name sort_array
+#' @family collection_funcs
+#' @export
+#' @examples
+#' \dontrun{
+#' sort_array(df$c)
+#' sort_array(df$c, FALSE)
+#' }
+setMethod("sort_array",
+          signature(x = "Column"),
+          function(x, asc = TRUE) {
+            jc <- callJStatic("org.apache.spark.sql.functions", "sort_array", x@jc, asc)
+            column(jc)
+          })
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index afdeffc2abd83..0dcd05438222b 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -644,6 +644,10 @@ setGeneric("add_months", function(y, x) { standardGeneric("add_months") })
 #' @export
 setGeneric("approxCountDistinct", function(x, ...) { standardGeneric("approxCountDistinct") })
 
+#' @rdname array_contains
+#' @export
+setGeneric("array_contains", function(x, value) { standardGeneric("array_contains") })
+
 #' @rdname ascii
 #' @export
 setGeneric("ascii", function(x) { standardGeneric("ascii") })
@@ -961,6 +965,10 @@ setGeneric("size", function(x) { standardGeneric("size") })
 #' @export
 setGeneric("skewness", function(x) { standardGeneric("skewness") })
 
+#' @rdname sort_array
+#' @export
+setGeneric("sort_array", function(x, asc = TRUE) { standardGeneric("sort_array") })
+
 #' @rdname soundex
 #' @export
 setGeneric("soundex", function(x) { standardGeneric("soundex") })
@@ -1076,4 +1084,4 @@ setGeneric("with")
 
 #' @rdname coltypes
 #' @export
-setGeneric("coltypes", function(x) { standardGeneric("coltypes") })
\ No newline at end of file
+setGeneric("coltypes", function(x) { standardGeneric("coltypes") })
diff --git a/R/pkg/R/utils.R b/R/pkg/R/utils.R
index db3b2c4bbd799..45c77a86c9582 100644
--- a/R/pkg/R/utils.R
+++ b/R/pkg/R/utils.R
@@ -635,4 +635,4 @@ assignNewEnv <- function(data) {
     assign(x = cols[i], value = data[, cols[i]], envir = env)
   }
   env
-}
\ No newline at end of file
+}
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 87ab33f6384b1..d9a94faff7ac0 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -878,6 +878,16 @@ test_that("column functions", {
 
   df4 <- createDataFrame(sqlContext, list(list(a = "010101")))
   expect_equal(collect(select(df4, conv(df4$a, 2, 16)))[1, 1], "15")
+
+  # Test array_contains() and sort_array()
+  df <- createDataFrame(sqlContext, list(list(list(1L, 2L, 3L)), list(list(6L, 5L, 4L))))
+  result <- collect(select(df, array_contains(df[[1]], 1L)))[[1]]
+  expect_equal(result, c(TRUE, FALSE))
+
+  result <- collect(select(df, sort_array(df[[1]], FALSE)))[[1]]
+  expect_equal(result, list(list(3L, 2L, 1L), list(6L, 5L, 4L)))
+  result <- collect(select(df, sort_array(df[[1]])))[[1]]
+  expect_equal(result, list(list(1L, 2L, 3L), list(4L, 5L, 6L)))
 })
 #
 test_that("column binary mathfunctions", {

From 3cca5ffb3d60d5de9a54bc71cf0b8279898936d2 Mon Sep 17 00:00:00 2001
From: Hurshal Patel <hpatel516@gmail.com>
Date: Wed, 18 Nov 2015 09:28:59 -0800
Subject: [PATCH 324/518] [SPARK-11195][CORE] Use correct classloader for
 TaskResultGetter

Make sure we are using the context classloader when deserializing failed TaskResults instead of the Spark classloader.

The issue is that `enqueueFailedTask` was using the incorrect classloader which results in `ClassNotFoundException`.

Adds a test in TaskResultGetterSuite that compiles a custom exception, throws it on the executor, and asserts that Spark handles the TaskResult deserialization instead of returning `UnknownReason`.

See #9367 for previous comments
See SPARK-11195 for a full repro

Author: Hurshal Patel <hpatel516@gmail.com>

Closes #9779 from choochootrain/spark-11195-master.
---
 .../scala/org/apache/spark/TestUtils.scala    | 11 ++--
 .../spark/scheduler/TaskResultGetter.scala    |  4 +-
 .../scheduler/TaskResultGetterSuite.scala     | 65 ++++++++++++++++++-
 3 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/TestUtils.scala b/core/src/main/scala/org/apache/spark/TestUtils.scala
index acfe751f6c746..43c89b258f2fa 100644
--- a/core/src/main/scala/org/apache/spark/TestUtils.scala
+++ b/core/src/main/scala/org/apache/spark/TestUtils.scala
@@ -20,6 +20,7 @@ package org.apache.spark
 import java.io.{ByteArrayInputStream, File, FileInputStream, FileOutputStream}
 import java.net.{URI, URL}
 import java.nio.charset.StandardCharsets
+import java.nio.file.Paths
 import java.util.Arrays
 import java.util.jar.{JarEntry, JarOutputStream}
 
@@ -83,15 +84,15 @@ private[spark] object TestUtils {
   }
 
   /**
-   * Create a jar file that contains this set of files. All files will be located at the root
-   * of the jar.
+   * Create a jar file that contains this set of files. All files will be located in the specified
+   * directory or at the root of the jar.
    */
-  def createJar(files: Seq[File], jarFile: File): URL = {
+  def createJar(files: Seq[File], jarFile: File, directoryPrefix: Option[String] = None): URL = {
     val jarFileStream = new FileOutputStream(jarFile)
     val jarStream = new JarOutputStream(jarFileStream, new java.util.jar.Manifest())
 
     for (file <- files) {
-      val jarEntry = new JarEntry(file.getName)
+      val jarEntry = new JarEntry(Paths.get(directoryPrefix.getOrElse(""), file.getName).toString)
       jarStream.putNextEntry(jarEntry)
 
       val in = new FileInputStream(file)
@@ -123,7 +124,7 @@ private[spark] object TestUtils {
       classpathUrls: Seq[URL]): File = {
     val compiler = ToolProvider.getSystemJavaCompiler
 
-    // Calling this outputs a class file in pwd. It's easier to just rename the file than
+    // Calling this outputs a class file in pwd. It's easier to just rename the files than
     // build a custom FileManager that controls the output location.
     val options = if (classpathUrls.nonEmpty) {
       Seq("-classpath", classpathUrls.map { _.getFile }.mkString(File.pathSeparator))
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskResultGetter.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskResultGetter.scala
index 46a6f6537e2ee..f4965994d8277 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskResultGetter.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskResultGetter.scala
@@ -103,16 +103,16 @@ private[spark] class TaskResultGetter(sparkEnv: SparkEnv, scheduler: TaskSchedul
     try {
       getTaskResultExecutor.execute(new Runnable {
         override def run(): Unit = Utils.logUncaughtExceptions {
+          val loader = Utils.getContextOrSparkClassLoader
           try {
             if (serializedData != null && serializedData.limit() > 0) {
               reason = serializer.get().deserialize[TaskEndReason](
-                serializedData, Utils.getSparkClassLoader)
+                serializedData, loader)
             }
           } catch {
             case cnd: ClassNotFoundException =>
               // Log an error but keep going here -- the task failed, so not catastrophic
               // if we can't deserialize the reason.
-              val loader = Utils.getContextOrSparkClassLoader
               logError(
                 "Could not deserialize TaskEndReason: ClassNotFound with classloader " + loader)
             case ex: Exception => {}
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskResultGetterSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskResultGetterSuite.scala
index 815caa79ff529..bc72c3685e8c1 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskResultGetterSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskResultGetterSuite.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.scheduler
 
+import java.io.File
+import java.net.URL
 import java.nio.ByteBuffer
 
 import scala.concurrent.duration._
@@ -26,8 +28,10 @@ import scala.util.control.NonFatal
 import org.scalatest.BeforeAndAfter
 import org.scalatest.concurrent.Eventually._
 
-import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkEnv, SparkFunSuite}
+import org.apache.spark._
 import org.apache.spark.storage.TaskResultBlockId
+import org.apache.spark.TestUtils.JavaSourceFromString
+import org.apache.spark.util.{MutableURLClassLoader, Utils}
 
 /**
  * Removes the TaskResult from the BlockManager before delegating to a normal TaskResultGetter.
@@ -119,5 +123,64 @@ class TaskResultGetterSuite extends SparkFunSuite with BeforeAndAfter with Local
     // Make sure two tasks were run (one failed one, and a second retried one).
     assert(scheduler.nextTaskId.get() === 2)
   }
+
+  /**
+   * Make sure we are using the context classloader when deserializing failed TaskResults instead
+   * of the Spark classloader.
+
+   * This test compiles a jar containing an exception and tests that when it is thrown on the
+   * executor, enqueueFailedTask can correctly deserialize the failure and identify the thrown
+   * exception as the cause.
+
+   * Before this fix, enqueueFailedTask would throw a ClassNotFoundException when deserializing
+   * the exception, resulting in an UnknownReason for the TaskEndResult.
+   */
+  test("failed task deserialized with the correct classloader (SPARK-11195)") {
+    // compile a small jar containing an exception that will be thrown on an executor.
+    val tempDir = Utils.createTempDir()
+    val srcDir = new File(tempDir, "repro/")
+    srcDir.mkdirs()
+    val excSource = new JavaSourceFromString(new File(srcDir, "MyException").getAbsolutePath,
+      """package repro;
+        |
+        |public class MyException extends Exception {
+        |}
+      """.stripMargin)
+    val excFile = TestUtils.createCompiledClass("MyException", srcDir, excSource, Seq.empty)
+    val jarFile = new File(tempDir, "testJar-%s.jar".format(System.currentTimeMillis()))
+    TestUtils.createJar(Seq(excFile), jarFile, directoryPrefix = Some("repro"))
+
+    // ensure we reset the classloader after the test completes
+    val originalClassLoader = Thread.currentThread.getContextClassLoader
+    try {
+      // load the exception from the jar
+      val loader = new MutableURLClassLoader(new Array[URL](0), originalClassLoader)
+      loader.addURL(jarFile.toURI.toURL)
+      Thread.currentThread().setContextClassLoader(loader)
+      val excClass: Class[_] = Utils.classForName("repro.MyException")
+
+      // NOTE: we must run the cluster with "local" so that the executor can load the compiled
+      // jar.
+      sc = new SparkContext("local", "test", conf)
+      val rdd = sc.parallelize(Seq(1), 1).map { _ =>
+        val exc = excClass.newInstance().asInstanceOf[Exception]
+        throw exc
+      }
+
+      // the driver should not have any problems resolving the exception class and determining
+      // why the task failed.
+      val exceptionMessage = intercept[SparkException] {
+        rdd.collect()
+      }.getMessage
+
+      val expectedFailure = """(?s).*Lost task.*: repro.MyException.*""".r
+      val unknownFailure = """(?s).*Lost task.*: UnknownReason.*""".r
+
+      assert(expectedFailure.findFirstMatchIn(exceptionMessage).isDefined)
+      assert(unknownFailure.findFirstMatchIn(exceptionMessage).isEmpty)
+    } finally {
+      Thread.currentThread.setContextClassLoader(originalClassLoader)
+    }
+  }
 }
 

From cffb899c4397ecccedbcc41e7cf3da91f953435a Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 18 Nov 2015 10:15:50 -0800
Subject: [PATCH 325/518] [SPARK-11803][SQL] fix Dataset self-join

When we resolve the join operator, we may change the output of right side if self-join is detected. So in `Dataset.joinWith`, we should resolve the join operator first, and then get the left output and right output from it, instead of using `left.output` and `right.output` directly.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9806 from cloud-fan/self-join.
---
 .../main/scala/org/apache/spark/sql/Dataset.scala  | 14 +++++++++-----
 .../scala/org/apache/spark/sql/DatasetSuite.scala  |  8 ++++----
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 817c20fdbb9f3..b644f6ad3096d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -498,13 +498,17 @@ class Dataset[T] private[sql](
     val left = this.logicalPlan
     val right = other.logicalPlan
 
+    val joined = sqlContext.executePlan(Join(left, right, Inner, Some(condition.expr)))
+    val leftOutput = joined.analyzed.output.take(left.output.length)
+    val rightOutput = joined.analyzed.output.takeRight(right.output.length)
+
     val leftData = this.unresolvedTEncoder match {
-      case e if e.flat => Alias(left.output.head, "_1")()
-      case _ => Alias(CreateStruct(left.output), "_1")()
+      case e if e.flat => Alias(leftOutput.head, "_1")()
+      case _ => Alias(CreateStruct(leftOutput), "_1")()
     }
     val rightData = other.unresolvedTEncoder match {
-      case e if e.flat => Alias(right.output.head, "_2")()
-      case _ => Alias(CreateStruct(right.output), "_2")()
+      case e if e.flat => Alias(rightOutput.head, "_2")()
+      case _ => Alias(CreateStruct(rightOutput), "_2")()
     }
 
 
@@ -513,7 +517,7 @@ class Dataset[T] private[sql](
     withPlan[(T, U)](other) { (left, right) =>
       Project(
         leftData :: rightData :: Nil,
-        Join(left, right, Inner, Some(condition.expr)))
+        joined.analyzed)
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index a522894c374f9..198962b8fb750 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -347,7 +347,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     checkAnswer(joined, ("2", 2))
   }
 
-  ignore("self join") {
+  test("self join") {
     val ds = Seq("1", "2").toDS().as("a")
     val joined = ds.joinWith(ds, lit(true))
     checkAnswer(joined, ("1", "1"), ("1", "2"), ("2", "1"), ("2", "2"))
@@ -360,15 +360,15 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
 
   test("kryo encoder") {
     implicit val kryoEncoder = Encoders.kryo[KryoData]
-    val ds = sqlContext.createDataset(Seq(KryoData(1), KryoData(2)))
+    val ds = Seq(KryoData(1), KryoData(2)).toDS()
 
     assert(ds.groupBy(p => p).count().collect().toSeq ==
       Seq((KryoData(1), 1L), (KryoData(2), 1L)))
   }
 
-  ignore("kryo encoder self join") {
+  test("kryo encoder self join") {
     implicit val kryoEncoder = Encoders.kryo[KryoData]
-    val ds = sqlContext.createDataset(Seq(KryoData(1), KryoData(2)))
+    val ds = Seq(KryoData(1), KryoData(2)).toDS()
     assert(ds.joinWith(ds, lit(true)).collect().toSet ==
       Set(
         (KryoData(1), KryoData(1)),

From 33b837333435ceb0c04d1f361a5383c4fe6a5a75 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 18 Nov 2015 10:23:12 -0800
Subject: [PATCH 326/518] [SPARK-11725][SQL] correctly handle null inputs for
 UDF

If user use primitive parameters in UDF, there is no way for him to do the null-check for primitive inputs, so we are assuming the primitive input is null-propagatable for this case and return null if the input is null.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9770 from cloud-fan/udf.
---
 .../spark/sql/catalyst/ScalaReflection.scala  |  9 ++++
 .../sql/catalyst/analysis/Analyzer.scala      | 32 +++++++++++++-
 .../sql/catalyst/expressions/ScalaUDF.scala   |  6 +++
 .../sql/catalyst/ScalaReflectionSuite.scala   | 17 +++++++
 .../sql/catalyst/analysis/AnalysisSuite.scala | 44 +++++++++++++++++++
 .../org/apache/spark/sql/DataFrameSuite.scala | 14 ++++++
 6 files changed, 121 insertions(+), 1 deletion(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 0b3dd351e38e8..38828e59a2152 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -719,6 +719,15 @@ trait ScalaReflection {
     }
   }
 
+  /**
+   * Returns classes of input parameters of scala function object.
+   */
+  def getParameterTypes(func: AnyRef): Seq[Class[_]] = {
+    val methods = func.getClass.getMethods.filter(m => m.getName == "apply" && !m.isBridge)
+    assert(methods.length == 1)
+    methods.head.getParameterTypes
+  }
+
   def typeOfObject: PartialFunction[Any, DataType] = {
     // The data type can be determined without ambiguity.
     case obj: Boolean => BooleanType
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index 2f4670b55bdba..f00c451b5981a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
 import org.apache.spark.sql.catalyst.trees.TreeNodeRef
-import org.apache.spark.sql.catalyst.{SimpleCatalystConf, CatalystConf}
+import org.apache.spark.sql.catalyst.{ScalaReflection, SimpleCatalystConf, CatalystConf}
 import org.apache.spark.sql.types._
 
 /**
@@ -85,6 +85,8 @@ class Analyzer(
       extendedResolutionRules : _*),
     Batch("Nondeterministic", Once,
       PullOutNondeterministic),
+    Batch("UDF", Once,
+      HandleNullInputsForUDF),
     Batch("Cleanup", fixedPoint,
       CleanupAliases)
   )
@@ -1063,6 +1065,34 @@ class Analyzer(
         Project(p.output, newPlan.withNewChildren(newChild :: Nil))
     }
   }
+
+  /**
+   * Correctly handle null primitive inputs for UDF by adding extra [[If]] expression to do the
+   * null check.  When user defines a UDF with primitive parameters, there is no way to tell if the
+   * primitive parameter is null or not, so here we assume the primitive input is null-propagatable
+   * and we should return null if the input is null.
+   */
+  object HandleNullInputsForUDF extends Rule[LogicalPlan] {
+    override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
+      case p if !p.resolved => p // Skip unresolved nodes.
+
+      case plan => plan transformExpressionsUp {
+
+        case udf @ ScalaUDF(func, _, inputs, _) =>
+          val parameterTypes = ScalaReflection.getParameterTypes(func)
+          assert(parameterTypes.length == inputs.length)
+
+          val inputsNullCheck = parameterTypes.zip(inputs)
+            // TODO: skip null handling for not-nullable primitive inputs after we can completely
+            // trust the `nullable` information.
+            // .filter { case (cls, expr) => cls.isPrimitive && expr.nullable }
+            .filter { case (cls, _) => cls.isPrimitive }
+            .map { case (_, expr) => IsNull(expr) }
+            .reduceLeftOption[Expression]((e1, e2) => Or(e1, e2))
+          inputsNullCheck.map(If(_, Literal.create(null, udf.dataType), udf)).getOrElse(udf)
+      }
+    }
+  }
 }
 
 /**
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
index 3388cc20a9803..03b89221ef2d3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala
@@ -24,7 +24,13 @@ import org.apache.spark.sql.types.DataType
 
 /**
  * User-defined function.
+ * @param function  The user defined scala function to run.
+ *                  Note that if you use primitive parameters, you are not able to check if it is
+ *                  null or not, and the UDF will return null for you if the primitive input is
+ *                  null. Use boxed type or [[Option]] if you wanna do the null-handling yourself.
  * @param dataType  Return type of function.
+ * @param children  The input expressions of this UDF.
+ * @param inputTypes  The expected input types of this UDF.
  */
 case class ScalaUDF(
     function: AnyRef,
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ScalaReflectionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ScalaReflectionSuite.scala
index 3b848cfdf737f..4ea410d492b01 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ScalaReflectionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ScalaReflectionSuite.scala
@@ -280,4 +280,21 @@ class ScalaReflectionSuite extends SparkFunSuite {
         assert(s.fields.map(_.dataType) === Seq(IntegerType, StringType, DoubleType))
     }
   }
+
+  test("get parameter type from a function object") {
+    val primitiveFunc = (i: Int, j: Long) => "x"
+    val primitiveTypes = getParameterTypes(primitiveFunc)
+    assert(primitiveTypes.forall(_.isPrimitive))
+    assert(primitiveTypes === Seq(classOf[Int], classOf[Long]))
+
+    val boxedFunc = (i: java.lang.Integer, j: java.lang.Long) => "x"
+    val boxedTypes = getParameterTypes(boxedFunc)
+    assert(boxedTypes.forall(!_.isPrimitive))
+    assert(boxedTypes === Seq(classOf[java.lang.Integer], classOf[java.lang.Long]))
+
+    val anyFunc = (i: Any, j: AnyRef) => "x"
+    val anyTypes = getParameterTypes(anyFunc)
+    assert(anyTypes.forall(!_.isPrimitive))
+    assert(anyTypes === Seq(classOf[java.lang.Object], classOf[java.lang.Object]))
+  }
 }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
index 65f09b46afae1..08586a97411ae 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
@@ -174,4 +174,48 @@ class AnalysisSuite extends AnalysisTest {
     )
     assertAnalysisError(plan, Seq("data type mismatch: Arguments must be same type"))
   }
+
+  test("SPARK-11725: correctly handle null inputs for ScalaUDF") {
+    val string = testRelation2.output(0)
+    val double = testRelation2.output(2)
+    val short = testRelation2.output(4)
+    val nullResult = Literal.create(null, StringType)
+
+    def checkUDF(udf: Expression, transformed: Expression): Unit = {
+      checkAnalysis(
+        Project(Alias(udf, "")() :: Nil, testRelation2),
+        Project(Alias(transformed, "")() :: Nil, testRelation2)
+      )
+    }
+
+    // non-primitive parameters do not need special null handling
+    val udf1 = ScalaUDF((s: String) => "x", StringType, string :: Nil)
+    val expected1 = udf1
+    checkUDF(udf1, expected1)
+
+    // only primitive parameter needs special null handling
+    val udf2 = ScalaUDF((s: String, d: Double) => "x", StringType, string :: double :: Nil)
+    val expected2 = If(IsNull(double), nullResult, udf2)
+    checkUDF(udf2, expected2)
+
+    // special null handling should apply to all primitive parameters
+    val udf3 = ScalaUDF((s: Short, d: Double) => "x", StringType, short :: double :: Nil)
+    val expected3 = If(
+      IsNull(short) || IsNull(double),
+      nullResult,
+      udf3)
+    checkUDF(udf3, expected3)
+
+    // we can skip special null handling for primitive parameters that are not nullable
+    // TODO: this is disabled for now as we can not completely trust `nullable`.
+    val udf4 = ScalaUDF(
+      (s: Short, d: Double) => "x",
+      StringType,
+      short :: double.withNullability(false) :: Nil)
+    val expected4 = If(
+      IsNull(short),
+      nullResult,
+      udf4)
+    // checkUDF(udf4, expected4)
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index 35cdab50bdec9..5a7f24684d1b7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -1115,4 +1115,18 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
     checkAnswer(df.select(df("*")), Row(1, "a"))
     checkAnswer(df.withColumnRenamed("d^'a.", "a"), Row(1, "a"))
   }
+
+  test("SPARK-11725: correctly handle null inputs for ScalaUDF") {
+    val df = Seq(
+      new java.lang.Integer(22) -> "John",
+      null.asInstanceOf[java.lang.Integer] -> "Lucy").toDF("age", "name")
+
+    val boxedUDF = udf[java.lang.Integer, java.lang.Integer] {
+      (i: java.lang.Integer) => if (i == null) null else i * 2
+    }
+    checkAnswer(df.select(boxedUDF($"age")), Row(44) :: Row(null) :: Nil)
+
+    val primitiveUDF = udf((i: Int) => i * 2)
+    checkAnswer(df.select(primitiveUDF($"age")), Row(44) :: Row(null) :: Nil)
+  }
 }

From dbf428c87ab34b6f76c75946043bdf5f60c9b1b3 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Wed, 18 Nov 2015 10:33:17 -0800
Subject: [PATCH 327/518] [SPARK-11795][SQL] combine grouping attributes into a
 single NamedExpression
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

we use `ExpressionEncoder.tuple` to build the result encoder, which assumes the input encoder should point to a struct type field if it’s non-flat.
However, our keyEncoder always point to a flat field/fields: `groupingAttributes`, we should combine them into a single `NamedExpression`.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9792 from cloud-fan/agg.
---
 .../main/scala/org/apache/spark/sql/GroupedDataset.scala | 9 +++++++--
 .../test/scala/org/apache/spark/sql/DatasetSuite.scala   | 5 ++---
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index c66162ee2148a..3f84e22a1025b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -22,7 +22,7 @@ import scala.collection.JavaConverters._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.function._
 import org.apache.spark.sql.catalyst.encoders.{FlatEncoder, ExpressionEncoder, encoderFor}
-import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.{Alias, CreateStruct, Attribute}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.QueryExecution
 
@@ -187,7 +187,12 @@ class GroupedDataset[K, T] private[sql](
     val namedColumns =
       columns.map(
         _.withInputType(resolvedTEncoder, dataAttributes).named)
-    val aggregate = Aggregate(groupingAttributes, groupingAttributes ++ namedColumns, logicalPlan)
+    val keyColumn = if (groupingAttributes.length > 1) {
+      Alias(CreateStruct(groupingAttributes), "key")()
+    } else {
+      groupingAttributes.head
+    }
+    val aggregate = Aggregate(groupingAttributes, keyColumn +: namedColumns, logicalPlan)
     val execution = new QueryExecution(sqlContext, aggregate)
 
     new Dataset(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 198962b8fb750..b6db583dfe01f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -84,8 +84,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       ("a", 2), ("b", 3), ("c", 4))
   }
 
-  ignore("Dataset should set the resolved encoders internally for maps") {
-    // TODO: Enable this once we fix SPARK-11793.
+  test("map and group by with class data") {
     // We inject a group by here to make sure this test case is future proof
     // when we implement better pipelining and local execution mode.
     val ds: Dataset[(ClassData, Long)] = Seq(ClassData("one", 1), ClassData("two", 2)).toDS()
@@ -94,7 +93,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
 
     checkAnswer(
       ds,
-      (ClassData("one", 1), 1L), (ClassData("two", 2), 1L))
+      (ClassData("one", 2), 1L), (ClassData("two", 3), 1L))
   }
 
   test("select") {

From 90a7519daaa7f4ee3be7c5a9aa244120811ff6eb Mon Sep 17 00:00:00 2001
From: Jakob Odersky <jodersky@gmail.com>
Date: Wed, 18 Nov 2015 11:35:41 -0800
Subject: [PATCH 328/518] [MINOR][BUILD] Ignore ensime cache

Using ENSIME, I often have `.ensime_cache` polluting my source tree. This PR simply adds the cache directory to `.gitignore`

Author: Jakob Odersky <jodersky@gmail.com>

Closes #9708 from jodersky/master.
---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index 08f2d8f7543f0..07524bc429e92 100644
--- a/.gitignore
+++ b/.gitignore
@@ -50,6 +50,7 @@ spark-tests.log
 streaming-tests.log
 dependency-reduced-pom.xml
 .ensime
+.ensime_cache/
 .ensime_lucene
 checkpoint
 derby.log

From 6f99522d13d8db9fcc767f7c3189557b9a53d283 Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Wed, 18 Nov 2015 11:49:12 -0800
Subject: [PATCH 329/518] [SPARK-11792] [SQL] [FOLLOW-UP] Change SizeEstimation
 to KnownSizeEstimation and make estimatedSize return Long instead of
 Option[Long]

https://issues.apache.org/jira/browse/SPARK-11792

The main changes include:
* Renaming `SizeEstimation` to `KnownSizeEstimation`. Hopefully this new name has more information.
* Making `estimatedSize` return `Long` instead of `Option[Long]`.
* In `UnsaveHashedRelation`, `estimatedSize` will delegate the work to `SizeEstimator` if we have not created a `BytesToBytesMap`.

Since we will put `UnsaveHashedRelation` to `BlockManager`, it is generally good to let it provide a more accurate size estimation. Also, if we do not put `BytesToBytesMap` directly into `BlockerManager`, I feel it is not really necessary to make `BytesToBytesMap` extends `KnownSizeEstimation`.

Author: Yin Huai <yhuai@databricks.com>

Closes #9813 from yhuai/SPARK-11792-followup.
---
 .../org/apache/spark/util/SizeEstimator.scala | 30 ++++++++++---------
 .../spark/util/SizeEstimatorSuite.scala       | 14 ++-------
 .../sql/execution/joins/HashedRelation.scala  | 12 +++++---
 3 files changed, 26 insertions(+), 30 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/util/SizeEstimator.scala b/core/src/main/scala/org/apache/spark/util/SizeEstimator.scala
index c3a2675ee5f45..09864e3f8392d 100644
--- a/core/src/main/scala/org/apache/spark/util/SizeEstimator.scala
+++ b/core/src/main/scala/org/apache/spark/util/SizeEstimator.scala
@@ -36,9 +36,14 @@ import org.apache.spark.util.collection.OpenHashSet
  * When a class extends it, [[SizeEstimator]] will query the `estimatedSize` first.
  * If `estimatedSize` does not return [[None]], [[SizeEstimator]] will use the returned size
  * as the size of the object. Otherwise, [[SizeEstimator]] will do the estimation work.
+ * The difference between a [[KnownSizeEstimation]] and
+ * [[org.apache.spark.util.collection.SizeTracker]] is that, a
+ * [[org.apache.spark.util.collection.SizeTracker]] still uses [[SizeEstimator]] to
+ * estimate the size. However, a [[KnownSizeEstimation]] can provide a better estimation without
+ * using [[SizeEstimator]].
  */
-private[spark] trait SizeEstimation {
-  def estimatedSize: Option[Long]
+private[spark] trait KnownSizeEstimation {
+  def estimatedSize: Long
 }
 
 /**
@@ -209,18 +214,15 @@ object SizeEstimator extends Logging {
       // the size estimator since it references the whole REPL. Do nothing in this case. In
       // general all ClassLoaders and Classes will be shared between objects anyway.
     } else {
-      val estimatedSize = obj match {
-        case s: SizeEstimation => s.estimatedSize
-        case _ => None
-      }
-      if (estimatedSize.isDefined) {
-        state.size += estimatedSize.get
-      } else {
-        val classInfo = getClassInfo(cls)
-        state.size += alignSize(classInfo.shellSize)
-        for (field <- classInfo.pointerFields) {
-          state.enqueue(field.get(obj))
-        }
+      obj match {
+        case s: KnownSizeEstimation =>
+          state.size += s.estimatedSize
+        case _ =>
+          val classInfo = getClassInfo(cls)
+          state.size += alignSize(classInfo.shellSize)
+          for (field <- classInfo.pointerFields) {
+            state.enqueue(field.get(obj))
+          }
       }
     }
   }
diff --git a/core/src/test/scala/org/apache/spark/util/SizeEstimatorSuite.scala b/core/src/test/scala/org/apache/spark/util/SizeEstimatorSuite.scala
index 9b6261af123e6..101610e38014e 100644
--- a/core/src/test/scala/org/apache/spark/util/SizeEstimatorSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/SizeEstimatorSuite.scala
@@ -60,16 +60,10 @@ class DummyString(val arr: Array[Char]) {
   @transient val hash32: Int = 0
 }
 
-class DummyClass8 extends SizeEstimation {
+class DummyClass8 extends KnownSizeEstimation {
   val x: Int = 0
 
-  override def estimatedSize: Option[Long] = Some(2015)
-}
-
-class DummyClass9 extends SizeEstimation {
-  val x: Int = 0
-
-  override def estimatedSize: Option[Long] = None
+  override def estimatedSize: Long = 2015
 }
 
 class SizeEstimatorSuite
@@ -231,9 +225,5 @@ class SizeEstimatorSuite
     // DummyClass8 provides its size estimation.
     assertResult(2015)(SizeEstimator.estimate(new DummyClass8))
     assertResult(20206)(SizeEstimator.estimate(Array.fill(10)(new DummyClass8)))
-
-    // DummyClass9 does not provide its size estimation.
-    assertResult(16)(SizeEstimator.estimate(new DummyClass9))
-    assertResult(216)(SizeEstimator.estimate(Array.fill(10)(new DummyClass9)))
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala
index 49ae09bf53782..aebfea5832402 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelation.scala
@@ -30,7 +30,7 @@ import org.apache.spark.sql.execution.metric.{LongSQLMetric, SQLMetrics}
 import org.apache.spark.unsafe.Platform
 import org.apache.spark.unsafe.map.BytesToBytesMap
 import org.apache.spark.unsafe.memory.MemoryLocation
-import org.apache.spark.util.{SizeEstimation, Utils}
+import org.apache.spark.util.{SizeEstimator, KnownSizeEstimation, Utils}
 import org.apache.spark.util.collection.CompactBuffer
 import org.apache.spark.{SparkConf, SparkEnv}
 
@@ -190,7 +190,7 @@ private[execution] object HashedRelation {
 private[joins] final class UnsafeHashedRelation(
     private var hashTable: JavaHashMap[UnsafeRow, CompactBuffer[UnsafeRow]])
   extends HashedRelation
-  with SizeEstimation
+  with KnownSizeEstimation
   with Externalizable {
 
   private[joins] def this() = this(null)  // Needed for serialization
@@ -217,8 +217,12 @@ private[joins] final class UnsafeHashedRelation(
     }
   }
 
-  override def estimatedSize: Option[Long] = {
-    Option(binaryMap).map(_.getTotalMemoryConsumption)
+  override def estimatedSize: Long = {
+    if (binaryMap != null) {
+      binaryMap.getTotalMemoryConsumption
+    } else {
+      SizeEstimator.estimate(hashTable)
+    }
   }
 
   override def get(key: InternalRow): Seq[InternalRow] = {

From 94624eacb0fdbbe210894151a956f8150cdf527e Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Wed, 18 Nov 2015 11:53:28 -0800
Subject: [PATCH 330/518] [SPARK-11739][SQL] clear the instantiated SQLContext

Currently, if the first SQLContext is not removed after stopping SparkContext, a SQLContext could set there forever. This patch make this more robust.

Author: Davies Liu <davies@databricks.com>

Closes #9706 from davies/clear_context.
---
 .../scala/org/apache/spark/sql/SQLContext.scala | 17 +++++++++++------
 .../spark/sql/MultiSQLContextsSuite.scala       |  5 ++---
 .../execution/ExchangeCoordinatorSuite.scala    |  2 +-
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index cd1fdc4edb39d..39471d2fb79a7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -1229,7 +1229,7 @@ class SQLContext private[sql](
   // construction of the instance.
   sparkContext.addSparkListener(new SparkListener {
     override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
-      SQLContext.clearInstantiatedContext(self)
+      SQLContext.clearInstantiatedContext()
     }
   })
 
@@ -1270,13 +1270,13 @@ object SQLContext {
    */
   def getOrCreate(sparkContext: SparkContext): SQLContext = {
     val ctx = activeContext.get()
-    if (ctx != null) {
+    if (ctx != null && !ctx.sparkContext.isStopped) {
       return ctx
     }
 
     synchronized {
       val ctx = instantiatedContext.get()
-      if (ctx == null) {
+      if (ctx == null || ctx.sparkContext.isStopped) {
         new SQLContext(sparkContext)
       } else {
         ctx
@@ -1284,12 +1284,17 @@ object SQLContext {
     }
   }
 
-  private[sql] def clearInstantiatedContext(sqlContext: SQLContext): Unit = {
-    instantiatedContext.compareAndSet(sqlContext, null)
+  private[sql] def clearInstantiatedContext(): Unit = {
+    instantiatedContext.set(null)
   }
 
   private[sql] def setInstantiatedContext(sqlContext: SQLContext): Unit = {
-    instantiatedContext.compareAndSet(null, sqlContext)
+    synchronized {
+      val ctx = instantiatedContext.get()
+      if (ctx == null || ctx.sparkContext.isStopped) {
+        instantiatedContext.set(sqlContext)
+      }
+    }
   }
 
   private[sql] def getInstantiatedContextOption(): Option[SQLContext] = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/MultiSQLContextsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/MultiSQLContextsSuite.scala
index 0e8fcb6a858b1..34c5c68fd1c18 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/MultiSQLContextsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/MultiSQLContextsSuite.scala
@@ -31,7 +31,7 @@ class MultiSQLContextsSuite extends SparkFunSuite with BeforeAndAfterAll {
     originalInstantiatedSQLContext = SQLContext.getInstantiatedContextOption()
 
     SQLContext.clearActive()
-    originalInstantiatedSQLContext.foreach(ctx => SQLContext.clearInstantiatedContext(ctx))
+    SQLContext.clearInstantiatedContext()
     sparkConf =
       new SparkConf(false)
         .setMaster("local[*]")
@@ -89,10 +89,9 @@ class MultiSQLContextsSuite extends SparkFunSuite with BeforeAndAfterAll {
         testNewSession(rootSQLContext)
         testNewSession(rootSQLContext)
         testCreatingNewSQLContext(allowMultipleSQLContexts)
-
-        SQLContext.clearInstantiatedContext(rootSQLContext)
       } finally {
         sc.stop()
+        SQLContext.clearInstantiatedContext()
       }
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala
index 25f2f5caeed15..b96d50a70b85c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala
@@ -34,7 +34,7 @@ class ExchangeCoordinatorSuite extends SparkFunSuite with BeforeAndAfterAll {
     originalInstantiatedSQLContext = SQLContext.getInstantiatedContextOption()
 
     SQLContext.clearActive()
-    originalInstantiatedSQLContext.foreach(ctx => SQLContext.clearInstantiatedContext(ctx))
+    SQLContext.clearInstantiatedContext()
   }
 
   override protected def afterAll(): Unit = {

From 31921e0f0bd559d042148d1ea32f865fb3068f38 Mon Sep 17 00:00:00 2001
From: Bryan Cutler <bjcutler@us.ibm.com>
Date: Wed, 18 Nov 2015 12:09:54 -0800
Subject: [PATCH 331/518] [SPARK-4557][STREAMING] Spark Streaming foreachRDD
 Java API method should accept a VoidFunction<...>

Currently streaming foreachRDD Java API uses a function prototype requiring a return value of null.  This PR deprecates the old method and uses VoidFunction to allow for more concise declaration.  Also added VoidFunction2 to Java API in order to use in Streaming methods.  Unit test is added for using foreachRDD with VoidFunction, and changes have been tested with Java 7 and Java 8 using lambdas.

Author: Bryan Cutler <bjcutler@us.ibm.com>

Closes #9488 from BryanCutler/foreachRDD-VoidFunction-SPARK-4557.
---
 .../api/java/function/VoidFunction2.java      | 27 ++++++++++++
 .../apache/spark/streaming/Java8APISuite.java | 26 ++++++++++++
 project/MimaExcludes.scala                    |  4 ++
 .../streaming/api/java/JavaDStreamLike.scala  | 24 ++++++++++-
 .../apache/spark/streaming/JavaAPISuite.java  | 41 ++++++++++++++++++-
 5 files changed, 120 insertions(+), 2 deletions(-)
 create mode 100644 core/src/main/java/org/apache/spark/api/java/function/VoidFunction2.java

diff --git a/core/src/main/java/org/apache/spark/api/java/function/VoidFunction2.java b/core/src/main/java/org/apache/spark/api/java/function/VoidFunction2.java
new file mode 100644
index 0000000000000..6c576ab678455
--- /dev/null
+++ b/core/src/main/java/org/apache/spark/api/java/function/VoidFunction2.java
@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+
+/**
+ * A two-argument function that takes arguments of type T1 and T2 with no return value.
+ */
+public interface VoidFunction2<T1, T2> extends Serializable {
+  public void call(T1 v1, T2 v2) throws Exception;
+}
diff --git a/extras/java8-tests/src/test/java/org/apache/spark/streaming/Java8APISuite.java b/extras/java8-tests/src/test/java/org/apache/spark/streaming/Java8APISuite.java
index 163ae92c12c6d..4eee97bc89613 100644
--- a/extras/java8-tests/src/test/java/org/apache/spark/streaming/Java8APISuite.java
+++ b/extras/java8-tests/src/test/java/org/apache/spark/streaming/Java8APISuite.java
@@ -28,6 +28,7 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import org.apache.spark.Accumulator;
 import org.apache.spark.HashPartitioner;
 import org.apache.spark.api.java.JavaPairRDD;
 import org.apache.spark.api.java.JavaRDD;
@@ -360,6 +361,31 @@ public void testFlatMap() {
     assertOrderInvariantEquals(expected, result);
   }
 
+  @Test
+  public void testForeachRDD() {
+    final Accumulator<Integer> accumRdd = ssc.sc().accumulator(0);
+    final Accumulator<Integer> accumEle = ssc.sc().accumulator(0);
+    List<List<Integer>> inputData = Arrays.asList(
+        Arrays.asList(1,1,1),
+        Arrays.asList(1,1,1));
+
+    JavaDStream<Integer> stream = JavaTestUtils.attachTestInputStream(ssc, inputData, 1);
+    JavaTestUtils.attachTestOutputStream(stream.count()); // dummy output
+
+    stream.foreachRDD(rdd -> {
+      accumRdd.add(1);
+      rdd.foreach(x -> accumEle.add(1));
+    });
+
+    // This is a test to make sure foreachRDD(VoidFunction2) can be called from Java
+    stream.foreachRDD((rdd, time) -> null);
+
+    JavaTestUtils.runStreams(ssc, 2, 2);
+
+    Assert.assertEquals(2, accumRdd.value().intValue());
+    Assert.assertEquals(6, accumEle.value().intValue());
+  }
+
   @Test
   public void testPairFlatMap() {
     List<List<String>> inputData = Arrays.asList(
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index eb70d27c34c20..bb45d1bb12146 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -142,6 +142,10 @@ object MimaExcludes {
           "org.apache.spark.streaming.kafka.KafkaUtilsPythonHelper.createDirectStream"),
         ProblemFilters.exclude[MissingMethodProblem](
           "org.apache.spark.streaming.kafka.KafkaUtilsPythonHelper.createRDD")
+      ) ++ Seq(
+        // SPARK-4557 Changed foreachRDD to use VoidFunction
+        ProblemFilters.exclude[MissingMethodProblem](
+          "org.apache.spark.streaming.api.java.JavaDStreamLike.foreachRDD")
       )
     case v if v.startsWith("1.5") =>
       Seq(
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaDStreamLike.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaDStreamLike.scala
index edfa474677f15..84acec7d8e330 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaDStreamLike.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/java/JavaDStreamLike.scala
@@ -27,7 +27,7 @@ import scala.reflect.ClassTag
 import org.apache.spark.api.java.{JavaPairRDD, JavaRDD, JavaRDDLike}
 import org.apache.spark.api.java.JavaPairRDD._
 import org.apache.spark.api.java.JavaSparkContext.fakeClassTag
-import org.apache.spark.api.java.function.{Function => JFunction, Function2 => JFunction2, Function3 => JFunction3, _}
+import org.apache.spark.api.java.function.{Function => JFunction, Function2 => JFunction2, Function3 => JFunction3, VoidFunction => JVoidFunction, VoidFunction2 => JVoidFunction2, _}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.streaming._
 import org.apache.spark.streaming.api.java.JavaDStream._
@@ -308,7 +308,10 @@ trait JavaDStreamLike[T, This <: JavaDStreamLike[T, This, R], R <: JavaRDDLike[T
   /**
    * Apply a function to each RDD in this DStream. This is an output operator, so
    * 'this' DStream will be registered as an output stream and therefore materialized.
+   *
+   * @deprecated  As of release 1.6.0, replaced by foreachRDD(JVoidFunction)
    */
+  @deprecated("Use foreachRDD(foreachFunc: JVoidFunction[R])", "1.6.0")
   def foreachRDD(foreachFunc: JFunction[R, Void]) {
     dstream.foreachRDD(rdd => foreachFunc.call(wrapRDD(rdd)))
   }
@@ -316,11 +319,30 @@ trait JavaDStreamLike[T, This <: JavaDStreamLike[T, This, R], R <: JavaRDDLike[T
   /**
    * Apply a function to each RDD in this DStream. This is an output operator, so
    * 'this' DStream will be registered as an output stream and therefore materialized.
+   *
+   * @deprecated  As of release 1.6.0, replaced by foreachRDD(JVoidFunction2)
    */
+  @deprecated("Use foreachRDD(foreachFunc: JVoidFunction2[R, Time])", "1.6.0")
   def foreachRDD(foreachFunc: JFunction2[R, Time, Void]) {
     dstream.foreachRDD((rdd, time) => foreachFunc.call(wrapRDD(rdd), time))
   }
 
+  /**
+   * Apply a function to each RDD in this DStream. This is an output operator, so
+   * 'this' DStream will be registered as an output stream and therefore materialized.
+   */
+  def foreachRDD(foreachFunc: JVoidFunction[R]) {
+    dstream.foreachRDD(rdd => foreachFunc.call(wrapRDD(rdd)))
+  }
+
+  /**
+   * Apply a function to each RDD in this DStream. This is an output operator, so
+   * 'this' DStream will be registered as an output stream and therefore materialized.
+   */
+  def foreachRDD(foreachFunc: JVoidFunction2[R, Time]) {
+    dstream.foreachRDD((rdd, time) => foreachFunc.call(wrapRDD(rdd), time))
+  }
+
   /**
    * Return a new DStream in which each RDD is generated by applying a function
    * on each RDD of 'this' DStream.
diff --git a/streaming/src/test/java/org/apache/spark/streaming/JavaAPISuite.java b/streaming/src/test/java/org/apache/spark/streaming/JavaAPISuite.java
index c5217149224e4..609bb4413b6b1 100644
--- a/streaming/src/test/java/org/apache/spark/streaming/JavaAPISuite.java
+++ b/streaming/src/test/java/org/apache/spark/streaming/JavaAPISuite.java
@@ -37,7 +37,9 @@
 import com.google.common.io.Files;
 import com.google.common.collect.Sets;
 
+import org.apache.spark.Accumulator;
 import org.apache.spark.HashPartitioner;
+import org.apache.spark.SparkConf;
 import org.apache.spark.api.java.JavaPairRDD;
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.api.java.JavaSparkContext;
@@ -45,7 +47,6 @@
 import org.apache.spark.storage.StorageLevel;
 import org.apache.spark.streaming.api.java.*;
 import org.apache.spark.util.Utils;
-import org.apache.spark.SparkConf;
 
 // The test suite itself is Serializable so that anonymous Function implementations can be
 // serialized, as an alternative to converting these anonymous classes to static inner classes;
@@ -768,6 +769,44 @@ public Iterable<String> call(String x) {
     assertOrderInvariantEquals(expected, result);
   }
 
+  @SuppressWarnings("unchecked")
+  @Test
+  public void testForeachRDD() {
+    final Accumulator<Integer> accumRdd = ssc.sc().accumulator(0);
+    final Accumulator<Integer> accumEle = ssc.sc().accumulator(0);
+    List<List<Integer>> inputData = Arrays.asList(
+        Arrays.asList(1,1,1),
+        Arrays.asList(1,1,1));
+
+    JavaDStream<Integer> stream = JavaTestUtils.attachTestInputStream(ssc, inputData, 1);
+    JavaTestUtils.attachTestOutputStream(stream.count()); // dummy output
+
+    stream.foreachRDD(new VoidFunction<JavaRDD<Integer>>() {
+      @Override
+      public void call(JavaRDD<Integer> rdd) {
+        accumRdd.add(1);
+        rdd.foreach(new VoidFunction<Integer>() {
+          @Override
+          public void call(Integer i) {
+            accumEle.add(1);
+          }
+        });
+      }
+    });
+
+    // This is a test to make sure foreachRDD(VoidFunction2) can be called from Java
+    stream.foreachRDD(new VoidFunction2<JavaRDD<Integer>, Time>() {
+      @Override
+      public void call(JavaRDD<Integer> rdd, Time time) {
+      }
+    });
+
+    JavaTestUtils.runStreams(ssc, 2, 2);
+
+    Assert.assertEquals(2, accumRdd.value().intValue());
+    Assert.assertEquals(6, accumEle.value().intValue());
+  }
+
   @SuppressWarnings("unchecked")
   @Test
   public void testPairFlatMap() {

From a416e41e285700f861559d710dbf857405bfddf6 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 18 Nov 2015 12:50:29 -0800
Subject: [PATCH 332/518] [SPARK-11809] Switch the default Mesos mode to
 coarse-grained mode

Based on my conversions with people, I believe the consensus is that the coarse-grained mode is more stable and easier to reason about. It is best to use that as the default rather than the more flaky fine-grained mode.

Author: Reynold Xin <rxin@databricks.com>

Closes #9795 from rxin/SPARK-11809.
---
 .../scala/org/apache/spark/SparkContext.scala |  2 +-
 docs/job-scheduling.md                        |  2 +-
 docs/running-on-mesos.md                      | 27 ++++++++++++-------
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index b5645b08f92d4..ab374cb71286a 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -2710,7 +2710,7 @@ object SparkContext extends Logging {
       case mesosUrl @ MESOS_REGEX(_) =>
         MesosNativeLibrary.load()
         val scheduler = new TaskSchedulerImpl(sc)
-        val coarseGrained = sc.conf.getBoolean("spark.mesos.coarse", false)
+        val coarseGrained = sc.conf.getBoolean("spark.mesos.coarse", defaultValue = true)
         val url = mesosUrl.stripPrefix("mesos://") // strip scheme from raw Mesos URLs
         val backend = if (coarseGrained) {
           new CoarseMesosSchedulerBackend(scheduler, sc, url, sc.env.securityManager)
diff --git a/docs/job-scheduling.md b/docs/job-scheduling.md
index a3c34cb6796fa..36327c6efeaf3 100644
--- a/docs/job-scheduling.md
+++ b/docs/job-scheduling.md
@@ -47,7 +47,7 @@ application is not running tasks on a machine, other applications may run tasks
 is useful when you expect large numbers of not overly active applications, such as shell sessions from
 separate users. However, it comes with a risk of less predictable latency, because it may take a while for
 an application to gain back cores on one node when it has work to do. To use this mode, simply use a
-`mesos://` URL without setting `spark.mesos.coarse` to true.
+`mesos://` URL and set `spark.mesos.coarse` to false.
 
 Note that none of the modes currently provide memory sharing across applications. If you would like to share
 data this way, we recommend running a single server application that can serve multiple requests by querying
diff --git a/docs/running-on-mesos.md b/docs/running-on-mesos.md
index 5be208cf3461e..a197d0e373027 100644
--- a/docs/running-on-mesos.md
+++ b/docs/running-on-mesos.md
@@ -161,21 +161,15 @@ Note that jars or python files that are passed to spark-submit should be URIs re
 
 # Mesos Run Modes
 
-Spark can run over Mesos in two modes: "fine-grained" (default) and "coarse-grained".
+Spark can run over Mesos in two modes: "coarse-grained" (default) and "fine-grained".
 
-In "fine-grained" mode (default), each Spark task runs as a separate Mesos task. This allows
-multiple instances of Spark (and other frameworks) to share machines at a very fine granularity,
-where each application gets more or fewer machines as it ramps up and down, but it comes with an
-additional overhead in launching each task. This mode may be inappropriate for low-latency
-requirements like interactive queries or serving web requests.
-
-The "coarse-grained" mode will instead launch only *one* long-running Spark task on each Mesos
+The "coarse-grained" mode will launch only *one* long-running Spark task on each Mesos
 machine, and dynamically schedule its own "mini-tasks" within it. The benefit is much lower startup
 overhead, but at the cost of reserving the Mesos resources for the complete duration of the
 application.
 
-To run in coarse-grained mode, set the `spark.mesos.coarse` property in your
-[SparkConf](configuration.html#spark-properties):
+Coarse-grained is the default mode. You can also set `spark.mesos.coarse` property to true
+to turn it on explictly in [SparkConf](configuration.html#spark-properties):
 
 {% highlight scala %}
 conf.set("spark.mesos.coarse", "true")
@@ -186,6 +180,19 @@ acquire. By default, it will acquire *all* cores in the cluster (that get offere
 only makes sense if you run just one application at a time. You can cap the maximum number of cores
 using `conf.set("spark.cores.max", "10")` (for example).
 
+In "fine-grained" mode, each Spark task runs as a separate Mesos task. This allows
+multiple instances of Spark (and other frameworks) to share machines at a very fine granularity,
+where each application gets more or fewer machines as it ramps up and down, but it comes with an
+additional overhead in launching each task. This mode may be inappropriate for low-latency
+requirements like interactive queries or serving web requests.
+
+To run in coarse-grained mode, set the `spark.mesos.coarse` property to false in your
+[SparkConf](configuration.html#spark-properties):
+
+{% highlight scala %}
+conf.set("spark.mesos.coarse", "false")
+{% endhighlight %}
+
 You may also make use of `spark.mesos.constraints` to set attribute based constraints on mesos resource offers. By default, all resource offers will be accepted.
 
 {% highlight scala %}

From 7c5b641808740ba5eed05ba8204cdbaf3fc579f5 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Wed, 18 Nov 2015 12:53:22 -0800
Subject: [PATCH 333/518] [SPARK-10745][CORE] Separate configs between shuffle
 and RPC

[SPARK-6028](https://issues.apache.org/jira/browse/SPARK-6028) uses network module to implement RPC. However, there are some configurations named with `spark.shuffle` prefix in the network module.

This PR refactors them to make sure the user can control them in shuffle and RPC separately. The user can use `spark.rpc.*` to set the configuration for netty RPC.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9481 from zsxwing/SPARK-10745.
---
 .../spark/deploy/ExternalShuffleService.scala |  3 +-
 .../netty/NettyBlockTransferService.scala     |  2 +-
 .../network/netty/SparkTransportConf.scala    | 12 ++--
 .../apache/spark/rpc/netty/NettyRpcEnv.scala  |  8 +--
 .../mesos/CoarseMesosSchedulerBackend.scala   |  2 +-
 .../shuffle/FileShuffleBlockResolver.scala    |  2 +-
 .../shuffle/IndexShuffleBlockResolver.scala   |  2 +-
 .../apache/spark/storage/BlockManager.scala   |  2 +-
 .../spark/ExternalShuffleServiceSuite.scala   |  2 +-
 .../spark/network/util/TransportConf.java     | 65 ++++++++++++++-----
 .../network/ChunkFetchIntegrationSuite.java   |  2 +-
 .../RequestTimeoutIntegrationSuite.java       |  2 +-
 .../spark/network/RpcIntegrationSuite.java    |  2 +-
 .../org/apache/spark/network/StreamSuite.java |  2 +-
 .../network/TransportClientFactorySuite.java  |  6 +-
 .../spark/network/sasl/SparkSaslSuite.java    |  6 +-
 .../network/sasl/SaslIntegrationSuite.java    |  2 +-
 .../ExternalShuffleBlockResolverSuite.java    |  2 +-
 .../shuffle/ExternalShuffleCleanupSuite.java  |  2 +-
 .../ExternalShuffleIntegrationSuite.java      |  2 +-
 .../shuffle/ExternalShuffleSecuritySuite.java |  2 +-
 .../shuffle/RetryingBlockFetcherSuite.java    |  2 +-
 .../network/yarn/YarnShuffleService.java      |  2 +-
 23 files changed, 84 insertions(+), 50 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala b/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala
index a039d543c35e7..e8a1e35c3fc48 100644
--- a/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/ExternalShuffleService.scala
@@ -45,7 +45,8 @@ class ExternalShuffleService(sparkConf: SparkConf, securityManager: SecurityMana
   private val port = sparkConf.getInt("spark.shuffle.service.port", 7337)
   private val useSasl: Boolean = securityManager.isAuthenticationEnabled()
 
-  private val transportConf = SparkTransportConf.fromSparkConf(sparkConf, numUsableCores = 0)
+  private val transportConf =
+    SparkTransportConf.fromSparkConf(sparkConf, "shuffle", numUsableCores = 0)
   private val blockHandler = newShuffleBlockHandler(transportConf)
   private val transportContext: TransportContext =
     new TransportContext(transportConf, blockHandler, true)
diff --git a/core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala b/core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala
index 70a42f9045e6b..b0694e3c6c8af 100644
--- a/core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala
+++ b/core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala
@@ -41,7 +41,7 @@ class NettyBlockTransferService(conf: SparkConf, securityManager: SecurityManage
   // TODO: Don't use Java serialization, use a more cross-version compatible serialization format.
   private val serializer = new JavaSerializer(conf)
   private val authEnabled = securityManager.isAuthenticationEnabled()
-  private val transportConf = SparkTransportConf.fromSparkConf(conf, numCores)
+  private val transportConf = SparkTransportConf.fromSparkConf(conf, "shuffle", numCores)
 
   private[this] var transportContext: TransportContext = _
   private[this] var server: TransportServer = _
diff --git a/core/src/main/scala/org/apache/spark/network/netty/SparkTransportConf.scala b/core/src/main/scala/org/apache/spark/network/netty/SparkTransportConf.scala
index cef203006d685..84833f59d7afe 100644
--- a/core/src/main/scala/org/apache/spark/network/netty/SparkTransportConf.scala
+++ b/core/src/main/scala/org/apache/spark/network/netty/SparkTransportConf.scala
@@ -40,23 +40,23 @@ object SparkTransportConf {
 
   /**
    * Utility for creating a [[TransportConf]] from a [[SparkConf]].
+   * @param _conf the [[SparkConf]]
+   * @param module the module name
    * @param numUsableCores if nonzero, this will restrict the server and client threads to only
    *                       use the given number of cores, rather than all of the machine's cores.
    *                       This restriction will only occur if these properties are not already set.
    */
-  def fromSparkConf(_conf: SparkConf, numUsableCores: Int = 0): TransportConf = {
+  def fromSparkConf(_conf: SparkConf, module: String, numUsableCores: Int = 0): TransportConf = {
     val conf = _conf.clone
 
     // Specify thread configuration based on our JVM's allocation of cores (rather than necessarily
     // assuming we have all the machine's cores).
     // NB: Only set if serverThreads/clientThreads not already set.
     val numThreads = defaultNumThreads(numUsableCores)
-    conf.set("spark.shuffle.io.serverThreads",
-      conf.get("spark.shuffle.io.serverThreads", numThreads.toString))
-    conf.set("spark.shuffle.io.clientThreads",
-      conf.get("spark.shuffle.io.clientThreads", numThreads.toString))
+    conf.setIfMissing(s"spark.$module.io.serverThreads", numThreads.toString)
+    conf.setIfMissing(s"spark.$module.io.clientThreads", numThreads.toString)
 
-    new TransportConf(new ConfigProvider {
+    new TransportConf(module, new ConfigProvider {
       override def get(name: String): String = conf.get(name)
     })
   }
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
index 09093819bb22c..3e0c497969502 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
@@ -22,16 +22,13 @@ import java.net.{InetSocketAddress, URI}
 import java.nio.ByteBuffer
 import java.util.concurrent._
 import java.util.concurrent.atomic.AtomicBoolean
-import javax.annotation.Nullable;
-import javax.annotation.concurrent.GuardedBy
+import javax.annotation.Nullable
 
-import scala.collection.mutable
 import scala.concurrent.{Future, Promise}
 import scala.reflect.ClassTag
 import scala.util.{DynamicVariable, Failure, Success}
 import scala.util.control.NonFatal
 
-import com.google.common.base.Preconditions
 import org.apache.spark.{Logging, SecurityManager, SparkConf}
 import org.apache.spark.network.TransportContext
 import org.apache.spark.network.client._
@@ -49,7 +46,8 @@ private[netty] class NettyRpcEnv(
     securityManager: SecurityManager) extends RpcEnv(conf) with Logging {
 
   private val transportConf = SparkTransportConf.fromSparkConf(
-    conf.clone.set("spark.shuffle.io.numConnectionsPerPeer", "1"),
+    conf.clone.set("spark.rpc.io.numConnectionsPerPeer", "1"),
+    "rpc",
     conf.getInt("spark.rpc.io.threads", 0))
 
   private val dispatcher: Dispatcher = new Dispatcher(this)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
index 2de9b6a651692..7d08eae0b4871 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/mesos/CoarseMesosSchedulerBackend.scala
@@ -109,7 +109,7 @@ private[spark] class CoarseMesosSchedulerBackend(
   private val mesosExternalShuffleClient: Option[MesosExternalShuffleClient] = {
     if (shuffleServiceEnabled) {
       Some(new MesosExternalShuffleClient(
-        SparkTransportConf.fromSparkConf(conf),
+        SparkTransportConf.fromSparkConf(conf, "shuffle"),
         securityManager,
         securityManager.isAuthenticationEnabled(),
         securityManager.isSaslEncryptionEnabled()))
diff --git a/core/src/main/scala/org/apache/spark/shuffle/FileShuffleBlockResolver.scala b/core/src/main/scala/org/apache/spark/shuffle/FileShuffleBlockResolver.scala
index 39fadd8783518..cc5f933393adf 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/FileShuffleBlockResolver.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/FileShuffleBlockResolver.scala
@@ -46,7 +46,7 @@ private[spark] trait ShuffleWriterGroup {
 private[spark] class FileShuffleBlockResolver(conf: SparkConf)
   extends ShuffleBlockResolver with Logging {
 
-  private val transportConf = SparkTransportConf.fromSparkConf(conf)
+  private val transportConf = SparkTransportConf.fromSparkConf(conf, "shuffle")
 
   private lazy val blockManager = SparkEnv.get.blockManager
 
diff --git a/core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala b/core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala
index 05b1eed7f3bef..fadb8fe7ed0ab 100644
--- a/core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala
+++ b/core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala
@@ -47,7 +47,7 @@ private[spark] class IndexShuffleBlockResolver(
 
   private lazy val blockManager = Option(_blockManager).getOrElse(SparkEnv.get.blockManager)
 
-  private val transportConf = SparkTransportConf.fromSparkConf(conf)
+  private val transportConf = SparkTransportConf.fromSparkConf(conf, "shuffle")
 
   def getDataFile(shuffleId: Int, mapId: Int): File = {
     blockManager.diskBlockManager.getFile(ShuffleDataBlockId(shuffleId, mapId, NOOP_REDUCE_ID))
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
index 661c706af32b1..ab0007fb78993 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
@@ -122,7 +122,7 @@ private[spark] class BlockManager(
   // Client to read other executors' shuffle files. This is either an external service, or just the
   // standard BlockTransferService to directly connect to other Executors.
   private[spark] val shuffleClient = if (externalShuffleServiceEnabled) {
-    val transConf = SparkTransportConf.fromSparkConf(conf, numUsableCores)
+    val transConf = SparkTransportConf.fromSparkConf(conf, "shuffle", numUsableCores)
     new ExternalShuffleClient(transConf, securityManager, securityManager.isAuthenticationEnabled(),
       securityManager.isSaslEncryptionEnabled())
   } else {
diff --git a/core/src/test/scala/org/apache/spark/ExternalShuffleServiceSuite.scala b/core/src/test/scala/org/apache/spark/ExternalShuffleServiceSuite.scala
index 231f4631e0a47..1c775bcb3d9c1 100644
--- a/core/src/test/scala/org/apache/spark/ExternalShuffleServiceSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ExternalShuffleServiceSuite.scala
@@ -35,7 +35,7 @@ class ExternalShuffleServiceSuite extends ShuffleSuite with BeforeAndAfterAll {
   var rpcHandler: ExternalShuffleBlockHandler = _
 
   override def beforeAll() {
-    val transportConf = SparkTransportConf.fromSparkConf(conf, numUsableCores = 2)
+    val transportConf = SparkTransportConf.fromSparkConf(conf, "shuffle", numUsableCores = 2)
     rpcHandler = new ExternalShuffleBlockHandler(transportConf, null)
     val transportContext = new TransportContext(transportConf, rpcHandler)
     server = transportContext.createServer()
diff --git a/network/common/src/main/java/org/apache/spark/network/util/TransportConf.java b/network/common/src/main/java/org/apache/spark/network/util/TransportConf.java
index 3b2eff377955a..115135d44adbd 100644
--- a/network/common/src/main/java/org/apache/spark/network/util/TransportConf.java
+++ b/network/common/src/main/java/org/apache/spark/network/util/TransportConf.java
@@ -23,18 +23,53 @@
  * A central location that tracks all the settings we expose to users.
  */
 public class TransportConf {
+
+  private final String SPARK_NETWORK_IO_MODE_KEY;
+  private final String SPARK_NETWORK_IO_PREFERDIRECTBUFS_KEY;
+  private final String SPARK_NETWORK_IO_CONNECTIONTIMEOUT_KEY;
+  private final String SPARK_NETWORK_IO_BACKLOG_KEY;
+  private final String SPARK_NETWORK_IO_NUMCONNECTIONSPERPEER_KEY;
+  private final String SPARK_NETWORK_IO_SERVERTHREADS_KEY;
+  private final String SPARK_NETWORK_IO_CLIENTTHREADS_KEY;
+  private final String SPARK_NETWORK_IO_RECEIVEBUFFER_KEY;
+  private final String SPARK_NETWORK_IO_SENDBUFFER_KEY;
+  private final String SPARK_NETWORK_SASL_TIMEOUT_KEY;
+  private final String SPARK_NETWORK_IO_MAXRETRIES_KEY;
+  private final String SPARK_NETWORK_IO_RETRYWAIT_KEY;
+  private final String SPARK_NETWORK_IO_LAZYFD_KEY;
+
   private final ConfigProvider conf;
 
-  public TransportConf(ConfigProvider conf) {
+  private final String module;
+
+  public TransportConf(String module, ConfigProvider conf) {
+    this.module = module;
     this.conf = conf;
+    SPARK_NETWORK_IO_MODE_KEY = getConfKey("io.mode");
+    SPARK_NETWORK_IO_PREFERDIRECTBUFS_KEY = getConfKey("io.preferDirectBufs");
+    SPARK_NETWORK_IO_CONNECTIONTIMEOUT_KEY = getConfKey("io.connectionTimeout");
+    SPARK_NETWORK_IO_BACKLOG_KEY = getConfKey("io.backLog");
+    SPARK_NETWORK_IO_NUMCONNECTIONSPERPEER_KEY =  getConfKey("io.numConnectionsPerPeer");
+    SPARK_NETWORK_IO_SERVERTHREADS_KEY = getConfKey("io.serverThreads");
+    SPARK_NETWORK_IO_CLIENTTHREADS_KEY = getConfKey("io.clientThreads");
+    SPARK_NETWORK_IO_RECEIVEBUFFER_KEY = getConfKey("io.receiveBuffer");
+    SPARK_NETWORK_IO_SENDBUFFER_KEY = getConfKey("io.sendBuffer");
+    SPARK_NETWORK_SASL_TIMEOUT_KEY = getConfKey("sasl.timeout");
+    SPARK_NETWORK_IO_MAXRETRIES_KEY = getConfKey("io.maxRetries");
+    SPARK_NETWORK_IO_RETRYWAIT_KEY = getConfKey("io.retryWait");
+    SPARK_NETWORK_IO_LAZYFD_KEY = getConfKey("io.lazyFD");
+  }
+
+  private String getConfKey(String suffix) {
+    return "spark." + module + "." + suffix;
   }
 
   /** IO mode: nio or epoll */
-  public String ioMode() { return conf.get("spark.shuffle.io.mode", "NIO").toUpperCase(); }
+  public String ioMode() { return conf.get(SPARK_NETWORK_IO_MODE_KEY, "NIO").toUpperCase(); }
 
   /** If true, we will prefer allocating off-heap byte buffers within Netty. */
   public boolean preferDirectBufs() {
-    return conf.getBoolean("spark.shuffle.io.preferDirectBufs", true);
+    return conf.getBoolean(SPARK_NETWORK_IO_PREFERDIRECTBUFS_KEY, true);
   }
 
   /** Connect timeout in milliseconds. Default 120 secs. */
@@ -42,23 +77,23 @@ public int connectionTimeoutMs() {
     long defaultNetworkTimeoutS = JavaUtils.timeStringAsSec(
       conf.get("spark.network.timeout", "120s"));
     long defaultTimeoutMs = JavaUtils.timeStringAsSec(
-      conf.get("spark.shuffle.io.connectionTimeout", defaultNetworkTimeoutS + "s")) * 1000;
+      conf.get(SPARK_NETWORK_IO_CONNECTIONTIMEOUT_KEY, defaultNetworkTimeoutS + "s")) * 1000;
     return (int) defaultTimeoutMs;
   }
 
   /** Number of concurrent connections between two nodes for fetching data. */
   public int numConnectionsPerPeer() {
-    return conf.getInt("spark.shuffle.io.numConnectionsPerPeer", 1);
+    return conf.getInt(SPARK_NETWORK_IO_NUMCONNECTIONSPERPEER_KEY, 1);
   }
 
   /** Requested maximum length of the queue of incoming connections. Default -1 for no backlog. */
-  public int backLog() { return conf.getInt("spark.shuffle.io.backLog", -1); }
+  public int backLog() { return conf.getInt(SPARK_NETWORK_IO_BACKLOG_KEY, -1); }
 
   /** Number of threads used in the server thread pool. Default to 0, which is 2x#cores. */
-  public int serverThreads() { return conf.getInt("spark.shuffle.io.serverThreads", 0); }
+  public int serverThreads() { return conf.getInt(SPARK_NETWORK_IO_SERVERTHREADS_KEY, 0); }
 
   /** Number of threads used in the client thread pool. Default to 0, which is 2x#cores. */
-  public int clientThreads() { return conf.getInt("spark.shuffle.io.clientThreads", 0); }
+  public int clientThreads() { return conf.getInt(SPARK_NETWORK_IO_CLIENTTHREADS_KEY, 0); }
 
   /**
    * Receive buffer size (SO_RCVBUF).
@@ -67,28 +102,28 @@ public int numConnectionsPerPeer() {
    * Assuming latency = 1ms, network_bandwidth = 10Gbps
    *  buffer size should be ~ 1.25MB
    */
-  public int receiveBuf() { return conf.getInt("spark.shuffle.io.receiveBuffer", -1); }
+  public int receiveBuf() { return conf.getInt(SPARK_NETWORK_IO_RECEIVEBUFFER_KEY, -1); }
 
   /** Send buffer size (SO_SNDBUF). */
-  public int sendBuf() { return conf.getInt("spark.shuffle.io.sendBuffer", -1); }
+  public int sendBuf() { return conf.getInt(SPARK_NETWORK_IO_SENDBUFFER_KEY, -1); }
 
   /** Timeout for a single round trip of SASL token exchange, in milliseconds. */
   public int saslRTTimeoutMs() {
-    return (int) JavaUtils.timeStringAsSec(conf.get("spark.shuffle.sasl.timeout", "30s")) * 1000;
+    return (int) JavaUtils.timeStringAsSec(conf.get(SPARK_NETWORK_SASL_TIMEOUT_KEY, "30s")) * 1000;
   }
 
   /**
    * Max number of times we will try IO exceptions (such as connection timeouts) per request.
    * If set to 0, we will not do any retries.
    */
-  public int maxIORetries() { return conf.getInt("spark.shuffle.io.maxRetries", 3); }
+  public int maxIORetries() { return conf.getInt(SPARK_NETWORK_IO_MAXRETRIES_KEY, 3); }
 
   /**
    * Time (in milliseconds) that we will wait in order to perform a retry after an IOException.
    * Only relevant if maxIORetries &gt; 0.
    */
   public int ioRetryWaitTimeMs() {
-    return (int) JavaUtils.timeStringAsSec(conf.get("spark.shuffle.io.retryWait", "5s")) * 1000;
+    return (int) JavaUtils.timeStringAsSec(conf.get(SPARK_NETWORK_IO_RETRYWAIT_KEY, "5s")) * 1000;
   }
 
   /**
@@ -101,11 +136,11 @@ public int memoryMapBytes() {
   }
 
   /**
-   * Whether to initialize shuffle FileDescriptor lazily or not. If true, file descriptors are
+   * Whether to initialize FileDescriptor lazily or not. If true, file descriptors are
    * created only when data is going to be transferred. This can reduce the number of open files.
    */
   public boolean lazyFileDescriptor() {
-    return conf.getBoolean("spark.shuffle.io.lazyFD", true);
+    return conf.getBoolean(SPARK_NETWORK_IO_LAZYFD_KEY, true);
   }
 
   /**
diff --git a/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java b/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
index dfb7740344ed0..dc5fa1cee69bc 100644
--- a/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
@@ -83,7 +83,7 @@ public static void setUp() throws Exception {
     fp.write(fileContent);
     fp.close();
 
-    final TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+    final TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     fileChunk = new FileSegmentManagedBuffer(conf, testFile, 10, testFile.length() - 25);
 
     streamManager = new StreamManager() {
diff --git a/network/common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java b/network/common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java
index 84ebb337e6d54..42955ef69235a 100644
--- a/network/common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java
@@ -60,7 +60,7 @@ public class RequestTimeoutIntegrationSuite {
   public void setUp() throws Exception {
     Map<String, String> configMap = Maps.newHashMap();
     configMap.put("spark.shuffle.io.connectionTimeout", "2s");
-    conf = new TransportConf(new MapConfigProvider(configMap));
+    conf = new TransportConf("shuffle", new MapConfigProvider(configMap));
 
     defaultManager = new StreamManager() {
       @Override
diff --git a/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java b/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
index 64b457b4b3f01..8eb56bdd9846f 100644
--- a/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
@@ -49,7 +49,7 @@ public class RpcIntegrationSuite {
 
   @BeforeClass
   public static void setUp() throws Exception {
-    TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+    TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     rpcHandler = new RpcHandler() {
       @Override
       public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
diff --git a/network/common/src/test/java/org/apache/spark/network/StreamSuite.java b/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
index 6dcec831dec71..00158fd081626 100644
--- a/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
@@ -89,7 +89,7 @@ public static void setUp() throws Exception {
       fp.close();
     }
 
-    final TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+    final TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     final StreamManager streamManager = new StreamManager() {
       @Override
       public ManagedBuffer getChunk(long streamId, int chunkIndex) {
diff --git a/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java b/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java
index f447137419306..dac7d4a5b0a09 100644
--- a/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/TransportClientFactorySuite.java
@@ -52,7 +52,7 @@ public class TransportClientFactorySuite {
 
   @Before
   public void setUp() {
-    conf = new TransportConf(new SystemPropertyConfigProvider());
+    conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     RpcHandler rpcHandler = new NoOpRpcHandler();
     context = new TransportContext(conf, rpcHandler);
     server1 = context.createServer();
@@ -76,7 +76,7 @@ private void testClientReuse(final int maxConnections, boolean concurrent)
 
     Map<String, String> configMap = Maps.newHashMap();
     configMap.put("spark.shuffle.io.numConnectionsPerPeer", Integer.toString(maxConnections));
-    TransportConf conf = new TransportConf(new MapConfigProvider(configMap));
+    TransportConf conf = new TransportConf("shuffle", new MapConfigProvider(configMap));
 
     RpcHandler rpcHandler = new NoOpRpcHandler();
     TransportContext context = new TransportContext(conf, rpcHandler);
@@ -182,7 +182,7 @@ public void closeBlockClientsWithFactory() throws IOException {
 
   @Test
   public void closeIdleConnectionForRequestTimeOut() throws IOException, InterruptedException {
-    TransportConf conf = new TransportConf(new ConfigProvider() {
+    TransportConf conf = new TransportConf("shuffle", new ConfigProvider() {
 
       @Override
       public String get(String name) {
diff --git a/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java b/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
index 3469e84e7f4da..b146899670180 100644
--- a/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
@@ -207,7 +207,7 @@ public void testEncryptedMessage() throws Exception {
   public void testEncryptedMessageChunking() throws Exception {
     File file = File.createTempFile("sasltest", ".txt");
     try {
-      TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+      TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
 
       byte[] data = new byte[8 * 1024];
       new Random().nextBytes(data);
@@ -242,7 +242,7 @@ public void testFileRegionEncryption() throws Exception {
     final File file = File.createTempFile("sasltest", ".txt");
     SaslTestCtx ctx = null;
     try {
-      final TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+      final TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
       StreamManager sm = mock(StreamManager.class);
       when(sm.getChunk(anyLong(), anyInt())).thenAnswer(new Answer<ManagedBuffer>() {
           @Override
@@ -368,7 +368,7 @@ private static class SaslTestCtx {
         boolean disableClientEncryption)
       throws Exception {
 
-      TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+      TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
 
       SecretKeyHolder keyHolder = mock(SecretKeyHolder.class);
       when(keyHolder.getSaslUser(anyString())).thenReturn("user");
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java
index c393a5e1e6810..1c2fa4d0d462c 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java
@@ -70,7 +70,7 @@ public class SaslIntegrationSuite {
 
   @BeforeClass
   public static void beforeAll() throws IOException {
-    conf = new TransportConf(new SystemPropertyConfigProvider());
+    conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     context = new TransportContext(conf, new TestRpcHandler());
 
     secretKeyHolder = mock(SecretKeyHolder.class);
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolverSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolverSuite.java
index 3c6cb367dea46..a9958232a1d28 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolverSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockResolverSuite.java
@@ -42,7 +42,7 @@ public class ExternalShuffleBlockResolverSuite {
 
   static TestShuffleDataContext dataContext;
 
-  static TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+  static TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
 
   @BeforeClass
   public static void beforeAll() throws IOException {
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleCleanupSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleCleanupSuite.java
index 2f4f1d0df478b..532d7ab8d01bd 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleCleanupSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleCleanupSuite.java
@@ -35,7 +35,7 @@ public class ExternalShuffleCleanupSuite {
 
   // Same-thread Executor used to ensure cleanup happens synchronously in test thread.
   Executor sameThreadExecutor = MoreExecutors.sameThreadExecutor();
-  TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+  TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
 
   @Test
   public void noCleanupAndCleanup() throws IOException {
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleIntegrationSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleIntegrationSuite.java
index a3f9a38b1aeb9..2095f41d79c16 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleIntegrationSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleIntegrationSuite.java
@@ -91,7 +91,7 @@ public static void beforeAll() throws IOException {
     dataContext1.create();
     dataContext1.insertHashShuffleData(1, 0, exec1Blocks);
 
-    conf = new TransportConf(new SystemPropertyConfigProvider());
+    conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     handler = new ExternalShuffleBlockHandler(conf, null);
     TransportContext transportContext = new TransportContext(conf, handler);
     server = transportContext.createServer();
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleSecuritySuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleSecuritySuite.java
index aa99efda94948..08ddb3755bd08 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleSecuritySuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleSecuritySuite.java
@@ -39,7 +39,7 @@
 
 public class ExternalShuffleSecuritySuite {
 
-  TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+  TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
   TransportServer server;
 
   @Before
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/RetryingBlockFetcherSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/RetryingBlockFetcherSuite.java
index 06e46f9241094..3a6ef0d3f8476 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/RetryingBlockFetcherSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/RetryingBlockFetcherSuite.java
@@ -254,7 +254,7 @@ private static void performInteractions(List<? extends Map<String, Object>> inte
                                           BlockFetchingListener listener)
     throws IOException {
 
-    TransportConf conf = new TransportConf(new SystemPropertyConfigProvider());
+    TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     BlockFetchStarter fetchStarter = mock(BlockFetchStarter.class);
 
     Stubber stub = null;
diff --git a/network/yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java b/network/yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java
index 11ea7f3fd3cfe..ba6d30a74c673 100644
--- a/network/yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java
+++ b/network/yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java
@@ -120,7 +120,7 @@ protected void serviceInit(Configuration conf) {
     registeredExecutorFile =
       findRegisteredExecutorFile(conf.getStrings("yarn.nodemanager.local-dirs"));
 
-    TransportConf transportConf = new TransportConf(new HadoopConfigProvider(conf));
+    TransportConf transportConf = new TransportConf("shuffle", new HadoopConfigProvider(conf));
     // If authentication is enabled, set up the shuffle server to use a
     // special RPC handler that filters out unauthenticated fetch requests
     boolean authEnabled = conf.getBoolean(SPARK_AUTHENTICATE_KEY, DEFAULT_SPARK_AUTHENTICATE);

From 09ad9533d5760652de59fa4830c24cb8667958ac Mon Sep 17 00:00:00 2001
From: JihongMa <linlin200605@gmail.com>
Date: Wed, 18 Nov 2015 13:03:37 -0800
Subject: [PATCH 334/518] [SPARK-11720][SQL][ML] Handle edge cases when count =
 0 or 1 for Stats function

return Double.NaN for mean/average when count == 0 for all numeric types that is converted to Double, Decimal type continue to return null.

Author: JihongMa <linlin200605@gmail.com>

Closes #9705 from JihongMA/SPARK-11720.
---
 python/pyspark/sql/dataframe.py                |  2 +-
 .../aggregate/CentralMomentAgg.scala           |  2 +-
 .../expressions/aggregate/Kurtosis.scala       |  9 +++++----
 .../expressions/aggregate/Skewness.scala       |  9 +++++----
 .../expressions/aggregate/Stddev.scala         | 18 ++++++++++++++----
 .../expressions/aggregate/Variance.scala       | 18 ++++++++++++++----
 .../spark/sql/DataFrameAggregateSuite.scala    | 18 ++++++++++++------
 .../org/apache/spark/sql/DataFrameSuite.scala  |  2 +-
 8 files changed, 53 insertions(+), 25 deletions(-)

diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index ad6ad0235a904..0dd75ba7ca820 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -761,7 +761,7 @@ def describe(self, *cols):
         +-------+------------------+-----+
         |  count|                 2|    2|
         |   mean|               3.5| null|
-        | stddev|2.1213203435596424|  NaN|
+        | stddev|2.1213203435596424| null|
         |    min|                 2|Alice|
         |    max|                 5|  Bob|
         +-------+------------------+-----+
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
index de5872ab11eb1..d07d4c338cdfe 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
@@ -206,7 +206,7 @@ abstract class CentralMomentAgg(child: Expression) extends ImperativeAggregate w
    * @param centralMoments Length `momentOrder + 1` array of central moments (un-normalized)
    *                       needed to compute the aggregate stat.
    */
-  def getStatistic(n: Double, mean: Double, centralMoments: Array[Double]): Double
+  def getStatistic(n: Double, mean: Double, centralMoments: Array[Double]): Any
 
   override final def eval(buffer: InternalRow): Any = {
     val n = buffer.getDouble(nOffset)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
index 8fa3aac9f1a51..c2bf2cb94116c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Kurtosis.scala
@@ -37,16 +37,17 @@ case class Kurtosis(child: Expression,
   override protected val momentOrder = 4
 
   // NOTE: this is the formula for excess kurtosis, which is default for R and SciPy
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Any = {
     require(moments.length == momentOrder + 1,
       s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
     val m2 = moments(2)
     val m4 = moments(4)
 
-    if (n == 0.0 || m2 == 0.0) {
+    if (n == 0.0) {
+      null
+    } else if (m2 == 0.0) {
       Double.NaN
-    }
-    else {
+    } else {
       n * m4 / (m2 * m2) - 3.0
     }
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
index e1c01a5b82781..9411bcea2539a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Skewness.scala
@@ -36,16 +36,17 @@ case class Skewness(child: Expression,
 
   override protected val momentOrder = 3
 
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Any = {
     require(moments.length == momentOrder + 1,
       s"$prettyName requires ${momentOrder + 1} central moments, received: ${moments.length}")
     val m2 = moments(2)
     val m3 = moments(3)
 
-    if (n == 0.0 || m2 == 0.0) {
+    if (n == 0.0) {
+      null
+    } else if (m2 == 0.0) {
       Double.NaN
-    }
-    else {
+    } else {
       math.sqrt(n) * m3 / math.sqrt(m2 * m2 * m2)
     }
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
index 05dd5e3b22543..eec79a9033e36 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Stddev.scala
@@ -36,11 +36,17 @@ case class StddevSamp(child: Expression,
 
   override protected val momentOrder = 2
 
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Any = {
     require(moments.length == momentOrder + 1,
       s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
 
-    if (n == 0.0 || n == 1.0) Double.NaN else math.sqrt(moments(2) / (n - 1.0))
+    if (n == 0.0) {
+      null
+    } else if (n == 1.0) {
+      Double.NaN
+    } else {
+      math.sqrt(moments(2) / (n - 1.0))
+    }
   }
 }
 
@@ -62,10 +68,14 @@ case class StddevPop(
 
   override protected val momentOrder = 2
 
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Any = {
     require(moments.length == momentOrder + 1,
       s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
 
-    if (n == 0.0) Double.NaN else math.sqrt(moments(2) / n)
+    if (n == 0.0) {
+      null
+    } else {
+      math.sqrt(moments(2) / n)
+    }
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala
index ede2da2805966..cf3a740305391 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Variance.scala
@@ -36,11 +36,17 @@ case class VarianceSamp(child: Expression,
 
   override protected val momentOrder = 2
 
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Any = {
     require(moments.length == momentOrder + 1,
       s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
 
-    if (n == 0.0 || n == 1.0) Double.NaN else moments(2) / (n - 1.0)
+    if (n == 0.0) {
+      null
+    } else if (n == 1.0) {
+      Double.NaN
+    } else {
+      moments(2) / (n - 1.0)
+    }
   }
 }
 
@@ -62,10 +68,14 @@ case class VariancePop(
 
   override protected val momentOrder = 2
 
-  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Double = {
+  override def getStatistic(n: Double, mean: Double, moments: Array[Double]): Any = {
     require(moments.length == momentOrder + 1,
       s"$prettyName requires ${momentOrder + 1} central moment, received: ${moments.length}")
 
-    if (n == 0.0) Double.NaN else moments(2) / n
+    if (n == 0.0) {
+      null
+    } else {
+      moments(2) / n
+    }
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index 432e8d17623a4..71adf2148a403 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -205,7 +205,7 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
     val emptyTableData = Seq.empty[(Int, Int)].toDF("a", "b")
     checkAnswer(
     emptyTableData.agg(stddev('a), stddev_pop('a), stddev_samp('a)),
-    Row(Double.NaN, Double.NaN, Double.NaN))
+    Row(null, null, null))
   }
 
   test("zero sum") {
@@ -244,17 +244,23 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
   test("zero moments") {
     val input = Seq((1, 2)).toDF("a", "b")
     checkAnswer(
-      input.agg(variance('a), var_samp('a), var_pop('a), skewness('a), kurtosis('a)),
-      Row(Double.NaN, Double.NaN, 0.0, Double.NaN, Double.NaN))
+      input.agg(stddev('a), stddev_samp('a), stddev_pop('a), variance('a),
+        var_samp('a), var_pop('a), skewness('a), kurtosis('a)),
+      Row(Double.NaN, Double.NaN, 0.0, Double.NaN, Double.NaN, 0.0,
+        Double.NaN, Double.NaN))
 
     checkAnswer(
       input.agg(
+        expr("stddev(a)"),
+        expr("stddev_samp(a)"),
+        expr("stddev_pop(a)"),
         expr("variance(a)"),
         expr("var_samp(a)"),
         expr("var_pop(a)"),
         expr("skewness(a)"),
         expr("kurtosis(a)")),
-      Row(Double.NaN, Double.NaN, 0.0, Double.NaN, Double.NaN))
+      Row(Double.NaN, Double.NaN, 0.0, Double.NaN, Double.NaN, 0.0,
+        Double.NaN, Double.NaN))
   }
 
   test("null moments") {
@@ -262,7 +268,7 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
 
     checkAnswer(
       emptyTableData.agg(variance('a), var_samp('a), var_pop('a), skewness('a), kurtosis('a)),
-      Row(Double.NaN, Double.NaN, Double.NaN, Double.NaN, Double.NaN))
+      Row(null, null, null, null, null))
 
     checkAnswer(
       emptyTableData.agg(
@@ -271,6 +277,6 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
         expr("var_pop(a)"),
         expr("skewness(a)"),
         expr("kurtosis(a)")),
-      Row(Double.NaN, Double.NaN, Double.NaN, Double.NaN, Double.NaN))
+      Row(null, null, null, null, null))
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index 5a7f24684d1b7..6399b0165c4c3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -459,7 +459,7 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
     val emptyDescribeResult = Seq(
       Row("count", "0", "0"),
       Row("mean", null, null),
-      Row("stddev", "NaN", "NaN"),
+      Row("stddev", null, null),
       Row("min", null, null),
       Row("max", null, null))
 

From 045a4f045821dcf60442f0600c2df1b79bddb536 Mon Sep 17 00:00:00 2001
From: Wenjian Huang <nextrush@163.com>
Date: Wed, 18 Nov 2015 13:06:25 -0800
Subject: [PATCH 335/518] [SPARK-6790][ML] Add spark.ml LinearRegression
 import/export

This replaces [https://github.com/apache/spark/pull/9656] with updates.

fayeshine should be the main author when this PR is committed.

CC: mengxr fayeshine

Author: Wenjian Huang <nextrush@163.com>
Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9814 from jkbradley/fayeshine-patch-6790.
---
 .../ml/regression/LinearRegression.scala      | 77 ++++++++++++++++++-
 .../ml/regression/LinearRegressionSuite.scala | 34 +++++++-
 2 files changed, 106 insertions(+), 5 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index 913140e581983..ca55d5915e688 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -22,6 +22,7 @@ import scala.collection.mutable
 import breeze.linalg.{DenseVector => BDV}
 import breeze.optimize.{CachedDiffFunction, DiffFunction, LBFGS => BreezeLBFGS, OWLQN => BreezeOWLQN}
 import breeze.stats.distributions.StudentsT
+import org.apache.hadoop.fs.Path
 
 import org.apache.spark.{Logging, SparkException}
 import org.apache.spark.ml.feature.Instance
@@ -30,7 +31,7 @@ import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.PredictorParams
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.evaluation.RegressionMetrics
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.linalg.BLAS._
@@ -65,7 +66,7 @@ private[regression] trait LinearRegressionParams extends PredictorParams
 @Experimental
 class LinearRegression @Since("1.3.0") (@Since("1.3.0") override val uid: String)
   extends Regressor[Vector, LinearRegression, LinearRegressionModel]
-  with LinearRegressionParams with Logging {
+  with LinearRegressionParams with Writable with Logging {
 
   @Since("1.4.0")
   def this() = this(Identifiable.randomUID("linReg"))
@@ -341,6 +342,19 @@ class LinearRegression @Since("1.3.0") (@Since("1.3.0") override val uid: String
 
   @Since("1.4.0")
   override def copy(extra: ParamMap): LinearRegression = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object LinearRegression extends Readable[LinearRegression] {
+
+  @Since("1.6.0")
+  override def read: Reader[LinearRegression] = new DefaultParamsReader[LinearRegression]
+
+  @Since("1.6.0")
+  override def load(path: String): LinearRegression = read.load(path)
 }
 
 /**
@@ -354,7 +368,7 @@ class LinearRegressionModel private[ml] (
     val coefficients: Vector,
     val intercept: Double)
   extends RegressionModel[Vector, LinearRegressionModel]
-  with LinearRegressionParams {
+  with LinearRegressionParams with Writable {
 
   private var trainingSummary: Option[LinearRegressionTrainingSummary] = None
 
@@ -422,6 +436,63 @@ class LinearRegressionModel private[ml] (
     if (trainingSummary.isDefined) newModel.setSummary(trainingSummary.get)
     newModel.setParent(parent)
   }
+
+  /**
+   * Returns a [[Writer]] instance for this ML instance.
+   *
+   * For [[LinearRegressionModel]], this does NOT currently save the training [[summary]].
+   * An option to save [[summary]] may be added in the future.
+   *
+   * This also does not save the [[parent]] currently.
+   */
+  @Since("1.6.0")
+  override def write: Writer = new LinearRegressionModel.LinearRegressionModelWriter(this)
+}
+
+@Since("1.6.0")
+object LinearRegressionModel extends Readable[LinearRegressionModel] {
+
+  @Since("1.6.0")
+  override def read: Reader[LinearRegressionModel] = new LinearRegressionModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): LinearRegressionModel = read.load(path)
+
+  /** [[Writer]] instance for [[LinearRegressionModel]] */
+  private[LinearRegressionModel] class LinearRegressionModelWriter(instance: LinearRegressionModel)
+    extends Writer with Logging {
+
+    private case class Data(intercept: Double, coefficients: Vector)
+
+    override protected def saveImpl(path: String): Unit = {
+      // Save metadata and Params
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      // Save model data: intercept, coefficients
+      val data = Data(instance.intercept, instance.coefficients)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).write.format("parquet").save(dataPath)
+    }
+  }
+
+  private class LinearRegressionModelReader extends Reader[LinearRegressionModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = "org.apache.spark.ml.regression.LinearRegressionModel"
+
+    override def load(path: String): LinearRegressionModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.format("parquet").load(dataPath)
+        .select("intercept", "coefficients").head()
+      val intercept = data.getDouble(0)
+      val coefficients = data.getAs[Vector](1)
+      val model = new LinearRegressionModel(metadata.uid, coefficients, intercept)
+
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
 }
 
 /**
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
index a1d86fe8fedad..2bdc0e184d734 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/LinearRegressionSuite.scala
@@ -22,14 +22,15 @@ import scala.util.Random
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.feature.Instance
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.regression.LabeledPoint
 import org.apache.spark.mllib.linalg.{Vector, DenseVector, Vectors}
 import org.apache.spark.mllib.util.{LinearDataGenerator, MLlibTestSparkContext}
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.{DataFrame, Row}
 
-class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
+class LinearRegressionSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   private val seed: Int = 42
   @transient var datasetWithDenseFeature: DataFrame = _
@@ -854,4 +855,33 @@ class LinearRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
     model.summary.tValues.zip(tValsR).foreach{ x => assert(x._1 ~== x._2 absTol 1E-3) }
     model.summary.pValues.zip(pValsR).foreach{ x => assert(x._1 ~== x._2 absTol 1E-3) }
   }
+
+  test("read/write") {
+    def checkModelData(model: LinearRegressionModel, model2: LinearRegressionModel): Unit = {
+      assert(model.intercept === model2.intercept)
+      assert(model.coefficients === model2.coefficients)
+    }
+    val lr = new LinearRegression()
+    testEstimatorAndModelReadWrite(lr, datasetWithWeight, LinearRegressionSuite.allParamSettings,
+      checkModelData)
+  }
+}
+
+object LinearRegressionSuite {
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = Map(
+    "predictionCol" -> "myPrediction",
+    "regParam" -> 0.01,
+    "elasticNetParam" -> 0.1,
+    "maxIter" -> 2,  // intentionally small
+    "fitIntercept" -> true,
+    "tol" -> 0.8,
+    "standardization" -> false,
+    "solver" -> "l-bfgs"
+  )
 }

From 2acdf10b1f3bb1242dba64efa798c672fde9f0d2 Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Wed, 18 Nov 2015 13:16:31 -0800
Subject: [PATCH 336/518] [SPARK-6789][ML] Add Readable, Writable support for
 spark.ml ALS, ALSModel

Also modifies DefaultParamsWriter.saveMetadata to take optional extra metadata.

CC: mengxr yanboliang

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9786 from jkbradley/als-io.
---
 .../apache/spark/ml/recommendation/ALS.scala  | 75 ++++++++++++++++--
 .../org/apache/spark/ml/util/ReadWrite.scala  | 14 +++-
 .../spark/ml/recommendation/ALSSuite.scala    | 78 ++++++++++++++++---
 3 files changed, 150 insertions(+), 17 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
index 535f266b9a944..d92514d2e239e 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
@@ -27,13 +27,16 @@ import scala.util.hashing.byteswap64
 
 import com.github.fommil.netlib.BLAS.{getInstance => blas}
 import org.apache.hadoop.fs.{FileSystem, Path}
+import org.json4s.{DefaultFormats, JValue}
+import org.json4s.JsonDSL._
+import org.json4s.jackson.JsonMethods._
 
 import org.apache.spark.{Logging, Partitioner}
-import org.apache.spark.annotation.{DeveloperApi, Experimental}
+import org.apache.spark.annotation.{Since, DeveloperApi, Experimental}
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg.CholeskyDecomposition
 import org.apache.spark.mllib.optimization.NNLS
 import org.apache.spark.rdd.RDD
@@ -182,7 +185,7 @@ class ALSModel private[ml] (
     val rank: Int,
     @transient val userFactors: DataFrame,
     @transient val itemFactors: DataFrame)
-  extends Model[ALSModel] with ALSModelParams {
+  extends Model[ALSModel] with ALSModelParams with Writable {
 
   /** @group setParam */
   def setUserCol(value: String): this.type = set(userCol, value)
@@ -220,8 +223,60 @@ class ALSModel private[ml] (
     val copied = new ALSModel(uid, rank, userFactors, itemFactors)
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: Writer = new ALSModel.ALSModelWriter(this)
 }
 
+@Since("1.6.0")
+object ALSModel extends Readable[ALSModel] {
+
+  @Since("1.6.0")
+  override def read: Reader[ALSModel] = new ALSModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): ALSModel = read.load(path)
+
+  private[recommendation] class ALSModelWriter(instance: ALSModel) extends Writer {
+
+    override protected def saveImpl(path: String): Unit = {
+      val extraMetadata = render("rank" -> instance.rank)
+      DefaultParamsWriter.saveMetadata(instance, path, sc, Some(extraMetadata))
+      val userPath = new Path(path, "userFactors").toString
+      instance.userFactors.write.format("parquet").save(userPath)
+      val itemPath = new Path(path, "itemFactors").toString
+      instance.itemFactors.write.format("parquet").save(itemPath)
+    }
+  }
+
+  private[recommendation] class ALSModelReader extends Reader[ALSModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = "org.apache.spark.ml.recommendation.ALSModel"
+
+    override def load(path: String): ALSModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      implicit val format = DefaultFormats
+      val rank: Int = metadata.extraMetadata match {
+        case Some(m: JValue) =>
+          (m \ "rank").extract[Int]
+        case None =>
+          throw new RuntimeException(s"ALSModel loader could not read rank from JSON metadata:" +
+            s" ${metadata.metadataStr}")
+      }
+
+      val userPath = new Path(path, "userFactors").toString
+      val userFactors = sqlContext.read.format("parquet").load(userPath)
+      val itemPath = new Path(path, "itemFactors").toString
+      val itemFactors = sqlContext.read.format("parquet").load(itemPath)
+
+      val model = new ALSModel(metadata.uid, rank, userFactors, itemFactors)
+
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+}
 
 /**
  * :: Experimental ::
@@ -254,7 +309,7 @@ class ALSModel private[ml] (
  * preferences rather than explicit ratings given to items.
  */
 @Experimental
-class ALS(override val uid: String) extends Estimator[ALSModel] with ALSParams {
+class ALS(override val uid: String) extends Estimator[ALSModel] with ALSParams with Writable {
 
   import org.apache.spark.ml.recommendation.ALS.Rating
 
@@ -336,8 +391,12 @@ class ALS(override val uid: String) extends Estimator[ALSModel] with ALSParams {
   }
 
   override def copy(extra: ParamMap): ALS = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
 }
 
+
 /**
  * :: DeveloperApi ::
  * An implementation of ALS that supports generic ID types, specialized for Int and Long. This is
@@ -347,7 +406,7 @@ class ALS(override val uid: String) extends Estimator[ALSModel] with ALSParams {
  * than 2 billion.
  */
 @DeveloperApi
-object ALS extends Logging {
+object ALS extends Readable[ALS] with Logging {
 
   /**
    * :: DeveloperApi ::
@@ -356,6 +415,12 @@ object ALS extends Logging {
   @DeveloperApi
   case class Rating[@specialized(Int, Long) ID](user: ID, item: ID, rating: Float)
 
+  @Since("1.6.0")
+  override def read: Reader[ALS] = new DefaultParamsReader[ALS]
+
+  @Since("1.6.0")
+  override def load(path: String): ALS = read.load(path)
+
   /** Trait for least squares solvers applied to the normal equation. */
   private[recommendation] trait LeastSquaresNESolver extends Serializable {
     /** Solves a least squares problem with regularization (possibly with other constraints). */
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index dddb72af5ba78..d8ce907af5323 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -194,7 +194,11 @@ private[ml] object DefaultParamsWriter {
    *  - uid
    *  - paramMap: These must be encodable using [[org.apache.spark.ml.param.Param.jsonEncode()]].
    */
-  def saveMetadata(instance: Params, path: String, sc: SparkContext): Unit = {
+  def saveMetadata(
+      instance: Params,
+      path: String,
+      sc: SparkContext,
+      extraMetadata: Option[JValue] = None): Unit = {
     val uid = instance.uid
     val cls = instance.getClass.getName
     val params = instance.extractParamMap().toSeq.asInstanceOf[Seq[ParamPair[Any]]]
@@ -205,7 +209,8 @@ private[ml] object DefaultParamsWriter {
       ("timestamp" -> System.currentTimeMillis()) ~
       ("sparkVersion" -> sc.version) ~
       ("uid" -> uid) ~
-      ("paramMap" -> jsonParams)
+      ("paramMap" -> jsonParams) ~
+      ("extraMetadata" -> extraMetadata)
     val metadataPath = new Path(path, "metadata").toString
     val metadataJson = compact(render(metadata))
     sc.parallelize(Seq(metadataJson), 1).saveAsTextFile(metadataPath)
@@ -236,6 +241,7 @@ private[ml] object DefaultParamsReader {
   /**
    * All info from metadata file.
    * @param params  paramMap, as a [[JValue]]
+   * @param extraMetadata  Extra metadata saved by [[DefaultParamsWriter.saveMetadata()]]
    * @param metadataStr  Full metadata file String (for debugging)
    */
   case class Metadata(
@@ -244,6 +250,7 @@ private[ml] object DefaultParamsReader {
       timestamp: Long,
       sparkVersion: String,
       params: JValue,
+      extraMetadata: Option[JValue],
       metadataStr: String)
 
   /**
@@ -262,12 +269,13 @@ private[ml] object DefaultParamsReader {
     val timestamp = (metadata \ "timestamp").extract[Long]
     val sparkVersion = (metadata \ "sparkVersion").extract[String]
     val params = metadata \ "paramMap"
+    val extraMetadata = (metadata \ "extraMetadata").extract[Option[JValue]]
     if (expectedClassName.nonEmpty) {
       require(className == expectedClassName, s"Error loading metadata: Expected class name" +
         s" $expectedClassName but found class name $className")
     }
 
-    Metadata(className, uid, timestamp, sparkVersion, params, metadataStr)
+    Metadata(className, uid, timestamp, sparkVersion, params, extraMetadata, metadataStr)
   }
 
   /**
diff --git a/mllib/src/test/scala/org/apache/spark/ml/recommendation/ALSSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/recommendation/ALSSuite.scala
index eadc80e0e62b1..2c3fb84160dcb 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/recommendation/ALSSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/recommendation/ALSSuite.scala
@@ -17,7 +17,6 @@
 
 package org.apache.spark.ml.recommendation
 
-import java.io.File
 import java.util.Random
 
 import scala.collection.mutable
@@ -26,28 +25,26 @@ import scala.language.existentials
 
 import com.github.fommil.netlib.BLAS.{getInstance => blas}
 
+import org.apache.spark.util.Utils
 import org.apache.spark.{Logging, SparkException, SparkFunSuite}
 import org.apache.spark.ml.recommendation.ALS._
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.Vectors
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{Row, SQLContext}
-import org.apache.spark.util.Utils
+import org.apache.spark.sql.{DataFrame, Row}
 
-class ALSSuite extends SparkFunSuite with MLlibTestSparkContext with Logging {
 
-  private var tempDir: File = _
+class ALSSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest with Logging {
 
   override def beforeAll(): Unit = {
     super.beforeAll()
-    tempDir = Utils.createTempDir()
     sc.setCheckpointDir(tempDir.getAbsolutePath)
   }
 
   override def afterAll(): Unit = {
-    Utils.deleteRecursively(tempDir)
     super.afterAll()
   }
 
@@ -186,7 +183,7 @@ class ALSSuite extends SparkFunSuite with MLlibTestSparkContext with Logging {
     assert(compressed.dstPtrs.toSeq === Seq(0, 2, 3, 4, 5))
     var decompressed = ArrayBuffer.empty[(Int, Int, Int, Float)]
     var i = 0
-    while (i < compressed.srcIds.size) {
+    while (i < compressed.srcIds.length) {
       var j = compressed.dstPtrs(i)
       while (j < compressed.dstPtrs(i + 1)) {
         val dstEncodedIndex = compressed.dstEncodedIndices(j)
@@ -483,4 +480,67 @@ class ALSSuite extends SparkFunSuite with MLlibTestSparkContext with Logging {
     ALS.train(ratings, rank = 1, maxIter = 50, numUserBlocks = 2, numItemBlocks = 2,
       implicitPrefs = true, seed = 0)
   }
+
+  test("read/write") {
+    import ALSSuite._
+    val (ratings, _) = genExplicitTestData(numUsers = 4, numItems = 4, rank = 1)
+    val als = new ALS()
+    allEstimatorParamSettings.foreach { case (p, v) =>
+      als.set(als.getParam(p), v)
+    }
+    val sqlContext = this.sqlContext
+    import sqlContext.implicits._
+    val model = als.fit(ratings.toDF())
+
+    // Test Estimator save/load
+    val als2 = testDefaultReadWrite(als)
+    allEstimatorParamSettings.foreach { case (p, v) =>
+      val param = als.getParam(p)
+      assert(als.get(param).get === als2.get(param).get)
+    }
+
+    // Test Model save/load
+    val model2 = testDefaultReadWrite(model)
+    allModelParamSettings.foreach { case (p, v) =>
+      val param = model.getParam(p)
+      assert(model.get(param).get === model2.get(param).get)
+    }
+    assert(model.rank === model2.rank)
+    def getFactors(df: DataFrame): Set[(Int, Array[Float])] = {
+      df.select("id", "features").collect().map { case r =>
+        (r.getInt(0), r.getAs[Array[Float]](1))
+      }.toSet
+    }
+    assert(getFactors(model.userFactors) === getFactors(model2.userFactors))
+    assert(getFactors(model.itemFactors) === getFactors(model2.itemFactors))
+  }
+}
+
+object ALSSuite {
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allModelParamSettings: Map[String, Any] = Map(
+    "predictionCol" -> "myPredictionCol"
+  )
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allEstimatorParamSettings: Map[String, Any] = allModelParamSettings ++ Map(
+    "maxIter" -> 1,
+    "rank" -> 1,
+    "regParam" -> 0.01,
+    "numUserBlocks" -> 2,
+    "numItemBlocks" -> 2,
+    "implicitPrefs" -> true,
+    "alpha" -> 0.9,
+    "nonnegative" -> true,
+    "checkpointInterval" -> 20
+  )
 }

From e391abdf2cb6098a35347bd123b815ee9ac5b689 Mon Sep 17 00:00:00 2001
From: Yuhao Yang <hhbyyh@gmail.com>
Date: Wed, 18 Nov 2015 13:25:15 -0800
Subject: [PATCH 337/518] [SPARK-11813][MLLIB] Avoid serialization of vocab in
 Word2Vec

jira: https://issues.apache.org/jira/browse/SPARK-11813

I found the problem during training a large corpus. Avoid serialization of vocab in Word2Vec has 2 benefits.
1. Performance improvement for less serialization.
2. Increase the capacity of Word2Vec a lot.
Currently in the fit of word2vec, the closure mainly includes serialization of Word2Vec and 2 global table.
the main part of Word2vec is the vocab of size: vocab * 40 * 2 * 4 = 320 vocab
2 global table: vocab * vectorSize * 8. If vectorSize = 20, that's 160 vocab.

Their sum cannot exceed Int.max due to the restriction of ByteArrayOutputStream. In any case, avoiding serialization of vocab helps decrease the size of the closure serialization, especially when vectorSize is small, thus to allow larger vocabulary.

Actually there's another possible fix, make local copy of fields to avoid including Word2Vec in the closure. Let me know if that's preferred.

Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #9803 from hhbyyh/w2vVocab.
---
 .../main/scala/org/apache/spark/mllib/feature/Word2Vec.scala  | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
index f3e4d346e358a..7ab0d89d23a3f 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
@@ -145,8 +145,8 @@ class Word2Vec extends Serializable with Logging {
 
   private var trainWordsCount = 0
   private var vocabSize = 0
-  private var vocab: Array[VocabWord] = null
-  private var vocabHash = mutable.HashMap.empty[String, Int]
+  @transient private var vocab: Array[VocabWord] = null
+  @transient private var vocabHash = mutable.HashMap.empty[String, Int]
 
   private def learnVocab(words: RDD[String]): Unit = {
     vocab = words.map(w => (w, 1))

From e222d758499ad2609046cc1a2cc8afb45c5bccbb Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Wed, 18 Nov 2015 13:30:29 -0800
Subject: [PATCH 338/518] [SPARK-11684][R][ML][DOC] Update SparkR glm API doc,
 user guide and example codes

This PR includes:
* Update SparkR:::glm, SparkR:::summary API docs.
* Update SparkR machine learning user guide and example codes to show:
  * supporting feature interaction in R formula.
  * summary for gaussian GLM model.
  * coefficients for binomial GLM model.

mengxr

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9727 from yanboliang/spark-11684.
---
 R/pkg/R/mllib.R                               | 18 +++++--
 docs/sparkr.md                                | 50 ++++++++++++++++---
 .../ml/regression/LinearRegression.scala      |  3 ++
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/R/pkg/R/mllib.R b/R/pkg/R/mllib.R
index f23e1c7f1fce4..8d3b4388ae575 100644
--- a/R/pkg/R/mllib.R
+++ b/R/pkg/R/mllib.R
@@ -32,6 +32,12 @@ setClass("PipelineModel", representation(model = "jobj"))
 #' @param family Error distribution. "gaussian" -> linear regression, "binomial" -> logistic reg.
 #' @param lambda Regularization parameter
 #' @param alpha Elastic-net mixing parameter (see glmnet's documentation for details)
+#' @param standardize Whether to standardize features before training
+#' @param solver The solver algorithm used for optimization, this can be "l-bfgs", "normal" and
+#'               "auto". "l-bfgs" denotes Limited-memory BFGS which is a limited-memory
+#'               quasi-Newton optimization method. "normal" denotes using Normal Equation as an
+#'               analytical solution to the linear regression problem. The default value is "auto"
+#'               which means that the solver algorithm is selected automatically.
 #' @return a fitted MLlib model
 #' @rdname glm
 #' @export
@@ -79,9 +85,15 @@ setMethod("predict", signature(object = "PipelineModel"),
 #'
 #' Returns the summary of a model produced by glm(), similarly to R's summary().
 #'
-#' @param x A fitted MLlib model
-#' @return a list with a 'coefficient' component, which is the matrix of coefficients. See
-#'         summary.glm for more information.
+#' @param object A fitted MLlib model
+#' @return a list with 'devianceResiduals' and 'coefficients' components for gaussian family
+#'         or a list with 'coefficients' component for binomial family. \cr
+#'         For gaussian family: the 'devianceResiduals' gives the min/max deviance residuals
+#'         of the estimation, the 'coefficients' gives the estimated coefficients and their
+#'         estimated standard errors, t values and p-values. (It only available when model
+#'         fitted by normal solver.) \cr
+#'         For binomial family: the 'coefficients' gives the estimated coefficients.
+#'         See summary.glm for more information. \cr
 #' @rdname summary
 #' @export
 #' @examples
diff --git a/docs/sparkr.md b/docs/sparkr.md
index 437bd4756c276..a744b76be7466 100644
--- a/docs/sparkr.md
+++ b/docs/sparkr.md
@@ -286,24 +286,37 @@ head(teenagers)
 
 # Machine Learning
 
-SparkR allows the fitting of generalized linear models over DataFrames using the [glm()](api/R/glm.html) function. Under the hood, SparkR uses MLlib to train a model of the specified family. Currently the gaussian and binomial families are supported. We support a subset of the available R formula operators for model fitting, including '~', '.', '+', and '-'. The example below shows the use of building a gaussian GLM model using SparkR.
+SparkR allows the fitting of generalized linear models over DataFrames using the [glm()](api/R/glm.html) function. Under the hood, SparkR uses MLlib to train a model of the specified family. Currently the gaussian and binomial families are supported. We support a subset of the available R formula operators for model fitting, including '~', '.', ':', '+', and '-'. 
+
+The [summary()](api/R/summary.html) function gives the summary of a model produced by [glm()](api/R/glm.html).
+
+* For gaussian GLM model, it returns a list with 'devianceResiduals' and 'coefficients' components. The 'devianceResiduals' gives the min/max deviance residuals of the estimation; the 'coefficients' gives the estimated coefficients and their estimated standard errors, t values and p-values. (It only available when model fitted by normal solver.)
+* For binomial GLM model, it returns a list with 'coefficients' component which gives the estimated coefficients.
+
+The examples below show the use of building gaussian GLM model and binomial GLM model using SparkR.
+
+## Gaussian GLM model
 
 <div data-lang="r"  markdown="1">
 {% highlight r %}
 # Create the DataFrame
 df <- createDataFrame(sqlContext, iris)
 
-# Fit a linear model over the dataset.
+# Fit a gaussian GLM model over the dataset.
 model <- glm(Sepal_Length ~ Sepal_Width + Species, data = df, family = "gaussian")
 
-# Model coefficients are returned in a similar format to R's native glm().
+# Model summary are returned in a similar format to R's native glm().
 summary(model)
+##$devianceResiduals
+## Min       Max     
+## -1.307112 1.412532
+##
 ##$coefficients
-##                    Estimate
-##(Intercept)        2.2513930
-##Sepal_Width        0.8035609
-##Species_versicolor 1.4587432
-##Species_virginica  1.9468169
+##                   Estimate  Std. Error t value  Pr(>|t|)    
+##(Intercept)        2.251393  0.3697543  6.08889  9.568102e-09
+##Sepal_Width        0.8035609 0.106339   7.556598 4.187317e-12
+##Species_versicolor 1.458743  0.1121079  13.01195 0           
+##Species_virginica  1.946817  0.100015   19.46525 0           
 
 # Make predictions based on the model.
 predictions <- predict(model, newData = df)
@@ -317,3 +330,24 @@ head(select(predictions, "Sepal_Length", "prediction"))
 ##6          5.4   5.385281
 {% endhighlight %}
 </div>
+
+## Binomial GLM model
+
+<div data-lang="r"  markdown="1">
+{% highlight r %}
+# Create the DataFrame
+df <- createDataFrame(sqlContext, iris)
+training <- filter(df, df$Species != "setosa")
+
+# Fit a binomial GLM model over the dataset.
+model <- glm(Species ~ Sepal_Length + Sepal_Width, data = training, family = "binomial")
+
+# Model coefficients are returned in a similar format to R's native glm().
+summary(model)
+##$coefficients
+##               Estimate
+##(Intercept)  -13.046005
+##Sepal_Length   1.902373
+##Sepal_Width    0.404655
+{% endhighlight %}
+</div>
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index ca55d5915e688..f7c44f0a51b8a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -145,6 +145,9 @@ class LinearRegression @Since("1.3.0") (@Since("1.3.0") override val uid: String
   /**
    * Set the solver algorithm used for optimization.
    * In case of linear regression, this can be "l-bfgs", "normal" and "auto".
+   * "l-bfgs" denotes Limited-memory BFGS which is a limited-memory quasi-Newton
+   * optimization method. "normal" denotes using Normal Equation as an analytical
+   * solution to the linear regression problem.
    * The default value is "auto" which means that the solver algorithm is
    * selected automatically.
    * @group setParam

From 603a721c21488e17c15c45ce1de893e6b3d02274 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Wed, 18 Nov 2015 13:32:06 -0800
Subject: [PATCH 339/518] [SPARK-11820][ML][PYSPARK] PySpark LiR & LoR should
 support weightCol

[SPARK-7685](https://issues.apache.org/jira/browse/SPARK-7685) and [SPARK-9642](https://issues.apache.org/jira/browse/SPARK-9642) have already supported setting weight column for ```LogisticRegression``` and ```LinearRegression```. It's a very important feature, PySpark should also support. mengxr

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9811 from yanboliang/spark-11820.
---
 python/pyspark/ml/classification.py | 17 +++++++++--------
 python/pyspark/ml/regression.py     | 16 ++++++++--------
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py
index 603f2c7f798dc..4a2982e2047ff 100644
--- a/python/pyspark/ml/classification.py
+++ b/python/pyspark/ml/classification.py
@@ -36,7 +36,8 @@
 @inherit_doc
 class LogisticRegression(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPredictionCol, HasMaxIter,
                          HasRegParam, HasTol, HasProbabilityCol, HasRawPredictionCol,
-                         HasElasticNetParam, HasFitIntercept, HasStandardization, HasThresholds):
+                         HasElasticNetParam, HasFitIntercept, HasStandardization, HasThresholds,
+                         HasWeightCol):
     """
     Logistic regression.
     Currently, this class only supports binary classification.
@@ -44,9 +45,9 @@ class LogisticRegression(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPredicti
     >>> from pyspark.sql import Row
     >>> from pyspark.mllib.linalg import Vectors
     >>> df = sc.parallelize([
-    ...     Row(label=1.0, features=Vectors.dense(1.0)),
-    ...     Row(label=0.0, features=Vectors.sparse(1, [], []))]).toDF()
-    >>> lr = LogisticRegression(maxIter=5, regParam=0.01)
+    ...     Row(label=1.0, weight=2.0, features=Vectors.dense(1.0)),
+    ...     Row(label=0.0, weight=2.0, features=Vectors.sparse(1, [], []))]).toDF()
+    >>> lr = LogisticRegression(maxIter=5, regParam=0.01, weightCol="weight")
     >>> model = lr.fit(df)
     >>> model.weights
     DenseVector([5.5...])
@@ -80,12 +81,12 @@ class LogisticRegression(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPredicti
     def __init__(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                  maxIter=100, regParam=0.1, elasticNetParam=0.0, tol=1e-6, fitIntercept=True,
                  threshold=0.5, thresholds=None, probabilityCol="probability",
-                 rawPredictionCol="rawPrediction", standardization=True):
+                 rawPredictionCol="rawPrediction", standardization=True, weightCol=None):
         """
         __init__(self, featuresCol="features", labelCol="label", predictionCol="prediction", \
                  maxIter=100, regParam=0.1, elasticNetParam=0.0, tol=1e-6, fitIntercept=True, \
                  threshold=0.5, thresholds=None, probabilityCol="probability", \
-                 rawPredictionCol="rawPrediction", standardization=True)
+                 rawPredictionCol="rawPrediction", standardization=True, weightCol=None)
         If the threshold and thresholds Params are both set, they must be equivalent.
         """
         super(LogisticRegression, self).__init__()
@@ -105,12 +106,12 @@ def __init__(self, featuresCol="features", labelCol="label", predictionCol="pred
     def setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                   maxIter=100, regParam=0.1, elasticNetParam=0.0, tol=1e-6, fitIntercept=True,
                   threshold=0.5, thresholds=None, probabilityCol="probability",
-                  rawPredictionCol="rawPrediction", standardization=True):
+                  rawPredictionCol="rawPrediction", standardization=True, weightCol=None):
         """
         setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction", \
                   maxIter=100, regParam=0.1, elasticNetParam=0.0, tol=1e-6, fitIntercept=True, \
                   threshold=0.5, thresholds=None, probabilityCol="probability", \
-                  rawPredictionCol="rawPrediction", standardization=True)
+                  rawPredictionCol="rawPrediction", standardization=True, weightCol=None)
         Sets params for logistic regression.
         If the threshold and thresholds Params are both set, they must be equivalent.
         """
diff --git a/python/pyspark/ml/regression.py b/python/pyspark/ml/regression.py
index 7648bf13266bf..944e648ec8801 100644
--- a/python/pyspark/ml/regression.py
+++ b/python/pyspark/ml/regression.py
@@ -35,7 +35,7 @@
 @inherit_doc
 class LinearRegression(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPredictionCol, HasMaxIter,
                        HasRegParam, HasTol, HasElasticNetParam, HasFitIntercept,
-                       HasStandardization, HasSolver):
+                       HasStandardization, HasSolver, HasWeightCol):
     """
     Linear regression.
 
@@ -50,9 +50,9 @@ class LinearRegression(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPrediction
 
     >>> from pyspark.mllib.linalg import Vectors
     >>> df = sqlContext.createDataFrame([
-    ...     (1.0, Vectors.dense(1.0)),
-    ...     (0.0, Vectors.sparse(1, [], []))], ["label", "features"])
-    >>> lr = LinearRegression(maxIter=5, regParam=0.0, solver="normal")
+    ...     (1.0, 2.0, Vectors.dense(1.0)),
+    ...     (0.0, 2.0, Vectors.sparse(1, [], []))], ["label", "weight", "features"])
+    >>> lr = LinearRegression(maxIter=5, regParam=0.0, solver="normal", weightCol="weight")
     >>> model = lr.fit(df)
     >>> test0 = sqlContext.createDataFrame([(Vectors.dense(-1.0),)], ["features"])
     >>> abs(model.transform(test0).head().prediction - (-1.0)) < 0.001
@@ -75,11 +75,11 @@ class LinearRegression(JavaEstimator, HasFeaturesCol, HasLabelCol, HasPrediction
     @keyword_only
     def __init__(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                  maxIter=100, regParam=0.0, elasticNetParam=0.0, tol=1e-6, fitIntercept=True,
-                 standardization=True, solver="auto"):
+                 standardization=True, solver="auto", weightCol=None):
         """
         __init__(self, featuresCol="features", labelCol="label", predictionCol="prediction", \
                  maxIter=100, regParam=0.0, elasticNetParam=0.0, tol=1e-6, fitIntercept=True, \
-                 standardization=True, solver="auto")
+                 standardization=True, solver="auto", weightCol=None)
         """
         super(LinearRegression, self).__init__()
         self._java_obj = self._new_java_obj(
@@ -92,11 +92,11 @@ def __init__(self, featuresCol="features", labelCol="label", predictionCol="pred
     @since("1.4.0")
     def setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction",
                   maxIter=100, regParam=0.0, elasticNetParam=0.0, tol=1e-6, fitIntercept=True,
-                  standardization=True, solver="auto"):
+                  standardization=True, solver="auto", weightCol=None):
         """
         setParams(self, featuresCol="features", labelCol="label", predictionCol="prediction", \
                   maxIter=100, regParam=0.0, elasticNetParam=0.0, tol=1e-6, fitIntercept=True, \
-                  standardization=True, solver="auto")
+                  standardization=True, solver="auto", weightCol=None)
         Sets params for linear regression.
         """
         kwargs = self.setParams._input_kwargs

From 54db79702513e11335c33bcf3a03c59e965e6f16 Mon Sep 17 00:00:00 2001
From: Dilip Biswal <dbiswal@us.ibm.com>
Date: Wed, 18 Nov 2015 14:05:18 -0800
Subject: [PATCH 340/518] [SPARK-11544][SQL] sqlContext doesn't use PathFilter

Apply the user supplied pathfilter while retrieving the files from fs.

Author: Dilip Biswal <dbiswal@us.ibm.com>

Closes #9652 from dilipbiswal/spark-11544.
---
 .../apache/spark/sql/sources/interfaces.scala | 25 ++++++++++---
 .../datasources/json/JsonSuite.scala          | 36 +++++++++++++++++--
 2 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index b3d3bdf50df63..f9465157c936d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -21,7 +21,8 @@ import scala.collection.mutable
 import scala.util.Try
 
 import org.apache.hadoop.conf.Configuration
-import org.apache.hadoop.fs.{FileStatus, FileSystem, Path}
+import org.apache.hadoop.fs.{PathFilter, FileStatus, FileSystem, Path}
+import org.apache.hadoop.mapred.{JobConf, FileInputFormat}
 import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext}
 
 import org.apache.spark.{Logging, SparkContext}
@@ -447,9 +448,15 @@ abstract class HadoopFsRelation private[sql](
           val hdfsPath = new Path(path)
           val fs = hdfsPath.getFileSystem(hadoopConf)
           val qualified = hdfsPath.makeQualified(fs.getUri, fs.getWorkingDirectory)
-
           logInfo(s"Listing $qualified on driver")
-          Try(fs.listStatus(qualified)).getOrElse(Array.empty)
+          // Dummy jobconf to get to the pathFilter defined in configuration
+          val jobConf = new JobConf(hadoopConf, this.getClass())
+          val pathFilter = FileInputFormat.getInputPathFilter(jobConf)
+          if (pathFilter != null) {
+            Try(fs.listStatus(qualified, pathFilter)).getOrElse(Array.empty)
+          } else {
+            Try(fs.listStatus(qualified)).getOrElse(Array.empty)
+          }
         }.filterNot { status =>
           val name = status.getPath.getName
           name.toLowerCase == "_temporary" || name.startsWith(".")
@@ -847,8 +854,16 @@ private[sql] object HadoopFsRelation extends Logging {
     if (name == "_temporary" || name.startsWith(".")) {
       Array.empty
     } else {
-      val (dirs, files) = fs.listStatus(status.getPath).partition(_.isDir)
-      files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
+      // Dummy jobconf to get to the pathFilter defined in configuration
+      val jobConf = new JobConf(fs.getConf, this.getClass())
+      val pathFilter = FileInputFormat.getInputPathFilter(jobConf)
+      if (pathFilter != null) {
+        val (dirs, files) = fs.listStatus(status.getPath, pathFilter).partition(_.isDir)
+        files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
+      } else {
+        val (dirs, files) = fs.listStatus(status.getPath).partition(_.isDir)
+        files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
+      }
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
index 6042b1178affe..f09b61e838159 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
@@ -19,19 +19,27 @@ package org.apache.spark.sql.execution.datasources.json
 
 import java.io.{File, StringWriter}
 import java.sql.{Date, Timestamp}
+import scala.collection.JavaConverters._
 
 import com.fasterxml.jackson.core.JsonFactory
-import org.apache.spark.rdd.RDD
+import org.apache.commons.io.FileUtils
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{Path, PathFilter}
 import org.scalactic.Tolerance._
 
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
-import org.apache.spark.sql.execution.datasources.{ResolvedDataSource, LogicalRelation}
+import org.apache.spark.sql.execution.datasources.{LogicalRelation, ResolvedDataSource}
 import org.apache.spark.sql.execution.datasources.json.InferSchema.compatibleType
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
 
+class TestFileFilter extends PathFilter {
+  override def accept(path: Path): Boolean = path.getParent.getName != "p=2"
+}
+
 class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
   import testImplicits._
 
@@ -1390,4 +1398,28 @@ class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
       )
     }
   }
+
+  test("SPARK-11544 test pathfilter") {
+    withTempPath { dir =>
+      val path = dir.getCanonicalPath
+
+      val df = sqlContext.range(2)
+      df.write.json(path + "/p=1")
+      df.write.json(path + "/p=2")
+      assert(sqlContext.read.json(path).count() === 4)
+
+      val clonedConf = new Configuration(hadoopConfiguration)
+      try {
+        hadoopConfiguration.setClass(
+          "mapreduce.input.pathFilter.class",
+          classOf[TestFileFilter],
+          classOf[PathFilter])
+        assert(sqlContext.read.json(path).count() === 2)
+      } finally {
+        // Hadoop 1 doesn't have `Configuration.unset`
+        hadoopConfiguration.clear()
+        clonedConf.asScala.foreach(entry => hadoopConfiguration.set(entry.getKey, entry.getValue))
+      }
+    }
+  }
 }

From 5df08949f5d9e5b4b0e9c2db50c1b4eb93383de3 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 18 Nov 2015 15:42:07 -0800
Subject: [PATCH 341/518] [SPARK-11810][SQL] Java-based encoder for opaque
 types in Datasets.

This patch refactors the existing Kryo encoder expressions and adds support for Java serialization.

Author: Reynold Xin <rxin@databricks.com>

Closes #9802 from rxin/SPARK-11810.
---
 .../scala/org/apache/spark/sql/Encoder.scala  | 41 +++++++++---
 .../sql/catalyst/expressions/objects.scala    | 67 ++++++++++++-------
 .../catalyst/encoders/FlatEncoderSuite.scala  | 27 ++++++--
 .../org/apache/spark/sql/DatasetSuite.scala   | 36 +++++++++-
 4 files changed, 130 insertions(+), 41 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index 79c2255641c06..1ed5111440c80 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -20,7 +20,7 @@ package org.apache.spark.sql
 import scala.reflect.{ClassTag, classTag}
 
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor}
-import org.apache.spark.sql.catalyst.expressions.{DeserializeWithKryo, BoundReference, SerializeWithKryo}
+import org.apache.spark.sql.catalyst.expressions.{DecodeUsingSerializer, BoundReference, EncodeUsingSerializer}
 import org.apache.spark.sql.types._
 
 /**
@@ -43,28 +43,49 @@ trait Encoder[T] extends Serializable {
  */
 object Encoders {
 
-  /**
-   * (Scala-specific) Creates an encoder that serializes objects of type T using Kryo.
-   * This encoder maps T into a single byte array (binary) field.
-   */
-  def kryo[T: ClassTag]: Encoder[T] = {
-    val ser = SerializeWithKryo(BoundReference(0, ObjectType(classOf[AnyRef]), nullable = true))
-    val deser = DeserializeWithKryo[T](BoundReference(0, BinaryType, nullable = true), classTag[T])
+  /** A way to construct encoders using generic serializers. */
+  private def genericSerializer[T: ClassTag](useKryo: Boolean): Encoder[T] = {
     ExpressionEncoder[T](
       schema = new StructType().add("value", BinaryType),
       flat = true,
-      toRowExpressions = Seq(ser),
-      fromRowExpression = deser,
+      toRowExpressions = Seq(
+        EncodeUsingSerializer(
+          BoundReference(0, ObjectType(classOf[AnyRef]), nullable = true), kryo = useKryo)),
+      fromRowExpression =
+        DecodeUsingSerializer[T](
+          BoundReference(0, BinaryType, nullable = true), classTag[T], kryo = useKryo),
       clsTag = classTag[T]
     )
   }
 
+  /**
+   * (Scala-specific) Creates an encoder that serializes objects of type T using Kryo.
+   * This encoder maps T into a single byte array (binary) field.
+   */
+  def kryo[T: ClassTag]: Encoder[T] = genericSerializer(useKryo = true)
+
   /**
    * Creates an encoder that serializes objects of type T using Kryo.
    * This encoder maps T into a single byte array (binary) field.
    */
   def kryo[T](clazz: Class[T]): Encoder[T] = kryo(ClassTag[T](clazz))
 
+  /**
+   * (Scala-specific) Creates an encoder that serializes objects of type T using generic Java
+   * serialization. This encoder maps T into a single byte array (binary) field.
+   *
+   * Note that this is extremely inefficient and should only be used as the last resort.
+   */
+  def javaSerialization[T: ClassTag]: Encoder[T] = genericSerializer(useKryo = false)
+
+  /**
+   * Creates an encoder that serializes objects of type T using generic Java serialization.
+   * This encoder maps T into a single byte array (binary) field.
+   *
+   * Note that this is extremely inefficient and should only be used as the last resort.
+   */
+  def javaSerialization[T](clazz: Class[T]): Encoder[T] = javaSerialization(ClassTag[T](clazz))
+
   def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder(flat = true)
   def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder(flat = true)
   def SHORT: Encoder[java.lang.Short] = ExpressionEncoder(flat = true)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index 489c6126f8cd3..acf0da240051e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -21,7 +21,7 @@ import scala.language.existentials
 import scala.reflect.ClassTag
 
 import org.apache.spark.SparkConf
-import org.apache.spark.serializer.{KryoSerializerInstance, KryoSerializer}
+import org.apache.spark.serializer._
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.analysis.SimpleAnalyzer
 import org.apache.spark.sql.catalyst.plans.logical.{Project, LocalRelation}
@@ -517,29 +517,39 @@ case class GetInternalRowField(child: Expression, ordinal: Int, dataType: DataTy
   }
 }
 
-/** Serializes an input object using Kryo serializer. */
-case class SerializeWithKryo(child: Expression) extends UnaryExpression {
+/**
+ * Serializes an input object using a generic serializer (Kryo or Java).
+ * @param kryo if true, use Kryo. Otherwise, use Java.
+ */
+case class EncodeUsingSerializer(child: Expression, kryo: Boolean) extends UnaryExpression {
 
   override def eval(input: InternalRow): Any =
     throw new UnsupportedOperationException("Only code-generated evaluation is supported")
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    val input = child.gen(ctx)
-    val kryo = ctx.freshName("kryoSerializer")
-    val kryoClass = classOf[KryoSerializer].getName
-    val kryoInstanceClass = classOf[KryoSerializerInstance].getName
-    val sparkConfClass = classOf[SparkConf].getName
+    // Code to initialize the serializer.
+    val serializer = ctx.freshName("serializer")
+    val (serializerClass, serializerInstanceClass) = {
+      if (kryo) {
+        (classOf[KryoSerializer].getName, classOf[KryoSerializerInstance].getName)
+      } else {
+        (classOf[JavaSerializer].getName, classOf[JavaSerializerInstance].getName)
+      }
+    }
+    val sparkConf = s"new ${classOf[SparkConf].getName}()"
     ctx.addMutableState(
-      kryoInstanceClass,
-      kryo,
-      s"$kryo = ($kryoInstanceClass) new $kryoClass(new $sparkConfClass()).newInstance();")
+      serializerInstanceClass,
+      serializer,
+      s"$serializer = ($serializerInstanceClass) new $serializerClass($sparkConf).newInstance();")
 
+    // Code to serialize.
+    val input = child.gen(ctx)
     s"""
       ${input.code}
       final boolean ${ev.isNull} = ${input.isNull};
       ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
       if (!${ev.isNull}) {
-        ${ev.value} = $kryo.serialize(${input.value}, null).array();
+        ${ev.value} = $serializer.serialize(${input.value}, null).array();
       }
      """
   }
@@ -548,29 +558,38 @@ case class SerializeWithKryo(child: Expression) extends UnaryExpression {
 }
 
 /**
- * Deserializes an input object using Kryo serializer. Note that the ClassTag is not an implicit
- * parameter because TreeNode cannot copy implicit parameters.
+ * Serializes an input object using a generic serializer (Kryo or Java).  Note that the ClassTag
+ * is not an implicit parameter because TreeNode cannot copy implicit parameters.
+ * @param kryo if true, use Kryo. Otherwise, use Java.
  */
-case class DeserializeWithKryo[T](child: Expression, tag: ClassTag[T]) extends UnaryExpression {
+case class DecodeUsingSerializer[T](child: Expression, tag: ClassTag[T], kryo: Boolean)
+  extends UnaryExpression {
 
   override protected def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    val input = child.gen(ctx)
-    val kryo = ctx.freshName("kryoSerializer")
-    val kryoClass = classOf[KryoSerializer].getName
-    val kryoInstanceClass = classOf[KryoSerializerInstance].getName
-    val sparkConfClass = classOf[SparkConf].getName
+    // Code to initialize the serializer.
+    val serializer = ctx.freshName("serializer")
+    val (serializerClass, serializerInstanceClass) = {
+      if (kryo) {
+        (classOf[KryoSerializer].getName, classOf[KryoSerializerInstance].getName)
+      } else {
+        (classOf[JavaSerializer].getName, classOf[JavaSerializerInstance].getName)
+      }
+    }
+    val sparkConf = s"new ${classOf[SparkConf].getName}()"
     ctx.addMutableState(
-      kryoInstanceClass,
-      kryo,
-      s"$kryo = ($kryoInstanceClass) new $kryoClass(new $sparkConfClass()).newInstance();")
+      serializerInstanceClass,
+      serializer,
+      s"$serializer = ($serializerInstanceClass) new $serializerClass($sparkConf).newInstance();")
 
+    // Code to serialize.
+    val input = child.gen(ctx)
     s"""
       ${input.code}
       final boolean ${ev.isNull} = ${input.isNull};
       ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
       if (!${ev.isNull}) {
         ${ev.value} = (${ctx.javaType(dataType)})
-          $kryo.deserialize(java.nio.ByteBuffer.wrap(${input.value}), null);
+          $serializer.deserialize(java.nio.ByteBuffer.wrap(${input.value}), null);
       }
      """
   }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
index 2729db84897a2..6e0322fb6e019 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
@@ -76,17 +76,34 @@ class FlatEncoderSuite extends ExpressionEncoderSuite {
   // Kryo encoders
   encodeDecodeTest(
     "hello",
-    Encoders.kryo[String].asInstanceOf[ExpressionEncoder[String]],
+    encoderFor(Encoders.kryo[String]),
     "kryo string")
   encodeDecodeTest(
-    new NotJavaSerializable(15),
-    Encoders.kryo[NotJavaSerializable].asInstanceOf[ExpressionEncoder[NotJavaSerializable]],
+    new KryoSerializable(15),
+    encoderFor(Encoders.kryo[KryoSerializable]),
     "kryo object serialization")
+
+  // Java encoders
+  encodeDecodeTest(
+    "hello",
+    encoderFor(Encoders.javaSerialization[String]),
+    "java string")
+  encodeDecodeTest(
+    new JavaSerializable(15),
+    encoderFor(Encoders.javaSerialization[JavaSerializable]),
+    "java object serialization")
 }
 
+/** For testing Kryo serialization based encoder. */
+class KryoSerializable(val value: Int) {
+  override def equals(other: Any): Boolean = {
+    this.value == other.asInstanceOf[KryoSerializable].value
+  }
+}
 
-class NotJavaSerializable(val value: Int) {
+/** For testing Java serialization based encoder. */
+class JavaSerializable(val value: Int) extends Serializable {
   override def equals(other: Any): Boolean = {
-    this.value == other.asInstanceOf[NotJavaSerializable].value
+    this.value == other.asInstanceOf[JavaSerializable].value
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index b6db583dfe01f..89d964aa3e469 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -357,7 +357,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     assert(ds.toString == "[_1: int, _2: int]")
   }
 
-  test("kryo encoder") {
+  test("Kryo encoder") {
     implicit val kryoEncoder = Encoders.kryo[KryoData]
     val ds = Seq(KryoData(1), KryoData(2)).toDS()
 
@@ -365,7 +365,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       Seq((KryoData(1), 1L), (KryoData(2), 1L)))
   }
 
-  test("kryo encoder self join") {
+  test("Kryo encoder self join") {
     implicit val kryoEncoder = Encoders.kryo[KryoData]
     val ds = Seq(KryoData(1), KryoData(2)).toDS()
     assert(ds.joinWith(ds, lit(true)).collect().toSet ==
@@ -375,6 +375,25 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
         (KryoData(2), KryoData(1)),
         (KryoData(2), KryoData(2))))
   }
+
+  test("Java encoder") {
+    implicit val kryoEncoder = Encoders.javaSerialization[JavaData]
+    val ds = Seq(JavaData(1), JavaData(2)).toDS()
+
+    assert(ds.groupBy(p => p).count().collect().toSeq ==
+      Seq((JavaData(1), 1L), (JavaData(2), 1L)))
+  }
+
+  ignore("Java encoder self join") {
+    implicit val kryoEncoder = Encoders.javaSerialization[JavaData]
+    val ds = Seq(JavaData(1), JavaData(2)).toDS()
+    assert(ds.joinWith(ds, lit(true)).collect().toSet ==
+      Set(
+        (JavaData(1), JavaData(1)),
+        (JavaData(1), JavaData(2)),
+        (JavaData(2), JavaData(1)),
+        (JavaData(2), JavaData(2))))
+  }
 }
 
 
@@ -406,3 +425,16 @@ class KryoData(val a: Int) {
 object KryoData {
   def apply(a: Int): KryoData = new KryoData(a)
 }
+
+/** Used to test Java encoder. */
+class JavaData(val a: Int) extends Serializable {
+  override def equals(other: Any): Boolean = {
+    a == other.asInstanceOf[JavaData].a
+  }
+  override def hashCode: Int = a
+  override def toString: String = s"JavaData($a)"
+}
+
+object JavaData {
+  def apply(a: Int): JavaData = new JavaData(a)
+}

From 7e987de1770f4ab3d54bc05db8de0a1ef035941d Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Wed, 18 Nov 2015 15:47:49 -0800
Subject: [PATCH 342/518] [SPARK-6787][ML] add read/write to estimators under
 ml.feature (1)

Add read/write support to the following estimators under spark.ml:

* CountVectorizer
* IDF
* MinMaxScaler
* StandardScaler (a little awkward because we store some params in spark.mllib model)
* StringIndexer

Added some necessary method for read/write. Maybe we should add `private[ml] trait DefaultParamsReadable` and `DefaultParamsWritable` to save some boilerplate code, though we still need to override `load` for Java compatibility.

jkbradley

Author: Xiangrui Meng <meng@databricks.com>

Closes #9798 from mengxr/SPARK-6787.
---
 .../spark/ml/feature/CountVectorizer.scala    | 72 +++++++++++++++--
 .../org/apache/spark/ml/feature/IDF.scala     | 71 ++++++++++++++++-
 .../spark/ml/feature/MinMaxScaler.scala       | 72 +++++++++++++++--
 .../spark/ml/feature/StandardScaler.scala     | 78 ++++++++++++++++++-
 .../spark/ml/feature/StringIndexer.scala      | 70 +++++++++++++++--
 .../ml/feature/CountVectorizerSuite.scala     | 24 +++++-
 .../apache/spark/ml/feature/IDFSuite.scala    | 19 ++++-
 .../spark/ml/feature/MinMaxScalerSuite.scala  | 25 +++++-
 .../ml/feature/StandardScalerSuite.scala      | 64 +++++++++++----
 .../spark/ml/feature/StringIndexerSuite.scala | 19 ++++-
 10 files changed, 467 insertions(+), 47 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
index 49028e4b85064..5ff9bfb7d1119 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
@@ -16,17 +16,19 @@
  */
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.broadcast.Broadcast
+import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
-import org.apache.spark.ml.{Estimator, Model}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg.{VectorUDT, Vectors}
 import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
-import org.apache.spark.sql.DataFrame
 import org.apache.spark.util.collection.OpenHashMap
 
 /**
@@ -105,7 +107,7 @@ private[feature] trait CountVectorizerParams extends Params with HasInputCol wit
  */
 @Experimental
 class CountVectorizer(override val uid: String)
-  extends Estimator[CountVectorizerModel] with CountVectorizerParams {
+  extends Estimator[CountVectorizerModel] with CountVectorizerParams with Writable {
 
   def this() = this(Identifiable.randomUID("cntVec"))
 
@@ -169,6 +171,19 @@ class CountVectorizer(override val uid: String)
   }
 
   override def copy(extra: ParamMap): CountVectorizer = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object CountVectorizer extends Readable[CountVectorizer] {
+
+  @Since("1.6.0")
+  override def read: Reader[CountVectorizer] = new DefaultParamsReader
+
+  @Since("1.6.0")
+  override def load(path: String): CountVectorizer = super.load(path)
 }
 
 /**
@@ -178,7 +193,9 @@ class CountVectorizer(override val uid: String)
  */
 @Experimental
 class CountVectorizerModel(override val uid: String, val vocabulary: Array[String])
-  extends Model[CountVectorizerModel] with CountVectorizerParams {
+  extends Model[CountVectorizerModel] with CountVectorizerParams with Writable {
+
+  import CountVectorizerModel._
 
   def this(vocabulary: Array[String]) = {
     this(Identifiable.randomUID("cntVecModel"), vocabulary)
@@ -232,4 +249,47 @@ class CountVectorizerModel(override val uid: String, val vocabulary: Array[Strin
     val copied = new CountVectorizerModel(uid, vocabulary).setParent(parent)
     copyValues(copied, extra)
   }
+
+  @Since("1.6.0")
+  override def write: Writer = new CountVectorizerModelWriter(this)
+}
+
+@Since("1.6.0")
+object CountVectorizerModel extends Readable[CountVectorizerModel] {
+
+  private[CountVectorizerModel]
+  class CountVectorizerModelWriter(instance: CountVectorizerModel) extends Writer {
+
+    private case class Data(vocabulary: Seq[String])
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = Data(instance.vocabulary)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class CountVectorizerModelReader extends Reader[CountVectorizerModel] {
+
+    private val className = "org.apache.spark.ml.feature.CountVectorizerModel"
+
+    override def load(path: String): CountVectorizerModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath)
+        .select("vocabulary")
+        .head()
+      val vocabulary = data.getAs[Seq[String]](0).toArray
+      val model = new CountVectorizerModel(metadata.uid, vocabulary)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: Reader[CountVectorizerModel] = new CountVectorizerModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): CountVectorizerModel = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
index 4c36df75d8aa0..53ad34ef12646 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
@@ -17,11 +17,13 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml._
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.feature
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT}
 import org.apache.spark.sql._
@@ -60,7 +62,7 @@ private[feature] trait IDFBase extends Params with HasInputCol with HasOutputCol
  * Compute the Inverse Document Frequency (IDF) given a collection of documents.
  */
 @Experimental
-final class IDF(override val uid: String) extends Estimator[IDFModel] with IDFBase {
+final class IDF(override val uid: String) extends Estimator[IDFModel] with IDFBase with Writable {
 
   def this() = this(Identifiable.randomUID("idf"))
 
@@ -85,6 +87,19 @@ final class IDF(override val uid: String) extends Estimator[IDFModel] with IDFBa
   }
 
   override def copy(extra: ParamMap): IDF = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object IDF extends Readable[IDF] {
+
+  @Since("1.6.0")
+  override def read: Reader[IDF] = new DefaultParamsReader
+
+  @Since("1.6.0")
+  override def load(path: String): IDF = super.load(path)
 }
 
 /**
@@ -95,7 +110,9 @@ final class IDF(override val uid: String) extends Estimator[IDFModel] with IDFBa
 class IDFModel private[ml] (
     override val uid: String,
     idfModel: feature.IDFModel)
-  extends Model[IDFModel] with IDFBase {
+  extends Model[IDFModel] with IDFBase with Writable {
+
+  import IDFModel._
 
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
@@ -117,4 +134,50 @@ class IDFModel private[ml] (
     val copied = new IDFModel(uid, idfModel)
     copyValues(copied, extra).setParent(parent)
   }
+
+  /** Returns the IDF vector. */
+  @Since("1.6.0")
+  def idf: Vector = idfModel.idf
+
+  @Since("1.6.0")
+  override def write: Writer = new IDFModelWriter(this)
+}
+
+@Since("1.6.0")
+object IDFModel extends Readable[IDFModel] {
+
+  private[IDFModel] class IDFModelWriter(instance: IDFModel) extends Writer {
+
+    private case class Data(idf: Vector)
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = Data(instance.idf)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class IDFModelReader extends Reader[IDFModel] {
+
+    private val className = "org.apache.spark.ml.feature.IDFModel"
+
+    override def load(path: String): IDFModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath)
+        .select("idf")
+        .head()
+      val idf = data.getAs[Vector](0)
+      val model = new IDFModel(metadata.uid, new feature.IDFModel(idf))
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: Reader[IDFModel] = new IDFModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): IDFModel = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
index 1b494ec8b1727..24d964fae834e 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
@@ -17,11 +17,14 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
-import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
-import org.apache.spark.ml.param.{ParamMap, DoubleParam, Params}
-import org.apache.spark.ml.util.Identifiable
+
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.{Estimator, Model}
+import org.apache.spark.ml.param.{DoubleParam, ParamMap, Params}
+import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT, Vectors}
 import org.apache.spark.mllib.stat.Statistics
 import org.apache.spark.sql._
@@ -85,7 +88,7 @@ private[feature] trait MinMaxScalerParams extends Params with HasInputCol with H
  */
 @Experimental
 class MinMaxScaler(override val uid: String)
-  extends Estimator[MinMaxScalerModel] with MinMaxScalerParams {
+  extends Estimator[MinMaxScalerModel] with MinMaxScalerParams with Writable {
 
   def this() = this(Identifiable.randomUID("minMaxScal"))
 
@@ -115,6 +118,19 @@ class MinMaxScaler(override val uid: String)
   }
 
   override def copy(extra: ParamMap): MinMaxScaler = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object MinMaxScaler extends Readable[MinMaxScaler] {
+
+  @Since("1.6.0")
+  override def read: Reader[MinMaxScaler] = new DefaultParamsReader
+
+  @Since("1.6.0")
+  override def load(path: String): MinMaxScaler = super.load(path)
 }
 
 /**
@@ -131,7 +147,9 @@ class MinMaxScalerModel private[ml] (
     override val uid: String,
     val originalMin: Vector,
     val originalMax: Vector)
-  extends Model[MinMaxScalerModel] with MinMaxScalerParams {
+  extends Model[MinMaxScalerModel] with MinMaxScalerParams with Writable {
+
+  import MinMaxScalerModel._
 
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
@@ -175,4 +193,46 @@ class MinMaxScalerModel private[ml] (
     val copied = new MinMaxScalerModel(uid, originalMin, originalMax)
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: Writer = new MinMaxScalerModelWriter(this)
+}
+
+@Since("1.6.0")
+object MinMaxScalerModel extends Readable[MinMaxScalerModel] {
+
+  private[MinMaxScalerModel]
+  class MinMaxScalerModelWriter(instance: MinMaxScalerModel) extends Writer {
+
+    private case class Data(originalMin: Vector, originalMax: Vector)
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = new Data(instance.originalMin, instance.originalMax)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class MinMaxScalerModelReader extends Reader[MinMaxScalerModel] {
+
+    private val className = "org.apache.spark.ml.feature.MinMaxScalerModel"
+
+    override def load(path: String): MinMaxScalerModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val Row(originalMin: Vector, originalMax: Vector) = sqlContext.read.parquet(dataPath)
+        .select("originalMin", "originalMax")
+        .head()
+      val model = new MinMaxScalerModel(metadata.uid, originalMin, originalMax)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: Reader[MinMaxScalerModel] = new MinMaxScalerModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): MinMaxScalerModel = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
index f6d0b0c0e9e75..ab04e5418dd4f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
@@ -17,11 +17,13 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml._
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.feature
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT}
 import org.apache.spark.sql._
@@ -57,7 +59,7 @@ private[feature] trait StandardScalerParams extends Params with HasInputCol with
  */
 @Experimental
 class StandardScaler(override val uid: String) extends Estimator[StandardScalerModel]
-  with StandardScalerParams {
+  with StandardScalerParams with Writable {
 
   def this() = this(Identifiable.randomUID("stdScal"))
 
@@ -94,6 +96,19 @@ class StandardScaler(override val uid: String) extends Estimator[StandardScalerM
   }
 
   override def copy(extra: ParamMap): StandardScaler = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object StandardScaler extends Readable[StandardScaler] {
+
+  @Since("1.6.0")
+  override def read: Reader[StandardScaler] = new DefaultParamsReader
+
+  @Since("1.6.0")
+  override def load(path: String): StandardScaler = super.load(path)
 }
 
 /**
@@ -104,7 +119,9 @@ class StandardScaler(override val uid: String) extends Estimator[StandardScalerM
 class StandardScalerModel private[ml] (
     override val uid: String,
     scaler: feature.StandardScalerModel)
-  extends Model[StandardScalerModel] with StandardScalerParams {
+  extends Model[StandardScalerModel] with StandardScalerParams with Writable {
+
+  import StandardScalerModel._
 
   /** Standard deviation of the StandardScalerModel */
   val std: Vector = scaler.std
@@ -112,6 +129,14 @@ class StandardScalerModel private[ml] (
   /** Mean of the StandardScalerModel */
   val mean: Vector = scaler.mean
 
+  /** Whether to scale to unit standard deviation. */
+  @Since("1.6.0")
+  def getWithStd: Boolean = scaler.withStd
+
+  /** Whether to center data with mean. */
+  @Since("1.6.0")
+  def getWithMean: Boolean = scaler.withMean
+
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
 
@@ -138,4 +163,49 @@ class StandardScalerModel private[ml] (
     val copied = new StandardScalerModel(uid, scaler)
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: Writer = new StandardScalerModelWriter(this)
+}
+
+@Since("1.6.0")
+object StandardScalerModel extends Readable[StandardScalerModel] {
+
+  private[StandardScalerModel]
+  class StandardScalerModelWriter(instance: StandardScalerModel) extends Writer {
+
+    private case class Data(std: Vector, mean: Vector, withStd: Boolean, withMean: Boolean)
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = Data(instance.std, instance.mean, instance.getWithStd, instance.getWithMean)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class StandardScalerModelReader extends Reader[StandardScalerModel] {
+
+    private val className = "org.apache.spark.ml.feature.StandardScalerModel"
+
+    override def load(path: String): StandardScalerModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val Row(std: Vector, mean: Vector, withStd: Boolean, withMean: Boolean) =
+        sqlContext.read.parquet(dataPath)
+          .select("std", "mean", "withStd", "withMean")
+          .head()
+      // This is very likely to change in the future because withStd and withMean should be params.
+      val oldModel = new feature.StandardScalerModel(std, mean, withStd, withMean)
+      val model = new StandardScalerModel(metadata.uid, oldModel)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: Reader[StandardScalerModel] = new StandardScalerModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): StandardScalerModel = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
index f782a272d11db..f16f6afc002d8 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
@@ -17,13 +17,14 @@
 
 package org.apache.spark.ml.feature
 
+import org.apache.hadoop.fs.Path
+
 import org.apache.spark.SparkException
-import org.apache.spark.annotation.{Since, Experimental}
-import org.apache.spark.ml.{Estimator, Model}
+import org.apache.spark.annotation.{Experimental, Since}
+import org.apache.spark.ml.{Estimator, Model, Transformer}
 import org.apache.spark.ml.attribute.{Attribute, NominalAttribute}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.Transformer
 import org.apache.spark.ml.util._
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions._
@@ -64,7 +65,7 @@ private[feature] trait StringIndexerBase extends Params with HasInputCol with Ha
  */
 @Experimental
 class StringIndexer(override val uid: String) extends Estimator[StringIndexerModel]
-  with StringIndexerBase {
+  with StringIndexerBase with Writable {
 
   def this() = this(Identifiable.randomUID("strIdx"))
 
@@ -92,6 +93,19 @@ class StringIndexer(override val uid: String) extends Estimator[StringIndexerMod
   }
 
   override def copy(extra: ParamMap): StringIndexer = defaultCopy(extra)
+
+  @Since("1.6.0")
+  override def write: Writer = new DefaultParamsWriter(this)
+}
+
+@Since("1.6.0")
+object StringIndexer extends Readable[StringIndexer] {
+
+  @Since("1.6.0")
+  override def read: Reader[StringIndexer] = new DefaultParamsReader
+
+  @Since("1.6.0")
+  override def load(path: String): StringIndexer = super.load(path)
 }
 
 /**
@@ -107,7 +121,10 @@ class StringIndexer(override val uid: String) extends Estimator[StringIndexerMod
 @Experimental
 class StringIndexerModel (
     override val uid: String,
-    val labels: Array[String]) extends Model[StringIndexerModel] with StringIndexerBase {
+    val labels: Array[String])
+  extends Model[StringIndexerModel] with StringIndexerBase with Writable {
+
+  import StringIndexerModel._
 
   def this(labels: Array[String]) = this(Identifiable.randomUID("strIdx"), labels)
 
@@ -176,6 +193,49 @@ class StringIndexerModel (
     val copied = new StringIndexerModel(uid, labels)
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: StringIndexModelWriter = new StringIndexModelWriter(this)
+}
+
+@Since("1.6.0")
+object StringIndexerModel extends Readable[StringIndexerModel] {
+
+  private[StringIndexerModel]
+  class StringIndexModelWriter(instance: StringIndexerModel) extends Writer {
+
+    private case class Data(labels: Array[String])
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = Data(instance.labels)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class StringIndexerModelReader extends Reader[StringIndexerModel] {
+
+    private val className = "org.apache.spark.ml.feature.StringIndexerModel"
+
+    override def load(path: String): StringIndexerModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath)
+        .select("labels")
+        .head()
+      val labels = data.getAs[Seq[String]](0).toArray
+      val model = new StringIndexerModel(metadata.uid, labels)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: Reader[StringIndexerModel] = new StringIndexerModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): StringIndexerModel = super.load(path)
 }
 
 /**
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala
index e192fa4850af0..9c9999017317d 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala
@@ -18,14 +18,17 @@ package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.Row
 
-class CountVectorizerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class CountVectorizerSuite extends SparkFunSuite with MLlibTestSparkContext
+  with DefaultReadWriteTest {
 
   test("params") {
+    ParamsSuite.checkParams(new CountVectorizer)
     ParamsSuite.checkParams(new CountVectorizerModel(Array("empty")))
   }
 
@@ -164,4 +167,23 @@ class CountVectorizerSuite extends SparkFunSuite with MLlibTestSparkContext {
         assert(features ~== expected absTol 1e-14)
     }
   }
+
+  test("CountVectorizer read/write") {
+    val t = new CountVectorizer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setMinDF(0.5)
+      .setMinTF(3.0)
+      .setVocabSize(10)
+    testDefaultReadWrite(t)
+  }
+
+  test("CountVectorizerModel read/write") {
+    val instance = new CountVectorizerModel("myCountVectorizerModel", Array("a", "b", "c"))
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setMinTF(3.0)
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.vocabulary === instance.vocabulary)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/IDFSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/IDFSuite.scala
index 08f80af03429b..bc958c15857ba 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/IDFSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/IDFSuite.scala
@@ -19,13 +19,14 @@ package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.feature.{IDFModel => OldIDFModel}
 import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.Row
 
-class IDFSuite extends SparkFunSuite with MLlibTestSparkContext {
+class IDFSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   def scaleDataWithIDF(dataSet: Array[Vector], model: Vector): Array[Vector] = {
     dataSet.map {
@@ -98,4 +99,20 @@ class IDFSuite extends SparkFunSuite with MLlibTestSparkContext {
         assert(x ~== y absTol 1e-5, "Transformed vector is different with expected vector.")
     }
   }
+
+  test("IDF read/write") {
+    val t = new IDF()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setMinDocFreq(5)
+    testDefaultReadWrite(t)
+  }
+
+  test("IDFModel read/write") {
+    val instance = new IDFModel("myIDFModel", new OldIDFModel(Vectors.dense(1.0, 2.0)))
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.idf === instance.idf)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/MinMaxScalerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/MinMaxScalerSuite.scala
index c04dda41eea34..09183fe65b722 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/MinMaxScalerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/MinMaxScalerSuite.scala
@@ -18,12 +18,12 @@
 package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{Row, SQLContext}
 
-class MinMaxScalerSuite extends SparkFunSuite with MLlibTestSparkContext {
+class MinMaxScalerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("MinMaxScaler fit basic case") {
     val sqlContext = new SQLContext(sc)
@@ -69,4 +69,25 @@ class MinMaxScalerSuite extends SparkFunSuite with MLlibTestSparkContext {
       }
     }
   }
+
+  test("MinMaxScaler read/write") {
+    val t = new MinMaxScaler()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setMax(1.0)
+      .setMin(-1.0)
+    testDefaultReadWrite(t)
+  }
+
+  test("MinMaxScalerModel read/write") {
+    val instance = new MinMaxScalerModel(
+        "myMinMaxScalerModel", Vectors.dense(-1.0, 0.0), Vectors.dense(1.0, 10.0))
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setMin(-1.0)
+      .setMax(1.0)
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.originalMin === instance.originalMin)
+    assert(newInstance.originalMax === instance.originalMax)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala
index 879a3ae875004..49a4b2efe0c29 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala
@@ -19,12 +19,16 @@ package org.apache.spark.ml.feature
 
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector, Vectors}
+import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
+import org.apache.spark.mllib.feature
+import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
-import org.apache.spark.sql.{DataFrame, Row, SQLContext}
+import org.apache.spark.sql.{DataFrame, Row}
 
-class StandardScalerSuite extends SparkFunSuite with MLlibTestSparkContext{
+class StandardScalerSuite extends SparkFunSuite with MLlibTestSparkContext
+  with DefaultReadWriteTest {
 
   @transient var data: Array[Vector] = _
   @transient var resWithStd: Array[Vector] = _
@@ -56,23 +60,29 @@ class StandardScalerSuite extends SparkFunSuite with MLlibTestSparkContext{
     )
   }
 
-  def assertResult(dataframe: DataFrame): Unit = {
-    dataframe.select("standarded_features", "expected").collect().foreach {
+  def assertResult(df: DataFrame): Unit = {
+    df.select("standardized_features", "expected").collect().foreach {
       case Row(vector1: Vector, vector2: Vector) =>
         assert(vector1 ~== vector2 absTol 1E-5,
           "The vector value is not correct after standardization.")
     }
   }
 
+  test("params") {
+    ParamsSuite.checkParams(new StandardScaler)
+    val oldModel = new feature.StandardScalerModel(Vectors.dense(1.0), Vectors.dense(2.0))
+    ParamsSuite.checkParams(new StandardScalerModel("empty", oldModel))
+  }
+
   test("Standardization with default parameter") {
     val df0 = sqlContext.createDataFrame(data.zip(resWithStd)).toDF("features", "expected")
 
-    val standardscaler0 = new StandardScaler()
+    val standardScaler0 = new StandardScaler()
       .setInputCol("features")
-      .setOutputCol("standarded_features")
+      .setOutputCol("standardized_features")
       .fit(df0)
 
-    assertResult(standardscaler0.transform(df0))
+    assertResult(standardScaler0.transform(df0))
   }
 
   test("Standardization with setter") {
@@ -80,29 +90,49 @@ class StandardScalerSuite extends SparkFunSuite with MLlibTestSparkContext{
     val df2 = sqlContext.createDataFrame(data.zip(resWithMean)).toDF("features", "expected")
     val df3 = sqlContext.createDataFrame(data.zip(data)).toDF("features", "expected")
 
-    val standardscaler1 = new StandardScaler()
+    val standardScaler1 = new StandardScaler()
       .setInputCol("features")
-      .setOutputCol("standarded_features")
+      .setOutputCol("standardized_features")
       .setWithMean(true)
       .setWithStd(true)
       .fit(df1)
 
-    val standardscaler2 = new StandardScaler()
+    val standardScaler2 = new StandardScaler()
       .setInputCol("features")
-      .setOutputCol("standarded_features")
+      .setOutputCol("standardized_features")
       .setWithMean(true)
       .setWithStd(false)
       .fit(df2)
 
-    val standardscaler3 = new StandardScaler()
+    val standardScaler3 = new StandardScaler()
       .setInputCol("features")
-      .setOutputCol("standarded_features")
+      .setOutputCol("standardized_features")
       .setWithMean(false)
       .setWithStd(false)
       .fit(df3)
 
-    assertResult(standardscaler1.transform(df1))
-    assertResult(standardscaler2.transform(df2))
-    assertResult(standardscaler3.transform(df3))
+    assertResult(standardScaler1.transform(df1))
+    assertResult(standardScaler2.transform(df2))
+    assertResult(standardScaler3.transform(df3))
+  }
+
+  test("StandardScaler read/write") {
+    val t = new StandardScaler()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setWithStd(false)
+      .setWithMean(true)
+    testDefaultReadWrite(t)
+  }
+
+  test("StandardScalerModel read/write") {
+    val oldModel = new feature.StandardScalerModel(
+      Vectors.dense(1.0, 2.0), Vectors.dense(3.0, 4.0), false, true)
+    val instance = new StandardScalerModel("myStandardScalerModel", oldModel)
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.std === instance.std)
+    assert(newInstance.mean === instance.mean)
+    assert(newInstance.getWithStd === instance.getWithStd)
+    assert(newInstance.getWithMean === instance.getWithMean)
   }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
index be37bfb438833..749bfac747826 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
@@ -118,6 +118,23 @@ class StringIndexerSuite
     assert(indexerModel.transform(df).eq(df))
   }
 
+  test("StringIndexer read/write") {
+    val t = new StringIndexer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setHandleInvalid("skip")
+    testDefaultReadWrite(t)
+  }
+
+  test("StringIndexerModel read/write") {
+    val instance = new StringIndexerModel("myStringIndexerModel", Array("a", "b", "c"))
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setHandleInvalid("skip")
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.labels === instance.labels)
+  }
+
   test("IndexToString params") {
     val idxToStr = new IndexToString()
     ParamsSuite.checkParams(idxToStr)
@@ -175,7 +192,7 @@ class StringIndexerSuite
     assert(outSchema("output").dataType === StringType)
   }
 
-  test("read/write") {
+  test("IndexToString read/write") {
     val t = new IndexToString()
       .setInputCol("myInputCol")
       .setOutputCol("myOutputCol")

From 3a9851936ddfe5bcb6a7f364d535fac977551f5d Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Wed, 18 Nov 2015 15:55:41 -0800
Subject: [PATCH 343/518] [SPARK-11649] Properly set Akka frame size in
 SparkListenerSuite test

SparkListenerSuite's _"onTaskGettingResult() called when result fetched remotely"_ test was extremely slow (1 to 4 minutes to run) and recently became extremely flaky, frequently failing with OutOfMemoryError.

The root cause was the fact that this was using `System.setProperty` to set the Akka frame size, which was not actually modifying the frame size. As a result, this test would allocate much more data than necessary. The fix here is to simply use SparkConf in order to configure the frame size.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9822 from JoshRosen/SPARK-11649.
---
 .../org/apache/spark/scheduler/SparkListenerSuite.scala      | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
index 53102b9f1c936..84e545851f49e 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
@@ -269,14 +269,15 @@ class SparkListenerSuite extends SparkFunSuite with LocalSparkContext with Match
   }
 
   test("onTaskGettingResult() called when result fetched remotely") {
-    sc = new SparkContext("local", "SparkListenerSuite")
+    val conf = new SparkConf().set("spark.akka.frameSize", "1")
+    sc = new SparkContext("local", "SparkListenerSuite", conf)
     val listener = new SaveTaskEvents
     sc.addSparkListener(listener)
 
     // Make a task whose result is larger than the akka frame size
-    System.setProperty("spark.akka.frameSize", "1")
     val akkaFrameSize =
       sc.env.actorSystem.settings.config.getBytes("akka.remote.netty.tcp.maximum-frame-size").toInt
+    assert(akkaFrameSize === 1024 * 1024)
     val result = sc.parallelize(Seq(1), 1)
       .map { x => 1.to(akkaFrameSize).toArray }
       .reduce { case (x, y) => x }

From c07a50b86254578625be777b1890ff95e832ac6e Mon Sep 17 00:00:00 2001
From: Derek Dagit <derekd@yahoo-inc.com>
Date: Wed, 18 Nov 2015 15:56:54 -0800
Subject: [PATCH 344/518] [SPARK-10930] History "Stages" page "duration" can be
 confusing

Author: Derek Dagit <derekd@yahoo-inc.com>

Closes #9051 from d2r/spark-10930-ui-max-task-dur.
---
 .../org/apache/spark/ui/jobs/StageTable.scala | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/ui/jobs/StageTable.scala b/core/src/main/scala/org/apache/spark/ui/jobs/StageTable.scala
index ea806d09b6009..2a1c3c1a50ec9 100644
--- a/core/src/main/scala/org/apache/spark/ui/jobs/StageTable.scala
+++ b/core/src/main/scala/org/apache/spark/ui/jobs/StageTable.scala
@@ -145,9 +145,22 @@ private[ui] class StageTableBase(
       case None => "Unknown"
     }
     val finishTime = s.completionTime.getOrElse(System.currentTimeMillis)
-    val duration = s.submissionTime.map { t =>
-      if (finishTime > t) finishTime - t else System.currentTimeMillis - t
-    }
+
+    // The submission time for a stage is misleading because it counts the time
+    // the stage waits to be launched. (SPARK-10930)
+    val taskLaunchTimes =
+      stageData.taskData.values.map(_.taskInfo.launchTime).filter(_ > 0)
+    val duration: Option[Long] =
+      if (taskLaunchTimes.nonEmpty) {
+        val startTime = taskLaunchTimes.min
+        if (finishTime > startTime) {
+          Some(finishTime - startTime)
+        } else {
+          Some(System.currentTimeMillis() - startTime)
+        }
+      } else {
+        None
+      }
     val formattedDuration = duration.map(d => UIUtils.formatDuration(d)).getOrElse("Unknown")
 
     val inputRead = stageData.inputBytes

From 4b117121900e5f242e7c8f46a69164385f0da7cc Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Wed, 18 Nov 2015 16:00:35 -0800
Subject: [PATCH 345/518] [SPARK-11495] Fix potential socket / file handle
 leaks that were found via static analysis

The HP Fortify Opens Source Review team (https://www.hpfod.com/open-source-review-project) reported a handful of potential resource leaks that were discovered using their static analysis tool. We should fix the issues identified by their scan.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9455 from JoshRosen/fix-potential-resource-leaks.
---
 .../spark/unsafe/map/BytesToBytesMap.java     |  7 ++++
 .../unsafe/sort/UnsafeSorterSpillReader.java  | 38 +++++++++++--------
 .../streaming/JavaCustomReceiver.java         | 31 +++++++--------
 .../network/ChunkFetchIntegrationSuite.java   | 15 ++++++--
 .../shuffle/TestShuffleDataContext.java       | 32 ++++++++++------
 .../spark/streaming/JavaReceiverAPISuite.java | 20 ++++++----
 6 files changed, 90 insertions(+), 53 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
index 04694dc54418c..3387f9a4177ce 100644
--- a/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
+++ b/core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java
@@ -24,6 +24,7 @@
 import java.util.LinkedList;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.io.Closeables;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -272,6 +273,7 @@ private void advanceToNextPage() {
             }
           }
           try {
+            Closeables.close(reader, /* swallowIOException = */ false);
             reader = spillWriters.getFirst().getReader(blockManager);
             recordsInPage = -1;
           } catch (IOException e) {
@@ -318,6 +320,11 @@ public Location next() {
         try {
           reader.loadNext();
         } catch (IOException e) {
+          try {
+            reader.close();
+          } catch(IOException e2) {
+            logger.error("Error while closing spill reader", e2);
+          }
           // Scala iterator does not handle exception
           Platform.throwException(e);
         }
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillReader.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillReader.java
index 039e940a357ea..dcb13e6581e54 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillReader.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeSorterSpillReader.java
@@ -20,8 +20,7 @@
 import java.io.*;
 
 import com.google.common.io.ByteStreams;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
+import com.google.common.io.Closeables;
 
 import org.apache.spark.storage.BlockId;
 import org.apache.spark.storage.BlockManager;
@@ -31,10 +30,8 @@
  * Reads spill files written by {@link UnsafeSorterSpillWriter} (see that class for a description
  * of the file format).
  */
-public final class UnsafeSorterSpillReader extends UnsafeSorterIterator {
-  private static final Logger logger = LoggerFactory.getLogger(UnsafeSorterSpillReader.class);
+public final class UnsafeSorterSpillReader extends UnsafeSorterIterator implements Closeable {
 
-  private final File file;
   private InputStream in;
   private DataInputStream din;
 
@@ -52,11 +49,15 @@ public UnsafeSorterSpillReader(
       File file,
       BlockId blockId) throws IOException {
     assert (file.length() > 0);
-    this.file = file;
     final BufferedInputStream bs = new BufferedInputStream(new FileInputStream(file));
-    this.in = blockManager.wrapForCompression(blockId, bs);
-    this.din = new DataInputStream(this.in);
-    numRecordsRemaining = din.readInt();
+    try {
+      this.in = blockManager.wrapForCompression(blockId, bs);
+      this.din = new DataInputStream(this.in);
+      numRecordsRemaining = din.readInt();
+    } catch (IOException e) {
+      Closeables.close(bs, /* swallowIOException = */ true);
+      throw e;
+    }
   }
 
   @Override
@@ -75,12 +76,7 @@ public void loadNext() throws IOException {
     ByteStreams.readFully(in, arr, 0, recordLength);
     numRecordsRemaining--;
     if (numRecordsRemaining == 0) {
-      in.close();
-      if (!file.delete() && file.exists()) {
-        logger.warn("Unable to delete spill file {}", file.getPath());
-      }
-      in = null;
-      din = null;
+      close();
     }
   }
 
@@ -103,4 +99,16 @@ public int getRecordLength() {
   public long getKeyPrefix() {
     return keyPrefix;
   }
+
+  @Override
+  public void close() throws IOException {
+   if (in != null) {
+     try {
+       in.close();
+     } finally {
+       in = null;
+       din = null;
+     }
+   }
+  }
 }
diff --git a/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java b/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java
index 99df259b4e8e6..4b50fbf59f80e 100644
--- a/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java
+++ b/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java
@@ -18,6 +18,7 @@
 package org.apache.spark.examples.streaming;
 
 import com.google.common.collect.Lists;
+import com.google.common.io.Closeables;
 
 import org.apache.spark.SparkConf;
 import org.apache.spark.api.java.function.FlatMapFunction;
@@ -121,23 +122,23 @@ public void onStop() {
 
   /** Create a socket connection and receive data until receiver is stopped */
   private void receive() {
-    Socket socket = null;
-    String userInput = null;
-
     try {
-      // connect to the server
-      socket = new Socket(host, port);
-
-      BufferedReader reader = new BufferedReader(new InputStreamReader(socket.getInputStream()));
-
-      // Until stopped or connection broken continue reading
-      while (!isStopped() && (userInput = reader.readLine()) != null) {
-        System.out.println("Received data '" + userInput + "'");
-        store(userInput);
+      Socket socket = null;
+      BufferedReader reader = null;
+      String userInput = null;
+      try {
+        // connect to the server
+        socket = new Socket(host, port);
+        reader = new BufferedReader(new InputStreamReader(socket.getInputStream()));
+        // Until stopped or connection broken continue reading
+        while (!isStopped() && (userInput = reader.readLine()) != null) {
+          System.out.println("Received data '" + userInput + "'");
+          store(userInput);
+        }
+      } finally {
+        Closeables.close(reader, /* swallowIOException = */ true);
+        Closeables.close(socket,  /* swallowIOException = */ true);
       }
-      reader.close();
-      socket.close();
-
       // Restart in an attempt to connect again when server is active again
       restart("Trying to connect again");
     } catch(ConnectException ce) {
diff --git a/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java b/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
index dc5fa1cee69bc..50a324e293386 100644
--- a/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
@@ -31,6 +31,7 @@
 
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
+import com.google.common.io.Closeables;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -78,10 +79,15 @@ public static void setUp() throws Exception {
     testFile = File.createTempFile("shuffle-test-file", "txt");
     testFile.deleteOnExit();
     RandomAccessFile fp = new RandomAccessFile(testFile, "rw");
-    byte[] fileContent = new byte[1024];
-    new Random().nextBytes(fileContent);
-    fp.write(fileContent);
-    fp.close();
+    boolean shouldSuppressIOException = true;
+    try {
+      byte[] fileContent = new byte[1024];
+      new Random().nextBytes(fileContent);
+      fp.write(fileContent);
+      shouldSuppressIOException = false;
+    } finally {
+      Closeables.close(fp, shouldSuppressIOException);
+    }
 
     final TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     fileChunk = new FileSegmentManagedBuffer(conf, testFile, 10, testFile.length() - 25);
@@ -117,6 +123,7 @@ public StreamManager getStreamManager() {
 
   @AfterClass
   public static void tearDown() {
+    bufferChunk.release();
     server.close();
     clientFactory.close();
     testFile.delete();
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/TestShuffleDataContext.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/TestShuffleDataContext.java
index 3fdde054ab6c7..7ac1ca128aed0 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/TestShuffleDataContext.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/TestShuffleDataContext.java
@@ -23,6 +23,7 @@
 import java.io.IOException;
 import java.io.OutputStream;
 
+import com.google.common.io.Closeables;
 import com.google.common.io.Files;
 
 import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo;
@@ -60,21 +61,28 @@ public void cleanup() {
   public void insertSortShuffleData(int shuffleId, int mapId, byte[][] blocks) throws IOException {
     String blockId = "shuffle_" + shuffleId + "_" + mapId + "_0";
 
-    OutputStream dataStream = new FileOutputStream(
-      ExternalShuffleBlockResolver.getFile(localDirs, subDirsPerLocalDir, blockId + ".data"));
-    DataOutputStream indexStream = new DataOutputStream(new FileOutputStream(
-      ExternalShuffleBlockResolver.getFile(localDirs, subDirsPerLocalDir, blockId + ".index")));
+    OutputStream dataStream = null;
+    DataOutputStream indexStream = null;
+    boolean suppressExceptionsDuringClose = true;
 
-    long offset = 0;
-    indexStream.writeLong(offset);
-    for (byte[] block : blocks) {
-      offset += block.length;
-      dataStream.write(block);
+    try {
+      dataStream = new FileOutputStream(
+        ExternalShuffleBlockResolver.getFile(localDirs, subDirsPerLocalDir, blockId + ".data"));
+      indexStream = new DataOutputStream(new FileOutputStream(
+        ExternalShuffleBlockResolver.getFile(localDirs, subDirsPerLocalDir, blockId + ".index")));
+
+      long offset = 0;
       indexStream.writeLong(offset);
+      for (byte[] block : blocks) {
+        offset += block.length;
+        dataStream.write(block);
+        indexStream.writeLong(offset);
+      }
+      suppressExceptionsDuringClose = false;
+    } finally {
+      Closeables.close(dataStream, suppressExceptionsDuringClose);
+      Closeables.close(indexStream, suppressExceptionsDuringClose);
     }
-
-    dataStream.close();
-    indexStream.close();
   }
 
   /** Creates reducer blocks in a hash-based data format within our local dirs. */
diff --git a/streaming/src/test/java/org/apache/spark/streaming/JavaReceiverAPISuite.java b/streaming/src/test/java/org/apache/spark/streaming/JavaReceiverAPISuite.java
index ec2bffd6a5b97..7a8ef9d14784c 100644
--- a/streaming/src/test/java/org/apache/spark/streaming/JavaReceiverAPISuite.java
+++ b/streaming/src/test/java/org/apache/spark/streaming/JavaReceiverAPISuite.java
@@ -23,6 +23,7 @@
 import org.apache.spark.streaming.api.java.JavaStreamingContext;
 import static org.junit.Assert.*;
 
+import com.google.common.io.Closeables;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -121,14 +122,19 @@ public void onStop() {
 
     private void receive() {
       try {
-        Socket socket = new Socket(host, port);
-        BufferedReader in = new BufferedReader(new InputStreamReader(socket.getInputStream()));
-        String userInput;
-        while ((userInput = in.readLine()) != null) {
-          store(userInput);
+        Socket socket = null;
+        BufferedReader in = null;
+        try {
+          socket = new Socket(host, port);
+          in = new BufferedReader(new InputStreamReader(socket.getInputStream()));
+          String userInput;
+          while ((userInput = in.readLine()) != null) {
+            store(userInput);
+          }
+        } finally {
+          Closeables.close(in, /* swallowIOException = */ true);
+          Closeables.close(socket,  /* swallowIOException = */ true);
         }
-        in.close();
-        socket.close();
       } catch(ConnectException ce) {
         ce.printStackTrace();
         restart("Could not connect", ce);

From a402c92c92b2e1c85d264f6077aec8f6d6a08270 Mon Sep 17 00:00:00 2001
From: Tathagata Das <tathagata.das1565@gmail.com>
Date: Wed, 18 Nov 2015 16:08:06 -0800
Subject: [PATCH 346/518] [SPARK-11814][STREAMING] Add better default
 checkpoint duration

DStream checkpoint interval is by default set at max(10 second, batch interval). That's bad for large batch intervals where the checkpoint interval = batch interval, and RDDs get checkpointed every batch.
This PR is to set the checkpoint interval of trackStateByKey to 10 * batch duration.

Author: Tathagata Das <tathagata.das1565@gmail.com>

Closes #9805 from tdas/SPARK-11814.
---
 .../streaming/dstream/TrackStateDStream.scala | 13 ++++++
 .../streaming/TrackStateByKeySuite.scala      | 44 ++++++++++++++++++-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
index 98e881e6ae115..0ada1111ce30a 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala
@@ -25,6 +25,7 @@ import org.apache.spark.rdd.{EmptyRDD, RDD}
 import org.apache.spark.storage.StorageLevel
 import org.apache.spark.streaming._
 import org.apache.spark.streaming.rdd.{TrackStateRDD, TrackStateRDDRecord}
+import org.apache.spark.streaming.dstream.InternalTrackStateDStream._
 
 /**
  * :: Experimental ::
@@ -120,6 +121,14 @@ class InternalTrackStateDStream[K: ClassTag, V: ClassTag, S: ClassTag, E: ClassT
   /** Enable automatic checkpointing */
   override val mustCheckpoint = true
 
+  /** Override the default checkpoint duration */
+  override def initialize(time: Time): Unit = {
+    if (checkpointDuration == null) {
+      checkpointDuration = slideDuration * DEFAULT_CHECKPOINT_DURATION_MULTIPLIER
+    }
+    super.initialize(time)
+  }
+
   /** Method that generates a RDD for the given time */
   override def compute(validTime: Time): Option[RDD[TrackStateRDDRecord[K, S, E]]] = {
     // Get the previous state or create a new empty state RDD
@@ -141,3 +150,7 @@ class InternalTrackStateDStream[K: ClassTag, V: ClassTag, S: ClassTag, E: ClassT
     }
   }
 }
+
+private[streaming] object InternalTrackStateDStream {
+  private val DEFAULT_CHECKPOINT_DURATION_MULTIPLIER = 10
+}
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala
index e3072b4442840..58aef74c0040f 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala
@@ -22,9 +22,10 @@ import java.io.File
 import scala.collection.mutable.{ArrayBuffer, SynchronizedBuffer}
 import scala.reflect.ClassTag
 
+import org.scalatest.PrivateMethodTester._
 import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll}
 
-import org.apache.spark.streaming.dstream.{TrackStateDStream, TrackStateDStreamImpl}
+import org.apache.spark.streaming.dstream.{InternalTrackStateDStream, TrackStateDStream, TrackStateDStreamImpl}
 import org.apache.spark.util.{ManualClock, Utils}
 import org.apache.spark.{SparkConf, SparkContext, SparkFunSuite}
 
@@ -57,6 +58,12 @@ class TrackStateByKeySuite extends SparkFunSuite with BeforeAndAfterAll with Bef
     sc = new SparkContext(conf)
   }
 
+  override def afterAll(): Unit = {
+    if (sc != null) {
+      sc.stop()
+    }
+  }
+
   test("state - get, exists, update, remove, ") {
     var state: StateImpl[Int] = null
 
@@ -436,6 +443,41 @@ class TrackStateByKeySuite extends SparkFunSuite with BeforeAndAfterAll with Bef
     assert(collectedStateSnapshots.last.toSet === Set(("a", 1)))
   }
 
+  test("trackStateByKey - checkpoint durations") {
+    val privateMethod = PrivateMethod[InternalTrackStateDStream[_, _, _, _]]('internalStream)
+
+    def testCheckpointDuration(
+        batchDuration: Duration,
+        expectedCheckpointDuration: Duration,
+        explicitCheckpointDuration: Option[Duration] = None
+      ): Unit = {
+      try {
+        ssc = new StreamingContext(sc, batchDuration)
+        val inputStream = new TestInputStream(ssc, Seq.empty[Seq[Int]], 2).map(_ -> 1)
+        val dummyFunc = (value: Option[Int], state: State[Int]) => 0
+        val trackStateStream = inputStream.trackStateByKey(StateSpec.function(dummyFunc))
+        val internalTrackStateStream = trackStateStream invokePrivate privateMethod()
+
+        explicitCheckpointDuration.foreach { d =>
+          trackStateStream.checkpoint(d)
+        }
+        trackStateStream.register()
+        ssc.start()  // should initialize all the checkpoint durations
+        assert(trackStateStream.checkpointDuration === null)
+        assert(internalTrackStateStream.checkpointDuration === expectedCheckpointDuration)
+      } finally {
+        StreamingContext.getActive().foreach { _.stop(stopSparkContext = false) }
+      }
+    }
+
+    testCheckpointDuration(Milliseconds(100), Seconds(1))
+    testCheckpointDuration(Seconds(1), Seconds(10))
+    testCheckpointDuration(Seconds(10), Seconds(100))
+
+    testCheckpointDuration(Milliseconds(100), Seconds(2), Some(Seconds(2)))
+    testCheckpointDuration(Seconds(1), Seconds(2), Some(Seconds(2)))
+    testCheckpointDuration(Seconds(10), Seconds(20), Some(Seconds(20)))
+  }
 
   private def testOperation[K: ClassTag, S: ClassTag, T: ClassTag](
       input: Seq[Seq[K]],

From 921900fd06362474f8caac675803d526a0986d70 Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Wed, 18 Nov 2015 16:19:00 -0800
Subject: [PATCH 347/518] [SPARK-11791] Fix flaky test in
 BatchedWriteAheadLogSuite

stack trace of failure:
```
org.scalatest.exceptions.TestFailedDueToTimeoutException: The code passed to eventually never returned normally. Attempted 62 times over 1.006322071 seconds. Last failure message:
Argument(s) are different! Wanted:
writeAheadLog.write(
    java.nio.HeapByteBuffer[pos=0 lim=124 cap=124],
    10
);
-> at org.apache.spark.streaming.util.BatchedWriteAheadLogSuite$$anonfun$23$$anonfun$apply$mcV$sp$15.apply(WriteAheadLogSuite.scala:518)
Actual invocation has different arguments:
writeAheadLog.write(
    java.nio.HeapByteBuffer[pos=0 lim=124 cap=124],
    10
);
-> at org.apache.spark.streaming.util.WriteAheadLogSuite$BlockingWriteAheadLog.write(WriteAheadLogSuite.scala:756)
```

I believe the issue was that due to a race condition, the ordering of the events could be messed up in the final ByteBuffer, therefore the comparison fails.

By adding eventually between the requests, we make sure the ordering is preserved. Note that in real life situations, the ordering across threads will not matter.

Another solution would be to implement a custom mockito matcher that sorts and then compares the results, but that kind of sounds like overkill to me. Let me know what you think tdas zsxwing

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9790 from brkyvz/fix-flaky-2.
---
 .../spark/streaming/util/WriteAheadLogSuite.scala    | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
index 7f80d6ecdbbb5..eaa88ea3cd380 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/util/WriteAheadLogSuite.scala
@@ -30,6 +30,7 @@ import scala.language.{implicitConversions, postfixOps}
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
+import org.mockito.ArgumentCaptor
 import org.mockito.Matchers.{eq => meq}
 import org.mockito.Matchers._
 import org.mockito.Mockito._
@@ -507,15 +508,18 @@ class BatchedWriteAheadLogSuite extends CommonWriteAheadLogTests(
     }
     blockingWal.allowWrite()
 
-    val buffer1 = wrapArrayArrayByte(Array(event1))
-    val buffer2 = wrapArrayArrayByte(Array(event2, event3, event4, event5))
+    val buffer = wrapArrayArrayByte(Array(event1))
+    val queuedEvents = Set(event2, event3, event4, event5)
 
     eventually(timeout(1 second)) {
       assert(batchedWal.invokePrivate(queueLength()) === 0)
-      verify(wal, times(1)).write(meq(buffer1), meq(3L))
+      verify(wal, times(1)).write(meq(buffer), meq(3L))
       // the file name should be the timestamp of the last record, as events should be naturally
       // in order of timestamp, and we need the last element.
-      verify(wal, times(1)).write(meq(buffer2), meq(10L))
+      val bufferCaptor = ArgumentCaptor.forClass(classOf[ByteBuffer])
+      verify(wal, times(1)).write(bufferCaptor.capture(), meq(10L))
+      val records = BatchedWriteAheadLog.deaggregate(bufferCaptor.getValue).map(byteBufferToString)
+      assert(records.toSet === queuedEvents)
     }
   }
 

From 59a501359a267fbdb7689058693aa788703e54b1 Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Wed, 18 Nov 2015 16:48:09 -0800
Subject: [PATCH 348/518] [SPARK-11636][SQL] Support classes defined in the
 REPL with Encoders

Before this PR there were two things that would blow up if you called `df.as[MyClass]` if `MyClass` was defined in the REPL:
 - [x] Because `classForName` doesn't work on the munged names returned by `tpe.erasure.typeSymbol.asClass.fullName`
 - [x] Because we don't have anything to pass into the constructor for the `$outer` pointer.

Note that this PR is just adding the infrastructure for working with inner classes in encoder and is not yet sufficient to make them work in the REPL.  Currently, the implementation show in https://github.com/marmbrus/spark/commit/95cec7d413b930b36420724fafd829bef8c732ab is causing a bug that breaks code gen due to some interaction between janino and the `ExecutorClassLoader`.  This will be addressed in a follow-up PR.

Author: Michael Armbrust <michael@databricks.com>

Closes #9602 from marmbrus/dataset-replClasses.
---
 .../spark/sql/catalyst/ScalaReflection.scala  | 81 ++++++++++---------
 .../catalyst/encoders/ExpressionEncoder.scala | 26 +++++-
 .../sql/catalyst/encoders/OuterScopes.scala   | 42 ++++++++++
 .../catalyst/encoders/ProductEncoder.scala    |  6 +-
 .../expressions/codegen/CodegenFallback.scala |  2 +-
 .../codegen/GenerateMutableProjection.scala   |  4 +-
 .../codegen/GenerateProjection.scala          | 10 +--
 .../codegen/GenerateSafeProjection.scala      |  4 +-
 .../codegen/GenerateUnsafeProjection.scala    |  4 +-
 .../codegen/GenerateUnsafeRowJoiner.scala     |  6 +-
 .../sql/catalyst/expressions/literals.scala   |  6 ++
 .../sql/catalyst/expressions/objects.scala    | 42 ++++++++--
 .../encoders/ExpressionEncoderSuite.scala     |  7 +-
 .../encoders/ProductEncoderSuite.scala        |  4 +
 .../scala/org/apache/spark/sql/Dataset.scala  |  4 +-
 .../org/apache/spark/sql/GroupedDataset.scala |  8 +-
 .../aggregate/TypedAggregateExpression.scala  | 19 ++---
 17 files changed, 193 insertions(+), 82 deletions(-)
 create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/OuterScopes.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 38828e59a2152..59ccf356f2c48 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -35,17 +35,6 @@ object ScalaReflection extends ScalaReflection {
   // class loader of the current thread.
   override def mirror: universe.Mirror =
     universe.runtimeMirror(Thread.currentThread().getContextClassLoader)
-}
-
-/**
- * Support for generating catalyst schemas for scala objects.
- */
-trait ScalaReflection {
-  /** The universe we work in (runtime or macro) */
-  val universe: scala.reflect.api.Universe
-
-  /** The mirror used to access types in the universe */
-  def mirror: universe.Mirror
 
   import universe._
 
@@ -53,30 +42,6 @@ trait ScalaReflection {
   // Since the map values can be mutable, we explicitly import scala.collection.Map at here.
   import scala.collection.Map
 
-  case class Schema(dataType: DataType, nullable: Boolean)
-
-  /** Returns a Sequence of attributes for the given case class type. */
-  def attributesFor[T: TypeTag]: Seq[Attribute] = schemaFor[T] match {
-    case Schema(s: StructType, _) =>
-      s.toAttributes
-  }
-
-  /** Returns a catalyst DataType and its nullability for the given Scala Type using reflection. */
-  def schemaFor[T: TypeTag]: Schema =
-    ScalaReflectionLock.synchronized { schemaFor(localTypeOf[T]) }
-
-  /**
-   * Return the Scala Type for `T` in the current classloader mirror.
-   *
-   * Use this method instead of the convenience method `universe.typeOf`, which
-   * assumes that all types can be found in the classloader that loaded scala-reflect classes.
-   * That's not necessarily the case when running using Eclipse launchers or even
-   * Sbt console or test (without `fork := true`).
-   *
-   * @see SPARK-5281
-   */
-  def localTypeOf[T: TypeTag]: `Type` = typeTag[T].in(mirror).tpe
-
   /**
    * Returns the Spark SQL DataType for a given scala type.  Where this is not an exact mapping
    * to a native type, an ObjectType is returned. Special handling is also used for Arrays including
@@ -114,7 +79,9 @@ trait ScalaReflection {
 
           }
           ObjectType(cls)
-        case other => ObjectType(Utils.classForName(className))
+        case other =>
+          val clazz = mirror.runtimeClass(tpe.erasure.typeSymbol.asClass)
+          ObjectType(clazz)
       }
   }
 
@@ -640,6 +607,48 @@ trait ScalaReflection {
       }
     }
   }
+}
+
+/**
+ * Support for generating catalyst schemas for scala objects.  Note that unlike its companion
+ * object, this trait able to work in both the runtime and the compile time (macro) universe.
+ */
+trait ScalaReflection {
+  /** The universe we work in (runtime or macro) */
+  val universe: scala.reflect.api.Universe
+
+  /** The mirror used to access types in the universe */
+  def mirror: universe.Mirror
+
+  import universe._
+
+  // The Predef.Map is scala.collection.immutable.Map.
+  // Since the map values can be mutable, we explicitly import scala.collection.Map at here.
+  import scala.collection.Map
+
+  case class Schema(dataType: DataType, nullable: Boolean)
+
+  /** Returns a Sequence of attributes for the given case class type. */
+  def attributesFor[T: TypeTag]: Seq[Attribute] = schemaFor[T] match {
+    case Schema(s: StructType, _) =>
+      s.toAttributes
+  }
+
+  /** Returns a catalyst DataType and its nullability for the given Scala Type using reflection. */
+  def schemaFor[T: TypeTag]: Schema =
+    ScalaReflectionLock.synchronized { schemaFor(localTypeOf[T]) }
+
+  /**
+   * Return the Scala Type for `T` in the current classloader mirror.
+   *
+   * Use this method instead of the convenience method `universe.typeOf`, which
+   * assumes that all types can be found in the classloader that loaded scala-reflect classes.
+   * That's not necessarily the case when running using Eclipse launchers or even
+   * Sbt console or test (without `fork := true`).
+   *
+   * @see SPARK-5281
+   */
+  def localTypeOf[T: TypeTag]: `Type` = typeTag[T].in(mirror).tpe
 
   /** Returns a catalyst DataType and its nullability for the given Scala Type using reflection. */
   def schemaFor(tpe: `Type`): Schema = ScalaReflectionLock.synchronized {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index b977f278c5b5c..456b595008479 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -17,11 +17,13 @@
 
 package org.apache.spark.sql.catalyst.encoders
 
+import java.util.concurrent.ConcurrentMap
+
 import scala.reflect.ClassTag
 import scala.reflect.runtime.universe.{typeTag, TypeTag}
 
 import org.apache.spark.util.Utils
-import org.apache.spark.sql.Encoder
+import org.apache.spark.sql.{AnalysisException, Encoder}
 import org.apache.spark.sql.catalyst.analysis.{SimpleAnalyzer, UnresolvedExtractValue, UnresolvedAttribute}
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project}
 import org.apache.spark.sql.catalyst.expressions._
@@ -211,7 +213,9 @@ case class ExpressionEncoder[T](
    * Returns a new copy of this encoder, where the expressions used by `fromRow` are resolved to the
    * given schema.
    */
-  def resolve(schema: Seq[Attribute]): ExpressionEncoder[T] = {
+  def resolve(
+      schema: Seq[Attribute],
+      outerScopes: ConcurrentMap[String, AnyRef]): ExpressionEncoder[T] = {
     val positionToAttribute = AttributeMap.toIndex(schema)
     val unbound = fromRowExpression transform {
       case b: BoundReference => positionToAttribute(b.ordinal)
@@ -219,7 +223,23 @@ case class ExpressionEncoder[T](
 
     val plan = Project(Alias(unbound, "")() :: Nil, LocalRelation(schema))
     val analyzedPlan = SimpleAnalyzer.execute(plan)
-    copy(fromRowExpression = analyzedPlan.expressions.head.children.head)
+
+    // In order to construct instances of inner classes (for example those declared in a REPL cell),
+    // we need an instance of the outer scope.  This rule substitues those outer objects into
+    // expressions that are missing them by looking up the name in the SQLContexts `outerScopes`
+    // registry.
+    copy(fromRowExpression = analyzedPlan.expressions.head.children.head transform {
+      case n: NewInstance if n.outerPointer.isEmpty && n.cls.isMemberClass =>
+        val outer = outerScopes.get(n.cls.getDeclaringClass.getName)
+        if (outer == null) {
+          throw new AnalysisException(
+            s"Unable to generate an encoder for inner class `${n.cls.getName}` without access " +
+            s"to the scope that this class was defined in. " + "" +
+             "Try moving this class out of its parent class.")
+        }
+
+        n.copy(outerPointer = Some(Literal.fromObject(outer)))
+    })
   }
 
   /**
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/OuterScopes.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/OuterScopes.scala
new file mode 100644
index 0000000000000..a753b187bcd32
--- /dev/null
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/OuterScopes.scala
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.encoders
+
+import java.util.concurrent.ConcurrentMap
+
+import com.google.common.collect.MapMaker
+
+object OuterScopes {
+  @transient
+  lazy val outerScopes: ConcurrentMap[String, AnyRef] =
+    new MapMaker().weakValues().makeMap()
+
+  /**
+   * Adds a new outer scope to this context that can be used when instantiating an `inner class`
+   * during deserialialization. Inner classes are created when a case class is defined in the
+   * Spark REPL and registering the outer scope that this class was defined in allows us to create
+   * new instances on the spark executors.  In normal use, users should not need to call this
+   * function.
+   *
+   * Warning: this function operates on the assumption that there is only ever one instance of any
+   * given wrapper class.
+   */
+  def addOuterScope(outer: AnyRef): Unit = {
+    outerScopes.putIfAbsent(outer.getClass.getName, outer)
+  }
+}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
index 55c4ee11b20f4..2914c6ee790ce 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
@@ -31,6 +31,7 @@ import scala.reflect.ClassTag
 
 object ProductEncoder {
   import ScalaReflection.universe._
+  import ScalaReflection.mirror
   import ScalaReflection.localTypeOf
   import ScalaReflection.dataTypeFor
   import ScalaReflection.Schema
@@ -420,8 +421,7 @@ object ProductEncoder {
           }
         }
 
-        val className: String = t.erasure.typeSymbol.asClass.fullName
-        val cls = Utils.classForName(className)
+        val cls = mirror.runtimeClass(tpe.erasure.typeSymbol.asClass)
 
         val arguments = params.head.zipWithIndex.map { case (p, i) =>
           val fieldName = p.name.toString
@@ -429,7 +429,7 @@ object ProductEncoder {
           val dataType = schemaFor(fieldType).dataType
 
           // For tuples, we based grab the inner fields by ordinal instead of name.
-          if (className startsWith "scala.Tuple") {
+          if (cls.getName startsWith "scala.Tuple") {
             constructorFor(fieldType, Some(addToPathOrdinal(i, dataType)))
           } else {
             constructorFor(fieldType, Some(addToPath(fieldName)))
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala
index d51a8dede7f34..a31574c251af5 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala
@@ -34,7 +34,7 @@ trait CodegenFallback extends Expression {
     val objectTerm = ctx.freshName("obj")
     s"""
       /* expression: ${this} */
-      Object $objectTerm = expressions[${ctx.references.size - 1}].eval(${ctx.INPUT_ROW});
+      java.lang.Object $objectTerm = expressions[${ctx.references.size - 1}].eval(${ctx.INPUT_ROW});
       boolean ${ev.isNull} = $objectTerm == null;
       ${ctx.javaType(this.dataType)} ${ev.value} = ${ctx.defaultValue(this.dataType)};
       if (!${ev.isNull}) {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala
index 4b66069b5f55a..40189f0877764 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala
@@ -82,7 +82,7 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], () => Mu
     val allUpdates = ctx.splitExpressions(ctx.INPUT_ROW, updates)
 
     val code = s"""
-      public Object generate($exprType[] expr) {
+      public java.lang.Object generate($exprType[] expr) {
         return new SpecificMutableProjection(expr);
       }
 
@@ -109,7 +109,7 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], () => Mu
           return (InternalRow) mutableRow;
         }
 
-        public Object apply(Object _i) {
+        public java.lang.Object apply(java.lang.Object _i) {
           InternalRow ${ctx.INPUT_ROW} = (InternalRow) _i;
           $allProjections
           // copy all the results into MutableRow
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateProjection.scala
index c0d313b2e1301..f229f2000d8e1 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateProjection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateProjection.scala
@@ -167,7 +167,7 @@ object GenerateProjection extends CodeGenerator[Seq[Expression], Projection] {
         ${initMutableStates(ctx)}
       }
 
-      public Object apply(Object r) {
+      public java.lang.Object apply(java.lang.Object r) {
         // GenerateProjection does not work with UnsafeRows.
         assert(!(r instanceof ${classOf[UnsafeRow].getName}));
         return new SpecificRow((InternalRow) r);
@@ -186,14 +186,14 @@ object GenerateProjection extends CodeGenerator[Seq[Expression], Projection] {
         public void setNullAt(int i) { nullBits[i] = true; }
         public boolean isNullAt(int i) { return nullBits[i]; }
 
-        public Object genericGet(int i) {
+        public java.lang.Object genericGet(int i) {
           if (isNullAt(i)) return null;
           switch (i) {
           $getCases
           }
           return null;
         }
-        public void update(int i, Object value) {
+        public void update(int i, java.lang.Object value) {
           if (value == null) {
             setNullAt(i);
             return;
@@ -212,7 +212,7 @@ object GenerateProjection extends CodeGenerator[Seq[Expression], Projection] {
           return result;
         }
 
-        public boolean equals(Object other) {
+        public boolean equals(java.lang.Object other) {
           if (other instanceof SpecificRow) {
             SpecificRow row = (SpecificRow) other;
             $columnChecks
@@ -222,7 +222,7 @@ object GenerateProjection extends CodeGenerator[Seq[Expression], Projection] {
         }
 
         public InternalRow copy() {
-          Object[] arr = new Object[${expressions.length}];
+          java.lang.Object[] arr = new java.lang.Object[${expressions.length}];
           ${copyColumns}
           return new ${classOf[GenericInternalRow].getName}(arr);
         }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateSafeProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateSafeProjection.scala
index f0ed8645d923f..b7926bda3de19 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateSafeProjection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateSafeProjection.scala
@@ -148,7 +148,7 @@ object GenerateSafeProjection extends CodeGenerator[Seq[Expression], Projection]
     }
     val allExpressions = ctx.splitExpressions(ctx.INPUT_ROW, expressionCodes)
     val code = s"""
-      public Object generate($exprType[] expr) {
+      public java.lang.Object generate($exprType[] expr) {
         return new SpecificSafeProjection(expr);
       }
 
@@ -165,7 +165,7 @@ object GenerateSafeProjection extends CodeGenerator[Seq[Expression], Projection]
           ${initMutableStates(ctx)}
         }
 
-        public Object apply(Object _i) {
+        public java.lang.Object apply(java.lang.Object _i) {
           InternalRow ${ctx.INPUT_ROW} = (InternalRow) _i;
           $allExpressions
           return mutableRow;
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
index 4c17d02a23725..7b6c9373ebe30 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
@@ -324,7 +324,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     val eval = createCode(ctx, expressions, subexpressionEliminationEnabled)
 
     val code = s"""
-      public Object generate($exprType[] exprs) {
+      public java.lang.Object generate($exprType[] exprs) {
         return new SpecificUnsafeProjection(exprs);
       }
 
@@ -342,7 +342,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
         }
 
         // Scala.Function1 need this
-        public Object apply(Object row) {
+        public java.lang.Object apply(java.lang.Object row) {
           return apply((InternalRow) row);
         }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeRowJoiner.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeRowJoiner.scala
index da91ff29537b3..da602d9b4bce1 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeRowJoiner.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeRowJoiner.scala
@@ -159,7 +159,7 @@ object GenerateUnsafeRowJoiner extends CodeGenerator[(StructType, StructType), U
 
     // ------------------------ Finally, put everything together  --------------------------- //
     val code = s"""
-       |public Object generate($exprType[] exprs) {
+       |public java.lang.Object generate($exprType[] exprs) {
        |  return new SpecificUnsafeRowJoiner();
        |}
        |
@@ -176,9 +176,9 @@ object GenerateUnsafeRowJoiner extends CodeGenerator[(StructType, StructType), U
        |      buf = new byte[sizeInBytes];
        |    }
        |
-       |    final Object obj1 = row1.getBaseObject();
+       |    final java.lang.Object obj1 = row1.getBaseObject();
        |    final long offset1 = row1.getBaseOffset();
-       |    final Object obj2 = row2.getBaseObject();
+       |    final java.lang.Object obj2 = row2.getBaseObject();
        |    final long offset2 = row2.getBaseOffset();
        |
        |    $copyBitset
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala
index 455fa2427c26d..e34fd49be8389 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala
@@ -48,6 +48,12 @@ object Literal {
       throw new RuntimeException("Unsupported literal type " + v.getClass + " " + v)
   }
 
+  /**
+   * Constructs a [[Literal]] of [[ObjectType]], for example when you need to pass an object
+   * into code generation.
+   */
+  def fromObject(obj: AnyRef): Literal = new Literal(obj, ObjectType(obj.getClass))
+
   def create(v: Any, dataType: DataType): Literal = {
     Literal(CatalystTypeConverters.convertToCatalyst(v), dataType)
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index acf0da240051e..f865a9408ef4e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -24,6 +24,7 @@ import org.apache.spark.SparkConf
 import org.apache.spark.serializer._
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.analysis.SimpleAnalyzer
+import org.apache.spark.sql.catalyst.encoders.ProductEncoder
 import org.apache.spark.sql.catalyst.plans.logical.{Project, LocalRelation}
 import org.apache.spark.sql.catalyst.util.GenericArrayData
 import org.apache.spark.sql.catalyst.InternalRow
@@ -178,6 +179,15 @@ case class Invoke(
   }
 }
 
+object NewInstance {
+  def apply(
+      cls: Class[_],
+      arguments: Seq[Expression],
+      propagateNull: Boolean = false,
+      dataType: DataType): NewInstance =
+    new NewInstance(cls, arguments, propagateNull, dataType, None)
+}
+
 /**
  * Constructs a new instance of the given class, using the result of evaluating the specified
  * expressions as arguments.
@@ -189,12 +199,15 @@ case class Invoke(
  * @param dataType The type of object being constructed, as a Spark SQL datatype.  This allows you
  *                 to manually specify the type when the object in question is a valid internal
  *                 representation (i.e. ArrayData) instead of an object.
+ * @param outerPointer If the object being constructed is an inner class the outerPointer must
+ *                     for the containing class must be specified.
  */
 case class NewInstance(
     cls: Class[_],
     arguments: Seq[Expression],
-    propagateNull: Boolean = true,
-    dataType: DataType) extends Expression {
+    propagateNull: Boolean,
+    dataType: DataType,
+    outerPointer: Option[Literal]) extends Expression {
   private val className = cls.getName
 
   override def nullable: Boolean = propagateNull
@@ -209,30 +222,43 @@ case class NewInstance(
     val argGen = arguments.map(_.gen(ctx))
     val argString = argGen.map(_.value).mkString(", ")
 
+    val outer = outerPointer.map(_.gen(ctx))
+
+    val setup =
+      s"""
+         ${argGen.map(_.code).mkString("\n")}
+         ${outer.map(_.code.mkString("")).getOrElse("")}
+       """.stripMargin
+
+    val constructorCall = outer.map { gen =>
+      s"""${gen.value}.new ${cls.getSimpleName}($argString)"""
+    }.getOrElse {
+      s"new $className($argString)"
+    }
+
     if (propagateNull) {
       val objNullCheck = if (ctx.defaultValue(dataType) == "null") {
         s"${ev.isNull} = ${ev.value} == null;"
       } else {
         ""
       }
-
       val argsNonNull = s"!(${argGen.map(_.isNull).mkString(" || ")})"
+
       s"""
-        ${argGen.map(_.code).mkString("\n")}
+        $setup
 
         boolean ${ev.isNull} = true;
         $javaType ${ev.value} = ${ctx.defaultValue(dataType)};
-
         if ($argsNonNull) {
-          ${ev.value} = new $className($argString);
+          ${ev.value} = $constructorCall;
           ${ev.isNull} = false;
         }
        """
     } else {
       s"""
-        ${argGen.map(_.code).mkString("\n")}
+        $setup
 
-        $javaType ${ev.value} = new $className($argString);
+        $javaType ${ev.value} = $constructorCall;
         final boolean ${ev.isNull} = ${ev.value} == null;
       """
     }
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
index 9fe64b4cf10e4..cde0364f3dd9d 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
@@ -18,6 +18,9 @@
 package org.apache.spark.sql.catalyst.encoders
 
 import java.util.Arrays
+import java.util.concurrent.ConcurrentMap
+
+import com.google.common.collect.MapMaker
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions.AttributeReference
@@ -25,6 +28,8 @@ import org.apache.spark.sql.catalyst.util.ArrayData
 import org.apache.spark.sql.types.ArrayType
 
 abstract class ExpressionEncoderSuite extends SparkFunSuite {
+  val outers: ConcurrentMap[String, AnyRef] = new MapMaker().weakValues().makeMap()
+
   protected def encodeDecodeTest[T](
       input: T,
       encoder: ExpressionEncoder[T],
@@ -32,7 +37,7 @@ abstract class ExpressionEncoderSuite extends SparkFunSuite {
     test(s"encode/decode for $testName: $input") {
       val row = encoder.toRow(input)
       val schema = encoder.schema.toAttributes
-      val boundEncoder = encoder.resolve(schema).bind(schema)
+      val boundEncoder = encoder.resolve(schema, outers).bind(schema)
       val convertedBack = try boundEncoder.fromRow(row) catch {
         case e: Exception =>
           fail(
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
index bc539d62c537d..1798514c5c38b 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
@@ -53,6 +53,10 @@ case class RepeatedData(
 case class SpecificCollection(l: List[Int])
 
 class ProductEncoderSuite extends ExpressionEncoderSuite {
+  outers.put(getClass.getName, this)
+
+  case class InnerClass(i: Int)
+  productTest(InnerClass(1))
 
   productTest(PrimitiveData(1, 1, 1, 1, 1, 1, true))
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index b644f6ad3096d..bdcdc5d47cbae 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -74,7 +74,7 @@ class Dataset[T] private[sql](
 
   /** The encoder for this [[Dataset]] that has been resolved to its output schema. */
   private[sql] val resolvedTEncoder: ExpressionEncoder[T] =
-    unresolvedTEncoder.resolve(queryExecution.analyzed.output)
+    unresolvedTEncoder.resolve(queryExecution.analyzed.output, OuterScopes.outerScopes)
 
   private implicit def classTag = resolvedTEncoder.clsTag
 
@@ -375,7 +375,7 @@ class Dataset[T] private[sql](
       sqlContext,
       Project(
         c1.withInputType(
-          resolvedTEncoder,
+          resolvedTEncoder.bind(queryExecution.analyzed.output),
           queryExecution.analyzed.output).named :: Nil,
         logicalPlan))
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 3f84e22a1025b..7e5acbe8517d1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -21,7 +21,7 @@ import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.function._
-import org.apache.spark.sql.catalyst.encoders.{FlatEncoder, ExpressionEncoder, encoderFor}
+import org.apache.spark.sql.catalyst.encoders.{FlatEncoder, ExpressionEncoder, encoderFor, OuterScopes}
 import org.apache.spark.sql.catalyst.expressions.{Alias, CreateStruct, Attribute}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.QueryExecution
@@ -52,8 +52,10 @@ class GroupedDataset[K, T] private[sql](
   private implicit val unresolvedKEncoder = encoderFor(kEncoder)
   private implicit val unresolvedTEncoder = encoderFor(tEncoder)
 
-  private val resolvedKEncoder = unresolvedKEncoder.resolve(groupingAttributes)
-  private val resolvedTEncoder = unresolvedTEncoder.resolve(dataAttributes)
+  private val resolvedKEncoder =
+    unresolvedKEncoder.resolve(groupingAttributes, OuterScopes.outerScopes)
+  private val resolvedTEncoder =
+    unresolvedTEncoder.resolve(dataAttributes, OuterScopes.outerScopes)
 
   private def logicalPlan = queryExecution.analyzed
   private def sqlContext = queryExecution.sqlContext
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
index 3f2775896bb8c..6ce41aaf01e27 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
@@ -52,8 +52,8 @@ object TypedAggregateExpression {
  */
 case class TypedAggregateExpression(
     aggregator: Aggregator[Any, Any, Any],
-    aEncoder: Option[ExpressionEncoder[Any]],
-    bEncoder: ExpressionEncoder[Any],
+    aEncoder: Option[ExpressionEncoder[Any]], // Should be bound.
+    bEncoder: ExpressionEncoder[Any], // Should be bound.
     cEncoder: ExpressionEncoder[Any],
     children: Seq[Attribute],
     mutableAggBufferOffset: Int,
@@ -92,9 +92,6 @@ case class TypedAggregateExpression(
   // We let the dataset do the binding for us.
   lazy val boundA = aEncoder.get
 
-  val bAttributes = bEncoder.schema.toAttributes
-  lazy val boundB = bEncoder.resolve(bAttributes).bind(bAttributes)
-
   private def updateBuffer(buffer: MutableRow, value: InternalRow): Unit = {
     // todo: need a more neat way to assign the value.
     var i = 0
@@ -114,24 +111,24 @@ case class TypedAggregateExpression(
 
   override def update(buffer: MutableRow, input: InternalRow): Unit = {
     val inputA = boundA.fromRow(input)
-    val currentB = boundB.shift(mutableAggBufferOffset).fromRow(buffer)
+    val currentB = bEncoder.shift(mutableAggBufferOffset).fromRow(buffer)
     val merged = aggregator.reduce(currentB, inputA)
-    val returned = boundB.toRow(merged)
+    val returned = bEncoder.toRow(merged)
 
     updateBuffer(buffer, returned)
   }
 
   override def merge(buffer1: MutableRow, buffer2: InternalRow): Unit = {
-    val b1 = boundB.shift(mutableAggBufferOffset).fromRow(buffer1)
-    val b2 = boundB.shift(inputAggBufferOffset).fromRow(buffer2)
+    val b1 = bEncoder.shift(mutableAggBufferOffset).fromRow(buffer1)
+    val b2 = bEncoder.shift(inputAggBufferOffset).fromRow(buffer2)
     val merged = aggregator.merge(b1, b2)
-    val returned = boundB.toRow(merged)
+    val returned = bEncoder.toRow(merged)
 
     updateBuffer(buffer1, returned)
   }
 
   override def eval(buffer: InternalRow): Any = {
-    val b = boundB.shift(mutableAggBufferOffset).fromRow(buffer)
+    val b = bEncoder.shift(mutableAggBufferOffset).fromRow(buffer)
     val result = cEncoder.toRow(aggregator.finish(b))
     dataType match {
       case _: StructType => result

From e99d3392068bc929c900a4cc7b50e9e2b437a23a Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Wed, 18 Nov 2015 18:34:01 -0800
Subject: [PATCH 349/518] [SPARK-11839][ML] refactor save/write traits

* add "ML" prefix to reader/writer/readable/writable to avoid name collision with java.util.*
* define `DefaultParamsReadable/Writable` and use them to save some code
* use `super.load` instead so people can jump directly to the doc of `Readable.load`, which documents the Java compatibility issues

jkbradley

Author: Xiangrui Meng <meng@databricks.com>

Closes #9827 from mengxr/SPARK-11839.
---
 .../scala/org/apache/spark/ml/Pipeline.scala  | 40 +++++++++----------
 .../classification/LogisticRegression.scala   | 29 +++++++-------
 .../apache/spark/ml/feature/Binarizer.scala   | 12 ++----
 .../apache/spark/ml/feature/Bucketizer.scala  | 12 ++----
 .../spark/ml/feature/CountVectorizer.scala    | 22 ++++------
 .../org/apache/spark/ml/feature/DCT.scala     | 12 ++----
 .../apache/spark/ml/feature/HashingTF.scala   | 12 ++----
 .../org/apache/spark/ml/feature/IDF.scala     | 23 +++++------
 .../apache/spark/ml/feature/Interaction.scala | 12 ++----
 .../spark/ml/feature/MinMaxScaler.scala       | 22 ++++------
 .../org/apache/spark/ml/feature/NGram.scala   | 12 ++----
 .../apache/spark/ml/feature/Normalizer.scala  | 12 ++----
 .../spark/ml/feature/OneHotEncoder.scala      | 12 ++----
 .../ml/feature/PolynomialExpansion.scala      | 12 ++----
 .../ml/feature/QuantileDiscretizer.scala      | 12 ++----
 .../spark/ml/feature/SQLTransformer.scala     | 13 ++----
 .../spark/ml/feature/StandardScaler.scala     | 22 ++++------
 .../spark/ml/feature/StopWordsRemover.scala   | 12 ++----
 .../spark/ml/feature/StringIndexer.scala      | 32 +++++----------
 .../apache/spark/ml/feature/Tokenizer.scala   | 24 +++--------
 .../spark/ml/feature/VectorAssembler.scala    | 12 ++----
 .../spark/ml/feature/VectorSlicer.scala       | 12 ++----
 .../apache/spark/ml/recommendation/ALS.scala  | 27 +++++--------
 .../ml/regression/LinearRegression.scala      | 30 ++++++--------
 .../org/apache/spark/ml/util/ReadWrite.scala  | 40 ++++++++++++-------
 .../org/apache/spark/ml/PipelineSuite.scala   | 14 +++----
 .../spark/ml/util/DefaultReadWriteTest.scala  | 17 ++++----
 27 files changed, 190 insertions(+), 321 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
index 25f0c696f42be..b0f22e042ec56 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
@@ -29,8 +29,8 @@ import org.json4s.jackson.JsonMethods._
 import org.apache.spark.{SparkContext, Logging}
 import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.ml.param.{Param, ParamMap, Params}
-import org.apache.spark.ml.util.Reader
-import org.apache.spark.ml.util.Writer
+import org.apache.spark.ml.util.MLReader
+import org.apache.spark.ml.util.MLWriter
 import org.apache.spark.ml.util._
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.types.StructType
@@ -89,7 +89,7 @@ abstract class PipelineStage extends Params with Logging {
  * an identity transformer.
  */
 @Experimental
-class Pipeline(override val uid: String) extends Estimator[PipelineModel] with Writable {
+class Pipeline(override val uid: String) extends Estimator[PipelineModel] with MLWritable {
 
   def this() = this(Identifiable.randomUID("pipeline"))
 
@@ -174,16 +174,16 @@ class Pipeline(override val uid: String) extends Estimator[PipelineModel] with W
     theStages.foldLeft(schema)((cur, stage) => stage.transformSchema(cur))
   }
 
-  override def write: Writer = new Pipeline.PipelineWriter(this)
+  override def write: MLWriter = new Pipeline.PipelineWriter(this)
 }
 
-object Pipeline extends Readable[Pipeline] {
+object Pipeline extends MLReadable[Pipeline] {
 
-  override def read: Reader[Pipeline] = new PipelineReader
+  override def read: MLReader[Pipeline] = new PipelineReader
 
-  override def load(path: String): Pipeline = read.load(path)
+  override def load(path: String): Pipeline = super.load(path)
 
-  private[ml] class PipelineWriter(instance: Pipeline) extends Writer {
+  private[ml] class PipelineWriter(instance: Pipeline) extends MLWriter {
 
     SharedReadWrite.validateStages(instance.getStages)
 
@@ -191,7 +191,7 @@ object Pipeline extends Readable[Pipeline] {
       SharedReadWrite.saveImpl(instance, instance.getStages, sc, path)
   }
 
-  private[ml] class PipelineReader extends Reader[Pipeline] {
+  private[ml] class PipelineReader extends MLReader[Pipeline] {
 
     /** Checked against metadata when loading model */
     private val className = "org.apache.spark.ml.Pipeline"
@@ -202,7 +202,7 @@ object Pipeline extends Readable[Pipeline] {
     }
   }
 
-  /** Methods for [[Reader]] and [[Writer]] shared between [[Pipeline]] and [[PipelineModel]] */
+  /** Methods for [[MLReader]] and [[MLWriter]] shared between [[Pipeline]] and [[PipelineModel]] */
   private[ml] object SharedReadWrite {
 
     import org.json4s.JsonDSL._
@@ -210,7 +210,7 @@ object Pipeline extends Readable[Pipeline] {
     /** Check that all stages are Writable */
     def validateStages(stages: Array[PipelineStage]): Unit = {
       stages.foreach {
-        case stage: Writable => // good
+        case stage: MLWritable => // good
         case other =>
           throw new UnsupportedOperationException("Pipeline write will fail on this Pipeline" +
             s" because it contains a stage which does not implement Writable. Non-Writable stage:" +
@@ -245,7 +245,7 @@ object Pipeline extends Readable[Pipeline] {
 
       // Save stages
       val stagesDir = new Path(path, "stages").toString
-      stages.zipWithIndex.foreach { case (stage: Writable, idx: Int) =>
+      stages.zipWithIndex.foreach { case (stage: MLWritable, idx: Int) =>
         stage.write.save(getStagePath(stage.uid, idx, stages.length, stagesDir))
       }
     }
@@ -285,7 +285,7 @@ object Pipeline extends Readable[Pipeline] {
         val stagePath = SharedReadWrite.getStagePath(stageUid, idx, stageUids.length, stagesDir)
         val stageMetadata = DefaultParamsReader.loadMetadata(stagePath, sc)
         val cls = Utils.classForName(stageMetadata.className)
-        cls.getMethod("read").invoke(null).asInstanceOf[Reader[PipelineStage]].load(stagePath)
+        cls.getMethod("read").invoke(null).asInstanceOf[MLReader[PipelineStage]].load(stagePath)
       }
       (metadata.uid, stages)
     }
@@ -308,7 +308,7 @@ object Pipeline extends Readable[Pipeline] {
 class PipelineModel private[ml] (
     override val uid: String,
     val stages: Array[Transformer])
-  extends Model[PipelineModel] with Writable with Logging {
+  extends Model[PipelineModel] with MLWritable with Logging {
 
   /** A Java/Python-friendly auxiliary constructor. */
   private[ml] def this(uid: String, stages: ju.List[Transformer]) = {
@@ -333,18 +333,18 @@ class PipelineModel private[ml] (
     new PipelineModel(uid, stages.map(_.copy(extra))).setParent(parent)
   }
 
-  override def write: Writer = new PipelineModel.PipelineModelWriter(this)
+  override def write: MLWriter = new PipelineModel.PipelineModelWriter(this)
 }
 
-object PipelineModel extends Readable[PipelineModel] {
+object PipelineModel extends MLReadable[PipelineModel] {
 
   import Pipeline.SharedReadWrite
 
-  override def read: Reader[PipelineModel] = new PipelineModelReader
+  override def read: MLReader[PipelineModel] = new PipelineModelReader
 
-  override def load(path: String): PipelineModel = read.load(path)
+  override def load(path: String): PipelineModel = super.load(path)
 
-  private[ml] class PipelineModelWriter(instance: PipelineModel) extends Writer {
+  private[ml] class PipelineModelWriter(instance: PipelineModel) extends MLWriter {
 
     SharedReadWrite.validateStages(instance.stages.asInstanceOf[Array[PipelineStage]])
 
@@ -352,7 +352,7 @@ object PipelineModel extends Readable[PipelineModel] {
       instance.stages.asInstanceOf[Array[PipelineStage]], sc, path)
   }
 
-  private[ml] class PipelineModelReader extends Reader[PipelineModel] {
+  private[ml] class PipelineModelReader extends MLReader[PipelineModel] {
 
     /** Checked against metadata when loading model */
     private val className = "org.apache.spark.ml.PipelineModel"
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
index 71c2533bcbf47..a3cc49f7f018c 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
@@ -29,9 +29,9 @@ import org.apache.spark.ml.feature.Instance
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
 import org.apache.spark.ml.util._
+import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics
 import org.apache.spark.mllib.linalg._
 import org.apache.spark.mllib.linalg.BLAS._
-import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics
 import org.apache.spark.mllib.stat.MultivariateOnlineSummarizer
 import org.apache.spark.mllib.util.MLUtils
 import org.apache.spark.rdd.RDD
@@ -157,7 +157,7 @@ private[classification] trait LogisticRegressionParams extends ProbabilisticClas
 @Experimental
 class LogisticRegression(override val uid: String)
   extends ProbabilisticClassifier[Vector, LogisticRegression, LogisticRegressionModel]
-  with LogisticRegressionParams with Writable with Logging {
+  with LogisticRegressionParams with DefaultParamsWritable with Logging {
 
   def this() = this(Identifiable.randomUID("logreg"))
 
@@ -385,12 +385,11 @@ class LogisticRegression(override val uid: String)
   }
 
   override def copy(extra: ParamMap): LogisticRegression = defaultCopy(extra)
-
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
-object LogisticRegression extends Readable[LogisticRegression] {
-  override def read: Reader[LogisticRegression] = new DefaultParamsReader[LogisticRegression]
+object LogisticRegression extends DefaultParamsReadable[LogisticRegression] {
+
+  override def load(path: String): LogisticRegression = super.load(path)
 }
 
 /**
@@ -403,7 +402,7 @@ class LogisticRegressionModel private[ml] (
     val coefficients: Vector,
     val intercept: Double)
   extends ProbabilisticClassificationModel[Vector, LogisticRegressionModel]
-  with LogisticRegressionParams with Writable {
+  with LogisticRegressionParams with MLWritable {
 
   @deprecated("Use coefficients instead.", "1.6.0")
   def weights: Vector = coefficients
@@ -519,26 +518,26 @@ class LogisticRegressionModel private[ml] (
   }
 
   /**
-   * Returns a [[Writer]] instance for this ML instance.
+   * Returns a [[MLWriter]] instance for this ML instance.
    *
    * For [[LogisticRegressionModel]], this does NOT currently save the training [[summary]].
    * An option to save [[summary]] may be added in the future.
    *
    * This also does not save the [[parent]] currently.
    */
-  override def write: Writer = new LogisticRegressionModel.LogisticRegressionModelWriter(this)
+  override def write: MLWriter = new LogisticRegressionModel.LogisticRegressionModelWriter(this)
 }
 
 
-object LogisticRegressionModel extends Readable[LogisticRegressionModel] {
+object LogisticRegressionModel extends MLReadable[LogisticRegressionModel] {
 
-  override def read: Reader[LogisticRegressionModel] = new LogisticRegressionModelReader
+  override def read: MLReader[LogisticRegressionModel] = new LogisticRegressionModelReader
 
-  override def load(path: String): LogisticRegressionModel = read.load(path)
+  override def load(path: String): LogisticRegressionModel = super.load(path)
 
-  /** [[Writer]] instance for [[LogisticRegressionModel]] */
+  /** [[MLWriter]] instance for [[LogisticRegressionModel]] */
   private[classification] class LogisticRegressionModelWriter(instance: LogisticRegressionModel)
-    extends Writer with Logging {
+    extends MLWriter with Logging {
 
     private case class Data(
         numClasses: Int,
@@ -558,7 +557,7 @@ object LogisticRegressionModel extends Readable[LogisticRegressionModel] {
   }
 
   private[classification] class LogisticRegressionModelReader
-    extends Reader[LogisticRegressionModel] {
+    extends MLReader[LogisticRegressionModel] {
 
     /** Checked against metadata when loading model */
     private val className = "org.apache.spark.ml.classification.LogisticRegressionModel"
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
index e2be6547d8f00..63c06581482ed 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala
@@ -33,7 +33,7 @@ import org.apache.spark.sql.types.{DoubleType, StructType}
  */
 @Experimental
 final class Binarizer(override val uid: String)
-  extends Transformer with Writable with HasInputCol with HasOutputCol {
+  extends Transformer with HasInputCol with HasOutputCol with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("binarizer"))
 
@@ -86,17 +86,11 @@ final class Binarizer(override val uid: String)
   }
 
   override def copy(extra: ParamMap): Binarizer = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object Binarizer extends Readable[Binarizer] {
-
-  @Since("1.6.0")
-  override def read: Reader[Binarizer] = new DefaultParamsReader[Binarizer]
+object Binarizer extends DefaultParamsReadable[Binarizer] {
 
   @Since("1.6.0")
-  override def load(path: String): Binarizer = read.load(path)
+  override def load(path: String): Binarizer = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala
index 7095fbd70aa07..324353a96afb3 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala
@@ -36,7 +36,7 @@ import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
  */
 @Experimental
 final class Bucketizer(override val uid: String)
-  extends Model[Bucketizer] with HasInputCol with HasOutputCol with Writable {
+  extends Model[Bucketizer] with HasInputCol with HasOutputCol with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("bucketizer"))
 
@@ -93,12 +93,9 @@ final class Bucketizer(override val uid: String)
   override def copy(extra: ParamMap): Bucketizer = {
     defaultCopy[Bucketizer](extra).setParent(parent)
   }
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
-object Bucketizer extends Readable[Bucketizer] {
+object Bucketizer extends DefaultParamsReadable[Bucketizer] {
 
   /** We require splits to be of length >= 3 and to be in strictly increasing order. */
   private[feature] def checkSplits(splits: Array[Double]): Boolean = {
@@ -140,8 +137,5 @@ object Bucketizer extends Readable[Bucketizer] {
   }
 
   @Since("1.6.0")
-  override def read: Reader[Bucketizer] = new DefaultParamsReader[Bucketizer]
-
-  @Since("1.6.0")
-  override def load(path: String): Bucketizer = read.load(path)
+  override def load(path: String): Bucketizer = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
index 5ff9bfb7d1119..4969cf42450d2 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
@@ -107,7 +107,7 @@ private[feature] trait CountVectorizerParams extends Params with HasInputCol wit
  */
 @Experimental
 class CountVectorizer(override val uid: String)
-  extends Estimator[CountVectorizerModel] with CountVectorizerParams with Writable {
+  extends Estimator[CountVectorizerModel] with CountVectorizerParams with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("cntVec"))
 
@@ -171,16 +171,10 @@ class CountVectorizer(override val uid: String)
   }
 
   override def copy(extra: ParamMap): CountVectorizer = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object CountVectorizer extends Readable[CountVectorizer] {
-
-  @Since("1.6.0")
-  override def read: Reader[CountVectorizer] = new DefaultParamsReader
+object CountVectorizer extends DefaultParamsReadable[CountVectorizer] {
 
   @Since("1.6.0")
   override def load(path: String): CountVectorizer = super.load(path)
@@ -193,7 +187,7 @@ object CountVectorizer extends Readable[CountVectorizer] {
  */
 @Experimental
 class CountVectorizerModel(override val uid: String, val vocabulary: Array[String])
-  extends Model[CountVectorizerModel] with CountVectorizerParams with Writable {
+  extends Model[CountVectorizerModel] with CountVectorizerParams with MLWritable {
 
   import CountVectorizerModel._
 
@@ -251,14 +245,14 @@ class CountVectorizerModel(override val uid: String, val vocabulary: Array[Strin
   }
 
   @Since("1.6.0")
-  override def write: Writer = new CountVectorizerModelWriter(this)
+  override def write: MLWriter = new CountVectorizerModelWriter(this)
 }
 
 @Since("1.6.0")
-object CountVectorizerModel extends Readable[CountVectorizerModel] {
+object CountVectorizerModel extends MLReadable[CountVectorizerModel] {
 
   private[CountVectorizerModel]
-  class CountVectorizerModelWriter(instance: CountVectorizerModel) extends Writer {
+  class CountVectorizerModelWriter(instance: CountVectorizerModel) extends MLWriter {
 
     private case class Data(vocabulary: Seq[String])
 
@@ -270,7 +264,7 @@ object CountVectorizerModel extends Readable[CountVectorizerModel] {
     }
   }
 
-  private class CountVectorizerModelReader extends Reader[CountVectorizerModel] {
+  private class CountVectorizerModelReader extends MLReader[CountVectorizerModel] {
 
     private val className = "org.apache.spark.ml.feature.CountVectorizerModel"
 
@@ -288,7 +282,7 @@ object CountVectorizerModel extends Readable[CountVectorizerModel] {
   }
 
   @Since("1.6.0")
-  override def read: Reader[CountVectorizerModel] = new CountVectorizerModelReader
+  override def read: MLReader[CountVectorizerModel] = new CountVectorizerModelReader
 
   @Since("1.6.0")
   override def load(path: String): CountVectorizerModel = super.load(path)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala
index 6ea5a616173ee..6bed72164a1da 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala
@@ -37,7 +37,7 @@ import org.apache.spark.sql.types.DataType
  */
 @Experimental
 class DCT(override val uid: String)
-  extends UnaryTransformer[Vector, Vector, DCT] with Writable {
+  extends UnaryTransformer[Vector, Vector, DCT] with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("dct"))
 
@@ -69,17 +69,11 @@ class DCT(override val uid: String)
   }
 
   override protected def outputDataType: DataType = new VectorUDT
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object DCT extends Readable[DCT] {
-
-  @Since("1.6.0")
-  override def read: Reader[DCT] = new DefaultParamsReader[DCT]
+object DCT extends DefaultParamsReadable[DCT] {
 
   @Since("1.6.0")
-  override def load(path: String): DCT = read.load(path)
+  override def load(path: String): DCT = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala
index 6d2ea675f5617..9e15835429a38 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala
@@ -34,7 +34,7 @@ import org.apache.spark.sql.types.{ArrayType, StructType}
  */
 @Experimental
 class HashingTF(override val uid: String)
-  extends Transformer with HasInputCol with HasOutputCol with Writable {
+  extends Transformer with HasInputCol with HasOutputCol with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("hashingTF"))
 
@@ -77,17 +77,11 @@ class HashingTF(override val uid: String)
   }
 
   override def copy(extra: ParamMap): HashingTF = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object HashingTF extends Readable[HashingTF] {
-
-  @Since("1.6.0")
-  override def read: Reader[HashingTF] = new DefaultParamsReader[HashingTF]
+object HashingTF extends DefaultParamsReadable[HashingTF] {
 
   @Since("1.6.0")
-  override def load(path: String): HashingTF = read.load(path)
+  override def load(path: String): HashingTF = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
index 53ad34ef12646..0e00ef6f2ee20 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
@@ -62,7 +62,8 @@ private[feature] trait IDFBase extends Params with HasInputCol with HasOutputCol
  * Compute the Inverse Document Frequency (IDF) given a collection of documents.
  */
 @Experimental
-final class IDF(override val uid: String) extends Estimator[IDFModel] with IDFBase with Writable {
+final class IDF(override val uid: String) extends Estimator[IDFModel] with IDFBase
+  with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("idf"))
 
@@ -87,16 +88,10 @@ final class IDF(override val uid: String) extends Estimator[IDFModel] with IDFBa
   }
 
   override def copy(extra: ParamMap): IDF = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object IDF extends Readable[IDF] {
-
-  @Since("1.6.0")
-  override def read: Reader[IDF] = new DefaultParamsReader
+object IDF extends DefaultParamsReadable[IDF] {
 
   @Since("1.6.0")
   override def load(path: String): IDF = super.load(path)
@@ -110,7 +105,7 @@ object IDF extends Readable[IDF] {
 class IDFModel private[ml] (
     override val uid: String,
     idfModel: feature.IDFModel)
-  extends Model[IDFModel] with IDFBase with Writable {
+  extends Model[IDFModel] with IDFBase with MLWritable {
 
   import IDFModel._
 
@@ -140,13 +135,13 @@ class IDFModel private[ml] (
   def idf: Vector = idfModel.idf
 
   @Since("1.6.0")
-  override def write: Writer = new IDFModelWriter(this)
+  override def write: MLWriter = new IDFModelWriter(this)
 }
 
 @Since("1.6.0")
-object IDFModel extends Readable[IDFModel] {
+object IDFModel extends MLReadable[IDFModel] {
 
-  private[IDFModel] class IDFModelWriter(instance: IDFModel) extends Writer {
+  private[IDFModel] class IDFModelWriter(instance: IDFModel) extends MLWriter {
 
     private case class Data(idf: Vector)
 
@@ -158,7 +153,7 @@ object IDFModel extends Readable[IDFModel] {
     }
   }
 
-  private class IDFModelReader extends Reader[IDFModel] {
+  private class IDFModelReader extends MLReader[IDFModel] {
 
     private val className = "org.apache.spark.ml.feature.IDFModel"
 
@@ -176,7 +171,7 @@ object IDFModel extends Readable[IDFModel] {
   }
 
   @Since("1.6.0")
-  override def read: Reader[IDFModel] = new IDFModelReader
+  override def read: MLReader[IDFModel] = new IDFModelReader
 
   @Since("1.6.0")
   override def load(path: String): IDFModel = super.load(path)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala
index 9df6b311cc9da..2181119f04a5d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala
@@ -45,7 +45,7 @@ import org.apache.spark.sql.types._
 @Since("1.6.0")
 @Experimental
 class Interaction @Since("1.6.0") (override val uid: String) extends Transformer
-  with HasInputCols with HasOutputCol with Writable {
+  with HasInputCols with HasOutputCol with DefaultParamsWritable {
 
   @Since("1.6.0")
   def this() = this(Identifiable.randomUID("interaction"))
@@ -224,19 +224,13 @@ class Interaction @Since("1.6.0") (override val uid: String) extends Transformer
     require($(inputCols).length > 0, "Input cols must have non-zero length.")
     require($(inputCols).distinct.length == $(inputCols).length, "Input cols must be distinct.")
   }
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object Interaction extends Readable[Interaction] {
-
-  @Since("1.6.0")
-  override def read: Reader[Interaction] = new DefaultParamsReader[Interaction]
+object Interaction extends DefaultParamsReadable[Interaction] {
 
   @Since("1.6.0")
-  override def load(path: String): Interaction = read.load(path)
+  override def load(path: String): Interaction = super.load(path)
 }
 
 /**
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
index 24d964fae834e..ed24eabb50444 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
@@ -88,7 +88,7 @@ private[feature] trait MinMaxScalerParams extends Params with HasInputCol with H
  */
 @Experimental
 class MinMaxScaler(override val uid: String)
-  extends Estimator[MinMaxScalerModel] with MinMaxScalerParams with Writable {
+  extends Estimator[MinMaxScalerModel] with MinMaxScalerParams with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("minMaxScal"))
 
@@ -118,16 +118,10 @@ class MinMaxScaler(override val uid: String)
   }
 
   override def copy(extra: ParamMap): MinMaxScaler = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object MinMaxScaler extends Readable[MinMaxScaler] {
-
-  @Since("1.6.0")
-  override def read: Reader[MinMaxScaler] = new DefaultParamsReader
+object MinMaxScaler extends DefaultParamsReadable[MinMaxScaler] {
 
   @Since("1.6.0")
   override def load(path: String): MinMaxScaler = super.load(path)
@@ -147,7 +141,7 @@ class MinMaxScalerModel private[ml] (
     override val uid: String,
     val originalMin: Vector,
     val originalMax: Vector)
-  extends Model[MinMaxScalerModel] with MinMaxScalerParams with Writable {
+  extends Model[MinMaxScalerModel] with MinMaxScalerParams with MLWritable {
 
   import MinMaxScalerModel._
 
@@ -195,14 +189,14 @@ class MinMaxScalerModel private[ml] (
   }
 
   @Since("1.6.0")
-  override def write: Writer = new MinMaxScalerModelWriter(this)
+  override def write: MLWriter = new MinMaxScalerModelWriter(this)
 }
 
 @Since("1.6.0")
-object MinMaxScalerModel extends Readable[MinMaxScalerModel] {
+object MinMaxScalerModel extends MLReadable[MinMaxScalerModel] {
 
   private[MinMaxScalerModel]
-  class MinMaxScalerModelWriter(instance: MinMaxScalerModel) extends Writer {
+  class MinMaxScalerModelWriter(instance: MinMaxScalerModel) extends MLWriter {
 
     private case class Data(originalMin: Vector, originalMax: Vector)
 
@@ -214,7 +208,7 @@ object MinMaxScalerModel extends Readable[MinMaxScalerModel] {
     }
   }
 
-  private class MinMaxScalerModelReader extends Reader[MinMaxScalerModel] {
+  private class MinMaxScalerModelReader extends MLReader[MinMaxScalerModel] {
 
     private val className = "org.apache.spark.ml.feature.MinMaxScalerModel"
 
@@ -231,7 +225,7 @@ object MinMaxScalerModel extends Readable[MinMaxScalerModel] {
   }
 
   @Since("1.6.0")
-  override def read: Reader[MinMaxScalerModel] = new MinMaxScalerModelReader
+  override def read: MLReader[MinMaxScalerModel] = new MinMaxScalerModelReader
 
   @Since("1.6.0")
   override def load(path: String): MinMaxScalerModel = super.load(path)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala
index 4a17acd95199f..65414ecbefbbd 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala
@@ -36,7 +36,7 @@ import org.apache.spark.sql.types.{ArrayType, DataType, StringType}
  */
 @Experimental
 class NGram(override val uid: String)
-  extends UnaryTransformer[Seq[String], Seq[String], NGram] with Writable {
+  extends UnaryTransformer[Seq[String], Seq[String], NGram] with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("ngram"))
 
@@ -66,17 +66,11 @@ class NGram(override val uid: String)
   }
 
   override protected def outputDataType: DataType = new ArrayType(StringType, false)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object NGram extends Readable[NGram] {
-
-  @Since("1.6.0")
-  override def read: Reader[NGram] = new DefaultParamsReader[NGram]
+object NGram extends DefaultParamsReadable[NGram] {
 
   @Since("1.6.0")
-  override def load(path: String): NGram = read.load(path)
+  override def load(path: String): NGram = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala
index 9df6a091d5058..c2d514fd9629e 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala
@@ -31,7 +31,7 @@ import org.apache.spark.sql.types.DataType
  */
 @Experimental
 class Normalizer(override val uid: String)
-  extends UnaryTransformer[Vector, Vector, Normalizer] with Writable {
+  extends UnaryTransformer[Vector, Vector, Normalizer] with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("normalizer"))
 
@@ -56,17 +56,11 @@ class Normalizer(override val uid: String)
   }
 
   override protected def outputDataType: DataType = new VectorUDT()
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object Normalizer extends Readable[Normalizer] {
-
-  @Since("1.6.0")
-  override def read: Reader[Normalizer] = new DefaultParamsReader[Normalizer]
+object Normalizer extends DefaultParamsReadable[Normalizer] {
 
   @Since("1.6.0")
-  override def load(path: String): Normalizer = read.load(path)
+  override def load(path: String): Normalizer = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
index 4e2adfaafa21e..d70164eaf0224 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
@@ -44,7 +44,7 @@ import org.apache.spark.sql.types.{DoubleType, StructType}
  */
 @Experimental
 class OneHotEncoder(override val uid: String) extends Transformer
-  with HasInputCol with HasOutputCol with Writable {
+  with HasInputCol with HasOutputCol with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("oneHot"))
 
@@ -165,17 +165,11 @@ class OneHotEncoder(override val uid: String) extends Transformer
   }
 
   override def copy(extra: ParamMap): OneHotEncoder = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object OneHotEncoder extends Readable[OneHotEncoder] {
-
-  @Since("1.6.0")
-  override def read: Reader[OneHotEncoder] = new DefaultParamsReader[OneHotEncoder]
+object OneHotEncoder extends DefaultParamsReadable[OneHotEncoder] {
 
   @Since("1.6.0")
-  override def load(path: String): OneHotEncoder = read.load(path)
+  override def load(path: String): OneHotEncoder = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala
index 49415398325fd..08610593fadda 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala
@@ -36,7 +36,7 @@ import org.apache.spark.sql.types.DataType
  */
 @Experimental
 class PolynomialExpansion(override val uid: String)
-  extends UnaryTransformer[Vector, Vector, PolynomialExpansion] with Writable {
+  extends UnaryTransformer[Vector, Vector, PolynomialExpansion] with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("poly"))
 
@@ -63,9 +63,6 @@ class PolynomialExpansion(override val uid: String)
   override protected def outputDataType: DataType = new VectorUDT()
 
   override def copy(extra: ParamMap): PolynomialExpansion = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 /**
@@ -81,7 +78,7 @@ class PolynomialExpansion(override val uid: String)
  * current index and increment it properly for sparse input.
  */
 @Since("1.6.0")
-object PolynomialExpansion extends Readable[PolynomialExpansion] {
+object PolynomialExpansion extends DefaultParamsReadable[PolynomialExpansion] {
 
   private def choose(n: Int, k: Int): Int = {
     Range(n, n - k, -1).product / Range(k, 1, -1).product
@@ -182,8 +179,5 @@ object PolynomialExpansion extends Readable[PolynomialExpansion] {
   }
 
   @Since("1.6.0")
-  override def read: Reader[PolynomialExpansion] = new DefaultParamsReader[PolynomialExpansion]
-
-  @Since("1.6.0")
-  override def load(path: String): PolynomialExpansion = read.load(path)
+  override def load(path: String): PolynomialExpansion = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/QuantileDiscretizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/QuantileDiscretizer.scala
index 2da5c966d2967..7bf67c6325a35 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/QuantileDiscretizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/QuantileDiscretizer.scala
@@ -60,7 +60,7 @@ private[feature] trait QuantileDiscretizerBase extends Params with HasInputCol w
  */
 @Experimental
 final class QuantileDiscretizer(override val uid: String)
-  extends Estimator[Bucketizer] with QuantileDiscretizerBase with Writable {
+  extends Estimator[Bucketizer] with QuantileDiscretizerBase with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("quantileDiscretizer"))
 
@@ -93,13 +93,10 @@ final class QuantileDiscretizer(override val uid: String)
   }
 
   override def copy(extra: ParamMap): QuantileDiscretizer = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object QuantileDiscretizer extends Readable[QuantileDiscretizer] with Logging {
+object QuantileDiscretizer extends DefaultParamsReadable[QuantileDiscretizer] with Logging {
   /**
    * Sampling from the given dataset to collect quantile statistics.
    */
@@ -179,8 +176,5 @@ object QuantileDiscretizer extends Readable[QuantileDiscretizer] with Logging {
   }
 
   @Since("1.6.0")
-  override def read: Reader[QuantileDiscretizer] = new DefaultParamsReader[QuantileDiscretizer]
-
-  @Since("1.6.0")
-  override def load(path: String): QuantileDiscretizer = read.load(path)
+  override def load(path: String): QuantileDiscretizer = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala
index c115064ff301a..3a735017ba836 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/SQLTransformer.scala
@@ -33,7 +33,8 @@ import org.apache.spark.sql.types.StructType
  */
 @Experimental
 @Since("1.6.0")
-class SQLTransformer @Since("1.6.0") (override val uid: String) extends Transformer with Writable {
+class SQLTransformer @Since("1.6.0") (override val uid: String) extends Transformer
+  with DefaultParamsWritable {
 
   @Since("1.6.0")
   def this() = this(Identifiable.randomUID("sql"))
@@ -77,17 +78,11 @@ class SQLTransformer @Since("1.6.0") (override val uid: String) extends Transfor
 
   @Since("1.6.0")
   override def copy(extra: ParamMap): SQLTransformer = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object SQLTransformer extends Readable[SQLTransformer] {
-
-  @Since("1.6.0")
-  override def read: Reader[SQLTransformer] = new DefaultParamsReader[SQLTransformer]
+object SQLTransformer extends DefaultParamsReadable[SQLTransformer] {
 
   @Since("1.6.0")
-  override def load(path: String): SQLTransformer = read.load(path)
+  override def load(path: String): SQLTransformer = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
index ab04e5418dd4f..1f689c1da1ba9 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
@@ -59,7 +59,7 @@ private[feature] trait StandardScalerParams extends Params with HasInputCol with
  */
 @Experimental
 class StandardScaler(override val uid: String) extends Estimator[StandardScalerModel]
-  with StandardScalerParams with Writable {
+  with StandardScalerParams with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("stdScal"))
 
@@ -96,16 +96,10 @@ class StandardScaler(override val uid: String) extends Estimator[StandardScalerM
   }
 
   override def copy(extra: ParamMap): StandardScaler = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object StandardScaler extends Readable[StandardScaler] {
-
-  @Since("1.6.0")
-  override def read: Reader[StandardScaler] = new DefaultParamsReader
+object StandardScaler extends DefaultParamsReadable[StandardScaler] {
 
   @Since("1.6.0")
   override def load(path: String): StandardScaler = super.load(path)
@@ -119,7 +113,7 @@ object StandardScaler extends Readable[StandardScaler] {
 class StandardScalerModel private[ml] (
     override val uid: String,
     scaler: feature.StandardScalerModel)
-  extends Model[StandardScalerModel] with StandardScalerParams with Writable {
+  extends Model[StandardScalerModel] with StandardScalerParams with MLWritable {
 
   import StandardScalerModel._
 
@@ -165,14 +159,14 @@ class StandardScalerModel private[ml] (
   }
 
   @Since("1.6.0")
-  override def write: Writer = new StandardScalerModelWriter(this)
+  override def write: MLWriter = new StandardScalerModelWriter(this)
 }
 
 @Since("1.6.0")
-object StandardScalerModel extends Readable[StandardScalerModel] {
+object StandardScalerModel extends MLReadable[StandardScalerModel] {
 
   private[StandardScalerModel]
-  class StandardScalerModelWriter(instance: StandardScalerModel) extends Writer {
+  class StandardScalerModelWriter(instance: StandardScalerModel) extends MLWriter {
 
     private case class Data(std: Vector, mean: Vector, withStd: Boolean, withMean: Boolean)
 
@@ -184,7 +178,7 @@ object StandardScalerModel extends Readable[StandardScalerModel] {
     }
   }
 
-  private class StandardScalerModelReader extends Reader[StandardScalerModel] {
+  private class StandardScalerModelReader extends MLReader[StandardScalerModel] {
 
     private val className = "org.apache.spark.ml.feature.StandardScalerModel"
 
@@ -204,7 +198,7 @@ object StandardScalerModel extends Readable[StandardScalerModel] {
   }
 
   @Since("1.6.0")
-  override def read: Reader[StandardScalerModel] = new StandardScalerModelReader
+  override def read: MLReader[StandardScalerModel] = new StandardScalerModelReader
 
   @Since("1.6.0")
   override def load(path: String): StandardScalerModel = super.load(path)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala
index f1146988dcc7c..318808596dc6a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala
@@ -86,7 +86,7 @@ private[spark] object StopWords {
  */
 @Experimental
 class StopWordsRemover(override val uid: String)
-  extends Transformer with HasInputCol with HasOutputCol with Writable {
+  extends Transformer with HasInputCol with HasOutputCol with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("stopWords"))
 
@@ -154,17 +154,11 @@ class StopWordsRemover(override val uid: String)
   }
 
   override def copy(extra: ParamMap): StopWordsRemover = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object StopWordsRemover extends Readable[StopWordsRemover] {
-
-  @Since("1.6.0")
-  override def read: Reader[StopWordsRemover] = new DefaultParamsReader[StopWordsRemover]
+object StopWordsRemover extends DefaultParamsReadable[StopWordsRemover] {
 
   @Since("1.6.0")
-  override def load(path: String): StopWordsRemover = read.load(path)
+  override def load(path: String): StopWordsRemover = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
index f16f6afc002d8..97a2e4f6d6ca4 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
@@ -65,7 +65,7 @@ private[feature] trait StringIndexerBase extends Params with HasInputCol with Ha
  */
 @Experimental
 class StringIndexer(override val uid: String) extends Estimator[StringIndexerModel]
-  with StringIndexerBase with Writable {
+  with StringIndexerBase with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("strIdx"))
 
@@ -93,16 +93,10 @@ class StringIndexer(override val uid: String) extends Estimator[StringIndexerMod
   }
 
   override def copy(extra: ParamMap): StringIndexer = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object StringIndexer extends Readable[StringIndexer] {
-
-  @Since("1.6.0")
-  override def read: Reader[StringIndexer] = new DefaultParamsReader
+object StringIndexer extends DefaultParamsReadable[StringIndexer] {
 
   @Since("1.6.0")
   override def load(path: String): StringIndexer = super.load(path)
@@ -122,7 +116,7 @@ object StringIndexer extends Readable[StringIndexer] {
 class StringIndexerModel (
     override val uid: String,
     val labels: Array[String])
-  extends Model[StringIndexerModel] with StringIndexerBase with Writable {
+  extends Model[StringIndexerModel] with StringIndexerBase with MLWritable {
 
   import StringIndexerModel._
 
@@ -199,10 +193,10 @@ class StringIndexerModel (
 }
 
 @Since("1.6.0")
-object StringIndexerModel extends Readable[StringIndexerModel] {
+object StringIndexerModel extends MLReadable[StringIndexerModel] {
 
   private[StringIndexerModel]
-  class StringIndexModelWriter(instance: StringIndexerModel) extends Writer {
+  class StringIndexModelWriter(instance: StringIndexerModel) extends MLWriter {
 
     private case class Data(labels: Array[String])
 
@@ -214,7 +208,7 @@ object StringIndexerModel extends Readable[StringIndexerModel] {
     }
   }
 
-  private class StringIndexerModelReader extends Reader[StringIndexerModel] {
+  private class StringIndexerModelReader extends MLReader[StringIndexerModel] {
 
     private val className = "org.apache.spark.ml.feature.StringIndexerModel"
 
@@ -232,7 +226,7 @@ object StringIndexerModel extends Readable[StringIndexerModel] {
   }
 
   @Since("1.6.0")
-  override def read: Reader[StringIndexerModel] = new StringIndexerModelReader
+  override def read: MLReader[StringIndexerModel] = new StringIndexerModelReader
 
   @Since("1.6.0")
   override def load(path: String): StringIndexerModel = super.load(path)
@@ -249,7 +243,7 @@ object StringIndexerModel extends Readable[StringIndexerModel] {
  */
 @Experimental
 class IndexToString private[ml] (override val uid: String)
-  extends Transformer with HasInputCol with HasOutputCol with Writable {
+  extends Transformer with HasInputCol with HasOutputCol with DefaultParamsWritable {
 
   def this() =
     this(Identifiable.randomUID("idxToStr"))
@@ -316,17 +310,11 @@ class IndexToString private[ml] (override val uid: String)
   override def copy(extra: ParamMap): IndexToString = {
     defaultCopy(extra)
   }
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object IndexToString extends Readable[IndexToString] {
-
-  @Since("1.6.0")
-  override def read: Reader[IndexToString] = new DefaultParamsReader[IndexToString]
+object IndexToString extends DefaultParamsReadable[IndexToString] {
 
   @Since("1.6.0")
-  override def load(path: String): IndexToString = read.load(path)
+  override def load(path: String): IndexToString = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
index 0e4445d1e2fa7..8ad7bbedaab5c 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Tokenizer.scala
@@ -31,7 +31,7 @@ import org.apache.spark.sql.types.{ArrayType, DataType, StringType}
  */
 @Experimental
 class Tokenizer(override val uid: String)
-  extends UnaryTransformer[String, Seq[String], Tokenizer] with Writable {
+  extends UnaryTransformer[String, Seq[String], Tokenizer] with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("tok"))
 
@@ -46,19 +46,13 @@ class Tokenizer(override val uid: String)
   override protected def outputDataType: DataType = new ArrayType(StringType, true)
 
   override def copy(extra: ParamMap): Tokenizer = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object Tokenizer extends Readable[Tokenizer] {
-
-  @Since("1.6.0")
-  override def read: Reader[Tokenizer] = new DefaultParamsReader[Tokenizer]
+object Tokenizer extends DefaultParamsReadable[Tokenizer] {
 
   @Since("1.6.0")
-  override def load(path: String): Tokenizer = read.load(path)
+  override def load(path: String): Tokenizer = super.load(path)
 }
 
 /**
@@ -70,7 +64,7 @@ object Tokenizer extends Readable[Tokenizer] {
  */
 @Experimental
 class RegexTokenizer(override val uid: String)
-  extends UnaryTransformer[String, Seq[String], RegexTokenizer] with Writable {
+  extends UnaryTransformer[String, Seq[String], RegexTokenizer] with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("regexTok"))
 
@@ -145,17 +139,11 @@ class RegexTokenizer(override val uid: String)
   override protected def outputDataType: DataType = new ArrayType(StringType, true)
 
   override def copy(extra: ParamMap): RegexTokenizer = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object RegexTokenizer extends Readable[RegexTokenizer] {
-
-  @Since("1.6.0")
-  override def read: Reader[RegexTokenizer] = new DefaultParamsReader[RegexTokenizer]
+object RegexTokenizer extends DefaultParamsReadable[RegexTokenizer] {
 
   @Since("1.6.0")
-  override def load(path: String): RegexTokenizer = read.load(path)
+  override def load(path: String): RegexTokenizer = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
index 7e54205292ca2..0feec0549852b 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
@@ -37,7 +37,7 @@ import org.apache.spark.sql.types._
  */
 @Experimental
 class VectorAssembler(override val uid: String)
-  extends Transformer with HasInputCols with HasOutputCol with Writable {
+  extends Transformer with HasInputCols with HasOutputCol with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("vecAssembler"))
 
@@ -120,19 +120,13 @@ class VectorAssembler(override val uid: String)
   }
 
   override def copy(extra: ParamMap): VectorAssembler = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object VectorAssembler extends Readable[VectorAssembler] {
-
-  @Since("1.6.0")
-  override def read: Reader[VectorAssembler] = new DefaultParamsReader[VectorAssembler]
+object VectorAssembler extends DefaultParamsReadable[VectorAssembler] {
 
   @Since("1.6.0")
-  override def load(path: String): VectorAssembler = read.load(path)
+  override def load(path: String): VectorAssembler = super.load(path)
 
   private[feature] def assemble(vv: Any*): Vector = {
     val indices = ArrayBuilder.make[Int]
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala
index 911582b55b574..5410a50bc2e47 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala
@@ -42,7 +42,7 @@ import org.apache.spark.sql.types.StructType
  */
 @Experimental
 final class VectorSlicer(override val uid: String)
-  extends Transformer with HasInputCol with HasOutputCol with Writable {
+  extends Transformer with HasInputCol with HasOutputCol with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("vectorSlicer"))
 
@@ -151,13 +151,10 @@ final class VectorSlicer(override val uid: String)
   }
 
   override def copy(extra: ParamMap): VectorSlicer = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object VectorSlicer extends Readable[VectorSlicer] {
+object VectorSlicer extends DefaultParamsReadable[VectorSlicer] {
 
   /** Return true if given feature indices are valid */
   private[feature] def validIndices(indices: Array[Int]): Boolean = {
@@ -174,8 +171,5 @@ object VectorSlicer extends Readable[VectorSlicer] {
   }
 
   @Since("1.6.0")
-  override def read: Reader[VectorSlicer] = new DefaultParamsReader[VectorSlicer]
-
-  @Since("1.6.0")
-  override def load(path: String): VectorSlicer = read.load(path)
+  override def load(path: String): VectorSlicer = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
index d92514d2e239e..795b73c4c2121 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
@@ -185,7 +185,7 @@ class ALSModel private[ml] (
     val rank: Int,
     @transient val userFactors: DataFrame,
     @transient val itemFactors: DataFrame)
-  extends Model[ALSModel] with ALSModelParams with Writable {
+  extends Model[ALSModel] with ALSModelParams with MLWritable {
 
   /** @group setParam */
   def setUserCol(value: String): this.type = set(userCol, value)
@@ -225,19 +225,19 @@ class ALSModel private[ml] (
   }
 
   @Since("1.6.0")
-  override def write: Writer = new ALSModel.ALSModelWriter(this)
+  override def write: MLWriter = new ALSModel.ALSModelWriter(this)
 }
 
 @Since("1.6.0")
-object ALSModel extends Readable[ALSModel] {
+object ALSModel extends MLReadable[ALSModel] {
 
   @Since("1.6.0")
-  override def read: Reader[ALSModel] = new ALSModelReader
+  override def read: MLReader[ALSModel] = new ALSModelReader
 
   @Since("1.6.0")
-  override def load(path: String): ALSModel = read.load(path)
+  override def load(path: String): ALSModel = super.load(path)
 
-  private[recommendation] class ALSModelWriter(instance: ALSModel) extends Writer {
+  private[recommendation] class ALSModelWriter(instance: ALSModel) extends MLWriter {
 
     override protected def saveImpl(path: String): Unit = {
       val extraMetadata = render("rank" -> instance.rank)
@@ -249,7 +249,7 @@ object ALSModel extends Readable[ALSModel] {
     }
   }
 
-  private[recommendation] class ALSModelReader extends Reader[ALSModel] {
+  private[recommendation] class ALSModelReader extends MLReader[ALSModel] {
 
     /** Checked against metadata when loading model */
     private val className = "org.apache.spark.ml.recommendation.ALSModel"
@@ -309,7 +309,8 @@ object ALSModel extends Readable[ALSModel] {
  * preferences rather than explicit ratings given to items.
  */
 @Experimental
-class ALS(override val uid: String) extends Estimator[ALSModel] with ALSParams with Writable {
+class ALS(override val uid: String) extends Estimator[ALSModel] with ALSParams
+  with DefaultParamsWritable {
 
   import org.apache.spark.ml.recommendation.ALS.Rating
 
@@ -391,9 +392,6 @@ class ALS(override val uid: String) extends Estimator[ALSModel] with ALSParams w
   }
 
   override def copy(extra: ParamMap): ALS = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 
@@ -406,7 +404,7 @@ class ALS(override val uid: String) extends Estimator[ALSModel] with ALSParams w
  * than 2 billion.
  */
 @DeveloperApi
-object ALS extends Readable[ALS] with Logging {
+object ALS extends DefaultParamsReadable[ALS] with Logging {
 
   /**
    * :: DeveloperApi ::
@@ -416,10 +414,7 @@ object ALS extends Readable[ALS] with Logging {
   case class Rating[@specialized(Int, Long) ID](user: ID, item: ID, rating: Float)
 
   @Since("1.6.0")
-  override def read: Reader[ALS] = new DefaultParamsReader[ALS]
-
-  @Since("1.6.0")
-  override def load(path: String): ALS = read.load(path)
+  override def load(path: String): ALS = super.load(path)
 
   /** Trait for least squares solvers applied to the normal equation. */
   private[recommendation] trait LeastSquaresNESolver extends Serializable {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index f7c44f0a51b8a..7ba1a60edaf71 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -66,7 +66,7 @@ private[regression] trait LinearRegressionParams extends PredictorParams
 @Experimental
 class LinearRegression @Since("1.3.0") (@Since("1.3.0") override val uid: String)
   extends Regressor[Vector, LinearRegression, LinearRegressionModel]
-  with LinearRegressionParams with Writable with Logging {
+  with LinearRegressionParams with DefaultParamsWritable with Logging {
 
   @Since("1.4.0")
   def this() = this(Identifiable.randomUID("linReg"))
@@ -345,19 +345,13 @@ class LinearRegression @Since("1.3.0") (@Since("1.3.0") override val uid: String
 
   @Since("1.4.0")
   override def copy(extra: ParamMap): LinearRegression = defaultCopy(extra)
-
-  @Since("1.6.0")
-  override def write: Writer = new DefaultParamsWriter(this)
 }
 
 @Since("1.6.0")
-object LinearRegression extends Readable[LinearRegression] {
-
-  @Since("1.6.0")
-  override def read: Reader[LinearRegression] = new DefaultParamsReader[LinearRegression]
+object LinearRegression extends DefaultParamsReadable[LinearRegression] {
 
   @Since("1.6.0")
-  override def load(path: String): LinearRegression = read.load(path)
+  override def load(path: String): LinearRegression = super.load(path)
 }
 
 /**
@@ -371,7 +365,7 @@ class LinearRegressionModel private[ml] (
     val coefficients: Vector,
     val intercept: Double)
   extends RegressionModel[Vector, LinearRegressionModel]
-  with LinearRegressionParams with Writable {
+  with LinearRegressionParams with MLWritable {
 
   private var trainingSummary: Option[LinearRegressionTrainingSummary] = None
 
@@ -441,7 +435,7 @@ class LinearRegressionModel private[ml] (
   }
 
   /**
-   * Returns a [[Writer]] instance for this ML instance.
+   * Returns a [[MLWriter]] instance for this ML instance.
    *
    * For [[LinearRegressionModel]], this does NOT currently save the training [[summary]].
    * An option to save [[summary]] may be added in the future.
@@ -449,21 +443,21 @@ class LinearRegressionModel private[ml] (
    * This also does not save the [[parent]] currently.
    */
   @Since("1.6.0")
-  override def write: Writer = new LinearRegressionModel.LinearRegressionModelWriter(this)
+  override def write: MLWriter = new LinearRegressionModel.LinearRegressionModelWriter(this)
 }
 
 @Since("1.6.0")
-object LinearRegressionModel extends Readable[LinearRegressionModel] {
+object LinearRegressionModel extends MLReadable[LinearRegressionModel] {
 
   @Since("1.6.0")
-  override def read: Reader[LinearRegressionModel] = new LinearRegressionModelReader
+  override def read: MLReader[LinearRegressionModel] = new LinearRegressionModelReader
 
   @Since("1.6.0")
-  override def load(path: String): LinearRegressionModel = read.load(path)
+  override def load(path: String): LinearRegressionModel = super.load(path)
 
-  /** [[Writer]] instance for [[LinearRegressionModel]] */
+  /** [[MLWriter]] instance for [[LinearRegressionModel]] */
   private[LinearRegressionModel] class LinearRegressionModelWriter(instance: LinearRegressionModel)
-    extends Writer with Logging {
+    extends MLWriter with Logging {
 
     private case class Data(intercept: Double, coefficients: Vector)
 
@@ -477,7 +471,7 @@ object LinearRegressionModel extends Readable[LinearRegressionModel] {
     }
   }
 
-  private class LinearRegressionModelReader extends Reader[LinearRegressionModel] {
+  private class LinearRegressionModelReader extends MLReader[LinearRegressionModel] {
 
     /** Checked against metadata when loading model */
     private val className = "org.apache.spark.ml.regression.LinearRegressionModel"
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index d8ce907af5323..ff9322dba122a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -31,7 +31,7 @@ import org.apache.spark.sql.SQLContext
 import org.apache.spark.util.Utils
 
 /**
- * Trait for [[Writer]] and [[Reader]].
+ * Trait for [[MLWriter]] and [[MLReader]].
  */
 private[util] sealed trait BaseReadWrite {
   private var optionSQLContext: Option[SQLContext] = None
@@ -64,7 +64,7 @@ private[util] sealed trait BaseReadWrite {
  */
 @Experimental
 @Since("1.6.0")
-abstract class Writer extends BaseReadWrite with Logging {
+abstract class MLWriter extends BaseReadWrite with Logging {
 
   protected var shouldOverwrite: Boolean = false
 
@@ -111,16 +111,16 @@ abstract class Writer extends BaseReadWrite with Logging {
 }
 
 /**
- * Trait for classes that provide [[Writer]].
+ * Trait for classes that provide [[MLWriter]].
  */
 @Since("1.6.0")
-trait Writable {
+trait MLWritable {
 
   /**
-   * Returns a [[Writer]] instance for this ML instance.
+   * Returns an [[MLWriter]] instance for this ML instance.
    */
   @Since("1.6.0")
-  def write: Writer
+  def write: MLWriter
 
   /**
    * Saves this ML instance to the input path, a shortcut of `write.save(path)`.
@@ -130,13 +130,18 @@ trait Writable {
   def save(path: String): Unit = write.save(path)
 }
 
+private[ml] trait DefaultParamsWritable extends MLWritable { self: Params =>
+
+  override def write: MLWriter = new DefaultParamsWriter(this)
+}
+
 /**
  * Abstract class for utility classes that can load ML instances.
  * @tparam T ML instance type
  */
 @Experimental
 @Since("1.6.0")
-abstract class Reader[T] extends BaseReadWrite {
+abstract class MLReader[T] extends BaseReadWrite {
 
   /**
    * Loads the ML component from the input path.
@@ -149,18 +154,18 @@ abstract class Reader[T] extends BaseReadWrite {
 }
 
 /**
- * Trait for objects that provide [[Reader]].
+ * Trait for objects that provide [[MLReader]].
  * @tparam T ML instance type
  */
 @Experimental
 @Since("1.6.0")
-trait Readable[T] {
+trait MLReadable[T] {
 
   /**
-   * Returns a [[Reader]] instance for this class.
+   * Returns an [[MLReader]] instance for this class.
    */
   @Since("1.6.0")
-  def read: Reader[T]
+  def read: MLReader[T]
 
   /**
    * Reads an ML instance from the input path, a shortcut of `read.load(path)`.
@@ -171,13 +176,18 @@ trait Readable[T] {
   def load(path: String): T = read.load(path)
 }
 
+private[ml] trait DefaultParamsReadable[T] extends MLReadable[T] {
+
+  override def read: MLReader[T] = new DefaultParamsReader
+}
+
 /**
- * Default [[Writer]] implementation for transformers and estimators that contain basic
+ * Default [[MLWriter]] implementation for transformers and estimators that contain basic
  * (json4s-serializable) params and no data. This will not handle more complex params or types with
  * data (e.g., models with coefficients).
  * @param instance object to save
  */
-private[ml] class DefaultParamsWriter(instance: Params) extends Writer {
+private[ml] class DefaultParamsWriter(instance: Params) extends MLWriter {
 
   override protected def saveImpl(path: String): Unit = {
     DefaultParamsWriter.saveMetadata(instance, path, sc)
@@ -218,13 +228,13 @@ private[ml] object DefaultParamsWriter {
 }
 
 /**
- * Default [[Reader]] implementation for transformers and estimators that contain basic
+ * Default [[MLReader]] implementation for transformers and estimators that contain basic
  * (json4s-serializable) params and no data. This will not handle more complex params or types with
  * data (e.g., models with coefficients).
  * @tparam T ML instance type
  * TODO: Consider adding check for correct class name.
  */
-private[ml] class DefaultParamsReader[T] extends Reader[T] {
+private[ml] class DefaultParamsReader[T] extends MLReader[T] {
 
   override def load(path: String): T = {
     val metadata = DefaultParamsReader.loadMetadata(path, sc)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
index 7f5c3895acb0c..12aba6bc6dbeb 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
@@ -179,8 +179,8 @@ class PipelineSuite extends SparkFunSuite with MLlibTestSparkContext with Defaul
 }
 
 
-/** Used to test [[Pipeline]] with [[Writable]] stages */
-class WritableStage(override val uid: String) extends Transformer with Writable {
+/** Used to test [[Pipeline]] with [[MLWritable]] stages */
+class WritableStage(override val uid: String) extends Transformer with MLWritable {
 
   final val intParam: IntParam = new IntParam(this, "intParam", "doc")
 
@@ -192,21 +192,21 @@ class WritableStage(override val uid: String) extends Transformer with Writable
 
   override def copy(extra: ParamMap): WritableStage = defaultCopy(extra)
 
-  override def write: Writer = new DefaultParamsWriter(this)
+  override def write: MLWriter = new DefaultParamsWriter(this)
 
   override def transform(dataset: DataFrame): DataFrame = dataset
 
   override def transformSchema(schema: StructType): StructType = schema
 }
 
-object WritableStage extends Readable[WritableStage] {
+object WritableStage extends MLReadable[WritableStage] {
 
-  override def read: Reader[WritableStage] = new DefaultParamsReader[WritableStage]
+  override def read: MLReader[WritableStage] = new DefaultParamsReader[WritableStage]
 
-  override def load(path: String): WritableStage = read.load(path)
+  override def load(path: String): WritableStage = super.load(path)
 }
 
-/** Used to test [[Pipeline]] with non-[[Writable]] stages */
+/** Used to test [[Pipeline]] with non-[[MLWritable]] stages */
 class UnWritableStage(override val uid: String) extends Transformer {
 
   final val intParam: IntParam = new IntParam(this, "intParam", "doc")
diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
index dd1e8acce9418..84d06b43d6224 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/DefaultReadWriteTest.scala
@@ -38,7 +38,7 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
    * @tparam T ML instance type
    * @return  Instance loaded from file
    */
-  def testDefaultReadWrite[T <: Params with Writable](
+  def testDefaultReadWrite[T <: Params with MLWritable](
       instance: T,
       testParams: Boolean = true): T = {
     val uid = instance.uid
@@ -52,7 +52,7 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
       instance.save(path)
     }
     instance.write.overwrite().save(path)
-    val loader = instance.getClass.getMethod("read").invoke(null).asInstanceOf[Reader[T]]
+    val loader = instance.getClass.getMethod("read").invoke(null).asInstanceOf[MLReader[T]]
     val newInstance = loader.load(path)
 
     assert(newInstance.uid === instance.uid)
@@ -92,7 +92,8 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
    * @tparam E  Type of [[Estimator]]
    * @tparam M  Type of [[Model]] produced by estimator
    */
-  def testEstimatorAndModelReadWrite[E <: Estimator[M] with Writable, M <: Model[M] with Writable](
+  def testEstimatorAndModelReadWrite[
+    E <: Estimator[M] with MLWritable, M <: Model[M] with MLWritable](
       estimator: E,
       dataset: DataFrame,
       testParams: Map[String, Any],
@@ -119,7 +120,7 @@ trait DefaultReadWriteTest extends TempDirectory { self: Suite =>
   }
 }
 
-class MyParams(override val uid: String) extends Params with Writable {
+class MyParams(override val uid: String) extends Params with MLWritable {
 
   final val intParamWithDefault: IntParam = new IntParam(this, "intParamWithDefault", "doc")
   final val intParam: IntParam = new IntParam(this, "intParam", "doc")
@@ -145,14 +146,14 @@ class MyParams(override val uid: String) extends Params with Writable {
 
   override def copy(extra: ParamMap): Params = defaultCopy(extra)
 
-  override def write: Writer = new DefaultParamsWriter(this)
+  override def write: MLWriter = new DefaultParamsWriter(this)
 }
 
-object MyParams extends Readable[MyParams] {
+object MyParams extends MLReadable[MyParams] {
 
-  override def read: Reader[MyParams] = new DefaultParamsReader[MyParams]
+  override def read: MLReader[MyParams] = new DefaultParamsReader[MyParams]
 
-  override def load(path: String): MyParams = read.load(path)
+  override def load(path: String): MyParams = super.load(path)
 }
 
 class DefaultReadWriteSuite extends SparkFunSuite with MLlibTestSparkContext

From e61367b9f9bfc8e123369d55d7ca5925568b98a7 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 18 Nov 2015 18:34:36 -0800
Subject: [PATCH 350/518] [SPARK-11833][SQL] Add Java tests for Kryo/Java
 Dataset encoders

Also added some nicer error messages for incompatible types (private types and primitive types) for Kryo/Java encoder.

Author: Reynold Xin <rxin@databricks.com>

Closes #9823 from rxin/SPARK-11833.
---
 .../scala/org/apache/spark/sql/Encoder.scala  | 69 +++++++++++------
 .../encoders/EncoderErrorMessageSuite.scala   | 40 ++++++++++
 .../catalyst/encoders/FlatEncoderSuite.scala  | 22 ++----
 .../apache/spark/sql/JavaDatasetSuite.java    | 75 ++++++++++++++++++-
 4 files changed, 166 insertions(+), 40 deletions(-)
 create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderErrorMessageSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index 1ed5111440c80..d54f2854fb33f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql
 
+import java.lang.reflect.Modifier
+
 import scala.reflect.{ClassTag, classTag}
 
 import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor}
@@ -43,30 +45,28 @@ trait Encoder[T] extends Serializable {
  */
 object Encoders {
 
-  /** A way to construct encoders using generic serializers. */
-  private def genericSerializer[T: ClassTag](useKryo: Boolean): Encoder[T] = {
-    ExpressionEncoder[T](
-      schema = new StructType().add("value", BinaryType),
-      flat = true,
-      toRowExpressions = Seq(
-        EncodeUsingSerializer(
-          BoundReference(0, ObjectType(classOf[AnyRef]), nullable = true), kryo = useKryo)),
-      fromRowExpression =
-        DecodeUsingSerializer[T](
-          BoundReference(0, BinaryType, nullable = true), classTag[T], kryo = useKryo),
-      clsTag = classTag[T]
-    )
-  }
+  def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder(flat = true)
+  def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder(flat = true)
+  def SHORT: Encoder[java.lang.Short] = ExpressionEncoder(flat = true)
+  def INT: Encoder[java.lang.Integer] = ExpressionEncoder(flat = true)
+  def LONG: Encoder[java.lang.Long] = ExpressionEncoder(flat = true)
+  def FLOAT: Encoder[java.lang.Float] = ExpressionEncoder(flat = true)
+  def DOUBLE: Encoder[java.lang.Double] = ExpressionEncoder(flat = true)
+  def STRING: Encoder[java.lang.String] = ExpressionEncoder(flat = true)
 
   /**
    * (Scala-specific) Creates an encoder that serializes objects of type T using Kryo.
    * This encoder maps T into a single byte array (binary) field.
+   *
+   * T must be publicly accessible.
    */
   def kryo[T: ClassTag]: Encoder[T] = genericSerializer(useKryo = true)
 
   /**
    * Creates an encoder that serializes objects of type T using Kryo.
    * This encoder maps T into a single byte array (binary) field.
+   *
+   * T must be publicly accessible.
    */
   def kryo[T](clazz: Class[T]): Encoder[T] = kryo(ClassTag[T](clazz))
 
@@ -75,6 +75,8 @@ object Encoders {
    * serialization. This encoder maps T into a single byte array (binary) field.
    *
    * Note that this is extremely inefficient and should only be used as the last resort.
+   *
+   * T must be publicly accessible.
    */
   def javaSerialization[T: ClassTag]: Encoder[T] = genericSerializer(useKryo = false)
 
@@ -83,17 +85,40 @@ object Encoders {
    * This encoder maps T into a single byte array (binary) field.
    *
    * Note that this is extremely inefficient and should only be used as the last resort.
+   *
+   * T must be publicly accessible.
    */
   def javaSerialization[T](clazz: Class[T]): Encoder[T] = javaSerialization(ClassTag[T](clazz))
 
-  def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder(flat = true)
-  def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder(flat = true)
-  def SHORT: Encoder[java.lang.Short] = ExpressionEncoder(flat = true)
-  def INT: Encoder[java.lang.Integer] = ExpressionEncoder(flat = true)
-  def LONG: Encoder[java.lang.Long] = ExpressionEncoder(flat = true)
-  def FLOAT: Encoder[java.lang.Float] = ExpressionEncoder(flat = true)
-  def DOUBLE: Encoder[java.lang.Double] = ExpressionEncoder(flat = true)
-  def STRING: Encoder[java.lang.String] = ExpressionEncoder(flat = true)
+  /** Throws an exception if T is not a public class. */
+  private def validatePublicClass[T: ClassTag](): Unit = {
+    if (!Modifier.isPublic(classTag[T].runtimeClass.getModifiers)) {
+      throw new UnsupportedOperationException(
+        s"${classTag[T].runtimeClass.getName} is not a public class. " +
+          "Only public classes are supported.")
+    }
+  }
+
+  /** A way to construct encoders using generic serializers. */
+  private def genericSerializer[T: ClassTag](useKryo: Boolean): Encoder[T] = {
+    if (classTag[T].runtimeClass.isPrimitive) {
+      throw new UnsupportedOperationException("Primitive types are not supported.")
+    }
+
+    validatePublicClass[T]()
+
+    ExpressionEncoder[T](
+      schema = new StructType().add("value", BinaryType),
+      flat = true,
+      toRowExpressions = Seq(
+        EncodeUsingSerializer(
+          BoundReference(0, ObjectType(classOf[AnyRef]), nullable = true), kryo = useKryo)),
+      fromRowExpression =
+        DecodeUsingSerializer[T](
+          BoundReference(0, BinaryType, nullable = true), classTag[T], kryo = useKryo),
+      clsTag = classTag[T]
+    )
+  }
 
   def tuple[T1, T2](
       e1: Encoder[T1],
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderErrorMessageSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderErrorMessageSuite.scala
new file mode 100644
index 0000000000000..0b2a10bb04c10
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderErrorMessageSuite.scala
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.encoders
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.Encoders
+
+
+class EncoderErrorMessageSuite extends SparkFunSuite {
+
+  // Note: we also test error messages for encoders for private classes in JavaDatasetSuite.
+  // That is done in Java because Scala cannot create truly private classes.
+
+  test("primitive types in encoders using Kryo serialization") {
+    intercept[UnsupportedOperationException] { Encoders.kryo[Int] }
+    intercept[UnsupportedOperationException] { Encoders.kryo[Long] }
+    intercept[UnsupportedOperationException] { Encoders.kryo[Char] }
+  }
+
+  test("primitive types in encoders using Java serialization") {
+    intercept[UnsupportedOperationException] { Encoders.javaSerialization[Int] }
+    intercept[UnsupportedOperationException] { Encoders.javaSerialization[Long] }
+    intercept[UnsupportedOperationException] { Encoders.javaSerialization[Char] }
+  }
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
index 6e0322fb6e019..07523d49f4266 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
@@ -74,24 +74,14 @@ class FlatEncoderSuite extends ExpressionEncoderSuite {
     FlatEncoder[Map[Int, Map[String, Int]]], "map of map")
 
   // Kryo encoders
-  encodeDecodeTest(
-    "hello",
-    encoderFor(Encoders.kryo[String]),
-    "kryo string")
-  encodeDecodeTest(
-    new KryoSerializable(15),
-    encoderFor(Encoders.kryo[KryoSerializable]),
-    "kryo object serialization")
+  encodeDecodeTest("hello", encoderFor(Encoders.kryo[String]), "kryo string")
+  encodeDecodeTest(new KryoSerializable(15),
+    encoderFor(Encoders.kryo[KryoSerializable]), "kryo object")
 
   // Java encoders
-  encodeDecodeTest(
-    "hello",
-    encoderFor(Encoders.javaSerialization[String]),
-    "java string")
-  encodeDecodeTest(
-    new JavaSerializable(15),
-    encoderFor(Encoders.javaSerialization[JavaSerializable]),
-    "java object serialization")
+  encodeDecodeTest("hello", encoderFor(Encoders.javaSerialization[String]), "java string")
+  encodeDecodeTest(new JavaSerializable(15),
+    encoderFor(Encoders.javaSerialization[JavaSerializable]), "java object")
 }
 
 /** For testing Kryo serialization based encoder. */
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index d9b22506fbd3b..ce40dd856f679 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -24,6 +24,7 @@
 import scala.Tuple3;
 import scala.Tuple4;
 import scala.Tuple5;
+
 import org.junit.*;
 
 import org.apache.spark.Accumulator;
@@ -410,8 +411,8 @@ public String call(Tuple2<String, Integer> value) throws Exception {
       .as(Encoders.tuple(Encoders.STRING(), Encoders.INT(), Encoders.LONG(), Encoders.LONG()));
     Assert.assertEquals(
       Arrays.asList(
-        new Tuple4<String, Integer, Long, Long>("a", 3, 3L, 2L),
-        new Tuple4<String, Integer, Long, Long>("b", 3, 3L, 1L)),
+        new Tuple4<>("a", 3, 3L, 2L),
+        new Tuple4<>("b", 3, 3L, 1L)),
       agged2.collectAsList());
   }
 
@@ -437,4 +438,74 @@ public Integer finish(Integer reduction) {
       return reduction;
     }
   }
+
+  public static class KryoSerializable {
+    String value;
+
+    KryoSerializable(String value) {
+      this.value = value;
+    }
+
+    @Override
+    public boolean equals(Object other) {
+      return this.value.equals(((KryoSerializable) other).value);
+    }
+
+    @Override
+    public int hashCode() {
+      return this.value.hashCode();
+    }
+  }
+
+  public static class JavaSerializable implements Serializable {
+    String value;
+
+    JavaSerializable(String value) {
+      this.value = value;
+    }
+
+    @Override
+    public boolean equals(Object other) {
+      return this.value.equals(((JavaSerializable) other).value);
+    }
+
+    @Override
+    public int hashCode() {
+      return this.value.hashCode();
+    }
+  }
+
+  @Test
+  public void testKryoEncoder() {
+    Encoder<KryoSerializable> encoder = Encoders.kryo(KryoSerializable.class);
+    List<KryoSerializable> data = Arrays.asList(
+      new KryoSerializable("hello"), new KryoSerializable("world"));
+    Dataset<KryoSerializable> ds = context.createDataset(data, encoder);
+    Assert.assertEquals(data, ds.collectAsList());
+  }
+
+  @Test
+  public void testJavaEncoder() {
+    Encoder<JavaSerializable> encoder = Encoders.javaSerialization(JavaSerializable.class);
+    List<JavaSerializable> data = Arrays.asList(
+      new JavaSerializable("hello"), new JavaSerializable("world"));
+    Dataset<JavaSerializable> ds = context.createDataset(data, encoder);
+    Assert.assertEquals(data, ds.collectAsList());
+  }
+
+  /**
+   * For testing error messages when creating an encoder on a private class. This is done
+   * here since we cannot create truly private classes in Scala.
+   */
+  private static class PrivateClassTest { }
+
+  @Test(expected = UnsupportedOperationException.class)
+  public void testJavaEncoderErrorMessageForPrivateClass() {
+    Encoders.javaSerialization(PrivateClassTest.class);
+  }
+
+  @Test(expected = UnsupportedOperationException.class)
+  public void testKryoEncoderErrorMessageForPrivateClass() {
+    Encoders.kryo(PrivateClassTest.class);
+  }
 }

From 6d0848b53bbe6c5acdcf5c033cd396b1ae6e293d Mon Sep 17 00:00:00 2001
From: Nong Li <nong@databricks.com>
Date: Wed, 18 Nov 2015 18:38:45 -0800
Subject: [PATCH 351/518] [SPARK-11787][SQL] Improve Parquet scan performance
 when using flat schemas.

This patch adds an alternate to the Parquet RecordReader from the parquet-mr project
that is much faster for flat schemas. Instead of using the general converter mechanism
from parquet-mr, this directly uses the lower level APIs from parquet-columnar and a
customer RecordReader that directly assembles into UnsafeRows.

This is optionally disabled and only used for supported schemas.

Using the tpcds store sales table and doing a sum of increasingly more columns, the results
are:

For 1 Column:
  Before: 11.3M rows/second
  After: 18.2M rows/second

For 2 Columns:
  Before: 7.2M rows/second
  After: 11.2M rows/second

For 5 Columns:
  Before: 2.9M rows/second
  After: 4.5M rows/second

Author: Nong Li <nong@databricks.com>

Closes #9774 from nongli/parquet.
---
 .../apache/spark/rdd/SqlNewHadoopRDD.scala    |  41 +-
 .../sql/catalyst/expressions/UnsafeRow.java   |   9 +
 .../expressions/codegen/BufferHolder.java     |  32 +-
 .../expressions/codegen/UnsafeRowWriter.java  |  20 +-
 .../SpecificParquetRecordReaderBase.java      | 240 +++++++
 .../parquet/UnsafeRowParquetRecordReader.java | 593 ++++++++++++++++++
 .../parquet/CatalystRowConverter.scala        |  48 +-
 .../parquet/ParquetFilterSuite.scala          |   4 +-
 8 files changed, 944 insertions(+), 43 deletions(-)
 create mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
 create mode 100644 sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/UnsafeRowParquetRecordReader.java

diff --git a/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDD.scala b/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDD.scala
index 264dae7f39085..4d176332b69ce 100644
--- a/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDD.scala
@@ -20,8 +20,6 @@ package org.apache.spark.rdd
 import java.text.SimpleDateFormat
 import java.util.Date
 
-import scala.reflect.ClassTag
-
 import org.apache.hadoop.conf.{Configurable, Configuration}
 import org.apache.hadoop.io.Writable
 import org.apache.hadoop.mapreduce._
@@ -30,10 +28,12 @@ import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.executor.DataReadMethod
 import org.apache.spark.mapreduce.SparkHadoopMapReduceUtil
+import org.apache.spark.storage.StorageLevel
 import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.util.{Utils, SerializableConfiguration, ShutdownHookManager}
 import org.apache.spark.{Partition => SparkPartition, _}
-import org.apache.spark.storage.StorageLevel
-import org.apache.spark.util.{SerializableConfiguration, ShutdownHookManager, Utils}
+
+import scala.reflect.ClassTag
 
 
 private[spark] class SqlNewHadoopPartition(
@@ -96,6 +96,11 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
 
   @transient protected val jobId = new JobID(jobTrackerId, id)
 
+  // If true, enable using the custom RecordReader for parquet. This only works for
+  // a subset of the types (no complex types).
+  protected val enableUnsafeRowParquetReader: Boolean =
+      sc.conf.getBoolean("spark.parquet.enableUnsafeRowRecordReader", true)
+
   override def getPartitions: Array[SparkPartition] = {
     val conf = getConf(isDriverSide = true)
     val inputFormat = inputFormatClass.newInstance
@@ -150,9 +155,31 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
           configurable.setConf(conf)
         case _ =>
       }
-      private[this] var reader = format.createRecordReader(
-        split.serializableHadoopSplit.value, hadoopAttemptContext)
-      reader.initialize(split.serializableHadoopSplit.value, hadoopAttemptContext)
+      private[this] var reader: RecordReader[Void, V] = null
+
+      /**
+        * If the format is ParquetInputFormat, try to create the optimized RecordReader. If this
+        * fails (for example, unsupported schema), try with the normal reader.
+        * TODO: plumb this through a different way?
+        */
+      if (enableUnsafeRowParquetReader &&
+          format.getClass.getName == "org.apache.parquet.hadoop.ParquetInputFormat") {
+        // TODO: move this class to sql.execution and remove this.
+        reader = Utils.classForName(
+          "org.apache.spark.sql.execution.datasources.parquet.UnsafeRowParquetRecordReader")
+          .newInstance().asInstanceOf[RecordReader[Void, V]]
+        try {
+          reader.initialize(split.serializableHadoopSplit.value, hadoopAttemptContext)
+        } catch {
+          case e: Exception => reader = null
+        }
+      }
+
+      if (reader == null) {
+        reader = format.createRecordReader(
+          split.serializableHadoopSplit.value, hadoopAttemptContext)
+        reader.initialize(split.serializableHadoopSplit.value, hadoopAttemptContext)
+      }
 
       // Register an on-task-completion callback to close the input stream.
       context.addTaskCompletionListener(context => close())
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
index 5ba14ebdb62a4..33769363a0ed5 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
@@ -178,6 +178,15 @@ public void pointTo(byte[] buf, int numFields, int sizeInBytes) {
     pointTo(buf, Platform.BYTE_ARRAY_OFFSET, numFields, sizeInBytes);
   }
 
+  /**
+   * Updates this UnsafeRow preserving the number of fields.
+   * @param buf byte array to point to
+   * @param sizeInBytes the number of bytes valid in the byte array
+   */
+  public void pointTo(byte[] buf, int sizeInBytes) {
+    pointTo(buf, numFields, sizeInBytes);
+  }
+
   @Override
   public void setNullAt(int i) {
     assertIndexIsValid(i);
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/BufferHolder.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/BufferHolder.java
index 9c9468678065d..d26b1b187c27b 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/BufferHolder.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/BufferHolder.java
@@ -17,19 +17,28 @@
 
 package org.apache.spark.sql.catalyst.expressions.codegen;
 
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
 import org.apache.spark.unsafe.Platform;
 
 /**
- * A helper class to manage the row buffer used in `GenerateUnsafeProjection`.
- *
- * Note that it is only used in `GenerateUnsafeProjection`, so it's safe to mark member variables
- * public for ease of use.
+ * A helper class to manage the row buffer when construct unsafe rows.
  */
 public class BufferHolder {
-  public byte[] buffer = new byte[64];
+  public byte[] buffer;
   public int cursor = Platform.BYTE_ARRAY_OFFSET;
 
-  public void grow(int neededSize) {
+  public BufferHolder() {
+    this(64);
+  }
+
+  public BufferHolder(int size) {
+    buffer = new byte[size];
+  }
+
+  /**
+   * Grows the buffer to at least neededSize. If row is non-null, points the row to the buffer.
+   */
+  public void grow(int neededSize, UnsafeRow row) {
     final int length = totalSize() + neededSize;
     if (buffer.length < length) {
       // This will not happen frequently, because the buffer is re-used.
@@ -41,12 +50,23 @@ public void grow(int neededSize) {
         Platform.BYTE_ARRAY_OFFSET,
         totalSize());
       buffer = tmp;
+      if (row != null) {
+        row.pointTo(buffer, length * 2);
+      }
     }
   }
 
+  public void grow(int neededSize) {
+    grow(neededSize, null);
+  }
+
   public void reset() {
     cursor = Platform.BYTE_ARRAY_OFFSET;
   }
+  public void resetTo(int offset) {
+    assert(offset <= buffer.length);
+    cursor = Platform.BYTE_ARRAY_OFFSET + offset;
+  }
 
   public int totalSize() {
     return cursor - Platform.BYTE_ARRAY_OFFSET;
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriter.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriter.java
index 048b7749d8fb4..e227c0dec9748 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriter.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/codegen/UnsafeRowWriter.java
@@ -35,6 +35,7 @@ public class UnsafeRowWriter {
   // The offset of the global buffer where we start to write this row.
   private int startingOffset;
   private int nullBitsSize;
+  private UnsafeRow row;
 
   public void initialize(BufferHolder holder, int numFields) {
     this.holder = holder;
@@ -43,7 +44,7 @@ public void initialize(BufferHolder holder, int numFields) {
 
     // grow the global buffer to make sure it has enough space to write fixed-length data.
     final int fixedSize = nullBitsSize + 8 * numFields;
-    holder.grow(fixedSize);
+    holder.grow(fixedSize, row);
     holder.cursor += fixedSize;
 
     // zero-out the null bits region
@@ -52,12 +53,19 @@ public void initialize(BufferHolder holder, int numFields) {
     }
   }
 
+  public void initialize(UnsafeRow row, BufferHolder holder, int numFields) {
+    initialize(holder, numFields);
+    this.row = row;
+  }
+
   private void zeroOutPaddingBytes(int numBytes) {
     if ((numBytes & 0x07) > 0) {
       Platform.putLong(holder.buffer, holder.cursor + ((numBytes >> 3) << 3), 0L);
     }
   }
 
+  public BufferHolder holder() { return holder; }
+
   public boolean isNullAt(int ordinal) {
     return BitSetMethods.isSet(holder.buffer, startingOffset, ordinal);
   }
@@ -90,7 +98,7 @@ public void alignToWords(int numBytes) {
 
     if (remainder > 0) {
       final int paddingBytes = 8 - remainder;
-      holder.grow(paddingBytes);
+      holder.grow(paddingBytes, row);
 
       for (int i = 0; i < paddingBytes; i++) {
         Platform.putByte(holder.buffer, holder.cursor, (byte) 0);
@@ -153,7 +161,7 @@ public void write(int ordinal, Decimal input, int precision, int scale) {
       }
     } else {
       // grow the global buffer before writing data.
-      holder.grow(16);
+      holder.grow(16, row);
 
       // zero-out the bytes
       Platform.putLong(holder.buffer, holder.cursor, 0L);
@@ -185,7 +193,7 @@ public void write(int ordinal, UTF8String input) {
     final int roundedSize = ByteArrayMethods.roundNumberOfBytesToNearestWord(numBytes);
 
     // grow the global buffer before writing data.
-    holder.grow(roundedSize);
+    holder.grow(roundedSize, row);
 
     zeroOutPaddingBytes(numBytes);
 
@@ -206,7 +214,7 @@ public void write(int ordinal, byte[] input, int offset, int numBytes) {
     final int roundedSize = ByteArrayMethods.roundNumberOfBytesToNearestWord(numBytes);
 
     // grow the global buffer before writing data.
-    holder.grow(roundedSize);
+    holder.grow(roundedSize, row);
 
     zeroOutPaddingBytes(numBytes);
 
@@ -222,7 +230,7 @@ public void write(int ordinal, byte[] input, int offset, int numBytes) {
 
   public void write(int ordinal, CalendarInterval input) {
     // grow the global buffer before writing data.
-    holder.grow(16);
+    holder.grow(16, row);
 
     // Write the months and microseconds fields of Interval to the variable length portion.
     Platform.putLong(holder.buffer, holder.cursor, input.months);
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
new file mode 100644
index 0000000000000..2ed30c1f5a8d9
--- /dev/null
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/SpecificParquetRecordReaderBase.java
@@ -0,0 +1,240 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+package org.apache.spark.sql.execution.datasources.parquet;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.apache.parquet.filter2.compat.RowGroupFilter.filterRowGroups;
+import static org.apache.parquet.format.converter.ParquetMetadataConverter.NO_FILTER;
+import static org.apache.parquet.format.converter.ParquetMetadataConverter.range;
+import static org.apache.parquet.hadoop.ParquetFileReader.readFooter;
+import static org.apache.parquet.hadoop.ParquetInputFormat.getFilter;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.mapreduce.InputSplit;
+import org.apache.hadoop.mapreduce.RecordReader;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+import org.apache.parquet.bytes.BytesInput;
+import org.apache.parquet.bytes.BytesUtils;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.values.ValuesReader;
+import org.apache.parquet.column.values.rle.RunLengthBitPackingHybridDecoder;
+import org.apache.parquet.filter2.compat.FilterCompat;
+import org.apache.parquet.hadoop.BadConfigurationException;
+import org.apache.parquet.hadoop.ParquetFileReader;
+import org.apache.parquet.hadoop.ParquetInputFormat;
+import org.apache.parquet.hadoop.ParquetInputSplit;
+import org.apache.parquet.hadoop.api.InitContext;
+import org.apache.parquet.hadoop.api.ReadSupport;
+import org.apache.parquet.hadoop.metadata.BlockMetaData;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+import org.apache.parquet.hadoop.util.ConfigurationUtil;
+import org.apache.parquet.schema.MessageType;
+
+/**
+ * Base class for custom RecordReaaders for Parquet that directly materialize to `T`.
+ * This class handles computing row groups, filtering on them, setting up the column readers,
+ * etc.
+ * This is heavily based on parquet-mr's RecordReader.
+ * TODO: move this to the parquet-mr project. There are performance benefits of doing it
+ * this way, albeit at a higher cost to implement. This base class is reusable.
+ */
+public abstract class SpecificParquetRecordReaderBase<T> extends RecordReader<Void, T> {
+  protected Path file;
+  protected MessageType fileSchema;
+  protected MessageType requestedSchema;
+  protected ReadSupport<T> readSupport;
+
+  /**
+   * The total number of rows this RecordReader will eventually read. The sum of the
+   * rows of all the row groups.
+   */
+  protected long totalRowCount;
+
+  protected ParquetFileReader reader;
+
+  public void initialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext)
+      throws IOException, InterruptedException {
+    Configuration configuration = taskAttemptContext.getConfiguration();
+    ParquetInputSplit split = (ParquetInputSplit)inputSplit;
+    this.file = split.getPath();
+    long[] rowGroupOffsets = split.getRowGroupOffsets();
+
+    ParquetMetadata footer;
+    List<BlockMetaData> blocks;
+
+    // if task.side.metadata is set, rowGroupOffsets is null
+    if (rowGroupOffsets == null) {
+      // then we need to apply the predicate push down filter
+      footer = readFooter(configuration, file, range(split.getStart(), split.getEnd()));
+      MessageType fileSchema = footer.getFileMetaData().getSchema();
+      FilterCompat.Filter filter = getFilter(configuration);
+      blocks = filterRowGroups(filter, footer.getBlocks(), fileSchema);
+    } else {
+      // otherwise we find the row groups that were selected on the client
+      footer = readFooter(configuration, file, NO_FILTER);
+      Set<Long> offsets = new HashSet<>();
+      for (long offset : rowGroupOffsets) {
+        offsets.add(offset);
+      }
+      blocks = new ArrayList<>();
+      for (BlockMetaData block : footer.getBlocks()) {
+        if (offsets.contains(block.getStartingPos())) {
+          blocks.add(block);
+        }
+      }
+      // verify we found them all
+      if (blocks.size() != rowGroupOffsets.length) {
+        long[] foundRowGroupOffsets = new long[footer.getBlocks().size()];
+        for (int i = 0; i < foundRowGroupOffsets.length; i++) {
+          foundRowGroupOffsets[i] = footer.getBlocks().get(i).getStartingPos();
+        }
+        // this should never happen.
+        // provide a good error message in case there's a bug
+        throw new IllegalStateException(
+            "All the offsets listed in the split should be found in the file."
+                + " expected: " + Arrays.toString(rowGroupOffsets)
+                + " found: " + blocks
+                + " out of: " + Arrays.toString(foundRowGroupOffsets)
+                + " in range " + split.getStart() + ", " + split.getEnd());
+      }
+    }
+    MessageType fileSchema = footer.getFileMetaData().getSchema();
+    Map<String, String> fileMetadata = footer.getFileMetaData().getKeyValueMetaData();
+    this.readSupport = getReadSupportInstance(
+        (Class<? extends ReadSupport<T>>) getReadSupportClass(configuration));
+    ReadSupport.ReadContext readContext = readSupport.init(new InitContext(
+        taskAttemptContext.getConfiguration(), toSetMultiMap(fileMetadata), fileSchema));
+    this.requestedSchema = readContext.getRequestedSchema();
+    this.fileSchema = fileSchema;
+    this.reader = new ParquetFileReader(configuration, file, blocks, requestedSchema.getColumns());
+    for (BlockMetaData block : blocks) {
+      this.totalRowCount += block.getRowCount();
+    }
+  }
+
+  @Override
+  public Void getCurrentKey() throws IOException, InterruptedException {
+    return null;
+  }
+
+  @Override
+  public void close() throws IOException {
+    if (reader != null) {
+      reader.close();
+      reader = null;
+    }
+  }
+
+  /**
+   * Utility classes to abstract over different way to read ints with different encodings.
+   * TODO: remove this layer of abstraction?
+   */
+  abstract static class IntIterator {
+    abstract int nextInt() throws IOException;
+  }
+
+  protected static final class ValuesReaderIntIterator extends IntIterator {
+    ValuesReader delegate;
+
+    public ValuesReaderIntIterator(ValuesReader delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    int nextInt() throws IOException {
+      return delegate.readInteger();
+    }
+  }
+
+  protected static final class RLEIntIterator extends IntIterator {
+    RunLengthBitPackingHybridDecoder delegate;
+
+    public RLEIntIterator(RunLengthBitPackingHybridDecoder delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    int nextInt() throws IOException {
+      return delegate.readInt();
+    }
+  }
+
+  protected static final class NullIntIterator extends IntIterator {
+    @Override
+    int nextInt() throws IOException { return 0; }
+  }
+
+  /**
+   * Creates a reader for definition and repetition levels, returning an optimized one if
+   * the levels are not needed.
+   */
+  static protected IntIterator createRLEIterator(int maxLevel, BytesInput bytes,
+                                              ColumnDescriptor descriptor) throws IOException {
+    try {
+      if (maxLevel == 0) return new NullIntIterator();
+      return new RLEIntIterator(
+          new RunLengthBitPackingHybridDecoder(
+              BytesUtils.getWidthFromMaxInt(maxLevel),
+              new ByteArrayInputStream(bytes.toByteArray())));
+    } catch (IOException e) {
+      throw new IOException("could not read levels in page for col " + descriptor, e);
+    }
+  }
+
+  private static <K, V> Map<K, Set<V>> toSetMultiMap(Map<K, V> map) {
+    Map<K, Set<V>> setMultiMap = new HashMap<>();
+    for (Map.Entry<K, V> entry : map.entrySet()) {
+      Set<V> set = new HashSet<>();
+      set.add(entry.getValue());
+      setMultiMap.put(entry.getKey(), Collections.unmodifiableSet(set));
+    }
+    return Collections.unmodifiableMap(setMultiMap);
+  }
+
+  private static Class<?> getReadSupportClass(Configuration configuration) {
+    return ConfigurationUtil.getClassFromConfig(configuration,
+        ParquetInputFormat.READ_SUPPORT_CLASS, ReadSupport.class);
+  }
+
+  /**
+   * @param readSupportClass to instantiate
+   * @return the configured read support
+   */
+  private static <T> ReadSupport<T> getReadSupportInstance(
+      Class<? extends ReadSupport<T>> readSupportClass){
+    try {
+      return readSupportClass.newInstance();
+    } catch (InstantiationException e) {
+      throw new BadConfigurationException("could not instantiate read support class", e);
+    } catch (IllegalAccessException e) {
+      throw new BadConfigurationException("could not instantiate read support class", e);
+    }
+  }
+}
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/UnsafeRowParquetRecordReader.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/UnsafeRowParquetRecordReader.java
new file mode 100644
index 0000000000000..8a92e489ccb7c
--- /dev/null
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/UnsafeRowParquetRecordReader.java
@@ -0,0 +1,593 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.parquet;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
+
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
+import org.apache.spark.sql.catalyst.expressions.codegen.BufferHolder;
+import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.types.UTF8String;
+
+import static org.apache.parquet.column.ValuesType.DEFINITION_LEVEL;
+import static org.apache.parquet.column.ValuesType.REPETITION_LEVEL;
+import static org.apache.parquet.column.ValuesType.VALUES;
+
+import org.apache.hadoop.mapreduce.InputSplit;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+import org.apache.parquet.Preconditions;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.Dictionary;
+import org.apache.parquet.column.Encoding;
+import org.apache.parquet.column.page.DataPage;
+import org.apache.parquet.column.page.DataPageV1;
+import org.apache.parquet.column.page.DataPageV2;
+import org.apache.parquet.column.page.DictionaryPage;
+import org.apache.parquet.column.page.PageReadStore;
+import org.apache.parquet.column.page.PageReader;
+import org.apache.parquet.column.values.ValuesReader;
+import org.apache.parquet.io.api.Binary;
+import org.apache.parquet.schema.OriginalType;
+import org.apache.parquet.schema.PrimitiveType;
+import org.apache.parquet.schema.Type;
+
+/**
+ * A specialized RecordReader that reads into UnsafeRows directly using the Parquet column APIs.
+ *
+ * This is somewhat based on parquet-mr's ColumnReader.
+ *
+ * TODO: handle complex types, decimal requiring more than 8 bytes, INT96. Schema mismatch.
+ * All of these can be handled efficiently and easily with codegen.
+ */
+public class UnsafeRowParquetRecordReader extends SpecificParquetRecordReaderBase<UnsafeRow> {
+  /**
+   * Batch of unsafe rows that we assemble and the current index we've returned. Everytime this
+   * batch is used up (batchIdx == numBatched), we populated the batch.
+   */
+  private UnsafeRow[] rows = new UnsafeRow[64];
+  private int batchIdx = 0;
+  private int numBatched = 0;
+
+  /**
+   * Used to write variable length columns. Same length as `rows`.
+   */
+  private UnsafeRowWriter[] rowWriters = null;
+  /**
+   * True if the row contains variable length fields.
+   */
+  private boolean containsVarLenFields;
+
+  /**
+   * The number of bytes in the fixed length portion of the row.
+   */
+  private int fixedSizeBytes;
+
+  /**
+   * For each request column, the reader to read this column.
+   * columnsReaders[i] populated the UnsafeRow's attribute at i.
+   */
+  private ColumnReader[] columnReaders;
+
+  /**
+   * The number of rows that have been returned.
+   */
+  private long rowsReturned;
+
+  /**
+   * The number of rows that have been reading, including the current in flight row group.
+   */
+  private long totalCountLoadedSoFar = 0;
+
+  /**
+   * For each column, the annotated original type.
+   */
+  private OriginalType[] originalTypes;
+
+  /**
+   * The default size for varlen columns. The row grows as necessary to accommodate the
+   * largest column.
+   */
+  private static final int DEFAULT_VAR_LEN_SIZE = 32;
+
+  /**
+   * Implementation of RecordReader API.
+   */
+  @Override
+  public void initialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext)
+      throws IOException, InterruptedException {
+    super.initialize(inputSplit, taskAttemptContext);
+
+    /**
+     * Check that the requested schema is supported.
+     */
+    if (requestedSchema.getFieldCount() == 0) {
+      // TODO: what does this mean?
+      throw new IOException("Empty request schema not supported.");
+    }
+    int numVarLenFields = 0;
+    originalTypes = new OriginalType[requestedSchema.getFieldCount()];
+    for (int i = 0; i < requestedSchema.getFieldCount(); ++i) {
+      Type t = requestedSchema.getFields().get(i);
+      if (!t.isPrimitive() || t.isRepetition(Type.Repetition.REPEATED)) {
+        throw new IOException("Complex types not supported.");
+      }
+      PrimitiveType primitiveType = t.asPrimitiveType();
+
+      originalTypes[i] = t.getOriginalType();
+
+      // TODO: Be extremely cautious in what is supported. Expand this.
+      if (originalTypes[i] != null && originalTypes[i] != OriginalType.DECIMAL &&
+          originalTypes[i] != OriginalType.UTF8 && originalTypes[i] != OriginalType.DATE) {
+        throw new IOException("Unsupported type: " + t);
+      }
+      if (originalTypes[i] == OriginalType.DECIMAL &&
+          primitiveType.getDecimalMetadata().getPrecision() >
+              CatalystSchemaConverter.MAX_PRECISION_FOR_INT64()) {
+        throw new IOException("Decimal with high precision is not supported.");
+      }
+      if (primitiveType.getPrimitiveTypeName() == PrimitiveType.PrimitiveTypeName.INT96) {
+        throw new IOException("Int96 not supported.");
+      }
+      ColumnDescriptor fd = fileSchema.getColumnDescription(requestedSchema.getPaths().get(i));
+      if (!fd.equals(requestedSchema.getColumns().get(i))) {
+        throw new IOException("Schema evolution not supported.");
+      }
+
+      if (primitiveType.getPrimitiveTypeName() == PrimitiveType.PrimitiveTypeName.BINARY) {
+        ++numVarLenFields;
+      }
+    }
+
+    /**
+     * Initialize rows and rowWriters. These objects are reused across all rows in the relation.
+     */
+    int rowByteSize = UnsafeRow.calculateBitSetWidthInBytes(requestedSchema.getFieldCount());
+    rowByteSize += 8 * requestedSchema.getFieldCount();
+    fixedSizeBytes = rowByteSize;
+    rowByteSize += numVarLenFields * DEFAULT_VAR_LEN_SIZE;
+    containsVarLenFields = numVarLenFields > 0;
+    rowWriters = new UnsafeRowWriter[rows.length];
+
+    for (int i = 0; i < rows.length; ++i) {
+      rows[i] = new UnsafeRow();
+      rowWriters[i] = new UnsafeRowWriter();
+      BufferHolder holder = new BufferHolder(rowByteSize);
+      rowWriters[i].initialize(rows[i], holder, requestedSchema.getFieldCount());
+      rows[i].pointTo(holder.buffer, Platform.BYTE_ARRAY_OFFSET, requestedSchema.getFieldCount(),
+          holder.buffer.length);
+    }
+  }
+
+  @Override
+  public boolean nextKeyValue() throws IOException, InterruptedException {
+    if (batchIdx >= numBatched) {
+      if (!loadBatch()) return false;
+    }
+    ++batchIdx;
+    return true;
+  }
+
+  @Override
+  public UnsafeRow getCurrentValue() throws IOException, InterruptedException {
+    return rows[batchIdx - 1];
+  }
+
+  @Override
+  public float getProgress() throws IOException, InterruptedException {
+    return (float) rowsReturned / totalRowCount;
+  }
+
+  /**
+   * Decodes a batch of values into `rows`. This function is the hot path.
+   */
+  private boolean loadBatch() throws IOException {
+    // no more records left
+    if (rowsReturned >= totalRowCount) { return false; }
+    checkEndOfRowGroup();
+
+    int num = (int)Math.min(rows.length, totalCountLoadedSoFar - rowsReturned);
+    rowsReturned += num;
+
+    if (containsVarLenFields) {
+      for (int i = 0; i < rowWriters.length; ++i) {
+        rowWriters[i].holder().resetTo(fixedSizeBytes);
+      }
+    }
+
+    for (int i = 0; i < columnReaders.length; ++i) {
+      switch (columnReaders[i].descriptor.getType()) {
+        case BOOLEAN:
+          decodeBooleanBatch(i, num);
+          break;
+        case INT32:
+          if (originalTypes[i] == OriginalType.DECIMAL) {
+            decodeIntAsDecimalBatch(i, num);
+          } else {
+            decodeIntBatch(i, num);
+          }
+          break;
+        case INT64:
+          Preconditions.checkState(originalTypes[i] == null
+              || originalTypes[i] == OriginalType.DECIMAL,
+              "Unexpected original type: " + originalTypes[i]);
+          decodeLongBatch(i, num);
+          break;
+        case FLOAT:
+          decodeFloatBatch(i, num);
+          break;
+        case DOUBLE:
+          decodeDoubleBatch(i, num);
+          break;
+        case BINARY:
+          decodeBinaryBatch(i, num);
+          break;
+        case FIXED_LEN_BYTE_ARRAY:
+          Preconditions.checkState(originalTypes[i] == OriginalType.DECIMAL,
+              "Unexpected original type: " + originalTypes[i]);
+          decodeFixedLenArrayAsDecimalBatch(i, num);
+          break;
+        case INT96:
+          throw new IOException("Unsupported " + columnReaders[i].descriptor.getType());
+      }
+      numBatched = num;
+      batchIdx = 0;
+    }
+    return true;
+  }
+
+  private void decodeBooleanBatch(int col, int num) throws IOException {
+    for (int n = 0; n < num; ++n) {
+      if (columnReaders[col].next()) {
+        rows[n].setBoolean(col, columnReaders[col].nextBoolean());
+      } else {
+        rows[n].setNullAt(col);
+      }
+    }
+  }
+
+  private void decodeIntBatch(int col, int num) throws IOException {
+    for (int n = 0; n < num; ++n) {
+      if (columnReaders[col].next()) {
+        rows[n].setInt(col, columnReaders[col].nextInt());
+      } else {
+        rows[n].setNullAt(col);
+      }
+    }
+  }
+
+  private void decodeIntAsDecimalBatch(int col, int num) throws IOException {
+    for (int n = 0; n < num; ++n) {
+      if (columnReaders[col].next()) {
+        // Since this is stored as an INT, it is always a compact decimal. Just set it as a long.
+        rows[n].setLong(col, columnReaders[col].nextInt());
+      } else {
+        rows[n].setNullAt(col);
+      }
+    }
+  }
+
+  private void decodeLongBatch(int col, int num) throws IOException {
+    for (int n = 0; n < num; ++n) {
+      if (columnReaders[col].next()) {
+        rows[n].setLong(col, columnReaders[col].nextLong());
+      } else {
+        rows[n].setNullAt(col);
+      }
+    }
+  }
+
+  private void decodeFloatBatch(int col, int num) throws IOException {
+    for (int n = 0; n < num; ++n) {
+      if (columnReaders[col].next()) {
+        rows[n].setFloat(col, columnReaders[col].nextFloat());
+      } else {
+        rows[n].setNullAt(col);
+      }
+    }
+  }
+
+  private void decodeDoubleBatch(int col, int num) throws IOException {
+    for (int n = 0; n < num; ++n) {
+      if (columnReaders[col].next()) {
+        rows[n].setDouble(col, columnReaders[col].nextDouble());
+      } else {
+        rows[n].setNullAt(col);
+      }
+    }
+  }
+
+  private void decodeBinaryBatch(int col, int num) throws IOException {
+    for (int n = 0; n < num; ++n) {
+      if (columnReaders[col].next()) {
+        ByteBuffer bytes = columnReaders[col].nextBinary().toByteBuffer();
+        int len = bytes.limit() - bytes.position();
+        if (originalTypes[col] == OriginalType.UTF8) {
+          UTF8String str = UTF8String.fromBytes(bytes.array(), bytes.position(), len);
+          rowWriters[n].write(col, str);
+        } else {
+          rowWriters[n].write(col, bytes.array(), bytes.position(), len);
+        }
+      } else {
+        rows[n].setNullAt(col);
+      }
+    }
+  }
+
+  private void decodeFixedLenArrayAsDecimalBatch(int col, int num) throws IOException {
+    PrimitiveType type = requestedSchema.getFields().get(col).asPrimitiveType();
+    int precision = type.getDecimalMetadata().getPrecision();
+    int scale = type.getDecimalMetadata().getScale();
+    Preconditions.checkState(precision <= CatalystSchemaConverter.MAX_PRECISION_FOR_INT64(),
+        "Unsupported precision.");
+
+    for (int n = 0; n < num; ++n) {
+      if (columnReaders[col].next()) {
+        Binary v = columnReaders[col].nextBinary();
+        // Constructs a `Decimal` with an unscaled `Long` value if possible.
+        long unscaled = CatalystRowConverter.binaryToUnscaledLong(v);
+        rows[n].setDecimal(col, Decimal.apply(unscaled, precision, scale), precision);
+      } else {
+        rows[n].setNullAt(col);
+      }
+    }
+  }
+
+  /**
+   *
+   * Decoder to return values from a single column.
+   */
+  private static final class ColumnReader {
+    /**
+     * Total number of values read.
+     */
+    private long valuesRead;
+
+    /**
+     * value that indicates the end of the current page. That is,
+     * if valuesRead == endOfPageValueCount, we are at the end of the page.
+     */
+    private long endOfPageValueCount;
+
+    /**
+     * The dictionary, if this column has dictionary encoding.
+     */
+    private final Dictionary dictionary;
+
+    /**
+     * If true, the current page is dictionary encoded.
+     */
+    private boolean useDictionary;
+
+    /**
+     * Maximum definition level for this column.
+     */
+    private final int maxDefLevel;
+
+    /**
+     * Repetition/Definition/Value readers.
+     */
+    private IntIterator repetitionLevelColumn;
+    private IntIterator definitionLevelColumn;
+    private ValuesReader dataColumn;
+
+    /**
+     * Total number of values in this column (in this row group).
+     */
+    private final long totalValueCount;
+
+    /**
+     * Total values in the current page.
+     */
+    private int pageValueCount;
+
+    private final PageReader pageReader;
+    private final ColumnDescriptor descriptor;
+
+    public ColumnReader(ColumnDescriptor descriptor, PageReader pageReader)
+        throws IOException {
+      this.descriptor = descriptor;
+      this.pageReader = pageReader;
+      this.maxDefLevel = descriptor.getMaxDefinitionLevel();
+
+      DictionaryPage dictionaryPage = pageReader.readDictionaryPage();
+      if (dictionaryPage != null) {
+        try {
+          this.dictionary = dictionaryPage.getEncoding().initDictionary(descriptor, dictionaryPage);
+          this.useDictionary = true;
+        } catch (IOException e) {
+          throw new IOException("could not decode the dictionary for " + descriptor, e);
+        }
+      } else {
+        this.dictionary = null;
+        this.useDictionary = false;
+      }
+      this.totalValueCount = pageReader.getTotalValueCount();
+      if (totalValueCount == 0) {
+        throw new IOException("totalValueCount == 0");
+      }
+    }
+
+    /**
+     * TODO: Hoist the useDictionary branch to decode*Batch and make the batch page aligned.
+     */
+    public boolean nextBoolean() {
+      if (!useDictionary) {
+        return dataColumn.readBoolean();
+      } else {
+        return dictionary.decodeToBoolean(dataColumn.readValueDictionaryId());
+      }
+    }
+
+    public int nextInt() {
+      if (!useDictionary) {
+        return dataColumn.readInteger();
+      } else {
+        return dictionary.decodeToInt(dataColumn.readValueDictionaryId());
+      }
+    }
+
+    public long nextLong() {
+      if (!useDictionary) {
+        return dataColumn.readLong();
+      } else {
+        return dictionary.decodeToLong(dataColumn.readValueDictionaryId());
+      }
+    }
+
+    public float nextFloat() {
+      if (!useDictionary) {
+        return dataColumn.readFloat();
+      } else {
+        return dictionary.decodeToFloat(dataColumn.readValueDictionaryId());
+      }
+    }
+
+    public double nextDouble() {
+      if (!useDictionary) {
+        return dataColumn.readDouble();
+      } else {
+        return dictionary.decodeToDouble(dataColumn.readValueDictionaryId());
+      }
+    }
+
+    public Binary nextBinary() {
+      if (!useDictionary) {
+        return dataColumn.readBytes();
+      } else {
+        return dictionary.decodeToBinary(dataColumn.readValueDictionaryId());
+      }
+    }
+
+    /**
+     * Advances to the next value. Returns true if the value is non-null.
+     */
+    private boolean next() throws IOException {
+      if (valuesRead >= endOfPageValueCount) {
+        if (valuesRead >= totalValueCount) {
+          // How do we get here? Throw end of stream exception?
+          return false;
+        }
+        readPage();
+      }
+      ++valuesRead;
+      // TODO: Don't read for flat schemas
+      //repetitionLevel = repetitionLevelColumn.nextInt();
+      return definitionLevelColumn.nextInt() == maxDefLevel;
+    }
+
+    private void readPage() throws IOException {
+      DataPage page = pageReader.readPage();
+      // TODO: Why is this a visitor?
+      page.accept(new DataPage.Visitor<Void>() {
+        @Override
+        public Void visit(DataPageV1 dataPageV1) {
+          try {
+            readPageV1(dataPageV1);
+            return null;
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        }
+
+        @Override
+        public Void visit(DataPageV2 dataPageV2) {
+          try {
+            readPageV2(dataPageV2);
+            return null;
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        }
+      });
+    }
+
+    private void initDataReader(Encoding dataEncoding, byte[] bytes, int offset, int valueCount)
+        throws IOException {
+      this.pageValueCount = valueCount;
+      this.endOfPageValueCount = valuesRead + pageValueCount;
+      if (dataEncoding.usesDictionary()) {
+        if (dictionary == null) {
+          throw new IOException(
+              "could not read page in col " + descriptor +
+                  " as the dictionary was missing for encoding " + dataEncoding);
+        }
+        this.dataColumn = dataEncoding.getDictionaryBasedValuesReader(
+            descriptor, VALUES, dictionary);
+        this.useDictionary = true;
+      } else {
+        this.dataColumn = dataEncoding.getValuesReader(descriptor, VALUES);
+        this.useDictionary = false;
+      }
+
+      try {
+        dataColumn.initFromPage(pageValueCount, bytes, offset);
+      } catch (IOException e) {
+        throw new IOException("could not read page in col " + descriptor, e);
+      }
+    }
+
+    private void readPageV1(DataPageV1 page) throws IOException {
+      ValuesReader rlReader = page.getRlEncoding().getValuesReader(descriptor, REPETITION_LEVEL);
+      ValuesReader dlReader = page.getDlEncoding().getValuesReader(descriptor, DEFINITION_LEVEL);
+      this.repetitionLevelColumn = new ValuesReaderIntIterator(rlReader);
+      this.definitionLevelColumn = new ValuesReaderIntIterator(dlReader);
+      try {
+        byte[] bytes = page.getBytes().toByteArray();
+        rlReader.initFromPage(pageValueCount, bytes, 0);
+        int next = rlReader.getNextOffset();
+        dlReader.initFromPage(pageValueCount, bytes, next);
+        next = dlReader.getNextOffset();
+        initDataReader(page.getValueEncoding(), bytes, next, page.getValueCount());
+      } catch (IOException e) {
+        throw new IOException("could not read page " + page + " in col " + descriptor, e);
+      }
+    }
+
+    private void readPageV2(DataPageV2 page) throws IOException {
+      this.repetitionLevelColumn = createRLEIterator(descriptor.getMaxRepetitionLevel(),
+          page.getRepetitionLevels(), descriptor);
+      this.definitionLevelColumn = createRLEIterator(descriptor.getMaxDefinitionLevel(),
+          page.getDefinitionLevels(), descriptor);
+      try {
+        initDataReader(page.getDataEncoding(), page.getData().toByteArray(), 0,
+            page.getValueCount());
+      } catch (IOException e) {
+        throw new IOException("could not read page " + page + " in col " + descriptor, e);
+      }
+    }
+  }
+
+  private void checkEndOfRowGroup() throws IOException {
+    if (rowsReturned != totalCountLoadedSoFar) return;
+    PageReadStore pages = reader.readNextRowGroup();
+    if (pages == null) {
+      throw new IOException("expecting more rows but reached last block. Read "
+          + rowsReturned + " out of " + totalRowCount);
+    }
+    List<ColumnDescriptor> columns = requestedSchema.getColumns();
+    columnReaders = new ColumnReader[columns.size()];
+    for (int i = 0; i < columns.size(); ++i) {
+      columnReaders[i] = new ColumnReader(columns.get(i), pages.getPageReader(columns.get(i)));
+    }
+    totalCountLoadedSoFar += pages.getRowCount();
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystRowConverter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystRowConverter.scala
index 1f653cd3d3cb1..94298fae2d69b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystRowConverter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/CatalystRowConverter.scala
@@ -370,35 +370,13 @@ private[parquet] class CatalystRowConverter(
     protected def decimalFromBinary(value: Binary): Decimal = {
       if (precision <= CatalystSchemaConverter.MAX_PRECISION_FOR_INT64) {
         // Constructs a `Decimal` with an unscaled `Long` value if possible.
-        val unscaled = binaryToUnscaledLong(value)
+        val unscaled = CatalystRowConverter.binaryToUnscaledLong(value)
         Decimal(unscaled, precision, scale)
       } else {
         // Otherwise, resorts to an unscaled `BigInteger` instead.
         Decimal(new BigDecimal(new BigInteger(value.getBytes), scale), precision, scale)
       }
     }
-
-    private def binaryToUnscaledLong(binary: Binary): Long = {
-      // The underlying `ByteBuffer` implementation is guaranteed to be `HeapByteBuffer`, so here
-      // we are using `Binary.toByteBuffer.array()` to steal the underlying byte array without
-      // copying it.
-      val buffer = binary.toByteBuffer
-      val bytes = buffer.array()
-      val start = buffer.position()
-      val end = buffer.limit()
-
-      var unscaled = 0L
-      var i = start
-
-      while (i < end) {
-        unscaled = (unscaled << 8) | (bytes(i) & 0xff)
-        i += 1
-      }
-
-      val bits = 8 * (end - start)
-      unscaled = (unscaled << (64 - bits)) >> (64 - bits)
-      unscaled
-    }
   }
 
   private class CatalystIntDictionaryAwareDecimalConverter(
@@ -658,3 +636,27 @@ private[parquet] class CatalystRowConverter(
     override def start(): Unit = elementConverter.start()
   }
 }
+
+private[parquet] object CatalystRowConverter {
+  def binaryToUnscaledLong(binary: Binary): Long = {
+    // The underlying `ByteBuffer` implementation is guaranteed to be `HeapByteBuffer`, so here
+    // we are using `Binary.toByteBuffer.array()` to steal the underlying byte array without
+    // copying it.
+    val buffer = binary.toByteBuffer
+    val bytes = buffer.array()
+    val start = buffer.position()
+    val end = buffer.limit()
+
+    var unscaled = 0L
+    var i = start
+
+    while (i < end) {
+      unscaled = (unscaled << 8) | (bytes(i) & 0xff)
+      i += 1
+    }
+
+    val bits = 8 * (end - start)
+    unscaled = (unscaled << (64 - bits)) >> (64 - bits)
+    unscaled
+  }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
index 458786f77af3f..c8028a5ef5528 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
@@ -337,7 +337,9 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
     }
   }
 
-  test("SPARK-11661 Still pushdown filters returned by unhandledFilters") {
+  // Renable when we can toggle custom ParquetRecordReader on/off. The custom reader does
+  // not do row by row filtering (and we probably don't want to push that).
+  ignore("SPARK-11661 Still pushdown filters returned by unhandledFilters") {
     import testImplicits._
     withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true") {
       withTempPath { dir =>

From 9c0654d36c6d171dd273850c2cc2f415cc2a5a6b Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Wed, 18 Nov 2015 18:41:40 -0800
Subject: [PATCH 352/518] Revert "[SPARK-11544][SQL] sqlContext doesn't use
 PathFilter"

This reverts commit 54db79702513e11335c33bcf3a03c59e965e6f16.
---
 .../apache/spark/sql/sources/interfaces.scala | 25 +++----------
 .../datasources/json/JsonSuite.scala          | 36 ++-----------------
 2 files changed, 7 insertions(+), 54 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index f9465157c936d..b3d3bdf50df63 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -21,8 +21,7 @@ import scala.collection.mutable
 import scala.util.Try
 
 import org.apache.hadoop.conf.Configuration
-import org.apache.hadoop.fs.{PathFilter, FileStatus, FileSystem, Path}
-import org.apache.hadoop.mapred.{JobConf, FileInputFormat}
+import org.apache.hadoop.fs.{FileStatus, FileSystem, Path}
 import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext}
 
 import org.apache.spark.{Logging, SparkContext}
@@ -448,15 +447,9 @@ abstract class HadoopFsRelation private[sql](
           val hdfsPath = new Path(path)
           val fs = hdfsPath.getFileSystem(hadoopConf)
           val qualified = hdfsPath.makeQualified(fs.getUri, fs.getWorkingDirectory)
+
           logInfo(s"Listing $qualified on driver")
-          // Dummy jobconf to get to the pathFilter defined in configuration
-          val jobConf = new JobConf(hadoopConf, this.getClass())
-          val pathFilter = FileInputFormat.getInputPathFilter(jobConf)
-          if (pathFilter != null) {
-            Try(fs.listStatus(qualified, pathFilter)).getOrElse(Array.empty)
-          } else {
-            Try(fs.listStatus(qualified)).getOrElse(Array.empty)
-          }
+          Try(fs.listStatus(qualified)).getOrElse(Array.empty)
         }.filterNot { status =>
           val name = status.getPath.getName
           name.toLowerCase == "_temporary" || name.startsWith(".")
@@ -854,16 +847,8 @@ private[sql] object HadoopFsRelation extends Logging {
     if (name == "_temporary" || name.startsWith(".")) {
       Array.empty
     } else {
-      // Dummy jobconf to get to the pathFilter defined in configuration
-      val jobConf = new JobConf(fs.getConf, this.getClass())
-      val pathFilter = FileInputFormat.getInputPathFilter(jobConf)
-      if (pathFilter != null) {
-        val (dirs, files) = fs.listStatus(status.getPath, pathFilter).partition(_.isDir)
-        files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
-      } else {
-        val (dirs, files) = fs.listStatus(status.getPath).partition(_.isDir)
-        files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
-      }
+      val (dirs, files) = fs.listStatus(status.getPath).partition(_.isDir)
+      files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
index f09b61e838159..6042b1178affe 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
@@ -19,27 +19,19 @@ package org.apache.spark.sql.execution.datasources.json
 
 import java.io.{File, StringWriter}
 import java.sql.{Date, Timestamp}
-import scala.collection.JavaConverters._
 
 import com.fasterxml.jackson.core.JsonFactory
-import org.apache.commons.io.FileUtils
-import org.apache.hadoop.conf.Configuration
-import org.apache.hadoop.fs.{Path, PathFilter}
+import org.apache.spark.rdd.RDD
 import org.scalactic.Tolerance._
 
-import org.apache.spark.rdd.RDD
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
-import org.apache.spark.sql.execution.datasources.{LogicalRelation, ResolvedDataSource}
+import org.apache.spark.sql.execution.datasources.{ResolvedDataSource, LogicalRelation}
 import org.apache.spark.sql.execution.datasources.json.InferSchema.compatibleType
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
 
-class TestFileFilter extends PathFilter {
-  override def accept(path: Path): Boolean = path.getParent.getName != "p=2"
-}
-
 class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
   import testImplicits._
 
@@ -1398,28 +1390,4 @@ class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
       )
     }
   }
-
-  test("SPARK-11544 test pathfilter") {
-    withTempPath { dir =>
-      val path = dir.getCanonicalPath
-
-      val df = sqlContext.range(2)
-      df.write.json(path + "/p=1")
-      df.write.json(path + "/p=2")
-      assert(sqlContext.read.json(path).count() === 4)
-
-      val clonedConf = new Configuration(hadoopConfiguration)
-      try {
-        hadoopConfiguration.setClass(
-          "mapreduce.input.pathFilter.class",
-          classOf[TestFileFilter],
-          classOf[PathFilter])
-        assert(sqlContext.read.json(path).count() === 2)
-      } finally {
-        // Hadoop 1 doesn't have `Configuration.unset`
-        hadoopConfiguration.clear()
-        clonedConf.asScala.foreach(entry => hadoopConfiguration.set(entry.getKey, entry.getValue))
-      }
-    }
-  }
 }

From 67c75828ff4df2e305bdf5d6be5a11201d1da3f3 Mon Sep 17 00:00:00 2001
From: Yuhao Yang <hhbyyh@gmail.com>
Date: Wed, 18 Nov 2015 18:49:46 -0800
Subject: [PATCH 353/518] [SPARK-11816][ML] fix some style issue in ML/MLlib
 examples

jira: https://issues.apache.org/jira/browse/SPARK-11816
Currently I only fixed some obvious comments issue like
// scalastyle:off println
on the bottom.

Yet the style in examples is not quite consistent, like only half of the examples  are with
// Example usage: ./bin/run-example mllib.FPGrowthExample \,

Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #9808 from hhbyyh/exampleStyle.
---
 .../java/org/apache/spark/examples/ml/JavaKMeansExample.java    | 2 +-
 .../apache/spark/examples/ml/AFTSurvivalRegressionExample.scala | 2 +-
 .../spark/examples/ml/DecisionTreeClassificationExample.scala   | 1 +
 .../spark/examples/ml/DecisionTreeRegressionExample.scala       | 1 +
 .../examples/ml/MultilayerPerceptronClassifierExample.scala     | 2 +-
 5 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaKMeansExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaKMeansExample.java
index be2bf0c7b465c..47665ff2b1f3c 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaKMeansExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaKMeansExample.java
@@ -41,7 +41,7 @@
  * An example demonstrating a k-means clustering.
  * Run with
  * <pre>
- * bin/run-example ml.JavaSimpleParamsExample <file> <k>
+ * bin/run-example ml.JavaKMeansExample <file> <k>
  * </pre>
  */
 public class JavaKMeansExample {
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala
index 5da285e83681f..f4b3613ccb94f 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/AFTSurvivalRegressionExample.scala
@@ -59,4 +59,4 @@ object AFTSurvivalRegressionExample {
     sc.stop()
   }
 }
-// scalastyle:off println
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
index ff8a0a90f1e44..db024b5cad935 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeClassificationExample.scala
@@ -90,3 +90,4 @@ object DecisionTreeClassificationExample {
     // $example off$
   }
 }
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala
index fc402724d2156..ad01f55df72b5 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/DecisionTreeRegressionExample.scala
@@ -78,3 +78,4 @@ object DecisionTreeRegressionExample {
     // $example off$
   }
 }
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala
index 146b83c8be490..9c98076bd24b1 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/MultilayerPerceptronClassifierExample.scala
@@ -66,4 +66,4 @@ object MultilayerPerceptronClassifierExample {
     sc.stop()
   }
 }
-// scalastyle:off println
+// scalastyle:on println

From fc3f77b42d62ca789d0ee07403795978961991c7 Mon Sep 17 00:00:00 2001
From: "navis.ryu" <navis@apache.org>
Date: Wed, 18 Nov 2015 19:37:14 -0800
Subject: [PATCH 354/518] [SPARK-11614][SQL] serde parameters should be set
 only when all params are ready

see HIVE-7975 and HIVE-12373

With changed semantic of setters in thrift objects in hive, setter should be called only after all parameters are set. It's not problem of current state but will be a problem in some day.

Author: navis.ryu <navis@apache.org>

Closes #9580 from navis/SPARK-11614.
---
 .../scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala
index f4d45714fae4e..9a981d02ad67c 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala
@@ -804,12 +804,13 @@ private[hive] case class MetastoreRelation
 
       val serdeInfo = new org.apache.hadoop.hive.metastore.api.SerDeInfo
       sd.setSerdeInfo(serdeInfo)
+      // maps and lists should be set only after all elements are ready (see HIVE-7975)
       serdeInfo.setSerializationLib(p.storage.serde)
 
       val serdeParameters = new java.util.HashMap[String, String]()
-      serdeInfo.setParameters(serdeParameters)
       table.serdeProperties.foreach { case (k, v) => serdeParameters.put(k, v) }
       p.storage.serdeProperties.foreach { case (k, v) => serdeParameters.put(k, v) }
+      serdeInfo.setParameters(serdeParameters)
 
       new Partition(hiveQlTable, tPartition)
     }

From d02d5b9295b169c3ebb0967453b2835edb8a121f Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Wed, 18 Nov 2015 21:44:01 -0800
Subject: [PATCH 355/518] [SPARK-11842][ML] Small cleanups to existing Readers
 and Writers

Updates:
* Add repartition(1) to save() methods' saving of data for LogisticRegressionModel, LinearRegressionModel.
* Strengthen privacy to class and companion object for Writers and Readers
* Change LogisticRegressionSuite read/write test to fit intercept
* Add Since versions for read/write methods in Pipeline, LogisticRegression
* Switch from hand-written class names in Readers to using getClass

CC: mengxr

CC: yanboliang Would you mind taking a look at this PR?  mengxr might not be able to soon.  Thank you!

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9829 from jkbradley/ml-io-cleanups.
---
 .../scala/org/apache/spark/ml/Pipeline.scala  | 22 +++++++++++++------
 .../classification/LogisticRegression.scala   | 19 ++++++++++------
 .../spark/ml/feature/CountVectorizer.scala    |  2 +-
 .../org/apache/spark/ml/feature/IDF.scala     |  2 +-
 .../spark/ml/feature/MinMaxScaler.scala       |  2 +-
 .../spark/ml/feature/StandardScaler.scala     |  2 +-
 .../spark/ml/feature/StringIndexer.scala      |  2 +-
 .../apache/spark/ml/recommendation/ALS.scala  |  6 ++---
 .../ml/regression/LinearRegression.scala      |  4 ++--
 .../LogisticRegressionSuite.scala             |  2 +-
 10 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
index b0f22e042ec56..6f15b37abcb30 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
@@ -27,7 +27,7 @@ import org.json4s._
 import org.json4s.jackson.JsonMethods._
 
 import org.apache.spark.{SparkContext, Logging}
-import org.apache.spark.annotation.{DeveloperApi, Experimental}
+import org.apache.spark.annotation.{Since, DeveloperApi, Experimental}
 import org.apache.spark.ml.param.{Param, ParamMap, Params}
 import org.apache.spark.ml.util.MLReader
 import org.apache.spark.ml.util.MLWriter
@@ -174,16 +174,20 @@ class Pipeline(override val uid: String) extends Estimator[PipelineModel] with M
     theStages.foldLeft(schema)((cur, stage) => stage.transformSchema(cur))
   }
 
+  @Since("1.6.0")
   override def write: MLWriter = new Pipeline.PipelineWriter(this)
 }
 
+@Since("1.6.0")
 object Pipeline extends MLReadable[Pipeline] {
 
+  @Since("1.6.0")
   override def read: MLReader[Pipeline] = new PipelineReader
 
+  @Since("1.6.0")
   override def load(path: String): Pipeline = super.load(path)
 
-  private[ml] class PipelineWriter(instance: Pipeline) extends MLWriter {
+  private[Pipeline] class PipelineWriter(instance: Pipeline) extends MLWriter {
 
     SharedReadWrite.validateStages(instance.getStages)
 
@@ -191,10 +195,10 @@ object Pipeline extends MLReadable[Pipeline] {
       SharedReadWrite.saveImpl(instance, instance.getStages, sc, path)
   }
 
-  private[ml] class PipelineReader extends MLReader[Pipeline] {
+  private class PipelineReader extends MLReader[Pipeline] {
 
     /** Checked against metadata when loading model */
-    private val className = "org.apache.spark.ml.Pipeline"
+    private val className = classOf[Pipeline].getName
 
     override def load(path: String): Pipeline = {
       val (uid: String, stages: Array[PipelineStage]) = SharedReadWrite.load(className, sc, path)
@@ -333,18 +337,22 @@ class PipelineModel private[ml] (
     new PipelineModel(uid, stages.map(_.copy(extra))).setParent(parent)
   }
 
+  @Since("1.6.0")
   override def write: MLWriter = new PipelineModel.PipelineModelWriter(this)
 }
 
+@Since("1.6.0")
 object PipelineModel extends MLReadable[PipelineModel] {
 
   import Pipeline.SharedReadWrite
 
+  @Since("1.6.0")
   override def read: MLReader[PipelineModel] = new PipelineModelReader
 
+  @Since("1.6.0")
   override def load(path: String): PipelineModel = super.load(path)
 
-  private[ml] class PipelineModelWriter(instance: PipelineModel) extends MLWriter {
+  private[PipelineModel] class PipelineModelWriter(instance: PipelineModel) extends MLWriter {
 
     SharedReadWrite.validateStages(instance.stages.asInstanceOf[Array[PipelineStage]])
 
@@ -352,10 +360,10 @@ object PipelineModel extends MLReadable[PipelineModel] {
       instance.stages.asInstanceOf[Array[PipelineStage]], sc, path)
   }
 
-  private[ml] class PipelineModelReader extends MLReader[PipelineModel] {
+  private class PipelineModelReader extends MLReader[PipelineModel] {
 
     /** Checked against metadata when loading model */
-    private val className = "org.apache.spark.ml.PipelineModel"
+    private val className = classOf[PipelineModel].getName
 
     override def load(path: String): PipelineModel = {
       val (uid: String, stages: Array[PipelineStage]) = SharedReadWrite.load(className, sc, path)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
index a3cc49f7f018c..418bbdc9a058f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
@@ -24,7 +24,7 @@ import breeze.optimize.{CachedDiffFunction, DiffFunction, LBFGS => BreezeLBFGS,
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.{Logging, SparkException}
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Since, Experimental}
 import org.apache.spark.ml.feature.Instance
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
@@ -525,18 +525,23 @@ class LogisticRegressionModel private[ml] (
    *
    * This also does not save the [[parent]] currently.
    */
+  @Since("1.6.0")
   override def write: MLWriter = new LogisticRegressionModel.LogisticRegressionModelWriter(this)
 }
 
 
+@Since("1.6.0")
 object LogisticRegressionModel extends MLReadable[LogisticRegressionModel] {
 
+  @Since("1.6.0")
   override def read: MLReader[LogisticRegressionModel] = new LogisticRegressionModelReader
 
+  @Since("1.6.0")
   override def load(path: String): LogisticRegressionModel = super.load(path)
 
   /** [[MLWriter]] instance for [[LogisticRegressionModel]] */
-  private[classification] class LogisticRegressionModelWriter(instance: LogisticRegressionModel)
+  private[LogisticRegressionModel]
+  class LogisticRegressionModelWriter(instance: LogisticRegressionModel)
     extends MLWriter with Logging {
 
     private case class Data(
@@ -552,15 +557,15 @@ object LogisticRegressionModel extends MLReadable[LogisticRegressionModel] {
       val data = Data(instance.numClasses, instance.numFeatures, instance.intercept,
         instance.coefficients)
       val dataPath = new Path(path, "data").toString
-      sqlContext.createDataFrame(Seq(data)).write.format("parquet").save(dataPath)
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
     }
   }
 
-  private[classification] class LogisticRegressionModelReader
+  private class LogisticRegressionModelReader
     extends MLReader[LogisticRegressionModel] {
 
     /** Checked against metadata when loading model */
-    private val className = "org.apache.spark.ml.classification.LogisticRegressionModel"
+    private val className = classOf[LogisticRegressionModel].getName
 
     override def load(path: String): LogisticRegressionModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
@@ -603,7 +608,7 @@ private[classification] class MultiClassSummarizer extends Serializable {
    * @return This MultilabelSummarizer
    */
   def add(label: Double, weight: Double = 1.0): this.type = {
-    require(weight >= 0.0, s"instance weight, ${weight} has to be >= 0.0")
+    require(weight >= 0.0, s"instance weight, $weight has to be >= 0.0")
 
     if (weight == 0.0) return this
 
@@ -839,7 +844,7 @@ private class LogisticAggregator(
     instance match { case Instance(label, weight, features) =>
       require(dim == features.size, s"Dimensions mismatch when adding new instance." +
         s" Expecting $dim but got ${features.size}.")
-      require(weight >= 0.0, s"instance weight, ${weight} has to be >= 0.0")
+      require(weight >= 0.0, s"instance weight, $weight has to be >= 0.0")
 
       if (weight == 0.0) return this
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
index 4969cf42450d2..b9e2144c0ad40 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
@@ -266,7 +266,7 @@ object CountVectorizerModel extends MLReadable[CountVectorizerModel] {
 
   private class CountVectorizerModelReader extends MLReader[CountVectorizerModel] {
 
-    private val className = "org.apache.spark.ml.feature.CountVectorizerModel"
+    private val className = classOf[CountVectorizerModel].getName
 
     override def load(path: String): CountVectorizerModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
index 0e00ef6f2ee20..f7b0f29a27c2d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
@@ -155,7 +155,7 @@ object IDFModel extends MLReadable[IDFModel] {
 
   private class IDFModelReader extends MLReader[IDFModel] {
 
-    private val className = "org.apache.spark.ml.feature.IDFModel"
+    private val className = classOf[IDFModel].getName
 
     override def load(path: String): IDFModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
index ed24eabb50444..c2866f5eceff3 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala
@@ -210,7 +210,7 @@ object MinMaxScalerModel extends MLReadable[MinMaxScalerModel] {
 
   private class MinMaxScalerModelReader extends MLReader[MinMaxScalerModel] {
 
-    private val className = "org.apache.spark.ml.feature.MinMaxScalerModel"
+    private val className = classOf[MinMaxScalerModel].getName
 
     override def load(path: String): MinMaxScalerModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
index 1f689c1da1ba9..6d545219ebf49 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
@@ -180,7 +180,7 @@ object StandardScalerModel extends MLReadable[StandardScalerModel] {
 
   private class StandardScalerModelReader extends MLReader[StandardScalerModel] {
 
-    private val className = "org.apache.spark.ml.feature.StandardScalerModel"
+    private val className = classOf[StandardScalerModel].getName
 
     override def load(path: String): StandardScalerModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
index 97a2e4f6d6ca4..5c40c35eeaa48 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
@@ -210,7 +210,7 @@ object StringIndexerModel extends MLReadable[StringIndexerModel] {
 
   private class StringIndexerModelReader extends MLReader[StringIndexerModel] {
 
-    private val className = "org.apache.spark.ml.feature.StringIndexerModel"
+    private val className = classOf[StringIndexerModel].getName
 
     override def load(path: String): StringIndexerModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
index 795b73c4c2121..4d35177ad9b0f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
@@ -237,7 +237,7 @@ object ALSModel extends MLReadable[ALSModel] {
   @Since("1.6.0")
   override def load(path: String): ALSModel = super.load(path)
 
-  private[recommendation] class ALSModelWriter(instance: ALSModel) extends MLWriter {
+  private[ALSModel] class ALSModelWriter(instance: ALSModel) extends MLWriter {
 
     override protected def saveImpl(path: String): Unit = {
       val extraMetadata = render("rank" -> instance.rank)
@@ -249,10 +249,10 @@ object ALSModel extends MLReadable[ALSModel] {
     }
   }
 
-  private[recommendation] class ALSModelReader extends MLReader[ALSModel] {
+  private class ALSModelReader extends MLReader[ALSModel] {
 
     /** Checked against metadata when loading model */
-    private val className = "org.apache.spark.ml.recommendation.ALSModel"
+    private val className = classOf[ALSModel].getName
 
     override def load(path: String): ALSModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index 7ba1a60edaf71..70ccec766c471 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -467,14 +467,14 @@ object LinearRegressionModel extends MLReadable[LinearRegressionModel] {
       // Save model data: intercept, coefficients
       val data = Data(instance.intercept, instance.coefficients)
       val dataPath = new Path(path, "data").toString
-      sqlContext.createDataFrame(Seq(data)).write.format("parquet").save(dataPath)
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
     }
   }
 
   private class LinearRegressionModelReader extends MLReader[LinearRegressionModel] {
 
     /** Checked against metadata when loading model */
-    private val className = "org.apache.spark.ml.regression.LinearRegressionModel"
+    private val className = classOf[LinearRegressionModel].getName
 
     override def load(path: String): LinearRegressionModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
index 48ce1bb630685..a9a6ff8a783d5 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
@@ -898,7 +898,7 @@ object LogisticRegressionSuite {
     "regParam" -> 0.01,
     "elasticNetParam" -> 0.1,
     "maxIter" -> 2,  // intentionally small
-    "fitIntercept" -> false,
+    "fitIntercept" -> true,
     "tol" -> 0.8,
     "standardization" -> false,
     "threshold" -> 0.6

From 1a93323c5bab18ed7e55bf6f7b13aae88cb9721c Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Wed, 18 Nov 2015 23:32:49 -0800
Subject: [PATCH 356/518] [SPARK-11339][SPARKR] Document the list of functions
 in R base package that are masked by functions with same name in SparkR

Added tests for function that are reported as masked, to make sure the base:: or stats:: function can be called.

For those we can't call, added them to SparkR programming guide.

It would seem to me `table, sample, subset, filter, cov` not working are not actually expected - I investigated/experimented with them but couldn't get them to work. It looks like as they are defined in base or stats they are missing the S3 generic, eg.
```
> methods("transform")
[1] transform,ANY-method       transform.data.frame
[3] transform,DataFrame-method transform.default
see '?methods' for accessing help and source code
> methods("subset")
[1] subset.data.frame       subset,DataFrame-method subset.default
[4] subset.matrix
see '?methods' for accessing help and source code
Warning message:
In .S3methods(generic.function, class, parent.frame()) :
  function 'subset' appears not to be S3 generic; found functions that look like S3 methods
```
Any idea?

More information on masking:
http://www.ats.ucla.edu/stat/r/faq/referencing_objects.htm
http://www.sfu.ca/~sweldon/howTo/guide4.pdf

This is what the output doc looks like (minus css):
![image](https://cloud.githubusercontent.com/assets/8969467/11229714/2946e5de-8d4d-11e5-94b0-dda9696b6fdd.png)

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9785 from felixcheung/rmasked.
---
 R/pkg/R/DataFrame.R              |  2 +-
 R/pkg/R/functions.R              |  2 +-
 R/pkg/R/generics.R               |  4 ++--
 R/pkg/inst/tests/test_mllib.R    |  5 +++++
 R/pkg/inst/tests/test_sparkSQL.R | 33 +++++++++++++++++++++++++++-
 docs/sparkr.md                   | 37 +++++++++++++++++++++++++++++++-
 6 files changed, 77 insertions(+), 6 deletions(-)

diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index 34177e3cdd94f..06b0108b1389e 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -2152,7 +2152,7 @@ setMethod("with",
           })
 
 #' Returns the column types of a DataFrame.
-#' 
+#'
 #' @name coltypes
 #' @title Get column types of a DataFrame
 #' @family dataframe_funcs
diff --git a/R/pkg/R/functions.R b/R/pkg/R/functions.R
index ff0f438045c14..25a1f22101494 100644
--- a/R/pkg/R/functions.R
+++ b/R/pkg/R/functions.R
@@ -2204,7 +2204,7 @@ setMethod("denseRank",
 #' @export
 #' @examples \dontrun{lag(df$c)}
 setMethod("lag",
-          signature(x = "characterOrColumn", offset = "numeric", defaultValue = "ANY"),
+          signature(x = "characterOrColumn"),
           function(x, offset, defaultValue = NULL) {
             col <- if (class(x) == "Column") {
               x@jc
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 0dcd05438222b..71004a05ba611 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -539,7 +539,7 @@ setGeneric("showDF", function(x,...) { standardGeneric("showDF") })
 
 # @rdname subset
 # @export
-setGeneric("subset", function(x, subset, select, ...) { standardGeneric("subset") })
+setGeneric("subset", function(x, ...) { standardGeneric("subset") })
 
 #' @rdname agg
 #' @export
@@ -790,7 +790,7 @@ setGeneric("kurtosis", function(x) { standardGeneric("kurtosis") })
 
 #' @rdname lag
 #' @export
-setGeneric("lag", function(x, offset, defaultValue = NULL) { standardGeneric("lag") })
+setGeneric("lag", function(x, ...) { standardGeneric("lag") })
 
 #' @rdname last
 #' @export
diff --git a/R/pkg/inst/tests/test_mllib.R b/R/pkg/inst/tests/test_mllib.R
index d497ad8c9daa3..e0667e5e22c18 100644
--- a/R/pkg/inst/tests/test_mllib.R
+++ b/R/pkg/inst/tests/test_mllib.R
@@ -31,6 +31,11 @@ test_that("glm and predict", {
   model <- glm(Sepal_Width ~ Sepal_Length, training, family = "gaussian")
   prediction <- predict(model, test)
   expect_equal(typeof(take(select(prediction, "prediction"), 1)$prediction), "double")
+
+  # Test stats::predict is working
+  x <- rnorm(15)
+  y <- x + rnorm(15)
+  expect_equal(length(predict(lm(y ~ x))), 15)
 })
 
 test_that("glm should work with long formula", {
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index d9a94faff7ac0..3f4f319fe745d 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -433,6 +433,10 @@ test_that("table() returns a new DataFrame", {
   expect_is(tabledf, "DataFrame")
   expect_equal(count(tabledf), 3)
   dropTempTable(sqlContext, "table1")
+
+  # Test base::table is working
+  #a <- letters[1:3]
+  #expect_equal(class(table(a, sample(a))), "table")
 })
 
 test_that("toRDD() returns an RRDD", {
@@ -673,6 +677,9 @@ test_that("sample on a DataFrame", {
   # Also test sample_frac
   sampled3 <- sample_frac(df, FALSE, 0.1, 0) # set seed for predictable result
   expect_true(count(sampled3) < 3)
+
+  # Test base::sample is working
+  #expect_equal(length(sample(1:12)), 12)
 })
 
 test_that("select operators", {
@@ -753,6 +760,9 @@ test_that("subsetting", {
   df6 <- subset(df, df$age %in% c(30), c(1,2))
   expect_equal(count(df6), 1)
   expect_equal(columns(df6), c("name", "age"))
+
+  # Test base::subset is working
+  expect_equal(nrow(subset(airquality, Temp > 80, select = c(Ozone, Temp))), 68)
 })
 
 test_that("selectExpr() on a DataFrame", {
@@ -888,6 +898,9 @@ test_that("column functions", {
   expect_equal(result, list(list(3L, 2L, 1L), list(6L, 5L, 4L)))
   result <- collect(select(df, sort_array(df[[1]])))[[1]]
   expect_equal(result, list(list(1L, 2L, 3L), list(4L, 5L, 6L)))
+
+  # Test that stats::lag is working
+  expect_equal(length(lag(ldeaths, 12)), 72)
 })
 #
 test_that("column binary mathfunctions", {
@@ -1086,7 +1099,7 @@ test_that("group by, agg functions", {
   gd3_local <- collect(agg(gd3, var(df8$age)))
   expect_equal(162, gd3_local[gd3_local$name == "Justin",][1, 2])
 
-  # make sure base:: or stats::sd, var are working
+  # Test stats::sd, stats::var are working
   expect_true(abs(sd(1:2) - 0.7071068) < 1e-6)
   expect_true(abs(var(1:5, 1:5) - 2.5) < 1e-6)
 
@@ -1138,6 +1151,9 @@ test_that("filter() on a DataFrame", {
   expect_equal(count(filtered5), 1)
   filtered6 <- where(df, df$age %in% c(19, 30))
   expect_equal(count(filtered6), 2)
+
+  # Test stats::filter is working
+  #expect_true(is.ts(filter(1:100, rep(1, 3))))
 })
 
 test_that("join() and merge() on a DataFrame", {
@@ -1284,6 +1300,12 @@ test_that("unionAll(), rbind(), except(), and intersect() on a DataFrame", {
   expect_is(unioned, "DataFrame")
   expect_equal(count(intersected), 1)
   expect_equal(first(intersected)$name, "Andy")
+
+  # Test base::rbind is working
+  expect_equal(length(rbind(1:4, c = 2, a = 10, 10, deparse.level = 0)), 16)
+
+  # Test base::intersect is working
+  expect_equal(length(intersect(1:20, 3:23)), 18)
 })
 
 test_that("withColumn() and withColumnRenamed()", {
@@ -1365,6 +1387,9 @@ test_that("describe() and summarize() on a DataFrame", {
   stats2 <- summary(df)
   expect_equal(collect(stats2)[4, "name"], "Andy")
   expect_equal(collect(stats2)[5, "age"], "30")
+
+  # Test base::summary is working
+  expect_equal(length(summary(attenu, digits = 4)), 35)
 })
 
 test_that("dropna() and na.omit() on a DataFrame", {
@@ -1448,6 +1473,9 @@ test_that("dropna() and na.omit() on a DataFrame", {
   expect_identical(expected, actual)
   actual <- collect(na.omit(df, minNonNulls = 3, cols = c("name", "age", "height")))
   expect_identical(expected, actual)
+
+  # Test stats::na.omit is working
+  expect_equal(nrow(na.omit(data.frame(x = c(0, 10, NA)))), 2)
 })
 
 test_that("fillna() on a DataFrame", {
@@ -1510,6 +1538,9 @@ test_that("cov() and corr() on a DataFrame", {
   expect_true(abs(result - 1.0) < 1e-12)
   result <- corr(df, "singles", "doubles", "pearson")
   expect_true(abs(result - 1.0) < 1e-12)
+
+  # Test stats::cov is working
+  #expect_true(abs(max(cov(swiss)) - 1739.295) < 1e-3)
 })
 
 test_that("freqItems() on a DataFrame", {
diff --git a/docs/sparkr.md b/docs/sparkr.md
index a744b76be7466..cfb9b41350f45 100644
--- a/docs/sparkr.md
+++ b/docs/sparkr.md
@@ -286,7 +286,7 @@ head(teenagers)
 
 # Machine Learning
 
-SparkR allows the fitting of generalized linear models over DataFrames using the [glm()](api/R/glm.html) function. Under the hood, SparkR uses MLlib to train a model of the specified family. Currently the gaussian and binomial families are supported. We support a subset of the available R formula operators for model fitting, including '~', '.', ':', '+', and '-'. 
+SparkR allows the fitting of generalized linear models over DataFrames using the [glm()](api/R/glm.html) function. Under the hood, SparkR uses MLlib to train a model of the specified family. Currently the gaussian and binomial families are supported. We support a subset of the available R formula operators for model fitting, including '~', '.', ':', '+', and '-'.
 
 The [summary()](api/R/summary.html) function gives the summary of a model produced by [glm()](api/R/glm.html).
 
@@ -351,3 +351,38 @@ summary(model)
 ##Sepal_Width    0.404655
 {% endhighlight %}
 </div>
+
+# R Function Name Conflicts
+
+When loading and attaching a new package in R, it is possible to have a name [conflict](https://stat.ethz.ch/R-manual/R-devel/library/base/html/library.html), where a
+function is masking another function.
+
+The following functions are masked by the SparkR package:
+
+<table class="table">
+  <tr><th>Masked function</th><th>How to Access</th></tr>
+  <tr>
+    <td><code>cov</code> in <code>package:stats</code></td>
+    <td><code><pre>stats::cov(x, y = NULL, use = "everything",
+           method = c("pearson", "kendall", "spearman"))</pre></code></td>
+  </tr>
+  <tr>
+    <td><code>filter</code> in <code>package:stats</code></td>
+    <td><code><pre>stats::filter(x, filter, method = c("convolution", "recursive"),
+              sides = 2, circular = FALSE, init)</pre></code></td>
+  </tr>
+  <tr>
+    <td><code>sample</code> in <code>package:base</code></td>
+    <td><code>base::sample(x, size, replace = FALSE, prob = NULL)</code></td>
+  </tr>
+  <tr>
+    <td><code>table</code> in <code>package:base</code></td>
+    <td><code><pre>base::table(...,
+            exclude = if (useNA == "no") c(NA, NaN),
+            useNA = c("no", "ifany", "always"),
+            dnn = list.names(...), deparse.level = 1)</pre></code></td>
+  </tr>
+</table>
+
+You can inspect the search path in R with [`search()`](https://stat.ethz.ch/R-manual/R-devel/library/base/html/search.html)
+

From f449992009becc8f7c7f06cda522b9beaa1e263c Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 19 Nov 2015 10:48:04 -0800
Subject: [PATCH 357/518] [SPARK-11849][SQL] Analyzer should replace
 current_date and current_timestamp with literals

We currently rely on the optimizer's constant folding to replace current_timestamp and current_date. However, this can still result in different values for different instances of current_timestamp/current_date if the optimizer is not running fast enough.

A better solution is to replace these functions in the analyzer in one shot.

Author: Reynold Xin <rxin@databricks.com>

Closes #9833 from rxin/SPARK-11849.
---
 .../sql/catalyst/analysis/Analyzer.scala      | 27 ++++++++++---
 .../sql/catalyst/analysis/AnalysisSuite.scala | 38 +++++++++++++++++++
 2 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index f00c451b5981a..84781cd57f3dc 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -65,9 +65,8 @@ class Analyzer(
 
   lazy val batches: Seq[Batch] = Seq(
     Batch("Substitution", fixedPoint,
-      CTESubstitution ::
-      WindowsSubstitution ::
-      Nil : _*),
+      CTESubstitution,
+      WindowsSubstitution),
     Batch("Resolution", fixedPoint,
       ResolveRelations ::
       ResolveReferences ::
@@ -84,7 +83,8 @@ class Analyzer(
       HiveTypeCoercion.typeCoercionRules ++
       extendedResolutionRules : _*),
     Batch("Nondeterministic", Once,
-      PullOutNondeterministic),
+      PullOutNondeterministic,
+      ComputeCurrentTime),
     Batch("UDF", Once,
       HandleNullInputsForUDF),
     Batch("Cleanup", fixedPoint,
@@ -1076,7 +1076,7 @@ class Analyzer(
     override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
       case p if !p.resolved => p // Skip unresolved nodes.
 
-      case plan => plan transformExpressionsUp {
+      case p => p transformExpressionsUp {
 
         case udf @ ScalaUDF(func, _, inputs, _) =>
           val parameterTypes = ScalaReflection.getParameterTypes(func)
@@ -1162,3 +1162,20 @@ object CleanupAliases extends Rule[LogicalPlan] {
       }
   }
 }
+
+/**
+ * Computes the current date and time to make sure we return the same result in a single query.
+ */
+object ComputeCurrentTime extends Rule[LogicalPlan] {
+  def apply(plan: LogicalPlan): LogicalPlan = {
+    val dateExpr = CurrentDate()
+    val timeExpr = CurrentTimestamp()
+    val currentDate = Literal.create(dateExpr.eval(EmptyRow), dateExpr.dataType)
+    val currentTime = Literal.create(timeExpr.eval(EmptyRow), timeExpr.dataType)
+
+    plan transformAllExpressions {
+      case CurrentDate() => currentDate
+      case CurrentTimestamp() => currentTime
+    }
+  }
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
index 08586a97411ae..e051069951887 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
@@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.types._
 
 class AnalysisSuite extends AnalysisTest {
@@ -218,4 +219,41 @@ class AnalysisSuite extends AnalysisTest {
       udf4)
     // checkUDF(udf4, expected4)
   }
+
+  test("analyzer should replace current_timestamp with literals") {
+    val in = Project(Seq(Alias(CurrentTimestamp(), "a")(), Alias(CurrentTimestamp(), "b")()),
+      LocalRelation())
+
+    val min = System.currentTimeMillis() * 1000
+    val plan = in.analyze.asInstanceOf[Project]
+    val max = (System.currentTimeMillis() + 1) * 1000
+
+    val lits = new scala.collection.mutable.ArrayBuffer[Long]
+    plan.transformAllExpressions { case e: Literal =>
+      lits += e.value.asInstanceOf[Long]
+      e
+    }
+    assert(lits.size == 2)
+    assert(lits(0) >= min && lits(0) <= max)
+    assert(lits(1) >= min && lits(1) <= max)
+    assert(lits(0) == lits(1))
+  }
+
+  test("analyzer should replace current_date with literals") {
+    val in = Project(Seq(Alias(CurrentDate(), "a")(), Alias(CurrentDate(), "b")()), LocalRelation())
+
+    val min = DateTimeUtils.millisToDays(System.currentTimeMillis())
+    val plan = in.analyze.asInstanceOf[Project]
+    val max = DateTimeUtils.millisToDays(System.currentTimeMillis())
+
+    val lits = new scala.collection.mutable.ArrayBuffer[Int]
+    plan.transformAllExpressions { case e: Literal =>
+      lits += e.value.asInstanceOf[Int]
+      e
+    }
+    assert(lits.size == 2)
+    assert(lits(0) >= min && lits(0) <= max)
+    assert(lits(1) >= min && lits(1) <= max)
+    assert(lits(0) == lits(1))
+  }
 }

From 962878843b611fa6229e3ee67bb22e2a4bc283cd Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Thu, 19 Nov 2015 11:02:17 -0800
Subject: [PATCH 358/518] [SPARK-11840][SQL] Restore the 1.5's behavior of
 planning a single distinct aggregation.

The impact of this change is for a query that has a single distinct column and does not have any grouping expression like
`SELECT COUNT(DISTINCT a) FROM table`
The plan will be changed from
```
AGG-2 (count distinct)
  Shuffle to a single reducer
    Partial-AGG-2 (count distinct)
      AGG-1 (grouping on a)
        Shuffle by a
          Partial-AGG-1 (grouping on 1)
```
to the following one (1.5 uses this)
```
AGG-2
  AGG-1 (grouping on a)
    Shuffle to a single reducer
      Partial-AGG-1(grouping on a)
```
The first plan is more robust. However, to better benchmark the impact of this change, we should use 1.5's plan and use the conf of `spark.sql.specializeSingleDistinctAggPlanning` to control the plan.

Author: Yin Huai <yhuai@databricks.com>

Closes #9828 from yhuai/distinctRewriter.
---
 .../sql/catalyst/analysis/DistinctAggregationRewriter.scala   | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
index c0c960471a61a..9c78f6d4cc71b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/DistinctAggregationRewriter.scala
@@ -126,8 +126,8 @@ case class DistinctAggregationRewriter(conf: CatalystConf) extends Rule[LogicalP
     val shouldRewrite = if (conf.specializeSingleDistinctAggPlanning) {
       // When the flag is set to specialize single distinct agg planning,
       // we will rely on our Aggregation strategy to handle queries with a single
-      // distinct column and this aggregate operator does have grouping expressions.
-      distinctAggGroups.size > 1 || (distinctAggGroups.size == 1 && a.groupingExpressions.isEmpty)
+      // distinct column.
+      distinctAggGroups.size > 1
     } else {
       distinctAggGroups.size >= 1
     }

From 72d150c271d2b206148fd0917a0def263445121b Mon Sep 17 00:00:00 2001
From: zsxwing <zsxwing@gmail.com>
Date: Thu, 19 Nov 2015 11:57:50 -0800
Subject: [PATCH 359/518] [SPARK-11830][CORE] Make NettyRpcEnv bind to the
 specified host

This PR includes the following change:

1. Bind NettyRpcEnv to the specified host
2. Fix the port information in the log for NettyRpcEnv.
3. Fix the service name of NettyRpcEnv.

Author: zsxwing <zsxwing@gmail.com>
Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9821 from zsxwing/SPARK-11830.
---
 .../src/main/scala/org/apache/spark/SparkEnv.scala |  9 ++++++++-
 .../org/apache/spark/rpc/netty/NettyRpcEnv.scala   |  7 +++----
 .../org/apache/spark/network/TransportContext.java |  8 +++++++-
 .../spark/network/server/TransportServer.java      | 14 ++++++++++----
 4 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala b/core/src/main/scala/org/apache/spark/SparkEnv.scala
index 4474a83bedbdb..88df27f733f2a 100644
--- a/core/src/main/scala/org/apache/spark/SparkEnv.scala
+++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala
@@ -258,8 +258,15 @@ object SparkEnv extends Logging {
       if (rpcEnv.isInstanceOf[AkkaRpcEnv]) {
         rpcEnv.asInstanceOf[AkkaRpcEnv].actorSystem
       } else {
+        val actorSystemPort = if (port == 0) 0 else rpcEnv.address.port + 1
         // Create a ActorSystem for legacy codes
-        AkkaUtils.createActorSystem(actorSystemName, hostname, port, conf, securityManager)._1
+        AkkaUtils.createActorSystem(
+          actorSystemName + "ActorSystem",
+          hostname,
+          actorSystemPort,
+          conf,
+          securityManager
+        )._1
       }
 
     // Figure out which port Akka actually bound to in case the original port is 0 or occupied.
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
index 3e0c497969502..3ce359868039b 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
@@ -102,7 +102,7 @@ private[netty] class NettyRpcEnv(
       } else {
         java.util.Collections.emptyList()
       }
-    server = transportContext.createServer(port, bootstraps)
+    server = transportContext.createServer(host, port, bootstraps)
     dispatcher.registerRpcEndpoint(
       RpcEndpointVerifier.NAME, new RpcEndpointVerifier(this, dispatcher))
   }
@@ -337,10 +337,10 @@ private[netty] class NettyRpcEnvFactory extends RpcEnvFactory with Logging {
     if (!config.clientMode) {
       val startNettyRpcEnv: Int => (NettyRpcEnv, Int) = { actualPort =>
         nettyEnv.startServer(actualPort)
-        (nettyEnv, actualPort)
+        (nettyEnv, nettyEnv.address.port)
       }
       try {
-        Utils.startServiceOnPort(config.port, startNettyRpcEnv, sparkConf, "NettyRpcEnv")._1
+        Utils.startServiceOnPort(config.port, startNettyRpcEnv, sparkConf, config.name)._1
       } catch {
         case NonFatal(e) =>
           nettyEnv.shutdown()
@@ -370,7 +370,6 @@ private[netty] class NettyRpcEnvFactory extends RpcEnvFactory with Logging {
  * @param conf Spark configuration.
  * @param endpointAddress The address where the endpoint is listening.
  * @param nettyEnv The RpcEnv associated with this ref.
- * @param local Whether the referenced endpoint lives in the same process.
  */
 private[netty] class NettyRpcEndpointRef(
     @transient private val conf: SparkConf,
diff --git a/network/common/src/main/java/org/apache/spark/network/TransportContext.java b/network/common/src/main/java/org/apache/spark/network/TransportContext.java
index 1b64b863a9fe5..238710d17249a 100644
--- a/network/common/src/main/java/org/apache/spark/network/TransportContext.java
+++ b/network/common/src/main/java/org/apache/spark/network/TransportContext.java
@@ -94,7 +94,13 @@ public TransportClientFactory createClientFactory() {
 
   /** Create a server which will attempt to bind to a specific port. */
   public TransportServer createServer(int port, List<TransportServerBootstrap> bootstraps) {
-    return new TransportServer(this, port, rpcHandler, bootstraps);
+    return new TransportServer(this, null, port, rpcHandler, bootstraps);
+  }
+
+  /** Create a server which will attempt to bind to a specific host and port. */
+  public TransportServer createServer(
+      String host, int port, List<TransportServerBootstrap> bootstraps) {
+    return new TransportServer(this, host, port, rpcHandler, bootstraps);
   }
 
   /** Creates a new server, binding to any available ephemeral port. */
diff --git a/network/common/src/main/java/org/apache/spark/network/server/TransportServer.java b/network/common/src/main/java/org/apache/spark/network/server/TransportServer.java
index f4fadb1ee3b8d..baae235e02205 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/TransportServer.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/TransportServer.java
@@ -55,9 +55,13 @@ public class TransportServer implements Closeable {
   private ChannelFuture channelFuture;
   private int port = -1;
 
-  /** Creates a TransportServer that binds to the given port, or to any available if 0. */
+  /**
+   * Creates a TransportServer that binds to the given host and the given port, or to any available
+   * if 0. If you don't want to bind to any special host, set "hostToBind" to null.
+   * */
   public TransportServer(
       TransportContext context,
+      String hostToBind,
       int portToBind,
       RpcHandler appRpcHandler,
       List<TransportServerBootstrap> bootstraps) {
@@ -67,7 +71,7 @@ public TransportServer(
     this.bootstraps = Lists.newArrayList(Preconditions.checkNotNull(bootstraps));
 
     try {
-      init(portToBind);
+      init(hostToBind, portToBind);
     } catch (RuntimeException e) {
       JavaUtils.closeQuietly(this);
       throw e;
@@ -81,7 +85,7 @@ public int getPort() {
     return port;
   }
 
-  private void init(int portToBind) {
+  private void init(String hostToBind, int portToBind) {
 
     IOMode ioMode = IOMode.valueOf(conf.ioMode());
     EventLoopGroup bossGroup =
@@ -120,7 +124,9 @@ protected void initChannel(SocketChannel ch) throws Exception {
       }
     });
 
-    channelFuture = bootstrap.bind(new InetSocketAddress(portToBind));
+    InetSocketAddress address = hostToBind == null ?
+        new InetSocketAddress(portToBind): new InetSocketAddress(hostToBind, portToBind);
+    channelFuture = bootstrap.bind(address);
     channelFuture.syncUninterruptibly();
 
     port = ((InetSocketAddress) channelFuture.channel().localAddress()).getPort();

From 276a7e130252c0e7aba702ae5570b3c4f424b23b Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Thu, 19 Nov 2015 12:45:04 -0800
Subject: [PATCH 360/518] [SPARK-11633][SQL] LogicalRDD throws TreeNode
 Exception : Failed to Copy Node

When handling self joins, the implementation did not consider the case insensitivity of HiveContext. It could cause an exception as shown in the JIRA:
```
TreeNodeException: Failed to copy node.
```

The fix is low risk. It avoids unnecessary attribute replacement. It should not affect the existing behavior of self joins. Also added the test case to cover this case.

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9762 from gatorsmile/joinMakeCopy.
---
 .../apache/spark/sql/execution/ExistingRDD.scala   |  4 ++++
 .../org/apache/spark/sql/DataFrameSuite.scala      | 14 ++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
index 62620ec642c78..623348f6768a4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
@@ -74,6 +74,10 @@ private[sql] case class LogicalRDD(
 
   override def children: Seq[LogicalPlan] = Nil
 
+  override protected final def otherCopyArgs: Seq[AnyRef] = {
+    sqlContext :: Nil
+  }
+
   override def newInstance(): LogicalRDD.this.type =
     LogicalRDD(output.map(_.newInstance()), rdd)(sqlContext).asInstanceOf[this.type]
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index 6399b0165c4c3..dd6d06512ff60 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -1110,6 +1110,20 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
     }
   }
 
+  // This test case is to verify a bug when making a new instance of LogicalRDD.
+  test("SPARK-11633: LogicalRDD throws TreeNode Exception: Failed to Copy Node") {
+    withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") {
+      val rdd = sparkContext.makeRDD(Seq(Row(1, 3), Row(2, 1)))
+      val df = sqlContext.createDataFrame(
+        rdd,
+        new StructType().add("f1", IntegerType).add("f2", IntegerType),
+        needsConversion = false).select($"F1", $"f2".as("f2"))
+      val df1 = df.as("a")
+      val df2 = df.as("b")
+      checkAnswer(df1.join(df2, $"a.f2" === $"b.f2"), Row(1, 3, 1, 3) :: Row(2, 1, 2, 1) :: Nil)
+    }
+  }
+
   test("SPARK-10656: completely support special chars") {
     val df = Seq(1 -> "a").toDF("i_$.a", "d^'a.")
     checkAnswer(df.select(df("*")), Row(1, "a"))

From 7d4aba18722727c85893ad8d8f07d4494665dcfc Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Thu, 19 Nov 2015 12:46:36 -0800
Subject: [PATCH 361/518] [SPARK-11848][SQL] Support EXPLAIN in DataSet APIs

When debugging DataSet API, I always need to print the logical and physical plans.

I am wondering if we should provide a simple API for EXPLAIN?

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9832 from gatorsmile/explainDS.
---
 .../org/apache/spark/sql/DataFrame.scala      | 23 +------------------
 .../spark/sql/execution/Queryable.scala       | 21 +++++++++++++++++
 2 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 3ba4ba18d2122..98358127e2709 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.{Inner, JoinType}
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, ScalaReflection, SqlParser}
-import org.apache.spark.sql.execution.{EvaluatePython, ExplainCommand, FileRelation, LogicalRDD, QueryExecution, Queryable, SQLExecution}
+import org.apache.spark.sql.execution.{EvaluatePython, FileRelation, LogicalRDD, QueryExecution, Queryable, SQLExecution}
 import org.apache.spark.sql.execution.datasources.{CreateTableUsingAsSelect, LogicalRelation}
 import org.apache.spark.sql.execution.datasources.json.JacksonGenerator
 import org.apache.spark.sql.sources.HadoopFsRelation
@@ -308,27 +308,6 @@ class DataFrame private[sql](
   def printSchema(): Unit = println(schema.treeString)
   // scalastyle:on println
 
-  /**
-   * Prints the plans (logical and physical) to the console for debugging purposes.
-   * @group basic
-   * @since 1.3.0
-   */
-  def explain(extended: Boolean): Unit = {
-    val explain = ExplainCommand(queryExecution.logical, extended = extended)
-    withPlan(explain).queryExecution.executedPlan.executeCollect().foreach {
-      // scalastyle:off println
-      r => println(r.getString(0))
-      // scalastyle:on println
-    }
-  }
-
-  /**
-   * Only prints the physical plan to the console for debugging purposes.
-   * @group basic
-   * @since 1.3.0
-   */
-  def explain(): Unit = explain(extended = false)
-
   /**
    * Returns true if the `collect` and `take` methods can be run locally
    * (without any Spark executors).
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
index 9ca383896a09b..e86a52c149a2f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.execution
 
+import org.apache.spark.sql.SQLContext
 import org.apache.spark.sql.types.StructType
 
 import scala.util.control.NonFatal
@@ -25,6 +26,7 @@ import scala.util.control.NonFatal
 private[sql] trait Queryable {
   def schema: StructType
   def queryExecution: QueryExecution
+  def sqlContext: SQLContext
 
   override def toString: String = {
     try {
@@ -34,4 +36,23 @@ private[sql] trait Queryable {
         s"Invalid tree; ${e.getMessage}:\n$queryExecution"
     }
   }
+
+  /**
+   * Prints the plans (logical and physical) to the console for debugging purposes.
+   * @since 1.3.0
+   */
+  def explain(extended: Boolean): Unit = {
+    val explain = ExplainCommand(queryExecution.logical, extended = extended)
+    sqlContext.executePlan(explain).executedPlan.executeCollect().foreach {
+      // scalastyle:off println
+      r => println(r.getString(0))
+      // scalastyle:on println
+    }
+  }
+
+  /**
+   * Only prints the physical plan to the console for debugging purposes.
+   * @since 1.3.0
+   */
+  def explain(): Unit = explain(extended = false)
 }

From 47d1c2325caaf9ffe31695b6fff529314b8582f7 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Thu, 19 Nov 2015 12:54:25 -0800
Subject: [PATCH 362/518] [SPARK-11750][SQL] revert SPARK-11727 and code clean
 up

After some experiment, I found it's not convenient to have separate encoder builders: `FlatEncoder` and `ProductEncoder`. For example, when create encoders for `ScalaUDF`, we have no idea if the type `T` is flat or not. So I revert the splitting change in https://github.com/apache/spark/pull/9693, while still keeping the bug fixes and tests.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9726 from cloud-fan/follow.
---
 .../scala/org/apache/spark/sql/Encoder.scala  |  16 +-
 .../spark/sql/catalyst/ScalaReflection.scala  | 354 +++++---------
 .../catalyst/encoders/ExpressionEncoder.scala |  19 +-
 .../sql/catalyst/encoders/FlatEncoder.scala   |  50 --
 .../catalyst/encoders/ProductEncoder.scala    | 452 ------------------
 .../sql/catalyst/encoders/RowEncoder.scala    |  12 +-
 .../sql/catalyst/expressions/objects.scala    |   7 +-
 .../sql/catalyst/ScalaReflectionSuite.scala   |  68 ---
 .../encoders/ExpressionEncoderSuite.scala     | 218 ++++++++-
 .../catalyst/encoders/FlatEncoderSuite.scala  |  99 ----
 .../encoders/ProductEncoderSuite.scala        | 156 ------
 .../org/apache/spark/sql/GroupedDataset.scala |   4 +-
 .../org/apache/spark/sql/SQLImplicits.scala   |  23 +-
 .../org/apache/spark/sql/functions.scala      |   4 +-
 14 files changed, 364 insertions(+), 1118 deletions(-)
 delete mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoder.scala
 delete mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
 delete mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
 delete mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index d54f2854fb33f..86bb536459035 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -45,14 +45,14 @@ trait Encoder[T] extends Serializable {
  */
 object Encoders {
 
-  def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder(flat = true)
-  def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder(flat = true)
-  def SHORT: Encoder[java.lang.Short] = ExpressionEncoder(flat = true)
-  def INT: Encoder[java.lang.Integer] = ExpressionEncoder(flat = true)
-  def LONG: Encoder[java.lang.Long] = ExpressionEncoder(flat = true)
-  def FLOAT: Encoder[java.lang.Float] = ExpressionEncoder(flat = true)
-  def DOUBLE: Encoder[java.lang.Double] = ExpressionEncoder(flat = true)
-  def STRING: Encoder[java.lang.String] = ExpressionEncoder(flat = true)
+  def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder()
+  def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder()
+  def SHORT: Encoder[java.lang.Short] = ExpressionEncoder()
+  def INT: Encoder[java.lang.Integer] = ExpressionEncoder()
+  def LONG: Encoder[java.lang.Long] = ExpressionEncoder()
+  def FLOAT: Encoder[java.lang.Float] = ExpressionEncoder()
+  def DOUBLE: Encoder[java.lang.Double] = ExpressionEncoder()
+  def STRING: Encoder[java.lang.String] = ExpressionEncoder()
 
   /**
    * (Scala-specific) Creates an encoder that serializes objects of type T using Kryo.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 59ccf356f2c48..33ae700706dae 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -50,39 +50,29 @@ object ScalaReflection extends ScalaReflection {
    * Unlike `schemaFor`, this function doesn't do any massaging of types into the Spark SQL type
    * system.  As a result, ObjectType will be returned for things like boxed Integers
    */
-  def dataTypeFor(tpe: `Type`): DataType = tpe match {
-    case t if t <:< definitions.IntTpe => IntegerType
-    case t if t <:< definitions.LongTpe => LongType
-    case t if t <:< definitions.DoubleTpe => DoubleType
-    case t if t <:< definitions.FloatTpe => FloatType
-    case t if t <:< definitions.ShortTpe => ShortType
-    case t if t <:< definitions.ByteTpe => ByteType
-    case t if t <:< definitions.BooleanTpe => BooleanType
-    case t if t <:< localTypeOf[Array[Byte]] => BinaryType
-    case _ =>
-      val className: String = tpe.erasure.typeSymbol.asClass.fullName
-      className match {
-        case "scala.Array" =>
-          val TypeRef(_, _, Seq(arrayType)) = tpe
-          val cls = arrayType match {
-            case t if t <:< definitions.IntTpe => classOf[Array[Int]]
-            case t if t <:< definitions.LongTpe => classOf[Array[Long]]
-            case t if t <:< definitions.DoubleTpe => classOf[Array[Double]]
-            case t if t <:< definitions.FloatTpe => classOf[Array[Float]]
-            case t if t <:< definitions.ShortTpe => classOf[Array[Short]]
-            case t if t <:< definitions.ByteTpe => classOf[Array[Byte]]
-            case t if t <:< definitions.BooleanTpe => classOf[Array[Boolean]]
-            case other =>
-              // There is probably a better way to do this, but I couldn't find it...
-              val elementType = dataTypeFor(other).asInstanceOf[ObjectType].cls
-              java.lang.reflect.Array.newInstance(elementType, 1).getClass
+  def dataTypeFor[T : TypeTag]: DataType = dataTypeFor(localTypeOf[T])
 
-          }
-          ObjectType(cls)
-        case other =>
-          val clazz = mirror.runtimeClass(tpe.erasure.typeSymbol.asClass)
-          ObjectType(clazz)
-      }
+  private def dataTypeFor(tpe: `Type`): DataType = ScalaReflectionLock.synchronized {
+    tpe match {
+      case t if t <:< definitions.IntTpe => IntegerType
+      case t if t <:< definitions.LongTpe => LongType
+      case t if t <:< definitions.DoubleTpe => DoubleType
+      case t if t <:< definitions.FloatTpe => FloatType
+      case t if t <:< definitions.ShortTpe => ShortType
+      case t if t <:< definitions.ByteTpe => ByteType
+      case t if t <:< definitions.BooleanTpe => BooleanType
+      case t if t <:< localTypeOf[Array[Byte]] => BinaryType
+      case _ =>
+        val className: String = tpe.erasure.typeSymbol.asClass.fullName
+        className match {
+          case "scala.Array" =>
+            val TypeRef(_, _, Seq(elementType)) = tpe
+            arrayClassFor(elementType)
+          case other =>
+            val clazz = mirror.runtimeClass(tpe.erasure.typeSymbol.asClass)
+            ObjectType(clazz)
+        }
+    }
   }
 
   /**
@@ -90,7 +80,7 @@ object ScalaReflection extends ScalaReflection {
    * Array[T].  Special handling is performed for primitive types to map them back to their raw
    * JVM form instead of the Scala Array that handles auto boxing.
    */
-  def arrayClassFor(tpe: `Type`): DataType = {
+  private def arrayClassFor(tpe: `Type`): DataType = ScalaReflectionLock.synchronized {
     val cls = tpe match {
       case t if t <:< definitions.IntTpe => classOf[Array[Int]]
       case t if t <:< definitions.LongTpe => classOf[Array[Long]]
@@ -108,6 +98,15 @@ object ScalaReflection extends ScalaReflection {
     ObjectType(cls)
   }
 
+  /**
+   * Returns true if the value of this data type is same between internal and external.
+   */
+  def isNativeType(dt: DataType): Boolean = dt match {
+    case BooleanType | ByteType | ShortType | IntegerType | LongType |
+         FloatType | DoubleType | BinaryType => true
+    case _ => false
+  }
+
   /**
    * Returns an expression that can be used to construct an object of type `T` given an input
    * row with a compatible schema.  Fields of the row will be extracted using UnresolvedAttributes
@@ -116,63 +115,33 @@ object ScalaReflection extends ScalaReflection {
    *
    * When used on a primitive type, the constructor will instead default to extracting the value
    * from ordinal 0 (since there are no names to map to).  The actual location can be moved by
-   * calling unbind/bind with a new schema.
+   * calling resolve/bind with a new schema.
    */
-  def constructorFor[T : TypeTag]: Expression = constructorFor(typeOf[T], None)
+  def constructorFor[T : TypeTag]: Expression = constructorFor(localTypeOf[T], None)
 
   private def constructorFor(
       tpe: `Type`,
       path: Option[Expression]): Expression = ScalaReflectionLock.synchronized {
 
     /** Returns the current path with a sub-field extracted. */
-    def addToPath(part: String): Expression =
-      path
-        .map(p => UnresolvedExtractValue(p, expressions.Literal(part)))
-        .getOrElse(UnresolvedAttribute(part))
+    def addToPath(part: String): Expression = path
+      .map(p => UnresolvedExtractValue(p, expressions.Literal(part)))
+      .getOrElse(UnresolvedAttribute(part))
 
     /** Returns the current path with a field at ordinal extracted. */
-    def addToPathOrdinal(ordinal: Int, dataType: DataType): Expression =
-      path
-        .map(p => GetStructField(p, StructField(s"_$ordinal", dataType), ordinal))
-        .getOrElse(BoundReference(ordinal, dataType, false))
+    def addToPathOrdinal(ordinal: Int, dataType: DataType): Expression = path
+      .map(p => GetInternalRowField(p, ordinal, dataType))
+      .getOrElse(BoundReference(ordinal, dataType, false))
 
-    /** Returns the current path or throws an error. */
-    def getPath = path.getOrElse(BoundReference(0, schemaFor(tpe).dataType, true))
+    /** Returns the current path or `BoundReference`. */
+    def getPath: Expression = path.getOrElse(BoundReference(0, schemaFor(tpe).dataType, true))
 
     tpe match {
-      case t if !dataTypeFor(t).isInstanceOf[ObjectType] =>
-        getPath
+      case t if !dataTypeFor(t).isInstanceOf[ObjectType] => getPath
 
       case t if t <:< localTypeOf[Option[_]] =>
         val TypeRef(_, _, Seq(optType)) = t
-        val boxedType = optType match {
-          // For primitive types we must manually box the primitive value.
-          case t if t <:< definitions.IntTpe => Some(classOf[java.lang.Integer])
-          case t if t <:< definitions.LongTpe => Some(classOf[java.lang.Long])
-          case t if t <:< definitions.DoubleTpe => Some(classOf[java.lang.Double])
-          case t if t <:< definitions.FloatTpe => Some(classOf[java.lang.Float])
-          case t if t <:< definitions.ShortTpe => Some(classOf[java.lang.Short])
-          case t if t <:< definitions.ByteTpe => Some(classOf[java.lang.Byte])
-          case t if t <:< definitions.BooleanTpe => Some(classOf[java.lang.Boolean])
-          case _ => None
-        }
-
-        boxedType.map { boxedType =>
-          val objectType = ObjectType(boxedType)
-          WrapOption(
-            objectType,
-            NewInstance(
-              boxedType,
-              getPath :: Nil,
-              propagateNull = true,
-              objectType))
-        }.getOrElse {
-          val className: String = optType.erasure.typeSymbol.asClass.fullName
-          val cls = Utils.classForName(className)
-          val objectType = ObjectType(cls)
-
-          WrapOption(objectType, constructorFor(optType, path))
-        }
+        WrapOption(constructorFor(optType, path))
 
       case t if t <:< localTypeOf[java.lang.Integer] =>
         val boxedType = classOf[java.lang.Integer]
@@ -231,11 +200,11 @@ object ScalaReflection extends ScalaReflection {
       case t if t <:< localTypeOf[java.math.BigDecimal] =>
         Invoke(getPath, "toJavaBigDecimal", ObjectType(classOf[java.math.BigDecimal]))
 
+      case t if t <:< localTypeOf[BigDecimal] =>
+        Invoke(getPath, "toBigDecimal", ObjectType(classOf[BigDecimal]))
+
       case t if t <:< localTypeOf[Array[_]] =>
         val TypeRef(_, _, Seq(elementType)) = t
-        val elementDataType = dataTypeFor(elementType)
-        val Schema(dataType, nullable) = schemaFor(elementType)
-
         val primitiveMethod = elementType match {
           case t if t <:< definitions.IntTpe => Some("toIntArray")
           case t if t <:< definitions.LongTpe => Some("toLongArray")
@@ -248,57 +217,52 @@ object ScalaReflection extends ScalaReflection {
         }
 
         primitiveMethod.map { method =>
-          Invoke(getPath, method, dataTypeFor(t))
+          Invoke(getPath, method, arrayClassFor(elementType))
         }.getOrElse {
-          val returnType = dataTypeFor(t)
           Invoke(
-            MapObjects(p => constructorFor(elementType, Some(p)), getPath, dataType),
+            MapObjects(
+              p => constructorFor(elementType, Some(p)),
+              getPath,
+              schemaFor(elementType).dataType),
             "array",
-            returnType)
+            arrayClassFor(elementType))
         }
 
+      case t if t <:< localTypeOf[Seq[_]] =>
+        val TypeRef(_, _, Seq(elementType)) = t
+        val arrayData =
+          Invoke(
+            MapObjects(
+              p => constructorFor(elementType, Some(p)),
+              getPath,
+              schemaFor(elementType).dataType),
+            "array",
+            ObjectType(classOf[Array[Any]]))
+
+        StaticInvoke(
+          scala.collection.mutable.WrappedArray,
+          ObjectType(classOf[Seq[_]]),
+          "make",
+          arrayData :: Nil)
+
       case t if t <:< localTypeOf[Map[_, _]] =>
         val TypeRef(_, _, Seq(keyType, valueType)) = t
-        val Schema(keyDataType, _) = schemaFor(keyType)
-        val Schema(valueDataType, valueNullable) = schemaFor(valueType)
-
-        val primitiveMethodKey = keyType match {
-          case t if t <:< definitions.IntTpe => Some("toIntArray")
-          case t if t <:< definitions.LongTpe => Some("toLongArray")
-          case t if t <:< definitions.DoubleTpe => Some("toDoubleArray")
-          case t if t <:< definitions.FloatTpe => Some("toFloatArray")
-          case t if t <:< definitions.ShortTpe => Some("toShortArray")
-          case t if t <:< definitions.ByteTpe => Some("toByteArray")
-          case t if t <:< definitions.BooleanTpe => Some("toBooleanArray")
-          case _ => None
-        }
 
         val keyData =
           Invoke(
             MapObjects(
               p => constructorFor(keyType, Some(p)),
-              Invoke(getPath, "keyArray", ArrayType(keyDataType)),
-              keyDataType),
+              Invoke(getPath, "keyArray", ArrayType(schemaFor(keyType).dataType)),
+              schemaFor(keyType).dataType),
             "array",
             ObjectType(classOf[Array[Any]]))
 
-        val primitiveMethodValue = valueType match {
-          case t if t <:< definitions.IntTpe => Some("toIntArray")
-          case t if t <:< definitions.LongTpe => Some("toLongArray")
-          case t if t <:< definitions.DoubleTpe => Some("toDoubleArray")
-          case t if t <:< definitions.FloatTpe => Some("toFloatArray")
-          case t if t <:< definitions.ShortTpe => Some("toShortArray")
-          case t if t <:< definitions.ByteTpe => Some("toByteArray")
-          case t if t <:< definitions.BooleanTpe => Some("toBooleanArray")
-          case _ => None
-        }
-
         val valueData =
           Invoke(
             MapObjects(
               p => constructorFor(valueType, Some(p)),
-              Invoke(getPath, "valueArray", ArrayType(valueDataType)),
-              valueDataType),
+              Invoke(getPath, "valueArray", ArrayType(schemaFor(valueType).dataType)),
+              schemaFor(valueType).dataType),
             "array",
             ObjectType(classOf[Array[Any]]))
 
@@ -308,40 +272,6 @@ object ScalaReflection extends ScalaReflection {
           "toScalaMap",
           keyData :: valueData :: Nil)
 
-      case t if t <:< localTypeOf[Seq[_]] =>
-        val TypeRef(_, _, Seq(elementType)) = t
-        val elementDataType = dataTypeFor(elementType)
-        val Schema(dataType, nullable) = schemaFor(elementType)
-
-        // Avoid boxing when possible by just wrapping a primitive array.
-        val primitiveMethod = elementType match {
-          case _ if nullable => None
-          case t if t <:< definitions.IntTpe => Some("toIntArray")
-          case t if t <:< definitions.LongTpe => Some("toLongArray")
-          case t if t <:< definitions.DoubleTpe => Some("toDoubleArray")
-          case t if t <:< definitions.FloatTpe => Some("toFloatArray")
-          case t if t <:< definitions.ShortTpe => Some("toShortArray")
-          case t if t <:< definitions.ByteTpe => Some("toByteArray")
-          case t if t <:< definitions.BooleanTpe => Some("toBooleanArray")
-          case _ => None
-        }
-
-        val arrayData = primitiveMethod.map { method =>
-          Invoke(getPath, method, arrayClassFor(elementType))
-        }.getOrElse {
-          Invoke(
-            MapObjects(p => constructorFor(elementType, Some(p)), getPath, dataType),
-            "array",
-            arrayClassFor(elementType))
-        }
-
-        StaticInvoke(
-          scala.collection.mutable.WrappedArray,
-          ObjectType(classOf[Seq[_]]),
-          "make",
-          arrayData :: Nil)
-
-
       case t if t <:< localTypeOf[Product] =>
         val formalTypeArgs = t.typeSymbol.asClass.typeParams
         val TypeRef(_, _, actualTypeArgs) = t
@@ -361,8 +291,7 @@ object ScalaReflection extends ScalaReflection {
           }
         }
 
-        val className: String = t.erasure.typeSymbol.asClass.fullName
-        val cls = Utils.classForName(className)
+        val cls = mirror.runtimeClass(tpe.erasure.typeSymbol.asClass)
 
         val arguments = params.head.zipWithIndex.map { case (p, i) =>
           val fieldName = p.name.toString
@@ -370,7 +299,7 @@ object ScalaReflection extends ScalaReflection {
           val dataType = schemaFor(fieldType).dataType
 
           // For tuples, we based grab the inner fields by ordinal instead of name.
-          if (className startsWith "scala.Tuple") {
+          if (cls.getName startsWith "scala.Tuple") {
             constructorFor(fieldType, Some(addToPathOrdinal(i, dataType)))
           } else {
             constructorFor(fieldType, Some(addToPath(fieldName)))
@@ -388,22 +317,19 @@ object ScalaReflection extends ScalaReflection {
         } else {
           newInstance
         }
-
     }
   }
 
   /** Returns expressions for extracting all the fields from the given type. */
   def extractorsFor[T : TypeTag](inputObject: Expression): CreateNamedStruct = {
-    ScalaReflectionLock.synchronized {
-      extractorFor(inputObject, typeTag[T].tpe) match {
-        case s: CreateNamedStruct => s
-        case o => CreateNamedStruct(expressions.Literal("value") :: o :: Nil)
-      }
+    extractorFor(inputObject, localTypeOf[T]) match {
+      case s: CreateNamedStruct => s
+      case other => CreateNamedStruct(expressions.Literal("value") :: other :: Nil)
     }
   }
 
   /** Helper for extracting internal fields from a case class. */
-  protected def extractorFor(
+  private def extractorFor(
       inputObject: Expression,
       tpe: `Type`): Expression = ScalaReflectionLock.synchronized {
     if (!inputObject.dataType.isInstanceOf[ObjectType]) {
@@ -491,51 +417,36 @@ object ScalaReflection extends ScalaReflection {
 
         case t if t <:< localTypeOf[Array[_]] =>
           val TypeRef(_, _, Seq(elementType)) = t
-          val elementDataType = dataTypeFor(elementType)
-          val Schema(dataType, nullable) = schemaFor(elementType)
-
-          if (!elementDataType.isInstanceOf[AtomicType]) {
-            MapObjects(extractorFor(_, elementType), inputObject, elementDataType)
-          } else {
-            NewInstance(
-              classOf[GenericArrayData],
-              inputObject :: Nil,
-              dataType = ArrayType(dataType, nullable))
-          }
+          toCatalystArray(inputObject, elementType)
 
         case t if t <:< localTypeOf[Seq[_]] =>
           val TypeRef(_, _, Seq(elementType)) = t
-          val elementDataType = dataTypeFor(elementType)
-          val Schema(dataType, nullable) = schemaFor(elementType)
-
-          if (dataType.isInstanceOf[AtomicType]) {
-            NewInstance(
-              classOf[GenericArrayData],
-              inputObject :: Nil,
-              dataType = ArrayType(dataType, nullable))
-          } else {
-            MapObjects(extractorFor(_, elementType), inputObject, elementDataType)
-          }
+          toCatalystArray(inputObject, elementType)
 
         case t if t <:< localTypeOf[Map[_, _]] =>
           val TypeRef(_, _, Seq(keyType, valueType)) = t
-          val Schema(keyDataType, _) = schemaFor(keyType)
-          val Schema(valueDataType, valueNullable) = schemaFor(valueType)
 
-          val rawMap = inputObject
           val keys =
-            NewInstance(
-              classOf[GenericArrayData],
-              Invoke(rawMap, "keys", ObjectType(classOf[scala.collection.GenIterable[_]])) :: Nil,
-              dataType = ObjectType(classOf[ArrayData]))
+            Invoke(
+              Invoke(inputObject, "keysIterator",
+                ObjectType(classOf[scala.collection.Iterator[_]])),
+              "toSeq",
+              ObjectType(classOf[scala.collection.Seq[_]]))
+          val convertedKeys = toCatalystArray(keys, keyType)
+
           val values =
-            NewInstance(
-              classOf[GenericArrayData],
-              Invoke(rawMap, "values", ObjectType(classOf[scala.collection.GenIterable[_]])) :: Nil,
-              dataType = ObjectType(classOf[ArrayData]))
+            Invoke(
+              Invoke(inputObject, "valuesIterator",
+                ObjectType(classOf[scala.collection.Iterator[_]])),
+              "toSeq",
+              ObjectType(classOf[scala.collection.Seq[_]]))
+          val convertedValues = toCatalystArray(values, valueType)
+
+          val Schema(keyDataType, _) = schemaFor(keyType)
+          val Schema(valueDataType, valueNullable) = schemaFor(valueType)
           NewInstance(
             classOf[ArrayBasedMapData],
-            keys :: values :: Nil,
+            convertedKeys :: convertedValues :: Nil,
             dataType = MapType(keyDataType, valueDataType, valueNullable))
 
         case t if t <:< localTypeOf[String] =>
@@ -558,6 +469,7 @@ object ScalaReflection extends ScalaReflection {
             DateType,
             "fromJavaDate",
             inputObject :: Nil)
+
         case t if t <:< localTypeOf[BigDecimal] =>
           StaticInvoke(
             Decimal,
@@ -587,26 +499,24 @@ object ScalaReflection extends ScalaReflection {
         case t if t <:< localTypeOf[java.lang.Boolean] =>
           Invoke(inputObject, "booleanValue", BooleanType)
 
-        case t if t <:< definitions.IntTpe =>
-          BoundReference(0, IntegerType, false)
-        case t if t <:< definitions.LongTpe =>
-          BoundReference(0, LongType, false)
-        case t if t <:< definitions.DoubleTpe =>
-          BoundReference(0, DoubleType, false)
-        case t if t <:< definitions.FloatTpe =>
-          BoundReference(0, FloatType, false)
-        case t if t <:< definitions.ShortTpe =>
-          BoundReference(0, ShortType, false)
-        case t if t <:< definitions.ByteTpe =>
-          BoundReference(0, ByteType, false)
-        case t if t <:< definitions.BooleanTpe =>
-          BoundReference(0, BooleanType, false)
-
         case other =>
           throw new UnsupportedOperationException(s"Extractor for type $other is not supported")
       }
     }
   }
+
+  private def toCatalystArray(input: Expression, elementType: `Type`): Expression = {
+    val externalDataType = dataTypeFor(elementType)
+    val Schema(catalystType, nullable) = schemaFor(elementType)
+    if (isNativeType(catalystType)) {
+      NewInstance(
+        classOf[GenericArrayData],
+        input :: Nil,
+        dataType = ArrayType(catalystType, nullable))
+    } else {
+      MapObjects(extractorFor(_, elementType), input, externalDataType)
+    }
+  }
 }
 
 /**
@@ -635,8 +545,7 @@ trait ScalaReflection {
   }
 
   /** Returns a catalyst DataType and its nullability for the given Scala Type using reflection. */
-  def schemaFor[T: TypeTag]: Schema =
-    ScalaReflectionLock.synchronized { schemaFor(localTypeOf[T]) }
+  def schemaFor[T: TypeTag]: Schema = schemaFor(localTypeOf[T])
 
   /**
    * Return the Scala Type for `T` in the current classloader mirror.
@@ -736,39 +645,4 @@ trait ScalaReflection {
     assert(methods.length == 1)
     methods.head.getParameterTypes
   }
-
-  def typeOfObject: PartialFunction[Any, DataType] = {
-    // The data type can be determined without ambiguity.
-    case obj: Boolean => BooleanType
-    case obj: Array[Byte] => BinaryType
-    case obj: String => StringType
-    case obj: UTF8String => StringType
-    case obj: Byte => ByteType
-    case obj: Short => ShortType
-    case obj: Int => IntegerType
-    case obj: Long => LongType
-    case obj: Float => FloatType
-    case obj: Double => DoubleType
-    case obj: java.sql.Date => DateType
-    case obj: java.math.BigDecimal => DecimalType.SYSTEM_DEFAULT
-    case obj: Decimal => DecimalType.SYSTEM_DEFAULT
-    case obj: java.sql.Timestamp => TimestampType
-    case null => NullType
-    // For other cases, there is no obvious mapping from the type of the given object to a
-    // Catalyst data type. A user should provide his/her specific rules
-    // (in a user-defined PartialFunction) to infer the Catalyst data type for other types of
-    // objects and then compose the user-defined PartialFunction with this one.
-  }
-
-  implicit class CaseClassRelation[A <: Product : TypeTag](data: Seq[A]) {
-
-    /**
-     * Implicitly added to Sequences of case class objects.  Returns a catalyst logical relation
-     * for the the data in the sequence.
-     */
-    def asRelation: LocalRelation = {
-      val output = attributesFor[A]
-      LocalRelation.fromProduct(output, data)
-    }
-  }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index 456b595008479..6eeba1442c1f3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -30,10 +30,10 @@ import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateSafeProjection, GenerateUnsafeProjection}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.ScalaReflection
-import org.apache.spark.sql.types.{NullType, StructField, ObjectType, StructType}
+import org.apache.spark.sql.types.{StructField, ObjectType, StructType}
 
 /**
- * A factory for constructing encoders that convert objects and primitves to and from the
+ * A factory for constructing encoders that convert objects and primitives to and from the
  * internal row format using catalyst expressions and code generation.  By default, the
  * expressions used to retrieve values from an input row when producing an object will be created as
  * follows:
@@ -44,20 +44,21 @@ import org.apache.spark.sql.types.{NullType, StructField, ObjectType, StructType
  *    to the name `value`.
  */
 object ExpressionEncoder {
-  def apply[T : TypeTag](flat: Boolean = false): ExpressionEncoder[T] = {
+  def apply[T : TypeTag](): ExpressionEncoder[T] = {
     // We convert the not-serializable TypeTag into StructType and ClassTag.
     val mirror = typeTag[T].mirror
     val cls = mirror.runtimeClass(typeTag[T].tpe)
+    val flat = !classOf[Product].isAssignableFrom(cls)
 
-    val inputObject = BoundReference(0, ObjectType(cls), nullable = true)
-    val extractExpression = ScalaReflection.extractorsFor[T](inputObject)
-    val constructExpression = ScalaReflection.constructorFor[T]
+    val inputObject = BoundReference(0, ScalaReflection.dataTypeFor[T], nullable = true)
+    val toRowExpression = ScalaReflection.extractorsFor[T](inputObject)
+    val fromRowExpression = ScalaReflection.constructorFor[T]
 
     new ExpressionEncoder[T](
-      extractExpression.dataType,
+      toRowExpression.dataType,
       flat,
-      extractExpression.flatten,
-      constructExpression,
+      toRowExpression.flatten,
+      fromRowExpression,
       ClassTag[T](cls))
   }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoder.scala
deleted file mode 100644
index 6d307ab13a9fc..0000000000000
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoder.scala
+++ /dev/null
@@ -1,50 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.catalyst.encoders
-
-import scala.reflect.ClassTag
-import scala.reflect.runtime.universe.{typeTag, TypeTag}
-
-import org.apache.spark.sql.types.StructType
-import org.apache.spark.sql.catalyst.expressions.{Literal, CreateNamedStruct, BoundReference}
-import org.apache.spark.sql.catalyst.ScalaReflection
-
-object FlatEncoder {
-  import ScalaReflection.schemaFor
-  import ScalaReflection.dataTypeFor
-
-  def apply[T : TypeTag]: ExpressionEncoder[T] = {
-    // We convert the not-serializable TypeTag into StructType and ClassTag.
-    val tpe = typeTag[T].tpe
-    val mirror = typeTag[T].mirror
-    val cls = mirror.runtimeClass(tpe)
-    assert(!schemaFor(tpe).dataType.isInstanceOf[StructType])
-
-    val input = BoundReference(0, dataTypeFor(tpe), nullable = true)
-    val toRowExpression = CreateNamedStruct(
-      Literal("value") :: ProductEncoder.extractorFor(input, tpe) :: Nil)
-    val fromRowExpression = ProductEncoder.constructorFor(tpe)
-
-    new ExpressionEncoder[T](
-      toRowExpression.dataType,
-      flat = true,
-      toRowExpression.flatten,
-      fromRowExpression,
-      ClassTag[T](cls))
-  }
-}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
deleted file mode 100644
index 2914c6ee790ce..0000000000000
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoder.scala
+++ /dev/null
@@ -1,452 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.catalyst.encoders
-
-import org.apache.spark.util.Utils
-import org.apache.spark.unsafe.types.UTF8String
-import org.apache.spark.sql.types._
-import org.apache.spark.sql.catalyst.ScalaReflectionLock
-import org.apache.spark.sql.catalyst.ScalaReflection
-import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedExtractValue}
-import org.apache.spark.sql.catalyst.expressions
-import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.util.{DateTimeUtils, ArrayBasedMapData, GenericArrayData}
-
-import scala.reflect.ClassTag
-
-object ProductEncoder {
-  import ScalaReflection.universe._
-  import ScalaReflection.mirror
-  import ScalaReflection.localTypeOf
-  import ScalaReflection.dataTypeFor
-  import ScalaReflection.Schema
-  import ScalaReflection.schemaFor
-  import ScalaReflection.arrayClassFor
-
-  def apply[T <: Product : TypeTag]: ExpressionEncoder[T] = {
-    // We convert the not-serializable TypeTag into StructType and ClassTag.
-    val tpe = typeTag[T].tpe
-    val mirror = typeTag[T].mirror
-    val cls = mirror.runtimeClass(tpe)
-
-    val inputObject = BoundReference(0, ObjectType(cls), nullable = true)
-    val toRowExpression = extractorFor(inputObject, tpe).asInstanceOf[CreateNamedStruct]
-    val fromRowExpression = constructorFor(tpe)
-
-    new ExpressionEncoder[T](
-      toRowExpression.dataType,
-      flat = false,
-      toRowExpression.flatten,
-      fromRowExpression,
-      ClassTag[T](cls))
-  }
-
-  // The Predef.Map is scala.collection.immutable.Map.
-  // Since the map values can be mutable, we explicitly import scala.collection.Map at here.
-  import scala.collection.Map
-
-  def extractorFor(
-      inputObject: Expression,
-      tpe: `Type`): Expression = ScalaReflectionLock.synchronized {
-    if (!inputObject.dataType.isInstanceOf[ObjectType]) {
-      inputObject
-    } else {
-      tpe match {
-        case t if t <:< localTypeOf[Option[_]] =>
-          val TypeRef(_, _, Seq(optType)) = t
-          optType match {
-            // For primitive types we must manually unbox the value of the object.
-            case t if t <:< definitions.IntTpe =>
-              Invoke(
-                UnwrapOption(ObjectType(classOf[java.lang.Integer]), inputObject),
-                "intValue",
-                IntegerType)
-            case t if t <:< definitions.LongTpe =>
-              Invoke(
-                UnwrapOption(ObjectType(classOf[java.lang.Long]), inputObject),
-                "longValue",
-                LongType)
-            case t if t <:< definitions.DoubleTpe =>
-              Invoke(
-                UnwrapOption(ObjectType(classOf[java.lang.Double]), inputObject),
-                "doubleValue",
-                DoubleType)
-            case t if t <:< definitions.FloatTpe =>
-              Invoke(
-                UnwrapOption(ObjectType(classOf[java.lang.Float]), inputObject),
-                "floatValue",
-                FloatType)
-            case t if t <:< definitions.ShortTpe =>
-              Invoke(
-                UnwrapOption(ObjectType(classOf[java.lang.Short]), inputObject),
-                "shortValue",
-                ShortType)
-            case t if t <:< definitions.ByteTpe =>
-              Invoke(
-                UnwrapOption(ObjectType(classOf[java.lang.Byte]), inputObject),
-                "byteValue",
-                ByteType)
-            case t if t <:< definitions.BooleanTpe =>
-              Invoke(
-                UnwrapOption(ObjectType(classOf[java.lang.Boolean]), inputObject),
-                "booleanValue",
-                BooleanType)
-
-            // For non-primitives, we can just extract the object from the Option and then recurse.
-            case other =>
-              val className: String = optType.erasure.typeSymbol.asClass.fullName
-              val classObj = Utils.classForName(className)
-              val optionObjectType = ObjectType(classObj)
-
-              val unwrapped = UnwrapOption(optionObjectType, inputObject)
-              expressions.If(
-                IsNull(unwrapped),
-                expressions.Literal.create(null, schemaFor(optType).dataType),
-                extractorFor(unwrapped, optType))
-          }
-
-        case t if t <:< localTypeOf[Product] =>
-          val formalTypeArgs = t.typeSymbol.asClass.typeParams
-          val TypeRef(_, _, actualTypeArgs) = t
-          val constructorSymbol = t.member(nme.CONSTRUCTOR)
-          val params = if (constructorSymbol.isMethod) {
-            constructorSymbol.asMethod.paramss
-          } else {
-            // Find the primary constructor, and use its parameter ordering.
-            val primaryConstructorSymbol: Option[Symbol] =
-              constructorSymbol.asTerm.alternatives.find(s =>
-                s.isMethod && s.asMethod.isPrimaryConstructor)
-
-            if (primaryConstructorSymbol.isEmpty) {
-              sys.error("Internal SQL error: Product object did not have a primary constructor.")
-            } else {
-              primaryConstructorSymbol.get.asMethod.paramss
-            }
-          }
-
-          CreateNamedStruct(params.head.flatMap { p =>
-            val fieldName = p.name.toString
-            val fieldType = p.typeSignature.substituteTypes(formalTypeArgs, actualTypeArgs)
-            val fieldValue = Invoke(inputObject, fieldName, dataTypeFor(fieldType))
-            expressions.Literal(fieldName) :: extractorFor(fieldValue, fieldType) :: Nil
-          })
-
-        case t if t <:< localTypeOf[Array[_]] =>
-          val TypeRef(_, _, Seq(elementType)) = t
-          toCatalystArray(inputObject, elementType)
-
-        case t if t <:< localTypeOf[Seq[_]] =>
-          val TypeRef(_, _, Seq(elementType)) = t
-          toCatalystArray(inputObject, elementType)
-
-        case t if t <:< localTypeOf[Map[_, _]] =>
-          val TypeRef(_, _, Seq(keyType, valueType)) = t
-
-          val keys =
-            Invoke(
-              Invoke(inputObject, "keysIterator",
-                ObjectType(classOf[scala.collection.Iterator[_]])),
-              "toSeq",
-              ObjectType(classOf[scala.collection.Seq[_]]))
-          val convertedKeys = toCatalystArray(keys, keyType)
-
-          val values =
-            Invoke(
-              Invoke(inputObject, "valuesIterator",
-                ObjectType(classOf[scala.collection.Iterator[_]])),
-              "toSeq",
-              ObjectType(classOf[scala.collection.Seq[_]]))
-          val convertedValues = toCatalystArray(values, valueType)
-
-          val Schema(keyDataType, _) = schemaFor(keyType)
-          val Schema(valueDataType, valueNullable) = schemaFor(valueType)
-          NewInstance(
-            classOf[ArrayBasedMapData],
-            convertedKeys :: convertedValues :: Nil,
-            dataType = MapType(keyDataType, valueDataType, valueNullable))
-
-        case t if t <:< localTypeOf[String] =>
-          StaticInvoke(
-            classOf[UTF8String],
-            StringType,
-            "fromString",
-            inputObject :: Nil)
-
-        case t if t <:< localTypeOf[java.sql.Timestamp] =>
-          StaticInvoke(
-            DateTimeUtils,
-            TimestampType,
-            "fromJavaTimestamp",
-            inputObject :: Nil)
-
-        case t if t <:< localTypeOf[java.sql.Date] =>
-          StaticInvoke(
-            DateTimeUtils,
-            DateType,
-            "fromJavaDate",
-            inputObject :: Nil)
-
-        case t if t <:< localTypeOf[BigDecimal] =>
-          StaticInvoke(
-            Decimal,
-            DecimalType.SYSTEM_DEFAULT,
-            "apply",
-            inputObject :: Nil)
-
-        case t if t <:< localTypeOf[java.math.BigDecimal] =>
-          StaticInvoke(
-            Decimal,
-            DecimalType.SYSTEM_DEFAULT,
-            "apply",
-            inputObject :: Nil)
-
-        case t if t <:< localTypeOf[java.lang.Integer] =>
-          Invoke(inputObject, "intValue", IntegerType)
-        case t if t <:< localTypeOf[java.lang.Long] =>
-          Invoke(inputObject, "longValue", LongType)
-        case t if t <:< localTypeOf[java.lang.Double] =>
-          Invoke(inputObject, "doubleValue", DoubleType)
-        case t if t <:< localTypeOf[java.lang.Float] =>
-          Invoke(inputObject, "floatValue", FloatType)
-        case t if t <:< localTypeOf[java.lang.Short] =>
-          Invoke(inputObject, "shortValue", ShortType)
-        case t if t <:< localTypeOf[java.lang.Byte] =>
-          Invoke(inputObject, "byteValue", ByteType)
-        case t if t <:< localTypeOf[java.lang.Boolean] =>
-          Invoke(inputObject, "booleanValue", BooleanType)
-
-        case other =>
-          throw new UnsupportedOperationException(s"Encoder for type $other is not supported")
-      }
-    }
-  }
-
-  private def toCatalystArray(input: Expression, elementType: `Type`): Expression = {
-    val externalDataType = dataTypeFor(elementType)
-    val Schema(catalystType, nullable) = schemaFor(elementType)
-    if (RowEncoder.isNativeType(catalystType)) {
-      NewInstance(
-        classOf[GenericArrayData],
-        input :: Nil,
-        dataType = ArrayType(catalystType, nullable))
-    } else {
-      MapObjects(extractorFor(_, elementType), input, externalDataType)
-    }
-  }
-
-  def constructorFor(
-      tpe: `Type`,
-      path: Option[Expression] = None): Expression = ScalaReflectionLock.synchronized {
-
-    /** Returns the current path with a sub-field extracted. */
-    def addToPath(part: String): Expression = path
-      .map(p => UnresolvedExtractValue(p, expressions.Literal(part)))
-      .getOrElse(UnresolvedAttribute(part))
-
-    /** Returns the current path with a field at ordinal extracted. */
-    def addToPathOrdinal(ordinal: Int, dataType: DataType): Expression = path
-      .map(p => GetInternalRowField(p, ordinal, dataType))
-      .getOrElse(BoundReference(ordinal, dataType, false))
-
-    /** Returns the current path or `BoundReference`. */
-    def getPath: Expression = path.getOrElse(BoundReference(0, schemaFor(tpe).dataType, true))
-
-    tpe match {
-      case t if !dataTypeFor(t).isInstanceOf[ObjectType] => getPath
-
-      case t if t <:< localTypeOf[Option[_]] =>
-        val TypeRef(_, _, Seq(optType)) = t
-        WrapOption(null, constructorFor(optType, path))
-
-      case t if t <:< localTypeOf[java.lang.Integer] =>
-        val boxedType = classOf[java.lang.Integer]
-        val objectType = ObjectType(boxedType)
-        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
-
-      case t if t <:< localTypeOf[java.lang.Long] =>
-        val boxedType = classOf[java.lang.Long]
-        val objectType = ObjectType(boxedType)
-        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
-
-      case t if t <:< localTypeOf[java.lang.Double] =>
-        val boxedType = classOf[java.lang.Double]
-        val objectType = ObjectType(boxedType)
-        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
-
-      case t if t <:< localTypeOf[java.lang.Float] =>
-        val boxedType = classOf[java.lang.Float]
-        val objectType = ObjectType(boxedType)
-        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
-
-      case t if t <:< localTypeOf[java.lang.Short] =>
-        val boxedType = classOf[java.lang.Short]
-        val objectType = ObjectType(boxedType)
-        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
-
-      case t if t <:< localTypeOf[java.lang.Byte] =>
-        val boxedType = classOf[java.lang.Byte]
-        val objectType = ObjectType(boxedType)
-        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
-
-      case t if t <:< localTypeOf[java.lang.Boolean] =>
-        val boxedType = classOf[java.lang.Boolean]
-        val objectType = ObjectType(boxedType)
-        NewInstance(boxedType, getPath :: Nil, propagateNull = true, objectType)
-
-      case t if t <:< localTypeOf[java.sql.Date] =>
-        StaticInvoke(
-          DateTimeUtils,
-          ObjectType(classOf[java.sql.Date]),
-          "toJavaDate",
-          getPath :: Nil,
-          propagateNull = true)
-
-      case t if t <:< localTypeOf[java.sql.Timestamp] =>
-        StaticInvoke(
-          DateTimeUtils,
-          ObjectType(classOf[java.sql.Timestamp]),
-          "toJavaTimestamp",
-          getPath :: Nil,
-          propagateNull = true)
-
-      case t if t <:< localTypeOf[java.lang.String] =>
-        Invoke(getPath, "toString", ObjectType(classOf[String]))
-
-      case t if t <:< localTypeOf[java.math.BigDecimal] =>
-        Invoke(getPath, "toJavaBigDecimal", ObjectType(classOf[java.math.BigDecimal]))
-
-      case t if t <:< localTypeOf[BigDecimal] =>
-        Invoke(getPath, "toBigDecimal", ObjectType(classOf[BigDecimal]))
-
-      case t if t <:< localTypeOf[Array[_]] =>
-        val TypeRef(_, _, Seq(elementType)) = t
-        val primitiveMethod = elementType match {
-          case t if t <:< definitions.IntTpe => Some("toIntArray")
-          case t if t <:< definitions.LongTpe => Some("toLongArray")
-          case t if t <:< definitions.DoubleTpe => Some("toDoubleArray")
-          case t if t <:< definitions.FloatTpe => Some("toFloatArray")
-          case t if t <:< definitions.ShortTpe => Some("toShortArray")
-          case t if t <:< definitions.ByteTpe => Some("toByteArray")
-          case t if t <:< definitions.BooleanTpe => Some("toBooleanArray")
-          case _ => None
-        }
-
-        primitiveMethod.map { method =>
-          Invoke(getPath, method, arrayClassFor(elementType))
-        }.getOrElse {
-          Invoke(
-            MapObjects(
-              p => constructorFor(elementType, Some(p)),
-              getPath,
-              schemaFor(elementType).dataType),
-            "array",
-            arrayClassFor(elementType))
-        }
-
-      case t if t <:< localTypeOf[Seq[_]] =>
-        val TypeRef(_, _, Seq(elementType)) = t
-        val arrayData =
-          Invoke(
-            MapObjects(
-              p => constructorFor(elementType, Some(p)),
-              getPath,
-              schemaFor(elementType).dataType),
-            "array",
-            ObjectType(classOf[Array[Any]]))
-
-        StaticInvoke(
-          scala.collection.mutable.WrappedArray,
-          ObjectType(classOf[Seq[_]]),
-          "make",
-          arrayData :: Nil)
-
-      case t if t <:< localTypeOf[Map[_, _]] =>
-        val TypeRef(_, _, Seq(keyType, valueType)) = t
-
-        val keyData =
-          Invoke(
-            MapObjects(
-              p => constructorFor(keyType, Some(p)),
-              Invoke(getPath, "keyArray", ArrayType(schemaFor(keyType).dataType)),
-              schemaFor(keyType).dataType),
-            "array",
-            ObjectType(classOf[Array[Any]]))
-
-        val valueData =
-          Invoke(
-            MapObjects(
-              p => constructorFor(valueType, Some(p)),
-              Invoke(getPath, "valueArray", ArrayType(schemaFor(valueType).dataType)),
-              schemaFor(valueType).dataType),
-            "array",
-            ObjectType(classOf[Array[Any]]))
-
-        StaticInvoke(
-          ArrayBasedMapData,
-          ObjectType(classOf[Map[_, _]]),
-          "toScalaMap",
-          keyData :: valueData :: Nil)
-
-      case t if t <:< localTypeOf[Product] =>
-        val formalTypeArgs = t.typeSymbol.asClass.typeParams
-        val TypeRef(_, _, actualTypeArgs) = t
-        val constructorSymbol = t.member(nme.CONSTRUCTOR)
-        val params = if (constructorSymbol.isMethod) {
-          constructorSymbol.asMethod.paramss
-        } else {
-          // Find the primary constructor, and use its parameter ordering.
-          val primaryConstructorSymbol: Option[Symbol] =
-            constructorSymbol.asTerm.alternatives.find(s =>
-              s.isMethod && s.asMethod.isPrimaryConstructor)
-
-          if (primaryConstructorSymbol.isEmpty) {
-            sys.error("Internal SQL error: Product object did not have a primary constructor.")
-          } else {
-            primaryConstructorSymbol.get.asMethod.paramss
-          }
-        }
-
-        val cls = mirror.runtimeClass(tpe.erasure.typeSymbol.asClass)
-
-        val arguments = params.head.zipWithIndex.map { case (p, i) =>
-          val fieldName = p.name.toString
-          val fieldType = p.typeSignature.substituteTypes(formalTypeArgs, actualTypeArgs)
-          val dataType = schemaFor(fieldType).dataType
-
-          // For tuples, we based grab the inner fields by ordinal instead of name.
-          if (cls.getName startsWith "scala.Tuple") {
-            constructorFor(fieldType, Some(addToPathOrdinal(i, dataType)))
-          } else {
-            constructorFor(fieldType, Some(addToPath(fieldName)))
-          }
-        }
-
-        val newInstance = NewInstance(cls, arguments, propagateNull = false, ObjectType(cls))
-
-        if (path.nonEmpty) {
-          expressions.If(
-            IsNull(getPath),
-            expressions.Literal.create(null, ObjectType(cls)),
-            newInstance
-          )
-        } else {
-          newInstance
-        }
-    }
-  }
-}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
index 9bb1602494b68..4cda4824acdc3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
@@ -23,6 +23,7 @@ import scala.reflect.ClassTag
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.util.{GenericArrayData, ArrayBasedMapData, DateTimeUtils}
+import org.apache.spark.sql.catalyst.ScalaReflection
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
@@ -132,17 +133,8 @@ object RowEncoder {
       CreateStruct(convertedFields)
   }
 
-  /**
-   * Returns true if the value of this data type is same between internal and external.
-   */
-  def isNativeType(dt: DataType): Boolean = dt match {
-    case BooleanType | ByteType | ShortType | IntegerType | LongType |
-         FloatType | DoubleType | BinaryType => true
-    case _ => false
-  }
-
   private def externalDataTypeFor(dt: DataType): DataType = dt match {
-    case _ if isNativeType(dt) => dt
+    case _ if ScalaReflection.isNativeType(dt) => dt
     case TimestampType => ObjectType(classOf[java.sql.Timestamp])
     case DateType => ObjectType(classOf[java.sql.Date])
     case _: DecimalType => ObjectType(classOf[java.math.BigDecimal])
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index f865a9408ef4e..ef7399e0196ab 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -24,7 +24,6 @@ import org.apache.spark.SparkConf
 import org.apache.spark.serializer._
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.analysis.SimpleAnalyzer
-import org.apache.spark.sql.catalyst.encoders.ProductEncoder
 import org.apache.spark.sql.catalyst.plans.logical.{Project, LocalRelation}
 import org.apache.spark.sql.catalyst.util.GenericArrayData
 import org.apache.spark.sql.catalyst.InternalRow
@@ -300,10 +299,9 @@ case class UnwrapOption(
 /**
  * Converts the result of evaluating `child` into an option, checking both the isNull bit and
  * (in the case of reference types) equality with null.
- * @param optionType The datatype to be held inside of the Option.
  * @param child The expression to evaluate and wrap.
  */
-case class WrapOption(optionType: DataType, child: Expression)
+case class WrapOption(child: Expression)
   extends UnaryExpression with ExpectsInputTypes {
 
   override def dataType: DataType = ObjectType(classOf[Option[_]])
@@ -316,14 +314,13 @@ case class WrapOption(optionType: DataType, child: Expression)
     throw new UnsupportedOperationException("Only code-generated evaluation is supported")
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    val javaType = ctx.javaType(optionType)
     val inputObject = child.gen(ctx)
 
     s"""
       ${inputObject.code}
 
       boolean ${ev.isNull} = false;
-      scala.Option<$javaType> ${ev.value} =
+      scala.Option ${ev.value} =
         ${inputObject.isNull} ?
         scala.Option$$.MODULE$$.apply(null) : new scala.Some(${inputObject.value});
     """
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ScalaReflectionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ScalaReflectionSuite.scala
index 4ea410d492b01..c2aace1ef238e 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ScalaReflectionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ScalaReflectionSuite.scala
@@ -186,74 +186,6 @@ class ScalaReflectionSuite extends SparkFunSuite {
       nullable = true))
   }
 
-  test("get data type of a value") {
-    // BooleanType
-    assert(BooleanType === typeOfObject(true))
-    assert(BooleanType === typeOfObject(false))
-
-    // BinaryType
-    assert(BinaryType === typeOfObject("string".getBytes))
-
-    // StringType
-    assert(StringType === typeOfObject("string"))
-
-    // ByteType
-    assert(ByteType === typeOfObject(127.toByte))
-
-    // ShortType
-    assert(ShortType === typeOfObject(32767.toShort))
-
-    // IntegerType
-    assert(IntegerType === typeOfObject(2147483647))
-
-    // LongType
-    assert(LongType === typeOfObject(9223372036854775807L))
-
-    // FloatType
-    assert(FloatType === typeOfObject(3.4028235E38.toFloat))
-
-    // DoubleType
-    assert(DoubleType === typeOfObject(1.7976931348623157E308))
-
-    // DecimalType
-    assert(DecimalType.SYSTEM_DEFAULT ===
-      typeOfObject(new java.math.BigDecimal("1.7976931348623157E318")))
-
-    // DateType
-    assert(DateType === typeOfObject(Date.valueOf("2014-07-25")))
-
-    // TimestampType
-    assert(TimestampType === typeOfObject(Timestamp.valueOf("2014-07-25 10:26:00")))
-
-    // NullType
-    assert(NullType === typeOfObject(null))
-
-    def typeOfObject1: PartialFunction[Any, DataType] = typeOfObject orElse {
-      case value: java.math.BigInteger => DecimalType.SYSTEM_DEFAULT
-      case value: java.math.BigDecimal => DecimalType.SYSTEM_DEFAULT
-      case _ => StringType
-    }
-
-    assert(DecimalType.SYSTEM_DEFAULT === typeOfObject1(
-      new BigInteger("92233720368547758070")))
-    assert(DecimalType.SYSTEM_DEFAULT === typeOfObject1(
-      new java.math.BigDecimal("1.7976931348623157E318")))
-    assert(StringType === typeOfObject1(BigInt("92233720368547758070")))
-
-    def typeOfObject2: PartialFunction[Any, DataType] = typeOfObject orElse {
-      case value: java.math.BigInteger => DecimalType.SYSTEM_DEFAULT
-    }
-
-    intercept[MatchError](typeOfObject2(BigInt("92233720368547758070")))
-
-    def typeOfObject3: PartialFunction[Any, DataType] = typeOfObject orElse {
-      case c: Seq[_] => ArrayType(typeOfObject3(c.head))
-    }
-
-    assert(ArrayType(IntegerType) === typeOfObject3(Seq(1, 2, 3)))
-    assert(ArrayType(ArrayType(IntegerType)) === typeOfObject3(Seq(Seq(1, 2, 3))))
-  }
-
   test("convert PrimitiveData to catalyst") {
     val data = PrimitiveData(1, 1, 1, 1, 1, 1, true)
     val convertedData = InternalRow(1, 1.toLong, 1.toDouble, 1.toFloat, 1.toShort, 1.toByte, true)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
index cde0364f3dd9d..76459b34a484f 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
@@ -17,24 +17,234 @@
 
 package org.apache.spark.sql.catalyst.encoders
 
+import java.sql.{Timestamp, Date}
 import java.util.Arrays
 import java.util.concurrent.ConcurrentMap
+import scala.collection.mutable.ArrayBuffer
+import scala.reflect.runtime.universe.TypeTag
 
 import com.google.common.collect.MapMaker
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.Encoders
 import org.apache.spark.sql.catalyst.expressions.AttributeReference
 import org.apache.spark.sql.catalyst.util.ArrayData
+import org.apache.spark.sql.catalyst.{OptionalData, PrimitiveData}
 import org.apache.spark.sql.types.ArrayType
 
-abstract class ExpressionEncoderSuite extends SparkFunSuite {
-  val outers: ConcurrentMap[String, AnyRef] = new MapMaker().weakValues().makeMap()
+case class RepeatedStruct(s: Seq[PrimitiveData])
 
-  protected def encodeDecodeTest[T](
+case class NestedArray(a: Array[Array[Int]]) {
+  override def equals(other: Any): Boolean = other match {
+    case NestedArray(otherArray) =>
+      java.util.Arrays.deepEquals(
+        a.asInstanceOf[Array[AnyRef]],
+        otherArray.asInstanceOf[Array[AnyRef]])
+    case _ => false
+  }
+}
+
+case class BoxedData(
+    intField: java.lang.Integer,
+    longField: java.lang.Long,
+    doubleField: java.lang.Double,
+    floatField: java.lang.Float,
+    shortField: java.lang.Short,
+    byteField: java.lang.Byte,
+    booleanField: java.lang.Boolean)
+
+case class RepeatedData(
+    arrayField: Seq[Int],
+    arrayFieldContainsNull: Seq[java.lang.Integer],
+    mapField: scala.collection.Map[Int, Long],
+    mapFieldNull: scala.collection.Map[Int, java.lang.Long],
+    structField: PrimitiveData)
+
+case class SpecificCollection(l: List[Int])
+
+/** For testing Kryo serialization based encoder. */
+class KryoSerializable(val value: Int) {
+  override def equals(other: Any): Boolean = {
+    this.value == other.asInstanceOf[KryoSerializable].value
+  }
+}
+
+/** For testing Java serialization based encoder. */
+class JavaSerializable(val value: Int) extends Serializable {
+  override def equals(other: Any): Boolean = {
+    this.value == other.asInstanceOf[JavaSerializable].value
+  }
+}
+
+class ExpressionEncoderSuite extends SparkFunSuite {
+  implicit def encoder[T : TypeTag]: ExpressionEncoder[T] = ExpressionEncoder()
+
+  // test flat encoders
+  encodeDecodeTest(false, "primitive boolean")
+  encodeDecodeTest(-3.toByte, "primitive byte")
+  encodeDecodeTest(-3.toShort, "primitive short")
+  encodeDecodeTest(-3, "primitive int")
+  encodeDecodeTest(-3L, "primitive long")
+  encodeDecodeTest(-3.7f, "primitive float")
+  encodeDecodeTest(-3.7, "primitive double")
+
+  encodeDecodeTest(new java.lang.Boolean(false), "boxed boolean")
+  encodeDecodeTest(new java.lang.Byte(-3.toByte), "boxed byte")
+  encodeDecodeTest(new java.lang.Short(-3.toShort), "boxed short")
+  encodeDecodeTest(new java.lang.Integer(-3), "boxed int")
+  encodeDecodeTest(new java.lang.Long(-3L), "boxed long")
+  encodeDecodeTest(new java.lang.Float(-3.7f), "boxed float")
+  encodeDecodeTest(new java.lang.Double(-3.7), "boxed double")
+
+  encodeDecodeTest(BigDecimal("32131413.211321313"), "scala decimal")
+  // encodeDecodeTest(new java.math.BigDecimal("231341.23123"), "java decimal")
+
+  encodeDecodeTest("hello", "string")
+  encodeDecodeTest(Date.valueOf("2012-12-23"), "date")
+  encodeDecodeTest(Timestamp.valueOf("2016-01-29 10:00:00"), "timestamp")
+  encodeDecodeTest(Array[Byte](13, 21, -23), "binary")
+
+  encodeDecodeTest(Seq(31, -123, 4), "seq of int")
+  encodeDecodeTest(Seq("abc", "xyz"), "seq of string")
+  encodeDecodeTest(Seq("abc", null, "xyz"), "seq of string with null")
+  encodeDecodeTest(Seq.empty[Int], "empty seq of int")
+  encodeDecodeTest(Seq.empty[String], "empty seq of string")
+
+  encodeDecodeTest(Seq(Seq(31, -123), null, Seq(4, 67)), "seq of seq of int")
+  encodeDecodeTest(Seq(Seq("abc", "xyz"), Seq[String](null), null, Seq("1", null, "2")),
+    "seq of seq of string")
+
+  encodeDecodeTest(Array(31, -123, 4), "array of int")
+  encodeDecodeTest(Array("abc", "xyz"), "array of string")
+  encodeDecodeTest(Array("a", null, "x"), "array of string with null")
+  encodeDecodeTest(Array.empty[Int], "empty array of int")
+  encodeDecodeTest(Array.empty[String], "empty array of string")
+
+  encodeDecodeTest(Array(Array(31, -123), null, Array(4, 67)), "array of array of int")
+  encodeDecodeTest(Array(Array("abc", "xyz"), Array[String](null), null, Array("1", null, "2")),
+    "array of array of string")
+
+  encodeDecodeTest(Map(1 -> "a", 2 -> "b"), "map")
+  encodeDecodeTest(Map(1 -> "a", 2 -> null), "map with null")
+  encodeDecodeTest(Map(1 -> Map("a" -> 1), 2 -> Map("b" -> 2)), "map of map")
+
+  // Kryo encoders
+  encodeDecodeTest("hello", "kryo string")(encoderFor(Encoders.kryo[String]))
+  encodeDecodeTest(new KryoSerializable(15), "kryo object")(
+    encoderFor(Encoders.kryo[KryoSerializable]))
+
+  // Java encoders
+  encodeDecodeTest("hello", "java string")(encoderFor(Encoders.javaSerialization[String]))
+  encodeDecodeTest(new JavaSerializable(15), "java object")(
+    encoderFor(Encoders.javaSerialization[JavaSerializable]))
+
+  // test product encoders
+  private def productTest[T <: Product : ExpressionEncoder](input: T): Unit = {
+    encodeDecodeTest(input, input.getClass.getSimpleName)
+  }
+
+  case class InnerClass(i: Int)
+  productTest(InnerClass(1))
+
+  productTest(PrimitiveData(1, 1, 1, 1, 1, 1, true))
+
+  productTest(
+    OptionalData(Some(2), Some(2), Some(2), Some(2), Some(2), Some(2), Some(true),
+      Some(PrimitiveData(1, 1, 1, 1, 1, 1, true))))
+
+  productTest(OptionalData(None, None, None, None, None, None, None, None))
+
+  productTest(BoxedData(1, 1L, 1.0, 1.0f, 1.toShort, 1.toByte, true))
+
+  productTest(BoxedData(null, null, null, null, null, null, null))
+
+  productTest(RepeatedStruct(PrimitiveData(1, 1, 1, 1, 1, 1, true) :: Nil))
+
+  productTest((1, "test", PrimitiveData(1, 1, 1, 1, 1, 1, true)))
+
+  productTest(
+    RepeatedData(
+      Seq(1, 2),
+      Seq(new Integer(1), null, new Integer(2)),
+      Map(1 -> 2L),
+      Map(1 -> null),
+      PrimitiveData(1, 1, 1, 1, 1, 1, true)))
+
+  productTest(NestedArray(Array(Array(1, -2, 3), null, Array(4, 5, -6))))
+
+  productTest(("Seq[(String, String)]",
+    Seq(("a", "b"))))
+  productTest(("Seq[(Int, Int)]",
+    Seq((1, 2))))
+  productTest(("Seq[(Long, Long)]",
+    Seq((1L, 2L))))
+  productTest(("Seq[(Float, Float)]",
+    Seq((1.toFloat, 2.toFloat))))
+  productTest(("Seq[(Double, Double)]",
+    Seq((1.toDouble, 2.toDouble))))
+  productTest(("Seq[(Short, Short)]",
+    Seq((1.toShort, 2.toShort))))
+  productTest(("Seq[(Byte, Byte)]",
+    Seq((1.toByte, 2.toByte))))
+  productTest(("Seq[(Boolean, Boolean)]",
+    Seq((true, false))))
+
+  productTest(("ArrayBuffer[(String, String)]",
+    ArrayBuffer(("a", "b"))))
+  productTest(("ArrayBuffer[(Int, Int)]",
+    ArrayBuffer((1, 2))))
+  productTest(("ArrayBuffer[(Long, Long)]",
+    ArrayBuffer((1L, 2L))))
+  productTest(("ArrayBuffer[(Float, Float)]",
+    ArrayBuffer((1.toFloat, 2.toFloat))))
+  productTest(("ArrayBuffer[(Double, Double)]",
+    ArrayBuffer((1.toDouble, 2.toDouble))))
+  productTest(("ArrayBuffer[(Short, Short)]",
+    ArrayBuffer((1.toShort, 2.toShort))))
+  productTest(("ArrayBuffer[(Byte, Byte)]",
+    ArrayBuffer((1.toByte, 2.toByte))))
+  productTest(("ArrayBuffer[(Boolean, Boolean)]",
+    ArrayBuffer((true, false))))
+
+  productTest(("Seq[Seq[(Int, Int)]]",
+    Seq(Seq((1, 2)))))
+
+  // test for ExpressionEncoder.tuple
+  encodeDecodeTest(
+    1 -> 10L,
+    "tuple with 2 flat encoders")(
+    ExpressionEncoder.tuple(ExpressionEncoder[Int], ExpressionEncoder[Long]))
+
+  encodeDecodeTest(
+    (PrimitiveData(1, 1, 1, 1, 1, 1, true), (3, 30L)),
+    "tuple with 2 product encoders")(
+    ExpressionEncoder.tuple(ExpressionEncoder[PrimitiveData], ExpressionEncoder[(Int, Long)]))
+
+  encodeDecodeTest(
+    (PrimitiveData(1, 1, 1, 1, 1, 1, true), 3),
+    "tuple with flat encoder and product encoder")(
+    ExpressionEncoder.tuple(ExpressionEncoder[PrimitiveData], ExpressionEncoder[Int]))
+
+  encodeDecodeTest(
+    (3, PrimitiveData(1, 1, 1, 1, 1, 1, true)),
+    "tuple with product encoder and flat encoder")(
+    ExpressionEncoder.tuple(ExpressionEncoder[Int], ExpressionEncoder[PrimitiveData]))
+
+  encodeDecodeTest(
+    (1, (10, 100L)),
+    "nested tuple encoder") {
+    val intEnc = ExpressionEncoder[Int]
+    val longEnc = ExpressionEncoder[Long]
+    ExpressionEncoder.tuple(intEnc, ExpressionEncoder.tuple(intEnc, longEnc))
+  }
+
+  private val outers: ConcurrentMap[String, AnyRef] = new MapMaker().weakValues().makeMap()
+  outers.put(getClass.getName, this)
+  private def encodeDecodeTest[T : ExpressionEncoder](
       input: T,
-      encoder: ExpressionEncoder[T],
       testName: String): Unit = {
     test(s"encode/decode for $testName: $input") {
+      val encoder = implicitly[ExpressionEncoder[T]]
       val row = encoder.toRow(input)
       val schema = encoder.schema.toAttributes
       val boundEncoder = encoder.resolve(schema, outers).bind(schema)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
deleted file mode 100644
index 07523d49f4266..0000000000000
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/FlatEncoderSuite.scala
+++ /dev/null
@@ -1,99 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.catalyst.encoders
-
-import java.sql.{Date, Timestamp}
-import org.apache.spark.sql.Encoders
-
-class FlatEncoderSuite extends ExpressionEncoderSuite {
-  encodeDecodeTest(false, FlatEncoder[Boolean], "primitive boolean")
-  encodeDecodeTest(-3.toByte, FlatEncoder[Byte], "primitive byte")
-  encodeDecodeTest(-3.toShort, FlatEncoder[Short], "primitive short")
-  encodeDecodeTest(-3, FlatEncoder[Int], "primitive int")
-  encodeDecodeTest(-3L, FlatEncoder[Long], "primitive long")
-  encodeDecodeTest(-3.7f, FlatEncoder[Float], "primitive float")
-  encodeDecodeTest(-3.7, FlatEncoder[Double], "primitive double")
-
-  encodeDecodeTest(new java.lang.Boolean(false), FlatEncoder[java.lang.Boolean], "boxed boolean")
-  encodeDecodeTest(new java.lang.Byte(-3.toByte), FlatEncoder[java.lang.Byte], "boxed byte")
-  encodeDecodeTest(new java.lang.Short(-3.toShort), FlatEncoder[java.lang.Short], "boxed short")
-  encodeDecodeTest(new java.lang.Integer(-3), FlatEncoder[java.lang.Integer], "boxed int")
-  encodeDecodeTest(new java.lang.Long(-3L), FlatEncoder[java.lang.Long], "boxed long")
-  encodeDecodeTest(new java.lang.Float(-3.7f), FlatEncoder[java.lang.Float], "boxed float")
-  encodeDecodeTest(new java.lang.Double(-3.7), FlatEncoder[java.lang.Double], "boxed double")
-
-  encodeDecodeTest(BigDecimal("32131413.211321313"), FlatEncoder[BigDecimal], "scala decimal")
-  type JDecimal = java.math.BigDecimal
-  // encodeDecodeTest(new JDecimal("231341.23123"), FlatEncoder[JDecimal], "java decimal")
-
-  encodeDecodeTest("hello", FlatEncoder[String], "string")
-  encodeDecodeTest(Date.valueOf("2012-12-23"), FlatEncoder[Date], "date")
-  encodeDecodeTest(Timestamp.valueOf("2016-01-29 10:00:00"), FlatEncoder[Timestamp], "timestamp")
-  encodeDecodeTest(Array[Byte](13, 21, -23), FlatEncoder[Array[Byte]], "binary")
-
-  encodeDecodeTest(Seq(31, -123, 4), FlatEncoder[Seq[Int]], "seq of int")
-  encodeDecodeTest(Seq("abc", "xyz"), FlatEncoder[Seq[String]], "seq of string")
-  encodeDecodeTest(Seq("abc", null, "xyz"), FlatEncoder[Seq[String]], "seq of string with null")
-  encodeDecodeTest(Seq.empty[Int], FlatEncoder[Seq[Int]], "empty seq of int")
-  encodeDecodeTest(Seq.empty[String], FlatEncoder[Seq[String]], "empty seq of string")
-
-  encodeDecodeTest(Seq(Seq(31, -123), null, Seq(4, 67)),
-    FlatEncoder[Seq[Seq[Int]]], "seq of seq of int")
-  encodeDecodeTest(Seq(Seq("abc", "xyz"), Seq[String](null), null, Seq("1", null, "2")),
-    FlatEncoder[Seq[Seq[String]]], "seq of seq of string")
-
-  encodeDecodeTest(Array(31, -123, 4), FlatEncoder[Array[Int]], "array of int")
-  encodeDecodeTest(Array("abc", "xyz"), FlatEncoder[Array[String]], "array of string")
-  encodeDecodeTest(Array("a", null, "x"), FlatEncoder[Array[String]], "array of string with null")
-  encodeDecodeTest(Array.empty[Int], FlatEncoder[Array[Int]], "empty array of int")
-  encodeDecodeTest(Array.empty[String], FlatEncoder[Array[String]], "empty array of string")
-
-  encodeDecodeTest(Array(Array(31, -123), null, Array(4, 67)),
-    FlatEncoder[Array[Array[Int]]], "array of array of int")
-  encodeDecodeTest(Array(Array("abc", "xyz"), Array[String](null), null, Array("1", null, "2")),
-    FlatEncoder[Array[Array[String]]], "array of array of string")
-
-  encodeDecodeTest(Map(1 -> "a", 2 -> "b"), FlatEncoder[Map[Int, String]], "map")
-  encodeDecodeTest(Map(1 -> "a", 2 -> null), FlatEncoder[Map[Int, String]], "map with null")
-  encodeDecodeTest(Map(1 -> Map("a" -> 1), 2 -> Map("b" -> 2)),
-    FlatEncoder[Map[Int, Map[String, Int]]], "map of map")
-
-  // Kryo encoders
-  encodeDecodeTest("hello", encoderFor(Encoders.kryo[String]), "kryo string")
-  encodeDecodeTest(new KryoSerializable(15),
-    encoderFor(Encoders.kryo[KryoSerializable]), "kryo object")
-
-  // Java encoders
-  encodeDecodeTest("hello", encoderFor(Encoders.javaSerialization[String]), "java string")
-  encodeDecodeTest(new JavaSerializable(15),
-    encoderFor(Encoders.javaSerialization[JavaSerializable]), "java object")
-}
-
-/** For testing Kryo serialization based encoder. */
-class KryoSerializable(val value: Int) {
-  override def equals(other: Any): Boolean = {
-    this.value == other.asInstanceOf[KryoSerializable].value
-  }
-}
-
-/** For testing Java serialization based encoder. */
-class JavaSerializable(val value: Int) extends Serializable {
-  override def equals(other: Any): Boolean = {
-    this.value == other.asInstanceOf[JavaSerializable].value
-  }
-}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
deleted file mode 100644
index 1798514c5c38b..0000000000000
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ProductEncoderSuite.scala
+++ /dev/null
@@ -1,156 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.catalyst.encoders
-
-import scala.collection.mutable.ArrayBuffer
-import scala.reflect.runtime.universe.TypeTag
-
-import org.apache.spark.sql.catalyst.{OptionalData, PrimitiveData}
-
-case class RepeatedStruct(s: Seq[PrimitiveData])
-
-case class NestedArray(a: Array[Array[Int]]) {
-  override def equals(other: Any): Boolean = other match {
-    case NestedArray(otherArray) =>
-      java.util.Arrays.deepEquals(
-        a.asInstanceOf[Array[AnyRef]],
-        otherArray.asInstanceOf[Array[AnyRef]])
-    case _ => false
-  }
-}
-
-case class BoxedData(
-    intField: java.lang.Integer,
-    longField: java.lang.Long,
-    doubleField: java.lang.Double,
-    floatField: java.lang.Float,
-    shortField: java.lang.Short,
-    byteField: java.lang.Byte,
-    booleanField: java.lang.Boolean)
-
-case class RepeatedData(
-    arrayField: Seq[Int],
-    arrayFieldContainsNull: Seq[java.lang.Integer],
-    mapField: scala.collection.Map[Int, Long],
-    mapFieldNull: scala.collection.Map[Int, java.lang.Long],
-    structField: PrimitiveData)
-
-case class SpecificCollection(l: List[Int])
-
-class ProductEncoderSuite extends ExpressionEncoderSuite {
-  outers.put(getClass.getName, this)
-
-  case class InnerClass(i: Int)
-  productTest(InnerClass(1))
-
-  productTest(PrimitiveData(1, 1, 1, 1, 1, 1, true))
-
-  productTest(
-    OptionalData(Some(2), Some(2), Some(2), Some(2), Some(2), Some(2), Some(true),
-      Some(PrimitiveData(1, 1, 1, 1, 1, 1, true))))
-
-  productTest(OptionalData(None, None, None, None, None, None, None, None))
-
-  productTest(BoxedData(1, 1L, 1.0, 1.0f, 1.toShort, 1.toByte, true))
-
-  productTest(BoxedData(null, null, null, null, null, null, null))
-
-  productTest(RepeatedStruct(PrimitiveData(1, 1, 1, 1, 1, 1, true) :: Nil))
-
-  productTest((1, "test", PrimitiveData(1, 1, 1, 1, 1, 1, true)))
-
-  productTest(
-    RepeatedData(
-      Seq(1, 2),
-      Seq(new Integer(1), null, new Integer(2)),
-      Map(1 -> 2L),
-      Map(1 -> null),
-      PrimitiveData(1, 1, 1, 1, 1, 1, true)))
-
-  productTest(NestedArray(Array(Array(1, -2, 3), null, Array(4, 5, -6))))
-
-  productTest(("Seq[(String, String)]",
-    Seq(("a", "b"))))
-  productTest(("Seq[(Int, Int)]",
-    Seq((1, 2))))
-  productTest(("Seq[(Long, Long)]",
-    Seq((1L, 2L))))
-  productTest(("Seq[(Float, Float)]",
-    Seq((1.toFloat, 2.toFloat))))
-  productTest(("Seq[(Double, Double)]",
-    Seq((1.toDouble, 2.toDouble))))
-  productTest(("Seq[(Short, Short)]",
-    Seq((1.toShort, 2.toShort))))
-  productTest(("Seq[(Byte, Byte)]",
-    Seq((1.toByte, 2.toByte))))
-  productTest(("Seq[(Boolean, Boolean)]",
-    Seq((true, false))))
-
-  productTest(("ArrayBuffer[(String, String)]",
-    ArrayBuffer(("a", "b"))))
-  productTest(("ArrayBuffer[(Int, Int)]",
-    ArrayBuffer((1, 2))))
-  productTest(("ArrayBuffer[(Long, Long)]",
-    ArrayBuffer((1L, 2L))))
-  productTest(("ArrayBuffer[(Float, Float)]",
-    ArrayBuffer((1.toFloat, 2.toFloat))))
-  productTest(("ArrayBuffer[(Double, Double)]",
-    ArrayBuffer((1.toDouble, 2.toDouble))))
-  productTest(("ArrayBuffer[(Short, Short)]",
-    ArrayBuffer((1.toShort, 2.toShort))))
-  productTest(("ArrayBuffer[(Byte, Byte)]",
-    ArrayBuffer((1.toByte, 2.toByte))))
-  productTest(("ArrayBuffer[(Boolean, Boolean)]",
-    ArrayBuffer((true, false))))
-
-  productTest(("Seq[Seq[(Int, Int)]]",
-    Seq(Seq((1, 2)))))
-
-  encodeDecodeTest(
-    1 -> 10L,
-    ExpressionEncoder.tuple(FlatEncoder[Int], FlatEncoder[Long]),
-    "tuple with 2 flat encoders")
-
-  encodeDecodeTest(
-    (PrimitiveData(1, 1, 1, 1, 1, 1, true), (3, 30L)),
-    ExpressionEncoder.tuple(ProductEncoder[PrimitiveData], ProductEncoder[(Int, Long)]),
-    "tuple with 2 product encoders")
-
-  encodeDecodeTest(
-    (PrimitiveData(1, 1, 1, 1, 1, 1, true), 3),
-    ExpressionEncoder.tuple(ProductEncoder[PrimitiveData], FlatEncoder[Int]),
-    "tuple with flat encoder and product encoder")
-
-  encodeDecodeTest(
-    (3, PrimitiveData(1, 1, 1, 1, 1, 1, true)),
-    ExpressionEncoder.tuple(FlatEncoder[Int], ProductEncoder[PrimitiveData]),
-    "tuple with product encoder and flat encoder")
-
-  encodeDecodeTest(
-    (1, (10, 100L)),
-    {
-      val intEnc = FlatEncoder[Int]
-      val longEnc = FlatEncoder[Long]
-      ExpressionEncoder.tuple(intEnc, ExpressionEncoder.tuple(intEnc, longEnc))
-    },
-    "nested tuple encoder")
-
-  private def productTest[T <: Product : TypeTag](input: T): Unit = {
-    encodeDecodeTest(input, ProductEncoder[T], input.getClass.getSimpleName)
-  }
-}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 7e5acbe8517d1..6de3dd626576a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -21,7 +21,7 @@ import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.function._
-import org.apache.spark.sql.catalyst.encoders.{FlatEncoder, ExpressionEncoder, encoderFor, OuterScopes}
+import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor, OuterScopes}
 import org.apache.spark.sql.catalyst.expressions.{Alias, CreateStruct, Attribute}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.QueryExecution
@@ -242,7 +242,7 @@ class GroupedDataset[K, T] private[sql](
    * Returns a [[Dataset]] that contains a tuple with each key and the number of items present
    * for that key.
    */
-  def count(): Dataset[(K, Long)] = agg(functions.count("*").as(FlatEncoder[Long]))
+  def count(): Dataset[(K, Long)] = agg(functions.count("*").as(ExpressionEncoder[Long]))
 
   /**
    * Applies the given function to each cogrouped data.  For each unique group, the function will
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
index 8471eea1b7d9c..25ffdcde17717 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
@@ -17,10 +17,6 @@
 
 package org.apache.spark.sql
 
-import org.apache.spark.sql.catalyst.encoders._
-import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
-import org.apache.spark.sql.execution.datasources.LogicalRelation
-
 import scala.language.implicitConversions
 import scala.reflect.runtime.universe.TypeTag
 
@@ -28,6 +24,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.SpecificMutableRow
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.types.StructField
 import org.apache.spark.unsafe.types.UTF8String
 
@@ -37,16 +34,16 @@ import org.apache.spark.unsafe.types.UTF8String
 abstract class SQLImplicits {
   protected def _sqlContext: SQLContext
 
-  implicit def newProductEncoder[T <: Product : TypeTag]: Encoder[T] = ProductEncoder[T]
+  implicit def newProductEncoder[T <: Product : TypeTag]: Encoder[T] = ExpressionEncoder()
 
-  implicit def newIntEncoder: Encoder[Int] = FlatEncoder[Int]
-  implicit def newLongEncoder: Encoder[Long] = FlatEncoder[Long]
-  implicit def newDoubleEncoder: Encoder[Double] = FlatEncoder[Double]
-  implicit def newFloatEncoder: Encoder[Float] = FlatEncoder[Float]
-  implicit def newByteEncoder: Encoder[Byte] = FlatEncoder[Byte]
-  implicit def newShortEncoder: Encoder[Short] = FlatEncoder[Short]
-  implicit def newBooleanEncoder: Encoder[Boolean] = FlatEncoder[Boolean]
-  implicit def newStringEncoder: Encoder[String] = FlatEncoder[String]
+  implicit def newIntEncoder: Encoder[Int] = ExpressionEncoder()
+  implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder()
+  implicit def newDoubleEncoder: Encoder[Double] = ExpressionEncoder()
+  implicit def newFloatEncoder: Encoder[Float] = ExpressionEncoder()
+  implicit def newByteEncoder: Encoder[Byte] = ExpressionEncoder()
+  implicit def newShortEncoder: Encoder[Short] = ExpressionEncoder()
+  implicit def newBooleanEncoder: Encoder[Boolean] = ExpressionEncoder()
+  implicit def newStringEncoder: Encoder[String] = ExpressionEncoder()
 
   /**
    * Creates a [[Dataset]] from an RDD.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 95158de710acf..b27b1340cce46 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -26,7 +26,7 @@ import scala.util.Try
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.sql.catalyst.{SqlParser, ScalaReflection}
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunction, Star}
-import org.apache.spark.sql.catalyst.encoders.FlatEncoder
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.BroadcastHint
@@ -267,7 +267,7 @@ object functions extends LegacyFunctions {
    * @since 1.3.0
    */
   def count(columnName: String): TypedColumn[Any, Long] =
-    count(Column(columnName)).as(FlatEncoder[Long])
+    count(Column(columnName)).as(ExpressionEncoder[Long])
 
   /**
    * Aggregate function: returns the number of distinct items in a group.

From 4700074530d9a398843e13f0ef514be97a237cea Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@oc0558782468.ibm.com>
Date: Thu, 19 Nov 2015 13:08:01 -0800
Subject: [PATCH 363/518] [SPARK-11778][SQL] parse table name before it is
 passed to lookupRelation

Fix a bug in DataFrameReader.table (table with schema name such as "db_name.table" doesn't work)
Use SqlParser.parseTableIdentifier to parse the table name before lookupRelation.

Author: Huaxin Gao <huaxing@oc0558782468.ibm.com>

Closes #9773 from huaxingao/spark-11778.
---
 .../scala/org/apache/spark/sql/DataFrameReader.scala   |  3 ++-
 .../spark/sql/hive/HiveDataFrameAnalyticsSuite.scala   | 10 ++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
index 5872fbded3833..dcb3737b70fbf 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
@@ -313,7 +313,8 @@ class DataFrameReader private[sql](sqlContext: SQLContext) extends Logging {
    * @since 1.4.0
    */
   def table(tableName: String): DataFrame = {
-    DataFrame(sqlContext, sqlContext.catalog.lookupRelation(TableIdentifier(tableName)))
+    DataFrame(sqlContext,
+      sqlContext.catalog.lookupRelation(SqlParser.parseTableIdentifier(tableName)))
   }
 
   /**
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
index 9864acf765265..f19a74d4b3724 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
@@ -34,10 +34,14 @@ class HiveDataFrameAnalyticsSuite extends QueryTest with TestHiveSingleton with
   override def beforeAll() {
     testData = Seq((1, 2), (2, 2), (3, 4)).toDF("a", "b")
     hiveContext.registerDataFrameAsTable(testData, "mytable")
+    hiveContext.sql("create schema usrdb")
+    hiveContext.sql("create table usrdb.test(c1 int)")
   }
 
   override def afterAll(): Unit = {
     hiveContext.dropTempTable("mytable")
+    hiveContext.sql("drop table usrdb.test")
+    hiveContext.sql("drop schema usrdb")
   }
 
   test("rollup") {
@@ -74,4 +78,10 @@ class HiveDataFrameAnalyticsSuite extends QueryTest with TestHiveSingleton with
       sql("select a, b, sum(b) from mytable group by a, b with cube").collect()
     )
   }
+
+  // There was a bug in DataFrameFrameReader.table and it has problem for table with schema name,
+  // Before fix, it throw Exceptionorg.apache.spark.sql.catalyst.analysis.NoSuchTableException
+  test("table name with schema") {
+    hiveContext.read.table("usrdb.test")
+  }
 }

From 599a8c6e2bf7da70b20ef3046f5ce099dfd637f8 Mon Sep 17 00:00:00 2001
From: David Tolpin <david.tolpin@gmail.com>
Date: Thu, 19 Nov 2015 13:57:23 -0800
Subject: [PATCH 364/518] [SPARK-11812][PYSPARK] invFunc=None works properly
 with python's reduceByKeyAndWindow

invFunc is optional and can be None. Instead of invFunc (the parameter) invReduceFunc (a local function) was checked for trueness (that is, not None, in this context). A local function is never None,
thus the case of invFunc=None (a common one when inverse reduction is not defined) was treated incorrectly, resulting in loss of data.

In addition, the docstring used wrong parameter names, also fixed.

Author: David Tolpin <david.tolpin@gmail.com>

Closes #9775 from dtolpin/master.
---
 python/pyspark/streaming/dstream.py |  6 +++---
 python/pyspark/streaming/tests.py   | 11 +++++++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/python/pyspark/streaming/dstream.py b/python/pyspark/streaming/dstream.py
index 698336cfce18d..acec850f02c2d 100644
--- a/python/pyspark/streaming/dstream.py
+++ b/python/pyspark/streaming/dstream.py
@@ -524,8 +524,8 @@ def reduceByKeyAndWindow(self, func, invFunc, windowDuration, slideDuration=None
         `invFunc` can be None, then it will reduce all the RDDs in window, could be slower
         than having `invFunc`.
 
-        @param reduceFunc:     associative reduce function
-        @param invReduceFunc:  inverse function of `reduceFunc`
+        @param func:           associative reduce function
+        @param invFunc:        inverse function of `reduceFunc`
         @param windowDuration: width of the window; must be a multiple of this DStream's
                               batching interval
         @param slideDuration:  sliding interval of the window (i.e., the interval after which
@@ -556,7 +556,7 @@ def invReduceFunc(t, a, b):
                                     if kv[1] is not None else kv[0])
 
         jreduceFunc = TransformFunction(self._sc, reduceFunc, reduced._jrdd_deserializer)
-        if invReduceFunc:
+        if invFunc:
             jinvReduceFunc = TransformFunction(self._sc, invReduceFunc, reduced._jrdd_deserializer)
         else:
             jinvReduceFunc = None
diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index 0bcd1f15532b5..3403f6d20d789 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -582,6 +582,17 @@ def test_reduce_by_invalid_window(self):
         self.assertRaises(ValueError, lambda: d1.reduceByKeyAndWindow(None, None, 0.1, 0.1))
         self.assertRaises(ValueError, lambda: d1.reduceByKeyAndWindow(None, None, 1, 0.1))
 
+    def test_reduce_by_key_and_window_with_none_invFunc(self):
+        input = [range(1), range(2), range(3), range(4), range(5), range(6)]
+
+        def func(dstream):
+            return dstream.map(lambda x: (x, 1))\
+                .reduceByKeyAndWindow(operator.add, None, 5, 1)\
+                .filter(lambda kv: kv[1] > 0).count()
+
+        expected = [[2], [4], [6], [6], [6], [6]]
+        self._test_func(input, func, expected)
+
 
 class StreamingContextTests(PySparkStreamingTestCase):
 

From 014c0f7a9dfdb1686fa9aeacaadb2a17a855a943 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 19 Nov 2015 14:48:18 -0800
Subject: [PATCH 365/518] [SPARK-11858][SQL] Move sql.columnar into
 sql.execution.

In addition, tightened visibility of a lot of classes in the columnar package from private[sql] to private[columnar].

Author: Reynold Xin <rxin@databricks.com>

Closes #9842 from rxin/SPARK-11858.
---
 .../spark/sql/execution/CacheManager.scala    |  2 +-
 .../spark/sql/execution/SparkStrategies.scala |  2 +-
 .../columnar/ColumnAccessor.scala             | 42 +++++++--------
 .../columnar/ColumnBuilder.scala              | 51 ++++++++++---------
 .../columnar/ColumnStats.scala                | 34 ++++++-------
 .../{ => execution}/columnar/ColumnType.scala | 48 ++++++++---------
 .../columnar/GenerateColumnAccessor.scala     |  4 +-
 .../columnar/InMemoryColumnarTableScan.scala  |  5 +-
 .../columnar/NullableColumnAccessor.scala     |  4 +-
 .../columnar/NullableColumnBuilder.scala      |  4 +-
 .../CompressibleColumnAccessor.scala          |  6 +--
 .../CompressibleColumnBuilder.scala           |  6 +--
 .../compression/CompressionScheme.scala       | 16 +++---
 .../compression/compressionSchemes.scala      | 16 +++---
 .../apache/spark/sql/execution/package.scala  |  2 +
 .../apache/spark/sql/CachedTableSuite.scala   |  4 +-
 .../org/apache/spark/sql/QueryTest.scala      |  2 +-
 .../columnar/ColumnStatsSuite.scala           |  6 +--
 .../columnar/ColumnTypeSuite.scala            |  4 +-
 .../columnar/ColumnarTestUtils.scala          |  2 +-
 .../columnar/InMemoryColumnarQuerySuite.scala |  2 +-
 .../NullableColumnAccessorSuite.scala         |  4 +-
 .../columnar/NullableColumnBuilderSuite.scala |  4 +-
 .../columnar/PartitionBatchPruningSuite.scala |  2 +-
 .../compression/BooleanBitSetSuite.scala      |  6 +--
 .../compression/DictionaryEncodingSuite.scala |  6 +--
 .../compression/IntegralDeltaSuite.scala      |  6 +--
 .../compression/RunLengthEncodingSuite.scala  |  6 +--
 .../TestCompressibleColumnBuilder.scala       |  4 +-
 .../spark/sql/hive/CachedTableSuite.scala     |  2 +-
 30 files changed, 155 insertions(+), 147 deletions(-)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/ColumnAccessor.scala (75%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/ColumnBuilder.scala (74%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/ColumnStats.scala (88%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/ColumnType.scala (93%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/GenerateColumnAccessor.scala (98%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/InMemoryColumnarTableScan.scala (98%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/NullableColumnAccessor.scala (94%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/NullableColumnBuilder.scala (95%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/compression/CompressibleColumnAccessor.scala (84%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/compression/CompressibleColumnBuilder.scala (94%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/compression/CompressionScheme.scala (83%)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/columnar/compression/compressionSchemes.scala (96%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/ColumnStatsSuite.scala (96%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/ColumnTypeSuite.scala (97%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/ColumnarTestUtils.scala (98%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/InMemoryColumnarQuerySuite.scala (99%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/NullableColumnAccessorSuite.scala (96%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/NullableColumnBuilderSuite.scala (96%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/PartitionBatchPruningSuite.scala (99%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/compression/BooleanBitSetSuite.scala (94%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/compression/DictionaryEncodingSuite.scala (96%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/compression/IntegralDeltaSuite.scala (96%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/compression/RunLengthEncodingSuite.scala (95%)
 rename sql/core/src/test/scala/org/apache/spark/sql/{ => execution}/columnar/compression/TestCompressibleColumnBuilder.scala (93%)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala
index f85aeb1b02694..293fcfe96e677 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala
@@ -22,7 +22,7 @@ import java.util.concurrent.locks.ReentrantReadWriteLock
 import org.apache.spark.Logging
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
-import org.apache.spark.sql.columnar.InMemoryRelation
+import org.apache.spark.sql.execution.columnar.InMemoryRelation
 import org.apache.spark.storage.StorageLevel
 import org.apache.spark.storage.StorageLevel.MEMORY_AND_DISK
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index 3d4ce633c07c9..f67c951bc0663 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.planning._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical.{BroadcastHint, LogicalPlan}
 import org.apache.spark.sql.catalyst.plans.physical._
-import org.apache.spark.sql.columnar.{InMemoryColumnarTableScan, InMemoryRelation}
+import org.apache.spark.sql.execution.columnar.{InMemoryColumnarTableScan, InMemoryRelation}
 import org.apache.spark.sql.execution.datasources.{CreateTableUsing, CreateTempTableUsing, DescribeCommand => LogicalDescribeCommand, _}
 import org.apache.spark.sql.execution.{DescribeCommand => RunnableDescribeCommand}
 import org.apache.spark.sql.{Strategy, execution}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnAccessor.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnAccessor.scala
similarity index 75%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnAccessor.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnAccessor.scala
index 42ec4d3433f16..fee36f6023895 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnAccessor.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnAccessor.scala
@@ -15,12 +15,12 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import java.nio.{ByteBuffer, ByteOrder}
 
 import org.apache.spark.sql.catalyst.expressions.{MutableRow, UnsafeArrayData, UnsafeMapData, UnsafeRow}
-import org.apache.spark.sql.columnar.compression.CompressibleColumnAccessor
+import org.apache.spark.sql.execution.columnar.compression.CompressibleColumnAccessor
 import org.apache.spark.sql.types._
 
 /**
@@ -29,7 +29,7 @@ import org.apache.spark.sql.types._
  * a [[MutableRow]]. In this way, boxing cost can be avoided by leveraging the setter methods
  * for primitive values provided by [[MutableRow]].
  */
-private[sql] trait ColumnAccessor {
+private[columnar] trait ColumnAccessor {
   initialize()
 
   protected def initialize()
@@ -41,7 +41,7 @@ private[sql] trait ColumnAccessor {
   protected def underlyingBuffer: ByteBuffer
 }
 
-private[sql] abstract class BasicColumnAccessor[JvmType](
+private[columnar] abstract class BasicColumnAccessor[JvmType](
     protected val buffer: ByteBuffer,
     protected val columnType: ColumnType[JvmType])
   extends ColumnAccessor {
@@ -61,65 +61,65 @@ private[sql] abstract class BasicColumnAccessor[JvmType](
   protected def underlyingBuffer = buffer
 }
 
-private[sql] class NullColumnAccessor(buffer: ByteBuffer)
+private[columnar] class NullColumnAccessor(buffer: ByteBuffer)
   extends BasicColumnAccessor[Any](buffer, NULL)
   with NullableColumnAccessor
 
-private[sql] abstract class NativeColumnAccessor[T <: AtomicType](
+private[columnar] abstract class NativeColumnAccessor[T <: AtomicType](
     override protected val buffer: ByteBuffer,
     override protected val columnType: NativeColumnType[T])
   extends BasicColumnAccessor(buffer, columnType)
   with NullableColumnAccessor
   with CompressibleColumnAccessor[T]
 
-private[sql] class BooleanColumnAccessor(buffer: ByteBuffer)
+private[columnar] class BooleanColumnAccessor(buffer: ByteBuffer)
   extends NativeColumnAccessor(buffer, BOOLEAN)
 
-private[sql] class ByteColumnAccessor(buffer: ByteBuffer)
+private[columnar] class ByteColumnAccessor(buffer: ByteBuffer)
   extends NativeColumnAccessor(buffer, BYTE)
 
-private[sql] class ShortColumnAccessor(buffer: ByteBuffer)
+private[columnar] class ShortColumnAccessor(buffer: ByteBuffer)
   extends NativeColumnAccessor(buffer, SHORT)
 
-private[sql] class IntColumnAccessor(buffer: ByteBuffer)
+private[columnar] class IntColumnAccessor(buffer: ByteBuffer)
   extends NativeColumnAccessor(buffer, INT)
 
-private[sql] class LongColumnAccessor(buffer: ByteBuffer)
+private[columnar] class LongColumnAccessor(buffer: ByteBuffer)
   extends NativeColumnAccessor(buffer, LONG)
 
-private[sql] class FloatColumnAccessor(buffer: ByteBuffer)
+private[columnar] class FloatColumnAccessor(buffer: ByteBuffer)
   extends NativeColumnAccessor(buffer, FLOAT)
 
-private[sql] class DoubleColumnAccessor(buffer: ByteBuffer)
+private[columnar] class DoubleColumnAccessor(buffer: ByteBuffer)
   extends NativeColumnAccessor(buffer, DOUBLE)
 
-private[sql] class StringColumnAccessor(buffer: ByteBuffer)
+private[columnar] class StringColumnAccessor(buffer: ByteBuffer)
   extends NativeColumnAccessor(buffer, STRING)
 
-private[sql] class BinaryColumnAccessor(buffer: ByteBuffer)
+private[columnar] class BinaryColumnAccessor(buffer: ByteBuffer)
   extends BasicColumnAccessor[Array[Byte]](buffer, BINARY)
   with NullableColumnAccessor
 
-private[sql] class CompactDecimalColumnAccessor(buffer: ByteBuffer, dataType: DecimalType)
+private[columnar] class CompactDecimalColumnAccessor(buffer: ByteBuffer, dataType: DecimalType)
   extends NativeColumnAccessor(buffer, COMPACT_DECIMAL(dataType))
 
-private[sql] class DecimalColumnAccessor(buffer: ByteBuffer, dataType: DecimalType)
+private[columnar] class DecimalColumnAccessor(buffer: ByteBuffer, dataType: DecimalType)
   extends BasicColumnAccessor[Decimal](buffer, LARGE_DECIMAL(dataType))
   with NullableColumnAccessor
 
-private[sql] class StructColumnAccessor(buffer: ByteBuffer, dataType: StructType)
+private[columnar] class StructColumnAccessor(buffer: ByteBuffer, dataType: StructType)
   extends BasicColumnAccessor[UnsafeRow](buffer, STRUCT(dataType))
   with NullableColumnAccessor
 
-private[sql] class ArrayColumnAccessor(buffer: ByteBuffer, dataType: ArrayType)
+private[columnar] class ArrayColumnAccessor(buffer: ByteBuffer, dataType: ArrayType)
   extends BasicColumnAccessor[UnsafeArrayData](buffer, ARRAY(dataType))
   with NullableColumnAccessor
 
-private[sql] class MapColumnAccessor(buffer: ByteBuffer, dataType: MapType)
+private[columnar] class MapColumnAccessor(buffer: ByteBuffer, dataType: MapType)
   extends BasicColumnAccessor[UnsafeMapData](buffer, MAP(dataType))
   with NullableColumnAccessor
 
-private[sql] object ColumnAccessor {
+private[columnar] object ColumnAccessor {
   def apply(dataType: DataType, buffer: ByteBuffer): ColumnAccessor = {
     val buf = buffer.order(ByteOrder.nativeOrder)
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnBuilder.scala
similarity index 74%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnBuilder.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnBuilder.scala
index 599f30f2d73b4..7e26f19bb7449 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnBuilder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnBuilder.scala
@@ -15,16 +15,16 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import java.nio.{ByteBuffer, ByteOrder}
 
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.columnar.ColumnBuilder._
-import org.apache.spark.sql.columnar.compression.{AllCompressionSchemes, CompressibleColumnBuilder}
+import org.apache.spark.sql.execution.columnar.ColumnBuilder._
+import org.apache.spark.sql.execution.columnar.compression.{AllCompressionSchemes, CompressibleColumnBuilder}
 import org.apache.spark.sql.types._
 
-private[sql] trait ColumnBuilder {
+private[columnar] trait ColumnBuilder {
   /**
    * Initializes with an approximate lower bound on the expected number of elements in this column.
    */
@@ -46,7 +46,7 @@ private[sql] trait ColumnBuilder {
   def build(): ByteBuffer
 }
 
-private[sql] class BasicColumnBuilder[JvmType](
+private[columnar] class BasicColumnBuilder[JvmType](
     val columnStats: ColumnStats,
     val columnType: ColumnType[JvmType])
   extends ColumnBuilder {
@@ -84,17 +84,17 @@ private[sql] class BasicColumnBuilder[JvmType](
   }
 }
 
-private[sql] class NullColumnBuilder
+private[columnar] class NullColumnBuilder
   extends BasicColumnBuilder[Any](new ObjectColumnStats(NullType), NULL)
   with NullableColumnBuilder
 
-private[sql] abstract class ComplexColumnBuilder[JvmType](
+private[columnar] abstract class ComplexColumnBuilder[JvmType](
     columnStats: ColumnStats,
     columnType: ColumnType[JvmType])
   extends BasicColumnBuilder[JvmType](columnStats, columnType)
   with NullableColumnBuilder
 
-private[sql] abstract class NativeColumnBuilder[T <: AtomicType](
+private[columnar] abstract class NativeColumnBuilder[T <: AtomicType](
     override val columnStats: ColumnStats,
     override val columnType: NativeColumnType[T])
   extends BasicColumnBuilder[T#InternalType](columnStats, columnType)
@@ -102,40 +102,45 @@ private[sql] abstract class NativeColumnBuilder[T <: AtomicType](
   with AllCompressionSchemes
   with CompressibleColumnBuilder[T]
 
-private[sql] class BooleanColumnBuilder extends NativeColumnBuilder(new BooleanColumnStats, BOOLEAN)
+private[columnar]
+class BooleanColumnBuilder extends NativeColumnBuilder(new BooleanColumnStats, BOOLEAN)
 
-private[sql] class ByteColumnBuilder extends NativeColumnBuilder(new ByteColumnStats, BYTE)
+private[columnar]
+class ByteColumnBuilder extends NativeColumnBuilder(new ByteColumnStats, BYTE)
 
-private[sql] class ShortColumnBuilder extends NativeColumnBuilder(new ShortColumnStats, SHORT)
+private[columnar] class ShortColumnBuilder extends NativeColumnBuilder(new ShortColumnStats, SHORT)
 
-private[sql] class IntColumnBuilder extends NativeColumnBuilder(new IntColumnStats, INT)
+private[columnar] class IntColumnBuilder extends NativeColumnBuilder(new IntColumnStats, INT)
 
-private[sql] class LongColumnBuilder extends NativeColumnBuilder(new LongColumnStats, LONG)
+private[columnar] class LongColumnBuilder extends NativeColumnBuilder(new LongColumnStats, LONG)
 
-private[sql] class FloatColumnBuilder extends NativeColumnBuilder(new FloatColumnStats, FLOAT)
+private[columnar] class FloatColumnBuilder extends NativeColumnBuilder(new FloatColumnStats, FLOAT)
 
-private[sql] class DoubleColumnBuilder extends NativeColumnBuilder(new DoubleColumnStats, DOUBLE)
+private[columnar]
+class DoubleColumnBuilder extends NativeColumnBuilder(new DoubleColumnStats, DOUBLE)
 
-private[sql] class StringColumnBuilder extends NativeColumnBuilder(new StringColumnStats, STRING)
+private[columnar]
+class StringColumnBuilder extends NativeColumnBuilder(new StringColumnStats, STRING)
 
-private[sql] class BinaryColumnBuilder extends ComplexColumnBuilder(new BinaryColumnStats, BINARY)
+private[columnar]
+class BinaryColumnBuilder extends ComplexColumnBuilder(new BinaryColumnStats, BINARY)
 
-private[sql] class CompactDecimalColumnBuilder(dataType: DecimalType)
+private[columnar] class CompactDecimalColumnBuilder(dataType: DecimalType)
   extends NativeColumnBuilder(new DecimalColumnStats(dataType), COMPACT_DECIMAL(dataType))
 
-private[sql] class DecimalColumnBuilder(dataType: DecimalType)
+private[columnar] class DecimalColumnBuilder(dataType: DecimalType)
   extends ComplexColumnBuilder(new DecimalColumnStats(dataType), LARGE_DECIMAL(dataType))
 
-private[sql] class StructColumnBuilder(dataType: StructType)
+private[columnar] class StructColumnBuilder(dataType: StructType)
   extends ComplexColumnBuilder(new ObjectColumnStats(dataType), STRUCT(dataType))
 
-private[sql] class ArrayColumnBuilder(dataType: ArrayType)
+private[columnar] class ArrayColumnBuilder(dataType: ArrayType)
   extends ComplexColumnBuilder(new ObjectColumnStats(dataType), ARRAY(dataType))
 
-private[sql] class MapColumnBuilder(dataType: MapType)
+private[columnar] class MapColumnBuilder(dataType: MapType)
   extends ComplexColumnBuilder(new ObjectColumnStats(dataType), MAP(dataType))
 
-private[sql] object ColumnBuilder {
+private[columnar] object ColumnBuilder {
   val DEFAULT_INITIAL_BUFFER_SIZE = 128 * 1024
   val MAX_BATCH_SIZE_IN_BYTE = 4 * 1024 * 1024L
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnStats.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnStats.scala
similarity index 88%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnStats.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnStats.scala
index 91a05650585cf..c52ee9ffd6d2a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnStats.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnStats.scala
@@ -15,14 +15,14 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, Attribute, AttributeMap, AttributeReference}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.types.UTF8String
 
-private[sql] class ColumnStatisticsSchema(a: Attribute) extends Serializable {
+private[columnar] class ColumnStatisticsSchema(a: Attribute) extends Serializable {
   val upperBound = AttributeReference(a.name + ".upperBound", a.dataType, nullable = true)()
   val lowerBound = AttributeReference(a.name + ".lowerBound", a.dataType, nullable = true)()
   val nullCount = AttributeReference(a.name + ".nullCount", IntegerType, nullable = false)()
@@ -32,7 +32,7 @@ private[sql] class ColumnStatisticsSchema(a: Attribute) extends Serializable {
   val schema = Seq(lowerBound, upperBound, nullCount, count, sizeInBytes)
 }
 
-private[sql] class PartitionStatistics(tableSchema: Seq[Attribute]) extends Serializable {
+private[columnar] class PartitionStatistics(tableSchema: Seq[Attribute]) extends Serializable {
   val (forAttribute, schema) = {
     val allStats = tableSchema.map(a => a -> new ColumnStatisticsSchema(a))
     (AttributeMap(allStats), allStats.map(_._2.schema).foldLeft(Seq.empty[Attribute])(_ ++ _))
@@ -45,10 +45,10 @@ private[sql] class PartitionStatistics(tableSchema: Seq[Attribute]) extends Seri
  * NOTE: we intentionally avoid using `Ordering[T]` to compare values here because `Ordering[T]`
  * brings significant performance penalty.
  */
-private[sql] sealed trait ColumnStats extends Serializable {
+private[columnar] sealed trait ColumnStats extends Serializable {
   protected var count = 0
   protected var nullCount = 0
-  private[sql] var sizeInBytes = 0L
+  private[columnar] var sizeInBytes = 0L
 
   /**
    * Gathers statistics information from `row(ordinal)`.
@@ -72,14 +72,14 @@ private[sql] sealed trait ColumnStats extends Serializable {
 /**
  * A no-op ColumnStats only used for testing purposes.
  */
-private[sql] class NoopColumnStats extends ColumnStats {
+private[columnar] class NoopColumnStats extends ColumnStats {
   override def gatherStats(row: InternalRow, ordinal: Int): Unit = super.gatherStats(row, ordinal)
 
   override def collectedStatistics: GenericInternalRow =
     new GenericInternalRow(Array[Any](null, null, nullCount, count, 0L))
 }
 
-private[sql] class BooleanColumnStats extends ColumnStats {
+private[columnar] class BooleanColumnStats extends ColumnStats {
   protected var upper = false
   protected var lower = true
 
@@ -97,7 +97,7 @@ private[sql] class BooleanColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class ByteColumnStats extends ColumnStats {
+private[columnar] class ByteColumnStats extends ColumnStats {
   protected var upper = Byte.MinValue
   protected var lower = Byte.MaxValue
 
@@ -115,7 +115,7 @@ private[sql] class ByteColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class ShortColumnStats extends ColumnStats {
+private[columnar] class ShortColumnStats extends ColumnStats {
   protected var upper = Short.MinValue
   protected var lower = Short.MaxValue
 
@@ -133,7 +133,7 @@ private[sql] class ShortColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class IntColumnStats extends ColumnStats {
+private[columnar] class IntColumnStats extends ColumnStats {
   protected var upper = Int.MinValue
   protected var lower = Int.MaxValue
 
@@ -151,7 +151,7 @@ private[sql] class IntColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class LongColumnStats extends ColumnStats {
+private[columnar] class LongColumnStats extends ColumnStats {
   protected var upper = Long.MinValue
   protected var lower = Long.MaxValue
 
@@ -169,7 +169,7 @@ private[sql] class LongColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class FloatColumnStats extends ColumnStats {
+private[columnar] class FloatColumnStats extends ColumnStats {
   protected var upper = Float.MinValue
   protected var lower = Float.MaxValue
 
@@ -187,7 +187,7 @@ private[sql] class FloatColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class DoubleColumnStats extends ColumnStats {
+private[columnar] class DoubleColumnStats extends ColumnStats {
   protected var upper = Double.MinValue
   protected var lower = Double.MaxValue
 
@@ -205,7 +205,7 @@ private[sql] class DoubleColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class StringColumnStats extends ColumnStats {
+private[columnar] class StringColumnStats extends ColumnStats {
   protected var upper: UTF8String = null
   protected var lower: UTF8String = null
 
@@ -223,7 +223,7 @@ private[sql] class StringColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class BinaryColumnStats extends ColumnStats {
+private[columnar] class BinaryColumnStats extends ColumnStats {
   override def gatherStats(row: InternalRow, ordinal: Int): Unit = {
     super.gatherStats(row, ordinal)
     if (!row.isNullAt(ordinal)) {
@@ -235,7 +235,7 @@ private[sql] class BinaryColumnStats extends ColumnStats {
     new GenericInternalRow(Array[Any](null, null, nullCount, count, sizeInBytes))
 }
 
-private[sql] class DecimalColumnStats(precision: Int, scale: Int) extends ColumnStats {
+private[columnar] class DecimalColumnStats(precision: Int, scale: Int) extends ColumnStats {
   def this(dt: DecimalType) = this(dt.precision, dt.scale)
 
   protected var upper: Decimal = null
@@ -256,7 +256,7 @@ private[sql] class DecimalColumnStats(precision: Int, scale: Int) extends Column
     new GenericInternalRow(Array[Any](lower, upper, nullCount, count, sizeInBytes))
 }
 
-private[sql] class ObjectColumnStats(dataType: DataType) extends ColumnStats {
+private[columnar] class ObjectColumnStats(dataType: DataType) extends ColumnStats {
   val columnType = ColumnType(dataType)
 
   override def gatherStats(row: InternalRow, ordinal: Int): Unit = {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnType.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnType.scala
similarity index 93%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnType.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnType.scala
index 68e509eb5047d..c9f2329db4b6d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/ColumnType.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ColumnType.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import java.math.{BigDecimal, BigInteger}
 import java.nio.ByteBuffer
@@ -41,7 +41,7 @@ import org.apache.spark.unsafe.types.UTF8String
  *
  * WARNNING: This only works with HeapByteBuffer
  */
-object ByteBufferHelper {
+private[columnar] object ByteBufferHelper {
   def getInt(buffer: ByteBuffer): Int = {
     val pos = buffer.position()
     buffer.position(pos + 4)
@@ -73,7 +73,7 @@ object ByteBufferHelper {
  *
  * @tparam JvmType Underlying Java type to represent the elements.
  */
-private[sql] sealed abstract class ColumnType[JvmType] {
+private[columnar] sealed abstract class ColumnType[JvmType] {
 
   // The catalyst data type of this column.
   def dataType: DataType
@@ -142,7 +142,7 @@ private[sql] sealed abstract class ColumnType[JvmType] {
   override def toString: String = getClass.getSimpleName.stripSuffix("$")
 }
 
-private[sql] object NULL extends ColumnType[Any] {
+private[columnar] object NULL extends ColumnType[Any] {
 
   override def dataType: DataType = NullType
   override def defaultSize: Int = 0
@@ -152,7 +152,7 @@ private[sql] object NULL extends ColumnType[Any] {
   override def getField(row: InternalRow, ordinal: Int): Any = null
 }
 
-private[sql] abstract class NativeColumnType[T <: AtomicType](
+private[columnar] abstract class NativeColumnType[T <: AtomicType](
     val dataType: T,
     val defaultSize: Int)
   extends ColumnType[T#InternalType] {
@@ -163,7 +163,7 @@ private[sql] abstract class NativeColumnType[T <: AtomicType](
   def scalaTag: TypeTag[dataType.InternalType] = dataType.tag
 }
 
-private[sql] object INT extends NativeColumnType(IntegerType, 4) {
+private[columnar] object INT extends NativeColumnType(IntegerType, 4) {
   override def append(v: Int, buffer: ByteBuffer): Unit = {
     buffer.putInt(v)
   }
@@ -192,7 +192,7 @@ private[sql] object INT extends NativeColumnType(IntegerType, 4) {
   }
 }
 
-private[sql] object LONG extends NativeColumnType(LongType, 8) {
+private[columnar] object LONG extends NativeColumnType(LongType, 8) {
   override def append(v: Long, buffer: ByteBuffer): Unit = {
     buffer.putLong(v)
   }
@@ -220,7 +220,7 @@ private[sql] object LONG extends NativeColumnType(LongType, 8) {
   }
 }
 
-private[sql] object FLOAT extends NativeColumnType(FloatType, 4) {
+private[columnar] object FLOAT extends NativeColumnType(FloatType, 4) {
   override def append(v: Float, buffer: ByteBuffer): Unit = {
     buffer.putFloat(v)
   }
@@ -248,7 +248,7 @@ private[sql] object FLOAT extends NativeColumnType(FloatType, 4) {
   }
 }
 
-private[sql] object DOUBLE extends NativeColumnType(DoubleType, 8) {
+private[columnar] object DOUBLE extends NativeColumnType(DoubleType, 8) {
   override def append(v: Double, buffer: ByteBuffer): Unit = {
     buffer.putDouble(v)
   }
@@ -276,7 +276,7 @@ private[sql] object DOUBLE extends NativeColumnType(DoubleType, 8) {
   }
 }
 
-private[sql] object BOOLEAN extends NativeColumnType(BooleanType, 1) {
+private[columnar] object BOOLEAN extends NativeColumnType(BooleanType, 1) {
   override def append(v: Boolean, buffer: ByteBuffer): Unit = {
     buffer.put(if (v) 1: Byte else 0: Byte)
   }
@@ -302,7 +302,7 @@ private[sql] object BOOLEAN extends NativeColumnType(BooleanType, 1) {
   }
 }
 
-private[sql] object BYTE extends NativeColumnType(ByteType, 1) {
+private[columnar] object BYTE extends NativeColumnType(ByteType, 1) {
   override def append(v: Byte, buffer: ByteBuffer): Unit = {
     buffer.put(v)
   }
@@ -330,7 +330,7 @@ private[sql] object BYTE extends NativeColumnType(ByteType, 1) {
   }
 }
 
-private[sql] object SHORT extends NativeColumnType(ShortType, 2) {
+private[columnar] object SHORT extends NativeColumnType(ShortType, 2) {
   override def append(v: Short, buffer: ByteBuffer): Unit = {
     buffer.putShort(v)
   }
@@ -362,7 +362,7 @@ private[sql] object SHORT extends NativeColumnType(ShortType, 2) {
  * A fast path to copy var-length bytes between ByteBuffer and UnsafeRow without creating wrapper
  * objects.
  */
-private[sql] trait DirectCopyColumnType[JvmType] extends ColumnType[JvmType] {
+private[columnar] trait DirectCopyColumnType[JvmType] extends ColumnType[JvmType] {
 
   // copy the bytes from ByteBuffer to UnsafeRow
   override def extract(buffer: ByteBuffer, row: MutableRow, ordinal: Int): Unit = {
@@ -387,7 +387,7 @@ private[sql] trait DirectCopyColumnType[JvmType] extends ColumnType[JvmType] {
   }
 }
 
-private[sql] object STRING
+private[columnar] object STRING
   extends NativeColumnType(StringType, 8) with DirectCopyColumnType[UTF8String] {
 
   override def actualSize(row: InternalRow, ordinal: Int): Int = {
@@ -425,7 +425,7 @@ private[sql] object STRING
   override def clone(v: UTF8String): UTF8String = v.clone()
 }
 
-private[sql] case class COMPACT_DECIMAL(precision: Int, scale: Int)
+private[columnar] case class COMPACT_DECIMAL(precision: Int, scale: Int)
   extends NativeColumnType(DecimalType(precision, scale), 8) {
 
   override def extract(buffer: ByteBuffer): Decimal = {
@@ -467,13 +467,13 @@ private[sql] case class COMPACT_DECIMAL(precision: Int, scale: Int)
   }
 }
 
-private[sql] object COMPACT_DECIMAL {
+private[columnar] object COMPACT_DECIMAL {
   def apply(dt: DecimalType): COMPACT_DECIMAL = {
     COMPACT_DECIMAL(dt.precision, dt.scale)
   }
 }
 
-private[sql] sealed abstract class ByteArrayColumnType[JvmType](val defaultSize: Int)
+private[columnar] sealed abstract class ByteArrayColumnType[JvmType](val defaultSize: Int)
   extends ColumnType[JvmType] with DirectCopyColumnType[JvmType] {
 
   def serialize(value: JvmType): Array[Byte]
@@ -492,7 +492,7 @@ private[sql] sealed abstract class ByteArrayColumnType[JvmType](val defaultSize:
   }
 }
 
-private[sql] object BINARY extends ByteArrayColumnType[Array[Byte]](16) {
+private[columnar] object BINARY extends ByteArrayColumnType[Array[Byte]](16) {
 
   def dataType: DataType = BinaryType
 
@@ -512,7 +512,7 @@ private[sql] object BINARY extends ByteArrayColumnType[Array[Byte]](16) {
   def deserialize(bytes: Array[Byte]): Array[Byte] = bytes
 }
 
-private[sql] case class LARGE_DECIMAL(precision: Int, scale: Int)
+private[columnar] case class LARGE_DECIMAL(precision: Int, scale: Int)
   extends ByteArrayColumnType[Decimal](12) {
 
   override val dataType: DataType = DecimalType(precision, scale)
@@ -539,13 +539,13 @@ private[sql] case class LARGE_DECIMAL(precision: Int, scale: Int)
   }
 }
 
-private[sql] object LARGE_DECIMAL {
+private[columnar] object LARGE_DECIMAL {
   def apply(dt: DecimalType): LARGE_DECIMAL = {
     LARGE_DECIMAL(dt.precision, dt.scale)
   }
 }
 
-private[sql] case class STRUCT(dataType: StructType)
+private[columnar] case class STRUCT(dataType: StructType)
   extends ColumnType[UnsafeRow] with DirectCopyColumnType[UnsafeRow] {
 
   private val numOfFields: Int = dataType.fields.size
@@ -586,7 +586,7 @@ private[sql] case class STRUCT(dataType: StructType)
   override def clone(v: UnsafeRow): UnsafeRow = v.copy()
 }
 
-private[sql] case class ARRAY(dataType: ArrayType)
+private[columnar] case class ARRAY(dataType: ArrayType)
   extends ColumnType[UnsafeArrayData] with DirectCopyColumnType[UnsafeArrayData] {
 
   override def defaultSize: Int = 16
@@ -625,7 +625,7 @@ private[sql] case class ARRAY(dataType: ArrayType)
   override def clone(v: UnsafeArrayData): UnsafeArrayData = v.copy()
 }
 
-private[sql] case class MAP(dataType: MapType)
+private[columnar] case class MAP(dataType: MapType)
   extends ColumnType[UnsafeMapData] with DirectCopyColumnType[UnsafeMapData] {
 
   override def defaultSize: Int = 32
@@ -663,7 +663,7 @@ private[sql] case class MAP(dataType: MapType)
   override def clone(v: UnsafeMapData): UnsafeMapData = v.copy()
 }
 
-private[sql] object ColumnType {
+private[columnar] object ColumnType {
   def apply(dataType: DataType): ColumnType[_] = {
     dataType match {
       case NullType => NULL
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/GenerateColumnAccessor.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/GenerateColumnAccessor.scala
similarity index 98%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/GenerateColumnAccessor.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/GenerateColumnAccessor.scala
index ff9393b465b7a..eaafc96e4d2e7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/GenerateColumnAccessor.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/GenerateColumnAccessor.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import org.apache.spark.Logging
 import org.apache.spark.sql.catalyst.InternalRow
@@ -121,7 +121,7 @@ object GenerateColumnAccessor extends CodeGenerator[Seq[DataType], ColumnarItera
       import org.apache.spark.sql.types.DataType;
       import org.apache.spark.sql.catalyst.expressions.codegen.BufferHolder;
       import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter;
-      import org.apache.spark.sql.columnar.MutableUnsafeRow;
+      import org.apache.spark.sql.execution.columnar.MutableUnsafeRow;
 
       public SpecificColumnarIterator generate($exprType[] expr) {
         return new SpecificColumnarIterator();
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarTableScan.scala
similarity index 98%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarTableScan.scala
index ae77298e6da2f..ce701fb3a7f28 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/InMemoryColumnarTableScan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarTableScan.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import scala.collection.mutable.ArrayBuffer
 
@@ -50,7 +50,8 @@ private[sql] object InMemoryRelation {
  * @param buffers The buffers for serialized columns
  * @param stats The stat of columns
  */
-private[sql] case class CachedBatch(numRows: Int, buffers: Array[Array[Byte]], stats: InternalRow)
+private[columnar]
+case class CachedBatch(numRows: Int, buffers: Array[Array[Byte]], stats: InternalRow)
 
 private[sql] case class InMemoryRelation(
     output: Seq[Attribute],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/NullableColumnAccessor.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/NullableColumnAccessor.scala
similarity index 94%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/NullableColumnAccessor.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/NullableColumnAccessor.scala
index 7eaecfe047c3f..8d99546924de1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/NullableColumnAccessor.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/NullableColumnAccessor.scala
@@ -15,13 +15,13 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import java.nio.{ByteOrder, ByteBuffer}
 
 import org.apache.spark.sql.catalyst.expressions.MutableRow
 
-private[sql] trait NullableColumnAccessor extends ColumnAccessor {
+private[columnar] trait NullableColumnAccessor extends ColumnAccessor {
   private var nullsBuffer: ByteBuffer = _
   private var nullCount: Int = _
   private var seenNulls: Int = 0
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/NullableColumnBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/NullableColumnBuilder.scala
similarity index 95%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/NullableColumnBuilder.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/NullableColumnBuilder.scala
index 76cfddf1cd01a..3a1931bfb5c84 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/NullableColumnBuilder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/NullableColumnBuilder.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import java.nio.{ByteBuffer, ByteOrder}
 
@@ -34,7 +34,7 @@ import org.apache.spark.sql.catalyst.InternalRow
  *   +---+-----+---------+
  * }}}
  */
-private[sql] trait NullableColumnBuilder extends ColumnBuilder {
+private[columnar] trait NullableColumnBuilder extends ColumnBuilder {
   protected var nulls: ByteBuffer = _
   protected var nullCount: Int = _
   private var pos: Int = _
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressibleColumnAccessor.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressibleColumnAccessor.scala
similarity index 84%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressibleColumnAccessor.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressibleColumnAccessor.scala
index cb205defbb1ad..6579b5068e65a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressibleColumnAccessor.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressibleColumnAccessor.scala
@@ -15,13 +15,13 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
 import org.apache.spark.sql.catalyst.expressions.MutableRow
-import org.apache.spark.sql.columnar.{ColumnAccessor, NativeColumnAccessor}
+import org.apache.spark.sql.execution.columnar.{ColumnAccessor, NativeColumnAccessor}
 import org.apache.spark.sql.types.AtomicType
 
-private[sql] trait CompressibleColumnAccessor[T <: AtomicType] extends ColumnAccessor {
+private[columnar] trait CompressibleColumnAccessor[T <: AtomicType] extends ColumnAccessor {
   this: NativeColumnAccessor[T] =>
 
   private var decoder: Decoder[T] = _
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressibleColumnBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressibleColumnBuilder.scala
similarity index 94%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressibleColumnBuilder.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressibleColumnBuilder.scala
index 161021ff96154..b0e216feb5595 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressibleColumnBuilder.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressibleColumnBuilder.scala
@@ -15,13 +15,13 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
 import java.nio.{ByteBuffer, ByteOrder}
 
 import org.apache.spark.Logging
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.columnar.{ColumnBuilder, NativeColumnBuilder}
+import org.apache.spark.sql.execution.columnar.{ColumnBuilder, NativeColumnBuilder}
 import org.apache.spark.sql.types.AtomicType
 
 /**
@@ -40,7 +40,7 @@ import org.apache.spark.sql.types.AtomicType
  *     header         body
  * }}}
  */
-private[sql] trait CompressibleColumnBuilder[T <: AtomicType]
+private[columnar] trait CompressibleColumnBuilder[T <: AtomicType]
   extends ColumnBuilder with Logging {
 
   this: NativeColumnBuilder[T] with WithCompressionSchemes =>
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressionScheme.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressionScheme.scala
similarity index 83%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressionScheme.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressionScheme.scala
index 9322b772fd898..920381f9c63d0 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/CompressionScheme.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressionScheme.scala
@@ -15,15 +15,15 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
 import java.nio.{ByteBuffer, ByteOrder}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.MutableRow
-import org.apache.spark.sql.columnar.{ColumnType, NativeColumnType}
+import org.apache.spark.sql.execution.columnar.{ColumnType, NativeColumnType}
 import org.apache.spark.sql.types.AtomicType
 
-private[sql] trait Encoder[T <: AtomicType] {
+private[columnar] trait Encoder[T <: AtomicType] {
   def gatherCompressibilityStats(row: InternalRow, ordinal: Int): Unit = {}
 
   def compressedSize: Int
@@ -37,13 +37,13 @@ private[sql] trait Encoder[T <: AtomicType] {
   def compress(from: ByteBuffer, to: ByteBuffer): ByteBuffer
 }
 
-private[sql] trait Decoder[T <: AtomicType] {
+private[columnar] trait Decoder[T <: AtomicType] {
   def next(row: MutableRow, ordinal: Int): Unit
 
   def hasNext: Boolean
 }
 
-private[sql] trait CompressionScheme {
+private[columnar] trait CompressionScheme {
   def typeId: Int
 
   def supports(columnType: ColumnType[_]): Boolean
@@ -53,15 +53,15 @@ private[sql] trait CompressionScheme {
   def decoder[T <: AtomicType](buffer: ByteBuffer, columnType: NativeColumnType[T]): Decoder[T]
 }
 
-private[sql] trait WithCompressionSchemes {
+private[columnar] trait WithCompressionSchemes {
   def schemes: Seq[CompressionScheme]
 }
 
-private[sql] trait AllCompressionSchemes extends WithCompressionSchemes {
+private[columnar] trait AllCompressionSchemes extends WithCompressionSchemes {
   override val schemes: Seq[CompressionScheme] = CompressionScheme.all
 }
 
-private[sql] object CompressionScheme {
+private[columnar] object CompressionScheme {
   val all: Seq[CompressionScheme] =
     Seq(PassThrough, RunLengthEncoding, DictionaryEncoding, BooleanBitSet, IntDelta, LongDelta)
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/compressionSchemes.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/compressionSchemes.scala
similarity index 96%
rename from sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/compressionSchemes.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/compressionSchemes.scala
index 41c9a284e3e4a..941f03b745a07 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/columnar/compression/compressionSchemes.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/compressionSchemes.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
 import java.nio.ByteBuffer
 
@@ -23,11 +23,11 @@ import scala.collection.mutable
 
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{MutableRow, SpecificMutableRow}
-import org.apache.spark.sql.columnar._
+import org.apache.spark.sql.execution.columnar._
 import org.apache.spark.sql.types._
 
 
-private[sql] case object PassThrough extends CompressionScheme {
+private[columnar] case object PassThrough extends CompressionScheme {
   override val typeId = 0
 
   override def supports(columnType: ColumnType[_]): Boolean = true
@@ -64,7 +64,7 @@ private[sql] case object PassThrough extends CompressionScheme {
   }
 }
 
-private[sql] case object RunLengthEncoding extends CompressionScheme {
+private[columnar] case object RunLengthEncoding extends CompressionScheme {
   override val typeId = 1
 
   override def encoder[T <: AtomicType](columnType: NativeColumnType[T]): Encoder[T] = {
@@ -172,7 +172,7 @@ private[sql] case object RunLengthEncoding extends CompressionScheme {
   }
 }
 
-private[sql] case object DictionaryEncoding extends CompressionScheme {
+private[columnar] case object DictionaryEncoding extends CompressionScheme {
   override val typeId = 2
 
   // 32K unique values allowed
@@ -281,7 +281,7 @@ private[sql] case object DictionaryEncoding extends CompressionScheme {
   }
 }
 
-private[sql] case object BooleanBitSet extends CompressionScheme {
+private[columnar] case object BooleanBitSet extends CompressionScheme {
   override val typeId = 3
 
   val BITS_PER_LONG = 64
@@ -371,7 +371,7 @@ private[sql] case object BooleanBitSet extends CompressionScheme {
   }
 }
 
-private[sql] case object IntDelta extends CompressionScheme {
+private[columnar] case object IntDelta extends CompressionScheme {
   override def typeId: Int = 4
 
   override def decoder[T <: AtomicType](buffer: ByteBuffer, columnType: NativeColumnType[T])
@@ -451,7 +451,7 @@ private[sql] case object IntDelta extends CompressionScheme {
   }
 }
 
-private[sql] case object LongDelta extends CompressionScheme {
+private[columnar] case object LongDelta extends CompressionScheme {
   override def typeId: Int = 5
 
   override def decoder[T <: AtomicType](buffer: ByteBuffer, columnType: NativeColumnType[T])
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/package.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/package.scala
index 28fa231e722d0..c912734bba9e3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/package.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/package.scala
@@ -19,5 +19,7 @@ package org.apache.spark.sql
 
 /**
  * The physical execution component of Spark SQL. Note that this is a private package.
+ * All classes in catalyst are considered an internal API to Spark SQL and are subject
+ * to change between minor releases.
  */
 package object execution
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
index bce94dafad755..d86df4cfb9b4d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala
@@ -27,7 +27,7 @@ import scala.language.postfixOps
 import org.scalatest.concurrent.Eventually._
 
 import org.apache.spark.Accumulators
-import org.apache.spark.sql.columnar._
+import org.apache.spark.sql.execution.columnar._
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.{SQLTestUtils, SharedSQLContext}
 import org.apache.spark.storage.{StorageLevel, RDDBlockId}
@@ -280,7 +280,7 @@ class CachedTableSuite extends QueryTest with SQLTestUtils with SharedSQLContext
     sql("CACHE TABLE testData")
     sqlContext.table("testData").queryExecution.withCachedData.collect {
       case cached: InMemoryRelation =>
-        val actualSizeInBytes = (1 to 100).map(i => INT.defaultSize + i.toString.length + 4).sum
+        val actualSizeInBytes = (1 to 100).map(i => 4 + i.toString.length + 4).sum
         assert(cached.statistics.sizeInBytes === actualSizeInBytes)
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
index b5417b195f396..6ea1fe4ccfd89 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
@@ -23,7 +23,7 @@ import scala.collection.JavaConverters._
 
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.util._
-import org.apache.spark.sql.columnar.InMemoryRelation
+import org.apache.spark.sql.execution.columnar.InMemoryRelation
 
 abstract class QueryTest extends PlanTest {
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnStatsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnStatsSuite.scala
similarity index 96%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnStatsSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnStatsSuite.scala
index 89a664001bdd2..b2d04f7c5a6e3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnStatsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnStatsSuite.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
@@ -50,7 +50,7 @@ class ColumnStatsSuite extends SparkFunSuite {
     }
 
     test(s"$columnStatsName: non-empty") {
-      import org.apache.spark.sql.columnar.ColumnarTestUtils._
+      import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
 
       val columnStats = columnStatsClass.newInstance()
       val rows = Seq.fill(10)(makeRandomRow(columnType)) ++ Seq.fill(10)(makeNullRow(1))
@@ -86,7 +86,7 @@ class ColumnStatsSuite extends SparkFunSuite {
     }
 
     test(s"$columnStatsName: non-empty") {
-      import org.apache.spark.sql.columnar.ColumnarTestUtils._
+      import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
 
       val columnStats = new DecimalColumnStats(15, 10)
       val rows = Seq.fill(10)(makeRandomRow(columnType)) ++ Seq.fill(10)(makeNullRow(1))
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnTypeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnTypeSuite.scala
similarity index 97%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnTypeSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnTypeSuite.scala
index 63bc39bfa0307..34dd96929e6c1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnTypeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnTypeSuite.scala
@@ -15,14 +15,14 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import java.nio.{ByteOrder, ByteBuffer}
 
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.CatalystTypeConverters
 import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, GenericMutableRow}
-import org.apache.spark.sql.columnar.ColumnarTestUtils._
+import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
 import org.apache.spark.sql.types._
 import org.apache.spark.{Logging, SparkFunSuite}
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnarTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnarTestUtils.scala
similarity index 98%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnarTestUtils.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnarTestUtils.scala
index a5882f7870e37..9cae65ef6f5dc 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/ColumnarTestUtils.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ColumnarTestUtils.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import scala.collection.immutable.HashSet
 import scala.util.Random
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala
similarity index 99%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/InMemoryColumnarQuerySuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala
index 6265e40a0a07b..25afed25c897b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/InMemoryColumnarQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import java.sql.{Date, Timestamp}
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/NullableColumnAccessorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/NullableColumnAccessorSuite.scala
similarity index 96%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/NullableColumnAccessorSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/NullableColumnAccessorSuite.scala
index aa1605fee8c73..35dc9a276cef7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/NullableColumnAccessorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/NullableColumnAccessorSuite.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import java.nio.ByteBuffer
 
@@ -38,7 +38,7 @@ object TestNullableColumnAccessor {
 }
 
 class NullableColumnAccessorSuite extends SparkFunSuite {
-  import org.apache.spark.sql.columnar.ColumnarTestUtils._
+  import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
 
   Seq(
     NULL, BOOLEAN, BYTE, SHORT, INT, LONG, FLOAT, DOUBLE,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/NullableColumnBuilderSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/NullableColumnBuilderSuite.scala
similarity index 96%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/NullableColumnBuilderSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/NullableColumnBuilderSuite.scala
index 91404577832a0..93be3e16a5ed9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/NullableColumnBuilderSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/NullableColumnBuilderSuite.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.CatalystTypeConverters
@@ -36,7 +36,7 @@ object TestNullableColumnBuilder {
 }
 
 class NullableColumnBuilderSuite extends SparkFunSuite {
-  import org.apache.spark.sql.columnar.ColumnarTestUtils._
+  import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
 
   Seq(
     BOOLEAN, BYTE, SHORT, INT, LONG, FLOAT, DOUBLE,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/PartitionBatchPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
similarity index 99%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/PartitionBatchPruningSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
index 6b7401464f46f..d762f7bfe914c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/PartitionBatchPruningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/PartitionBatchPruningSuite.scala
@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar
+package org.apache.spark.sql.execution.columnar
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/BooleanBitSetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/BooleanBitSetSuite.scala
similarity index 94%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/BooleanBitSetSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/BooleanBitSetSuite.scala
index 9a2948c59ba42..ccbddef0fad3a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/BooleanBitSetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/BooleanBitSetSuite.scala
@@ -15,13 +15,13 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.GenericMutableRow
-import org.apache.spark.sql.columnar.ColumnarTestUtils._
-import org.apache.spark.sql.columnar.{BOOLEAN, NoopColumnStats}
+import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
+import org.apache.spark.sql.execution.columnar.{BOOLEAN, NoopColumnStats}
 
 class BooleanBitSetSuite extends SparkFunSuite {
   import BooleanBitSet._
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/DictionaryEncodingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/DictionaryEncodingSuite.scala
similarity index 96%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/DictionaryEncodingSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/DictionaryEncodingSuite.scala
index acfab6586c0d1..830ca0294e1b8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/DictionaryEncodingSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/DictionaryEncodingSuite.scala
@@ -15,14 +15,14 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
 import java.nio.ByteBuffer
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions.GenericMutableRow
-import org.apache.spark.sql.columnar._
-import org.apache.spark.sql.columnar.ColumnarTestUtils._
+import org.apache.spark.sql.execution.columnar._
+import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
 import org.apache.spark.sql.types.AtomicType
 
 class DictionaryEncodingSuite extends SparkFunSuite {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/IntegralDeltaSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/IntegralDeltaSuite.scala
similarity index 96%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/IntegralDeltaSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/IntegralDeltaSuite.scala
index 2111e9fbe62cb..988a577a7b4d0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/IntegralDeltaSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/IntegralDeltaSuite.scala
@@ -15,12 +15,12 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions.GenericMutableRow
-import org.apache.spark.sql.columnar._
-import org.apache.spark.sql.columnar.ColumnarTestUtils._
+import org.apache.spark.sql.execution.columnar._
+import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
 import org.apache.spark.sql.types.IntegralType
 
 class IntegralDeltaSuite extends SparkFunSuite {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/RunLengthEncodingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/RunLengthEncodingSuite.scala
similarity index 95%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/RunLengthEncodingSuite.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/RunLengthEncodingSuite.scala
index 67ec08f594a43..ce3affba55c71 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/RunLengthEncodingSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/RunLengthEncodingSuite.scala
@@ -15,12 +15,12 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions.GenericMutableRow
-import org.apache.spark.sql.columnar._
-import org.apache.spark.sql.columnar.ColumnarTestUtils._
+import org.apache.spark.sql.execution.columnar._
+import org.apache.spark.sql.execution.columnar.ColumnarTestUtils._
 import org.apache.spark.sql.types.AtomicType
 
 class RunLengthEncodingSuite extends SparkFunSuite {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/TestCompressibleColumnBuilder.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/TestCompressibleColumnBuilder.scala
similarity index 93%
rename from sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/TestCompressibleColumnBuilder.scala
rename to sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/TestCompressibleColumnBuilder.scala
index 5268dfe0aa03e..5e078f251375a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/columnar/compression/TestCompressibleColumnBuilder.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/compression/TestCompressibleColumnBuilder.scala
@@ -15,9 +15,9 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql.columnar.compression
+package org.apache.spark.sql.execution.columnar.compression
 
-import org.apache.spark.sql.columnar._
+import org.apache.spark.sql.execution.columnar._
 import org.apache.spark.sql.types.AtomicType
 
 class TestCompressibleColumnBuilder[T <: AtomicType](
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/CachedTableSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/CachedTableSuite.scala
index 5c2fc7d82ffbd..99478e82d419f 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/CachedTableSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/CachedTableSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.hive
 
 import java.io.File
 
-import org.apache.spark.sql.columnar.InMemoryColumnarTableScan
+import org.apache.spark.sql.execution.columnar.InMemoryColumnarTableScan
 import org.apache.spark.sql.execution.datasources.parquet.ParquetRelation
 import org.apache.spark.sql.hive.test.TestHiveSingleton
 import org.apache.spark.sql.{AnalysisException, QueryTest, SaveMode}

From 90d384dcbc1d1a3466cf8bae570a26f23012c102 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Thu, 19 Nov 2015 14:49:25 -0800
Subject: [PATCH 366/518] [SPARK-11831][CORE][TESTS] Use port 0 to avoid port
 conflicts in tests

Use port 0 to fix port-contention-related flakiness

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9841 from zsxwing/SPARK-11831.
---
 .../org/apache/spark/rpc/RpcEnvSuite.scala    | 24 +++++++++----------
 .../spark/rpc/akka/AkkaRpcEnvSuite.scala      |  4 ++--
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
index 834e4743df866..2f55006420ce1 100644
--- a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
@@ -39,7 +39,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
 
   override def beforeAll(): Unit = {
     val conf = new SparkConf()
-    env = createRpcEnv(conf, "local", 12345)
+    env = createRpcEnv(conf, "local", 0)
   }
 
   override def afterAll(): Unit = {
@@ -76,7 +76,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 0, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef("local", env.address, "send-remotely")
     try {
@@ -130,7 +130,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 0, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef("local", env.address, "ask-remotely")
     try {
@@ -158,7 +158,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     val shortProp = "spark.rpc.short.timeout"
     conf.set("spark.rpc.retry.wait", "0")
     conf.set("spark.rpc.numRetries", "1")
-    val anotherEnv = createRpcEnv(conf, "remote", 13345, clientMode = true)
+    val anotherEnv = createRpcEnv(conf, "remote", 0, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef("local", env.address, "ask-timeout")
     try {
@@ -417,7 +417,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 0, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef("local", env.address, "sendWithReply-remotely")
     try {
@@ -457,7 +457,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 0, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef(
       "local", env.address, "sendWithReply-remotely-error")
@@ -497,7 +497,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
 
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 0, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef(
       "local", env.address, "network-events")
@@ -543,7 +543,7 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
       }
     })
 
-    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 13345, clientMode = true)
+    val anotherEnv = createRpcEnv(new SparkConf(), "remote", 0, clientMode = true)
     // Use anotherEnv to find out the RpcEndpointRef
     val rpcEndpointRef = anotherEnv.setupEndpointRef(
       "local", env.address, "sendWithReply-unserializable-error")
@@ -571,8 +571,8 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     conf.set("spark.authenticate", "true")
     conf.set("spark.authenticate.secret", "good")
 
-    val localEnv = createRpcEnv(conf, "authentication-local", 13345)
-    val remoteEnv = createRpcEnv(conf, "authentication-remote", 14345, clientMode = true)
+    val localEnv = createRpcEnv(conf, "authentication-local", 0)
+    val remoteEnv = createRpcEnv(conf, "authentication-remote", 0, clientMode = true)
 
     try {
       @volatile var message: String = null
@@ -602,8 +602,8 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     conf.set("spark.authenticate", "true")
     conf.set("spark.authenticate.secret", "good")
 
-    val localEnv = createRpcEnv(conf, "authentication-local", 13345)
-    val remoteEnv = createRpcEnv(conf, "authentication-remote", 14345, clientMode = true)
+    val localEnv = createRpcEnv(conf, "authentication-local", 0)
+    val remoteEnv = createRpcEnv(conf, "authentication-remote", 0, clientMode = true)
 
     try {
       localEnv.setupEndpoint("ask-authentication", new RpcEndpoint {
diff --git a/core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala b/core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala
index 6478ab51c4da2..7aac02775e1bf 100644
--- a/core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/akka/AkkaRpcEnvSuite.scala
@@ -40,7 +40,7 @@ class AkkaRpcEnvSuite extends RpcEnvSuite {
     })
     val conf = new SparkConf()
     val newRpcEnv = new AkkaRpcEnvFactory().create(
-      RpcEnvConfig(conf, "test", "localhost", 12346, new SecurityManager(conf), false))
+      RpcEnvConfig(conf, "test", "localhost", 0, new SecurityManager(conf), false))
     try {
       val newRef = newRpcEnv.setupEndpointRef("local", ref.address, "test_endpoint")
       assert(s"akka.tcp://local@${env.address}/user/test_endpoint" ===
@@ -59,7 +59,7 @@ class AkkaRpcEnvSuite extends RpcEnvSuite {
     val conf = SSLSampleConfigs.sparkSSLConfig()
     val securityManager = new SecurityManager(conf)
     val rpcEnv = new AkkaRpcEnvFactory().create(
-      RpcEnvConfig(conf, "test", "localhost", 12346, securityManager, false))
+      RpcEnvConfig(conf, "test", "localhost", 0, securityManager, false))
     try {
       val uri = rpcEnv.uriOf("local", RpcAddress("1.2.3.4", 12345), "test_endpoint")
       assert("akka.ssl.tcp://local@1.2.3.4:12345/user/test_endpoint" === uri)

From 3bd77b213a9cd177c3ea3c61d37e5098e55f75a5 Mon Sep 17 00:00:00 2001
From: Srinivasa Reddy Vundela <vsr@cloudera.com>
Date: Thu, 19 Nov 2015 14:51:40 -0800
Subject: [PATCH 367/518] =?UTF-8?q?[SPARK-11799][CORE]=20Make=20it=20expli?=
 =?UTF-8?q?cit=20in=20executor=20logs=20that=20uncaught=20e=E2=80=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…xceptions are thrown during executor shutdown

This commit will make sure that when uncaught exceptions are prepended with [Container in shutdown] when JVM is shutting down.

Author: Srinivasa Reddy Vundela <vsr@cloudera.com>

Closes #9809 from vundela/master_11799.
---
 .../apache/spark/util/SparkUncaughtExceptionHandler.scala   | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/util/SparkUncaughtExceptionHandler.scala b/core/src/main/scala/org/apache/spark/util/SparkUncaughtExceptionHandler.scala
index 7248187247330..5e322557e9649 100644
--- a/core/src/main/scala/org/apache/spark/util/SparkUncaughtExceptionHandler.scala
+++ b/core/src/main/scala/org/apache/spark/util/SparkUncaughtExceptionHandler.scala
@@ -29,7 +29,11 @@ private[spark] object SparkUncaughtExceptionHandler
 
   override def uncaughtException(thread: Thread, exception: Throwable) {
     try {
-      logError("Uncaught exception in thread " + thread, exception)
+      // Make it explicit that uncaught exceptions are thrown when container is shutting down.
+      // It will help users when they analyze the executor logs
+      val inShutdownMsg = if (ShutdownHookManager.inShutdown()) "[Container in shutdown] " else ""
+      val errMsg = "Uncaught exception in thread "
+      logError(inShutdownMsg + errMsg + thread, exception)
 
       // We may have been called from a shutdown hook. If so, we must not call System.exit().
       // (If we do, we will deadlock.)

From f7135ed7194d4f936f6f58e14f02b1ed93f68ad1 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Thu, 19 Nov 2015 14:53:58 -0800
Subject: [PATCH 368/518] [SPARK-11828][CORE] Register DAGScheduler metrics
 source after app id is known.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9820 from vanzin/SPARK-11828.
---
 core/src/main/scala/org/apache/spark/SparkContext.scala       | 1 +
 .../main/scala/org/apache/spark/scheduler/DAGScheduler.scala  | 4 +---
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index ab374cb71286a..af4456c05b0a1 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -581,6 +581,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
 
     // Post init
     _taskScheduler.postStartHook()
+    _env.metricsSystem.registerSource(_dagScheduler.metricsSource)
     _env.metricsSystem.registerSource(new BlockManagerSource(_env.blockManager))
     _executorAllocationManager.foreach { e =>
       _env.metricsSystem.registerSource(e.executorAllocationManagerSource)
diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index 4a9518fff4e7b..ae725b467d8c4 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -130,7 +130,7 @@ class DAGScheduler(
 
   def this(sc: SparkContext) = this(sc, sc.taskScheduler)
 
-  private[scheduler] val metricsSource: DAGSchedulerSource = new DAGSchedulerSource(this)
+  private[spark] val metricsSource: DAGSchedulerSource = new DAGSchedulerSource(this)
 
   private[scheduler] val nextJobId = new AtomicInteger(0)
   private[scheduler] def numTotalJobs: Int = nextJobId.get()
@@ -1580,8 +1580,6 @@ class DAGScheduler(
     taskScheduler.stop()
   }
 
-  // Start the event thread and register the metrics source at the end of the constructor
-  env.metricsSystem.registerSource(metricsSource)
   eventProcessLoop.start()
 }
 

From 01403aa97b6aaab9b86ae806b5ea9e82690a741f Mon Sep 17 00:00:00 2001
From: hushan <hushan@xiaomi.com>
Date: Thu, 19 Nov 2015 14:56:00 -0800
Subject: [PATCH 369/518] [SPARK-11746][CORE] Use cache-aware method
 dependencies

a small change

Author: hushan <hushan@xiaomi.com>

Closes #9691 from suyanNone/unify-getDependency.
---
 .../main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala b/core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala
index d6a37e8cc5dac..0c6ddda52cee9 100644
--- a/core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala
@@ -65,7 +65,7 @@ class PartitionPruningRDD[T: ClassTag](
   }
 
   override protected def getPartitions: Array[Partition] =
-    getDependencies.head.asInstanceOf[PruneDependency[T]].partitions
+    dependencies.head.asInstanceOf[PruneDependency[T]].partitions
 }
 
 

From 37cff1b1a79cad11277612cb9bc8bc2365cf5ff2 Mon Sep 17 00:00:00 2001
From: Andrew Ray <ray.andrew@gmail.com>
Date: Thu, 19 Nov 2015 15:11:30 -0800
Subject: [PATCH 370/518] [SPARK-11275][SQL] Incorrect results when using
 rollup/cube

Fixes bug with grouping sets (including cube/rollup) where aggregates that included grouping expressions would return the wrong (null) result.

Also simplifies the analyzer rule a bit and leaves column pruning to the optimizer.

Added multiple unit tests to DataFrameAggregateSuite and verified it passes hive compatibility suite:
```
build/sbt -Phive -Dspark.hive.whitelist='groupby.*_grouping.*' 'test-only org.apache.spark.sql.hive.execution.HiveCompatibilitySuite'
```

This is an alternative to pr https://github.com/apache/spark/pull/9419 but I think its better as it simplifies the analyzer rule instead of adding another special case to it.

Author: Andrew Ray <ray.andrew@gmail.com>

Closes #9815 from aray/groupingset-agg-fix.
---
 .../sql/catalyst/analysis/Analyzer.scala      | 58 +++++++----------
 .../plans/logical/basicOperators.scala        |  4 ++
 .../spark/sql/DataFrameAggregateSuite.scala   | 62 +++++++++++++++++++
 3 files changed, 90 insertions(+), 34 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index 84781cd57f3dc..47962ebe6ef82 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -213,45 +213,35 @@ class Analyzer(
         GroupingSets(bitmasks(a), a.groupByExprs, a.child, a.aggregations)
       case x: GroupingSets =>
         val gid = AttributeReference(VirtualColumn.groupingIdName, IntegerType, false)()
-        // We will insert another Projection if the GROUP BY keys contains the
-        // non-attribute expressions. And the top operators can references those
-        // expressions by its alias.
-        // e.g. SELECT key%5 as c1 FROM src GROUP BY key%5 ==>
-        //      SELECT a as c1 FROM (SELECT key%5 AS a FROM src) GROUP BY a
-
-        // find all of the non-attribute expressions in the GROUP BY keys
-        val nonAttributeGroupByExpressions = new ArrayBuffer[Alias]()
-
-        // The pair of (the original GROUP BY key, associated attribute)
-        val groupByExprPairs = x.groupByExprs.map(_ match {
-          case e: NamedExpression => (e, e.toAttribute)
-          case other => {
-            val alias = Alias(other, other.toString)()
-            nonAttributeGroupByExpressions += alias // add the non-attributes expression alias
-            (other, alias.toAttribute)
-          }
-        })
-
-        // substitute the non-attribute expressions for aggregations.
-        val aggregation = x.aggregations.map(expr => expr.transformDown {
-          case e => groupByExprPairs.find(_._1.semanticEquals(e)).map(_._2).getOrElse(e)
-        }.asInstanceOf[NamedExpression])
 
-        // substitute the group by expressions.
-        val newGroupByExprs = groupByExprPairs.map(_._2)
+        // Expand works by setting grouping expressions to null as determined by the bitmasks. To
+        // prevent these null values from being used in an aggregate instead of the original value
+        // we need to create new aliases for all group by expressions that will only be used for
+        // the intended purpose.
+        val groupByAliases: Seq[Alias] = x.groupByExprs.map {
+          case e: NamedExpression => Alias(e, e.name)()
+          case other => Alias(other, other.toString)()
+        }
 
-        val child = if (nonAttributeGroupByExpressions.length > 0) {
-          // insert additional projection if contains the
-          // non-attribute expressions in the GROUP BY keys
-          Project(x.child.output ++ nonAttributeGroupByExpressions, x.child)
-        } else {
-          x.child
+        val aggregations: Seq[NamedExpression] = x.aggregations.map {
+          // If an expression is an aggregate (contains a AggregateExpression) then we dont change
+          // it so that the aggregation is computed on the unmodified value of its argument
+          // expressions.
+          case expr if expr.find(_.isInstanceOf[AggregateExpression]).nonEmpty => expr
+          // If not then its a grouping expression and we need to use the modified (with nulls from
+          // Expand) value of the expression.
+          case expr => expr.transformDown {
+            case e => groupByAliases.find(_.child.semanticEquals(e)).map(_.toAttribute).getOrElse(e)
+          }.asInstanceOf[NamedExpression]
         }
 
+        val child = Project(x.child.output ++ groupByAliases, x.child)
+        val groupByAttributes = groupByAliases.map(_.toAttribute)
+
         Aggregate(
-          newGroupByExprs :+ VirtualColumn.groupingIdAttribute,
-          aggregation,
-          Expand(x.bitmasks, newGroupByExprs, gid, child))
+          groupByAttributes :+ VirtualColumn.groupingIdAttribute,
+          aggregations,
+          Expand(x.bitmasks, groupByAttributes, gid, child))
     }
   }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index 45630a591d349..0c444482c5e4c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -323,6 +323,10 @@ trait GroupingAnalytics extends UnaryNode {
 
   override def output: Seq[Attribute] = aggregations.map(_.toAttribute)
 
+  // Needs to be unresolved before its translated to Aggregate + Expand because output attributes
+  // will change in analysis.
+  override lazy val resolved: Boolean = false
+
   def withNewAggs(aggs: Seq[NamedExpression]): GroupingAnalytics
 }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index 71adf2148a403..9c42f65bb6f52 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -60,6 +60,68 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
     )
   }
 
+  test("rollup") {
+    checkAnswer(
+      courseSales.rollup("course", "year").sum("earnings"),
+      Row("Java", 2012, 20000.0) ::
+        Row("Java", 2013, 30000.0) ::
+        Row("Java", null, 50000.0) ::
+        Row("dotNET", 2012, 15000.0) ::
+        Row("dotNET", 2013, 48000.0) ::
+        Row("dotNET", null, 63000.0) ::
+        Row(null, null, 113000.0) :: Nil
+    )
+  }
+
+  test("cube") {
+    checkAnswer(
+      courseSales.cube("course", "year").sum("earnings"),
+      Row("Java", 2012, 20000.0) ::
+        Row("Java", 2013, 30000.0) ::
+        Row("Java", null, 50000.0) ::
+        Row("dotNET", 2012, 15000.0) ::
+        Row("dotNET", 2013, 48000.0) ::
+        Row("dotNET", null, 63000.0) ::
+        Row(null, 2012, 35000.0) ::
+        Row(null, 2013, 78000.0) ::
+        Row(null, null, 113000.0) :: Nil
+    )
+  }
+
+  test("rollup overlapping columns") {
+    checkAnswer(
+      testData2.rollup($"a" + $"b" as "foo", $"b" as "bar").agg(sum($"a" - $"b") as "foo"),
+      Row(2, 1, 0) :: Row(3, 2, -1) :: Row(3, 1, 1) :: Row(4, 2, 0) :: Row(4, 1, 2) :: Row(5, 2, 1)
+        :: Row(2, null, 0) :: Row(3, null, 0) :: Row(4, null, 2) :: Row(5, null, 1)
+        :: Row(null, null, 3) :: Nil
+    )
+
+    checkAnswer(
+      testData2.rollup("a", "b").agg(sum("b")),
+      Row(1, 1, 1) :: Row(1, 2, 2) :: Row(2, 1, 1) :: Row(2, 2, 2) :: Row(3, 1, 1) :: Row(3, 2, 2)
+        :: Row(1, null, 3) :: Row(2, null, 3) :: Row(3, null, 3)
+        :: Row(null, null, 9) :: Nil
+    )
+  }
+
+  test("cube overlapping columns") {
+    checkAnswer(
+      testData2.cube($"a" + $"b", $"b").agg(sum($"a" - $"b")),
+      Row(2, 1, 0) :: Row(3, 2, -1) :: Row(3, 1, 1) :: Row(4, 2, 0) :: Row(4, 1, 2) :: Row(5, 2, 1)
+        :: Row(2, null, 0) :: Row(3, null, 0) :: Row(4, null, 2) :: Row(5, null, 1)
+        :: Row(null, 1, 3) :: Row(null, 2, 0)
+        :: Row(null, null, 3) :: Nil
+    )
+
+    checkAnswer(
+      testData2.cube("a", "b").agg(sum("b")),
+      Row(1, 1, 1) :: Row(1, 2, 2) :: Row(2, 1, 1) :: Row(2, 2, 2) :: Row(3, 1, 1) :: Row(3, 2, 2)
+        :: Row(1, null, 3) :: Row(2, null, 3) :: Row(3, null, 3)
+        :: Row(null, 1, 3) :: Row(null, 2, 6)
+        :: Row(null, null, 9) :: Nil
+    )
+  }
+
   test("spark.sql.retainGroupColumns config") {
     checkAnswer(
       testData2.groupBy("a").agg(sum($"b")),

From 880128f37e1bc0b9d98d1786670be62a06c648f2 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Thu, 19 Nov 2015 16:49:18 -0800
Subject: [PATCH 371/518] [SPARK-4134][CORE] Lower severity of some executor
 loss logs.

Don't log ERROR messages when executors are explicitly killed or when
the exit reason is not yet known.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9780 from vanzin/SPARK-11789.
---
 .../spark/scheduler/ExecutorLossReason.scala  |  2 +
 .../spark/scheduler/TaskSchedulerImpl.scala   | 44 ++++++++++++-------
 .../spark/scheduler/TaskSetManager.scala      |  1 +
 .../CoarseGrainedSchedulerBackend.scala       | 18 +++++---
 .../spark/deploy/yarn/YarnAllocator.scala     |  4 +-
 5 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
index 47a5cbff4930b..7e1197d742802 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/ExecutorLossReason.scala
@@ -40,6 +40,8 @@ private[spark] object ExecutorExited {
   }
 }
 
+private[spark] object ExecutorKilled extends ExecutorLossReason("Executor killed by driver.")
+
 /**
  * A loss reason that means we don't yet know why the executor exited.
  *
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index bf0419db1f75e..bdf19f9f277d9 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -470,25 +470,25 @@ private[spark] class TaskSchedulerImpl(
     synchronized {
       if (executorIdToTaskCount.contains(executorId)) {
         val hostPort = executorIdToHost(executorId)
-        logError("Lost executor %s on %s: %s".format(executorId, hostPort, reason))
+        logExecutorLoss(executorId, hostPort, reason)
         removeExecutor(executorId, reason)
         failedExecutor = Some(executorId)
       } else {
-         executorIdToHost.get(executorId) match {
-           case Some(_) =>
-             // If the host mapping still exists, it means we don't know the loss reason for the
-             // executor. So call removeExecutor() to update tasks running on that executor when
-             // the real loss reason is finally known.
-             logError(s"Actual reason for lost executor $executorId: ${reason.message}")
-             removeExecutor(executorId, reason)
-
-           case None =>
-             // We may get multiple executorLost() calls with different loss reasons. For example,
-             // one may be triggered by a dropped connection from the slave while another may be a
-             // report of executor termination from Mesos. We produce log messages for both so we
-             // eventually report the termination reason.
-             logError("Lost an executor " + executorId + " (already removed): " + reason)
-         }
+        executorIdToHost.get(executorId) match {
+          case Some(hostPort) =>
+            // If the host mapping still exists, it means we don't know the loss reason for the
+            // executor. So call removeExecutor() to update tasks running on that executor when
+            // the real loss reason is finally known.
+            logExecutorLoss(executorId, hostPort, reason)
+            removeExecutor(executorId, reason)
+
+          case None =>
+            // We may get multiple executorLost() calls with different loss reasons. For example,
+            // one may be triggered by a dropped connection from the slave while another may be a
+            // report of executor termination from Mesos. We produce log messages for both so we
+            // eventually report the termination reason.
+            logError(s"Lost an executor $executorId (already removed): $reason")
+        }
       }
     }
     // Call dagScheduler.executorLost without holding the lock on this to prevent deadlock
@@ -498,6 +498,18 @@ private[spark] class TaskSchedulerImpl(
     }
   }
 
+  private def logExecutorLoss(
+      executorId: String,
+      hostPort: String,
+      reason: ExecutorLossReason): Unit = reason match {
+    case LossReasonPending =>
+      logDebug(s"Executor $executorId on $hostPort lost, but reason not yet known.")
+    case ExecutorKilled =>
+      logInfo(s"Executor $executorId on $hostPort killed by driver.")
+    case _ =>
+      logError(s"Lost executor $executorId on $hostPort: $reason")
+  }
+
   /**
    * Remove an executor from all our data structures and mark it as lost. If the executor's loss
    * reason is not yet known, do not yet remove its association with its host nor update the status
diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
index 114468c48c44c..a02f3017cb6e9 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala
@@ -800,6 +800,7 @@ private[spark] class TaskSetManager(
     for ((tid, info) <- taskInfos if info.running && info.executorId == execId) {
       val exitCausedByApp: Boolean = reason match {
         case exited: ExecutorExited => exited.exitCausedByApp
+        case ExecutorKilled => false
         case _ => true
       }
       handleFailedTask(tid, TaskState.FAILED, ExecutorLostFailure(info.executorId, exitCausedByApp,
diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
index 6f0c910c009a5..505c161141c88 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala
@@ -64,8 +64,10 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
 
   private val listenerBus = scheduler.sc.listenerBus
 
-  // Executors we have requested the cluster manager to kill that have not died yet
-  private val executorsPendingToRemove = new HashSet[String]
+  // Executors we have requested the cluster manager to kill that have not died yet; maps
+  // the executor ID to whether it was explicitly killed by the driver (and thus shouldn't
+  // be considered an app-related failure).
+  private val executorsPendingToRemove = new HashMap[String, Boolean]
 
   // A map to store hostname with its possible task number running on it
   protected var hostToLocalTaskCount: Map[String, Int] = Map.empty
@@ -250,15 +252,15 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
         case Some(executorInfo) =>
           // This must be synchronized because variables mutated
           // in this block are read when requesting executors
-          CoarseGrainedSchedulerBackend.this.synchronized {
+          val killed = CoarseGrainedSchedulerBackend.this.synchronized {
             addressToExecutorId -= executorInfo.executorAddress
             executorDataMap -= executorId
-            executorsPendingToRemove -= executorId
             executorsPendingLossReason -= executorId
+            executorsPendingToRemove.remove(executorId).getOrElse(false)
           }
           totalCoreCount.addAndGet(-executorInfo.totalCores)
           totalRegisteredExecutors.addAndGet(-1)
-          scheduler.executorLost(executorId, reason)
+          scheduler.executorLost(executorId, if (killed) ExecutorKilled else reason)
           listenerBus.post(
             SparkListenerExecutorRemoved(System.currentTimeMillis(), executorId, reason.toString))
         case None => logInfo(s"Asked to remove non-existent executor $executorId")
@@ -459,6 +461,10 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
   /**
    * Request that the cluster manager kill the specified executors.
    *
+   * When asking the executor to be replaced, the executor loss is considered a failure, and
+   * killed tasks that are running on the executor will count towards the failure limits. If no
+   * replacement is being requested, then the tasks will not count towards the limit.
+   *
    * @param executorIds identifiers of executors to kill
    * @param replace whether to replace the killed executors with new ones
    * @param force whether to force kill busy executors
@@ -479,7 +485,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp
     val executorsToKill = knownExecutors
       .filter { id => !executorsPendingToRemove.contains(id) }
       .filter { id => force || !scheduler.isExecutorBusy(id) }
-    executorsPendingToRemove ++= executorsToKill
+    executorsToKill.foreach { id => executorsPendingToRemove(id) = !replace }
 
     // If we do not wish to replace the executors we kill, sync the target number of executors
     // with the cluster manager to avoid allocating new ones. When computing the new target,
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
index 7e39c3ea56af3..73cd9031f0250 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocator.scala
@@ -481,7 +481,7 @@ private[yarn] class YarnAllocator(
             (true, memLimitExceededLogMessage(
               completedContainer.getDiagnostics,
               PMEM_EXCEEDED_PATTERN))
-          case unknown =>
+          case _ =>
             numExecutorsFailed += 1
             (true, "Container marked as failed: " + containerId + onHostStr +
               ". Exit status: " + completedContainer.getExitStatus +
@@ -493,7 +493,7 @@ private[yarn] class YarnAllocator(
         } else {
           logInfo(containerExitReason)
         }
-        ExecutorExited(0, exitCausedByApp, containerExitReason)
+        ExecutorExited(exitStatus, exitCausedByApp, containerExitReason)
       } else {
         // If we have already released this container, then it must mean
         // that the driver has explicitly requested it to be killed

From b2cecb80ece59a1c086d4ae7aeebef445a4e7299 Mon Sep 17 00:00:00 2001
From: Tathagata Das <tathagata.das1565@gmail.com>
Date: Thu, 19 Nov 2015 16:50:08 -0800
Subject: [PATCH 372/518] [SPARK-11845][STREAMING][TEST] Added unit test to
 verify TrackStateRDD is correctly checkpointed

To make sure that all lineage is correctly truncated for TrackStateRDD when checkpointed.

Author: Tathagata Das <tathagata.das1565@gmail.com>

Closes #9831 from tdas/SPARK-11845.
---
 .../org/apache/spark/CheckpointSuite.scala    | 411 +++++++++---------
 .../streaming/rdd/TrackStateRDDSuite.scala    |  60 ++-
 2 files changed, 267 insertions(+), 204 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
index 119e5fc28e412..ab23326c6c25d 100644
--- a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
+++ b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
@@ -21,17 +21,223 @@ import java.io.File
 
 import scala.reflect.ClassTag
 
+import org.apache.spark.CheckpointSuite._
 import org.apache.spark.rdd._
 import org.apache.spark.storage.{BlockId, StorageLevel, TestBlockId}
 import org.apache.spark.util.Utils
 
+trait RDDCheckpointTester { self: SparkFunSuite =>
+
+  protected val partitioner = new HashPartitioner(2)
+
+  private def defaultCollectFunc[T](rdd: RDD[T]): Any = rdd.collect()
+
+  /** Implementations of this trait must implement this method */
+  protected def sparkContext: SparkContext
+
+  /**
+   * Test checkpointing of the RDD generated by the given operation. It tests whether the
+   * serialized size of the RDD is reduce after checkpointing or not. This function should be called
+   * on all RDDs that have a parent RDD (i.e., do not call on ParallelCollection, BlockRDD, etc.).
+   *
+   * @param op an operation to run on the RDD
+   * @param reliableCheckpoint if true, use reliable checkpoints, otherwise use local checkpoints
+   * @param collectFunc a function for collecting the values in the RDD, in case there are
+   *                    non-comparable types like arrays that we want to convert to something
+   *                    that supports ==
+   */
+  protected def testRDD[U: ClassTag](
+      op: (RDD[Int]) => RDD[U],
+      reliableCheckpoint: Boolean,
+      collectFunc: RDD[U] => Any = defaultCollectFunc[U] _): Unit = {
+    // Generate the final RDD using given RDD operation
+    val baseRDD = generateFatRDD()
+    val operatedRDD = op(baseRDD)
+    val parentRDD = operatedRDD.dependencies.headOption.orNull
+    val rddType = operatedRDD.getClass.getSimpleName
+    val numPartitions = operatedRDD.partitions.length
+
+    // Force initialization of all the data structures in RDDs
+    // Without this, serializing the RDD will give a wrong estimate of the size of the RDD
+    initializeRdd(operatedRDD)
+
+    val partitionsBeforeCheckpoint = operatedRDD.partitions
+
+    // Find serialized sizes before and after the checkpoint
+    logInfo("RDD before checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
+    val (rddSizeBeforeCheckpoint, partitionSizeBeforeCheckpoint) = getSerializedSizes(operatedRDD)
+    checkpoint(operatedRDD, reliableCheckpoint)
+    val result = collectFunc(operatedRDD)
+    operatedRDD.collect() // force re-initialization of post-checkpoint lazy variables
+    val (rddSizeAfterCheckpoint, partitionSizeAfterCheckpoint) = getSerializedSizes(operatedRDD)
+    logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
+
+    // Test whether the checkpoint file has been created
+    if (reliableCheckpoint) {
+      assert(
+        collectFunc(sparkContext.checkpointFile[U](operatedRDD.getCheckpointFile.get)) === result)
+    }
+
+    // Test whether dependencies have been changed from its earlier parent RDD
+    assert(operatedRDD.dependencies.head.rdd != parentRDD)
+
+    // Test whether the partitions have been changed from its earlier partitions
+    assert(operatedRDD.partitions.toList != partitionsBeforeCheckpoint.toList)
+
+    // Test whether the partitions have been changed to the new Hadoop partitions
+    assert(operatedRDD.partitions.toList === operatedRDD.checkpointData.get.getPartitions.toList)
+
+    // Test whether the number of partitions is same as before
+    assert(operatedRDD.partitions.length === numPartitions)
+
+    // Test whether the data in the checkpointed RDD is same as original
+    assert(collectFunc(operatedRDD) === result)
+
+    // Test whether serialized size of the RDD has reduced.
+    logInfo("Size of " + rddType +
+      " [" + rddSizeBeforeCheckpoint + " --> " + rddSizeAfterCheckpoint + "]")
+    assert(
+      rddSizeAfterCheckpoint < rddSizeBeforeCheckpoint,
+      "Size of " + rddType + " did not reduce after checkpointing " +
+        " [" + rddSizeBeforeCheckpoint + " --> " + rddSizeAfterCheckpoint + "]"
+    )
+  }
+
+  /**
+   * Test whether checkpointing of the parent of the generated RDD also
+   * truncates the lineage or not. Some RDDs like CoGroupedRDD hold on to its parent
+   * RDDs partitions. So even if the parent RDD is checkpointed and its partitions changed,
+   * the generated RDD will remember the partitions and therefore potentially the whole lineage.
+   * This function should be called only those RDD whose partitions refer to parent RDD's
+   * partitions (i.e., do not call it on simple RDD like MappedRDD).
+   *
+   * @param op an operation to run on the RDD
+   * @param reliableCheckpoint if true, use reliable checkpoints, otherwise use local checkpoints
+   * @param collectFunc a function for collecting the values in the RDD, in case there are
+   *                    non-comparable types like arrays that we want to convert to something
+   *                    that supports ==
+   */
+  protected def testRDDPartitions[U: ClassTag](
+      op: (RDD[Int]) => RDD[U],
+      reliableCheckpoint: Boolean,
+      collectFunc: RDD[U] => Any = defaultCollectFunc[U] _): Unit = {
+    // Generate the final RDD using given RDD operation
+    val baseRDD = generateFatRDD()
+    val operatedRDD = op(baseRDD)
+    val parentRDDs = operatedRDD.dependencies.map(_.rdd)
+    val rddType = operatedRDD.getClass.getSimpleName
+
+    // Force initialization of all the data structures in RDDs
+    // Without this, serializing the RDD will give a wrong estimate of the size of the RDD
+    initializeRdd(operatedRDD)
+
+    // Find serialized sizes before and after the checkpoint
+    logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
+    val (rddSizeBeforeCheckpoint, partitionSizeBeforeCheckpoint) = getSerializedSizes(operatedRDD)
+    // checkpoint the parent RDD, not the generated one
+    parentRDDs.foreach { rdd =>
+      checkpoint(rdd, reliableCheckpoint)
+    }
+    val result = collectFunc(operatedRDD) // force checkpointing
+    operatedRDD.collect() // force re-initialization of post-checkpoint lazy variables
+    val (rddSizeAfterCheckpoint, partitionSizeAfterCheckpoint) = getSerializedSizes(operatedRDD)
+    logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
+
+    // Test whether the data in the checkpointed RDD is same as original
+    assert(collectFunc(operatedRDD) === result)
+
+    // Test whether serialized size of the partitions has reduced
+    logInfo("Size of partitions of " + rddType +
+      " [" + partitionSizeBeforeCheckpoint + " --> " + partitionSizeAfterCheckpoint + "]")
+    assert(
+      partitionSizeAfterCheckpoint < partitionSizeBeforeCheckpoint,
+      "Size of " + rddType + " partitions did not reduce after checkpointing parent RDDs" +
+        " [" + partitionSizeBeforeCheckpoint + " --> " + partitionSizeAfterCheckpoint + "]"
+    )
+  }
+
+  /**
+   * Get serialized sizes of the RDD and its partitions, in order to test whether the size shrinks
+   * upon checkpointing. Ignores the checkpointData field, which may grow when we checkpoint.
+   */
+  private def getSerializedSizes(rdd: RDD[_]): (Int, Int) = {
+    val rddSize = Utils.serialize(rdd).size
+    val rddCpDataSize = Utils.serialize(rdd.checkpointData).size
+    val rddPartitionSize = Utils.serialize(rdd.partitions).size
+    val rddDependenciesSize = Utils.serialize(rdd.dependencies).size
+
+    // Print detailed size, helps in debugging
+    logInfo("Serialized sizes of " + rdd +
+      ": RDD = " + rddSize +
+      ", RDD checkpoint data = " + rddCpDataSize +
+      ", RDD partitions = " + rddPartitionSize +
+      ", RDD dependencies = " + rddDependenciesSize
+    )
+    // this makes sure that serializing the RDD's checkpoint data does not
+    // serialize the whole RDD as well
+    assert(
+      rddSize > rddCpDataSize,
+      "RDD's checkpoint data (" + rddCpDataSize + ") is equal or larger than the " +
+        "whole RDD with checkpoint data (" + rddSize + ")"
+    )
+    (rddSize - rddCpDataSize, rddPartitionSize)
+  }
+
+  /**
+   * Serialize and deserialize an object. This is useful to verify the objects
+   * contents after deserialization (e.g., the contents of an RDD split after
+   * it is sent to a slave along with a task)
+   */
+  protected def serializeDeserialize[T](obj: T): T = {
+    val bytes = Utils.serialize(obj)
+    Utils.deserialize[T](bytes)
+  }
+
+  /**
+   * Recursively force the initialization of the all members of an RDD and it parents.
+   */
+  private def initializeRdd(rdd: RDD[_]): Unit = {
+    rdd.partitions // forces the initialization of the partitions
+    rdd.dependencies.map(_.rdd).foreach(initializeRdd)
+  }
+
+  /** Checkpoint the RDD either locally or reliably. */
+  protected def checkpoint(rdd: RDD[_], reliableCheckpoint: Boolean): Unit = {
+    if (reliableCheckpoint) {
+      rdd.checkpoint()
+    } else {
+      rdd.localCheckpoint()
+    }
+  }
+
+  /** Run a test twice, once for local checkpointing and once for reliable checkpointing. */
+  protected def runTest(name: String)(body: Boolean => Unit): Unit = {
+    test(name + " [reliable checkpoint]")(body(true))
+    test(name + " [local checkpoint]")(body(false))
+  }
+
+  /**
+   * Generate an RDD such that both the RDD and its partitions have large size.
+   */
+  protected def generateFatRDD(): RDD[Int] = {
+    new FatRDD(sparkContext.makeRDD(1 to 100, 4)).map(x => x)
+  }
+
+  /**
+   * Generate an pair RDD (with partitioner) such that both the RDD and its partitions
+   * have large size.
+   */
+  protected def generateFatPairRDD(): RDD[(Int, Int)] = {
+    new FatPairRDD(sparkContext.makeRDD(1 to 100, 4), partitioner).mapValues(x => x)
+  }
+}
+
 /**
  * Test suite for end-to-end checkpointing functionality.
  * This tests both reliable checkpoints and local checkpoints.
  */
-class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging {
+class CheckpointSuite extends SparkFunSuite with RDDCheckpointTester with LocalSparkContext {
   private var checkpointDir: File = _
-  private val partitioner = new HashPartitioner(2)
 
   override def beforeEach(): Unit = {
     super.beforeEach()
@@ -46,6 +252,8 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     Utils.deleteRecursively(checkpointDir)
   }
 
+  override def sparkContext: SparkContext = sc
+
   runTest("basic checkpointing") { reliableCheckpoint: Boolean =>
     val parCollection = sc.makeRDD(1 to 4)
     val flatMappedRDD = parCollection.flatMap(x => 1 to x)
@@ -250,204 +458,6 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     assert(rdd.isCheckpointedAndMaterialized === true)
     assert(rdd.partitions.size === 0)
   }
-
-  // Utility test methods
-
-  /** Checkpoint the RDD either locally or reliably. */
-  private def checkpoint(rdd: RDD[_], reliableCheckpoint: Boolean): Unit = {
-    if (reliableCheckpoint) {
-      rdd.checkpoint()
-    } else {
-      rdd.localCheckpoint()
-    }
-  }
-
-  /** Run a test twice, once for local checkpointing and once for reliable checkpointing. */
-  private def runTest(name: String)(body: Boolean => Unit): Unit = {
-    test(name + " [reliable checkpoint]")(body(true))
-    test(name + " [local checkpoint]")(body(false))
-  }
-
-  private def defaultCollectFunc[T](rdd: RDD[T]): Any = rdd.collect()
-
-  /**
-   * Test checkpointing of the RDD generated by the given operation. It tests whether the
-   * serialized size of the RDD is reduce after checkpointing or not. This function should be called
-   * on all RDDs that have a parent RDD (i.e., do not call on ParallelCollection, BlockRDD, etc.).
-   *
-   * @param op an operation to run on the RDD
-   * @param reliableCheckpoint if true, use reliable checkpoints, otherwise use local checkpoints
-   * @param collectFunc a function for collecting the values in the RDD, in case there are
-   *   non-comparable types like arrays that we want to convert to something that supports ==
-   */
-  private def testRDD[U: ClassTag](
-      op: (RDD[Int]) => RDD[U],
-      reliableCheckpoint: Boolean,
-      collectFunc: RDD[U] => Any = defaultCollectFunc[U] _): Unit = {
-    // Generate the final RDD using given RDD operation
-    val baseRDD = generateFatRDD()
-    val operatedRDD = op(baseRDD)
-    val parentRDD = operatedRDD.dependencies.headOption.orNull
-    val rddType = operatedRDD.getClass.getSimpleName
-    val numPartitions = operatedRDD.partitions.length
-
-    // Force initialization of all the data structures in RDDs
-    // Without this, serializing the RDD will give a wrong estimate of the size of the RDD
-    initializeRdd(operatedRDD)
-
-    val partitionsBeforeCheckpoint = operatedRDD.partitions
-
-    // Find serialized sizes before and after the checkpoint
-    logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
-    val (rddSizeBeforeCheckpoint, partitionSizeBeforeCheckpoint) = getSerializedSizes(operatedRDD)
-    checkpoint(operatedRDD, reliableCheckpoint)
-    val result = collectFunc(operatedRDD)
-    operatedRDD.collect() // force re-initialization of post-checkpoint lazy variables
-    val (rddSizeAfterCheckpoint, partitionSizeAfterCheckpoint) = getSerializedSizes(operatedRDD)
-    logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
-
-    // Test whether the checkpoint file has been created
-    if (reliableCheckpoint) {
-      assert(collectFunc(sc.checkpointFile[U](operatedRDD.getCheckpointFile.get)) === result)
-    }
-
-    // Test whether dependencies have been changed from its earlier parent RDD
-    assert(operatedRDD.dependencies.head.rdd != parentRDD)
-
-    // Test whether the partitions have been changed from its earlier partitions
-    assert(operatedRDD.partitions.toList != partitionsBeforeCheckpoint.toList)
-
-    // Test whether the partitions have been changed to the new Hadoop partitions
-    assert(operatedRDD.partitions.toList === operatedRDD.checkpointData.get.getPartitions.toList)
-
-    // Test whether the number of partitions is same as before
-    assert(operatedRDD.partitions.length === numPartitions)
-
-    // Test whether the data in the checkpointed RDD is same as original
-    assert(collectFunc(operatedRDD) === result)
-
-    // Test whether serialized size of the RDD has reduced.
-    logInfo("Size of " + rddType +
-      " [" + rddSizeBeforeCheckpoint + " --> " + rddSizeAfterCheckpoint + "]")
-    assert(
-      rddSizeAfterCheckpoint < rddSizeBeforeCheckpoint,
-      "Size of " + rddType + " did not reduce after checkpointing " +
-        " [" + rddSizeBeforeCheckpoint + " --> " + rddSizeAfterCheckpoint + "]"
-    )
-  }
-
-  /**
-   * Test whether checkpointing of the parent of the generated RDD also
-   * truncates the lineage or not. Some RDDs like CoGroupedRDD hold on to its parent
-   * RDDs partitions. So even if the parent RDD is checkpointed and its partitions changed,
-   * the generated RDD will remember the partitions and therefore potentially the whole lineage.
-   * This function should be called only those RDD whose partitions refer to parent RDD's
-   * partitions (i.e., do not call it on simple RDD like MappedRDD).
-   *
-   * @param op an operation to run on the RDD
-   * @param reliableCheckpoint if true, use reliable checkpoints, otherwise use local checkpoints
-   * @param collectFunc a function for collecting the values in the RDD, in case there are
-   *   non-comparable types like arrays that we want to convert to something that supports ==
-   */
-  private def testRDDPartitions[U: ClassTag](
-      op: (RDD[Int]) => RDD[U],
-      reliableCheckpoint: Boolean,
-      collectFunc: RDD[U] => Any = defaultCollectFunc[U] _): Unit = {
-    // Generate the final RDD using given RDD operation
-    val baseRDD = generateFatRDD()
-    val operatedRDD = op(baseRDD)
-    val parentRDDs = operatedRDD.dependencies.map(_.rdd)
-    val rddType = operatedRDD.getClass.getSimpleName
-
-    // Force initialization of all the data structures in RDDs
-    // Without this, serializing the RDD will give a wrong estimate of the size of the RDD
-    initializeRdd(operatedRDD)
-
-    // Find serialized sizes before and after the checkpoint
-    logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
-    val (rddSizeBeforeCheckpoint, partitionSizeBeforeCheckpoint) = getSerializedSizes(operatedRDD)
-    // checkpoint the parent RDD, not the generated one
-    parentRDDs.foreach { rdd =>
-      checkpoint(rdd, reliableCheckpoint)
-    }
-    val result = collectFunc(operatedRDD)  // force checkpointing
-    operatedRDD.collect() // force re-initialization of post-checkpoint lazy variables
-    val (rddSizeAfterCheckpoint, partitionSizeAfterCheckpoint) = getSerializedSizes(operatedRDD)
-    logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
-
-    // Test whether the data in the checkpointed RDD is same as original
-    assert(collectFunc(operatedRDD) === result)
-
-    // Test whether serialized size of the partitions has reduced
-    logInfo("Size of partitions of " + rddType +
-      " [" + partitionSizeBeforeCheckpoint + " --> " + partitionSizeAfterCheckpoint + "]")
-    assert(
-      partitionSizeAfterCheckpoint < partitionSizeBeforeCheckpoint,
-      "Size of " + rddType + " partitions did not reduce after checkpointing parent RDDs" +
-        " [" + partitionSizeBeforeCheckpoint + " --> " + partitionSizeAfterCheckpoint + "]"
-    )
-  }
-
-  /**
-   * Generate an RDD such that both the RDD and its partitions have large size.
-   */
-  private def generateFatRDD(): RDD[Int] = {
-    new FatRDD(sc.makeRDD(1 to 100, 4)).map(x => x)
-  }
-
-  /**
-   * Generate an pair RDD (with partitioner) such that both the RDD and its partitions
-   * have large size.
-   */
-  private def generateFatPairRDD(): RDD[(Int, Int)] = {
-    new FatPairRDD(sc.makeRDD(1 to 100, 4), partitioner).mapValues(x => x)
-  }
-
-  /**
-   * Get serialized sizes of the RDD and its partitions, in order to test whether the size shrinks
-   * upon checkpointing. Ignores the checkpointData field, which may grow when we checkpoint.
-   */
-  private def getSerializedSizes(rdd: RDD[_]): (Int, Int) = {
-    val rddSize = Utils.serialize(rdd).size
-    val rddCpDataSize = Utils.serialize(rdd.checkpointData).size
-    val rddPartitionSize = Utils.serialize(rdd.partitions).size
-    val rddDependenciesSize = Utils.serialize(rdd.dependencies).size
-
-    // Print detailed size, helps in debugging
-    logInfo("Serialized sizes of " + rdd +
-      ": RDD = " + rddSize +
-      ", RDD checkpoint data = " + rddCpDataSize +
-      ", RDD partitions = " + rddPartitionSize +
-      ", RDD dependencies = " + rddDependenciesSize
-    )
-    // this makes sure that serializing the RDD's checkpoint data does not
-    // serialize the whole RDD as well
-    assert(
-      rddSize > rddCpDataSize,
-      "RDD's checkpoint data (" + rddCpDataSize  + ") is equal or larger than the " +
-        "whole RDD with checkpoint data (" + rddSize + ")"
-    )
-    (rddSize - rddCpDataSize, rddPartitionSize)
-  }
-
-  /**
-   * Serialize and deserialize an object. This is useful to verify the objects
-   * contents after deserialization (e.g., the contents of an RDD split after
-   * it is sent to a slave along with a task)
-   */
-  private def serializeDeserialize[T](obj: T): T = {
-    val bytes = Utils.serialize(obj)
-    Utils.deserialize[T](bytes)
-  }
-
-  /**
-   * Recursively force the initialization of the all members of an RDD and it parents.
-   */
-  private def initializeRdd(rdd: RDD[_]): Unit = {
-    rdd.partitions // forces the
-    rdd.dependencies.map(_.rdd).foreach(initializeRdd)
-  }
-
 }
 
 /** RDD partition that has large serialized size. */
@@ -494,5 +504,4 @@ object CheckpointSuite {
       part
     ).asInstanceOf[RDD[(K, Array[Iterable[V]])]]
   }
-
 }
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
index 19ef5a14f8ab4..0feb3af1abb0f 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
@@ -17,31 +17,40 @@
 
 package org.apache.spark.streaming.rdd
 
+import java.io.File
+
 import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
 
 import org.scalatest.BeforeAndAfterAll
 
+import org.apache.spark._
 import org.apache.spark.rdd.RDD
 import org.apache.spark.streaming.util.OpenHashMapBasedStateMap
-import org.apache.spark.streaming.{Time, State}
-import org.apache.spark.{HashPartitioner, SparkConf, SparkContext, SparkFunSuite}
+import org.apache.spark.streaming.{State, Time}
+import org.apache.spark.util.Utils
 
-class TrackStateRDDSuite extends SparkFunSuite with BeforeAndAfterAll {
+class TrackStateRDDSuite extends SparkFunSuite with RDDCheckpointTester with BeforeAndAfterAll {
 
   private var sc: SparkContext = null
+  private var checkpointDir: File = _
 
   override def beforeAll(): Unit = {
     sc = new SparkContext(
       new SparkConf().setMaster("local").setAppName("TrackStateRDDSuite"))
+    checkpointDir = Utils.createTempDir()
+    sc.setCheckpointDir(checkpointDir.toString)
   }
 
   override def afterAll(): Unit = {
     if (sc != null) {
       sc.stop()
     }
+    Utils.deleteRecursively(checkpointDir)
   }
 
+  override def sparkContext: SparkContext = sc
+
   test("creation from pair RDD") {
     val data = Seq((1, "1"), (2, "2"), (3, "3"))
     val partitioner = new HashPartitioner(10)
@@ -278,6 +287,51 @@ class TrackStateRDDSuite extends SparkFunSuite with BeforeAndAfterAll {
       rdd7, Seq(("k3", 2)), Set())
   }
 
+  test("checkpointing") {
+    /**
+     * This tests whether the TrackStateRDD correctly truncates any references to its parent RDDs -
+     * the data RDD and the parent TrackStateRDD.
+     */
+    def rddCollectFunc(rdd: RDD[TrackStateRDDRecord[Int, Int, Int]])
+      : Set[(List[(Int, Int, Long)], List[Int])] = {
+      rdd.map { record => (record.stateMap.getAll().toList, record.emittedRecords.toList) }
+         .collect.toSet
+    }
+
+    /** Generate TrackStateRDD with data RDD having a long lineage */
+    def makeStateRDDWithLongLineageDataRDD(longLineageRDD: RDD[Int])
+      : TrackStateRDD[Int, Int, Int, Int] = {
+      TrackStateRDD.createFromPairRDD(longLineageRDD.map { _ -> 1}, partitioner, Time(0))
+    }
+
+    testRDD(
+      makeStateRDDWithLongLineageDataRDD, reliableCheckpoint = true, rddCollectFunc _)
+    testRDDPartitions(
+      makeStateRDDWithLongLineageDataRDD, reliableCheckpoint = true, rddCollectFunc _)
+
+    /** Generate TrackStateRDD with parent state RDD having a long lineage */
+    def makeStateRDDWithLongLineageParenttateRDD(
+        longLineageRDD: RDD[Int]): TrackStateRDD[Int, Int, Int, Int] = {
+
+      // Create a TrackStateRDD that has a long lineage using the data RDD with a long lineage
+      val stateRDDWithLongLineage = makeStateRDDWithLongLineageDataRDD(longLineageRDD)
+
+      // Create a new TrackStateRDD, with the lineage lineage TrackStateRDD as the parent
+      new TrackStateRDD[Int, Int, Int, Int](
+        stateRDDWithLongLineage,
+        stateRDDWithLongLineage.sparkContext.emptyRDD[(Int, Int)].partitionBy(partitioner),
+        (time: Time, key: Int, value: Option[Int], state: State[Int]) => None,
+        Time(10),
+        None
+      )
+    }
+
+    testRDD(
+      makeStateRDDWithLongLineageParenttateRDD, reliableCheckpoint = true, rddCollectFunc _)
+    testRDDPartitions(
+      makeStateRDDWithLongLineageParenttateRDD, reliableCheckpoint = true, rddCollectFunc _)
+  }
+
   /** Assert whether the `trackStateByKey` operation generates expected results */
   private def assertOperation[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
       testStateRDD: TrackStateRDD[K, V, S, T],

From ee21407747fb00db2f26d1119446ccbb20c19232 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Thu, 19 Nov 2015 17:14:10 -0800
Subject: [PATCH 373/518] [SPARK-11864][SQL] Improve performance of max/min

This PR has the following optimization:

1) The greatest/least already does the null-check, so the `If` and `IsNull` are not necessary.

2) In greatest/least, it should initialize the result using the first child (removing one block).

3) For primitive types, the generated greater expression is too complicated (`a > b ? 1 : (a < b) ? -1 : 0) > 0`), should be as simple as `a > b`

Combine these optimization, this could improve the performance of `ss_max` query by 30%.

Author: Davies Liu <davies@databricks.com>

Closes #9846 from davies/improve_max.
---
 .../catalyst/expressions/aggregate/Max.scala  |  5 +--
 .../catalyst/expressions/aggregate/Min.scala  |  5 +--
 .../expressions/codegen/CodeGenerator.scala   | 12 ++++++
 .../expressions/conditionalExpressions.scala  | 38 +++++++++++--------
 .../expressions/nullExpressions.scala         | 10 +++--
 5 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
index 61cae44cd0f5b..906003188d4ff 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala
@@ -46,13 +46,12 @@ case class Max(child: Expression) extends DeclarativeAggregate {
   )
 
   override lazy val updateExpressions: Seq[Expression] = Seq(
-    /* max = */ If(IsNull(child), max, If(IsNull(max), child, Greatest(Seq(max, child))))
+    /* max = */ Greatest(Seq(max, child))
   )
 
   override lazy val mergeExpressions: Seq[Expression] = {
-    val greatest = Greatest(Seq(max.left, max.right))
     Seq(
-      /* max = */ If(IsNull(max.right), max.left, If(IsNull(max.left), max.right, greatest))
+      /* max = */ Greatest(Seq(max.left, max.right))
     )
   }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
index 242456d9e2e18..39f7afbd081cd 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala
@@ -47,13 +47,12 @@ case class Min(child: Expression) extends DeclarativeAggregate {
   )
 
   override lazy val updateExpressions: Seq[Expression] = Seq(
-    /* min = */ If(IsNull(child), min, If(IsNull(min), child, Least(Seq(min, child))))
+    /* min = */ Least(Seq(min, child))
   )
 
   override lazy val mergeExpressions: Seq[Expression] = {
-    val least = Least(Seq(min.left, min.right))
     Seq(
-      /* min = */ If(IsNull(min.right), min.left, If(IsNull(min.left), min.right, least))
+      /* min = */ Least(Seq(min.left, min.right))
     )
   }
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
index 1718cfbd35332..1b7260cdfe515 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
@@ -329,6 +329,18 @@ class CodeGenContext {
       throw new IllegalArgumentException("cannot generate compare code for un-comparable type")
   }
 
+  /**
+    * Generates code for greater of two expressions.
+    *
+    * @param dataType data type of the expressions
+    * @param c1 name of the variable of expression 1's output
+    * @param c2 name of the variable of expression 2's output
+    */
+  def genGreater(dataType: DataType, c1: String, c2: String): String = javaType(dataType) match {
+    case JAVA_BYTE | JAVA_SHORT | JAVA_INT | JAVA_LONG => s"$c1 > $c2"
+    case _ => s"(${genComp(dataType, c1, c2)}) > 0"
+  }
+
   /**
    * List of java data types that have special accessors and setters in [[InternalRow]].
    */
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
index 0d4af43978ea1..694a2a7c54a90 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
@@ -348,19 +348,22 @@ case class Least(children: Seq[Expression]) extends Expression {
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val evalChildren = children.map(_.gen(ctx))
-    def updateEval(i: Int): String =
+    val first = evalChildren(0)
+    val rest = evalChildren.drop(1)
+    def updateEval(eval: GeneratedExpressionCode): String =
       s"""
-        if (!${evalChildren(i).isNull} && (${ev.isNull} ||
-          ${ctx.genComp(dataType, evalChildren(i).value, ev.value)} < 0)) {
+        ${eval.code}
+        if (!${eval.isNull} && (${ev.isNull} ||
+          ${ctx.genGreater(dataType, ev.value, eval.value)})) {
           ${ev.isNull} = false;
-          ${ev.value} = ${evalChildren(i).value};
+          ${ev.value} = ${eval.value};
         }
       """
     s"""
-      ${evalChildren.map(_.code).mkString("\n")}
-      boolean ${ev.isNull} = true;
-      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
-      ${children.indices.map(updateEval).mkString("\n")}
+      ${first.code}
+      boolean ${ev.isNull} = ${first.isNull};
+      ${ctx.javaType(dataType)} ${ev.value} = ${first.value};
+      ${rest.map(updateEval).mkString("\n")}
     """
   }
 }
@@ -403,19 +406,22 @@ case class Greatest(children: Seq[Expression]) extends Expression {
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     val evalChildren = children.map(_.gen(ctx))
-    def updateEval(i: Int): String =
+    val first = evalChildren(0)
+    val rest = evalChildren.drop(1)
+    def updateEval(eval: GeneratedExpressionCode): String =
       s"""
-        if (!${evalChildren(i).isNull} && (${ev.isNull} ||
-          ${ctx.genComp(dataType, evalChildren(i).value, ev.value)} > 0)) {
+        ${eval.code}
+        if (!${eval.isNull} && (${ev.isNull} ||
+          ${ctx.genGreater(dataType, eval.value, ev.value)})) {
           ${ev.isNull} = false;
-          ${ev.value} = ${evalChildren(i).value};
+          ${ev.value} = ${eval.value};
         }
       """
     s"""
-      ${evalChildren.map(_.code).mkString("\n")}
-      boolean ${ev.isNull} = true;
-      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
-      ${children.indices.map(updateEval).mkString("\n")}
+      ${first.code}
+      boolean ${ev.isNull} = ${first.isNull};
+      ${ctx.javaType(dataType)} ${ev.value} = ${first.value};
+      ${rest.map(updateEval).mkString("\n")}
     """
   }
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala
index 94deafb75b69c..df4747d4e6f7a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala
@@ -62,11 +62,15 @@ case class Coalesce(children: Seq[Expression]) extends Expression {
   }
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
+    val first = children(0)
+    val rest = children.drop(1)
+    val firstEval = first.gen(ctx)
     s"""
-      boolean ${ev.isNull} = true;
-      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
+      ${firstEval.code}
+      boolean ${ev.isNull} = ${firstEval.isNull};
+      ${ctx.javaType(dataType)} ${ev.value} = ${firstEval.value};
     """ +
-    children.map { e =>
+      rest.map { e =>
       val eval = e.gen(ctx)
       s"""
         if (${ev.isNull}) {

From 7ee7d5a3c4ff77d2cee2afce36ff41f6302e6315 Mon Sep 17 00:00:00 2001
From: Dilip Biswal <dbiswal@us.ibm.com>
Date: Thu, 19 Nov 2015 19:46:10 -0800
Subject: [PATCH 374/518] [SPARK-11544][SQL][TEST-HADOOP1.0] sqlContext doesn't
 use PathFilter

Apply the user supplied pathfilter while retrieving the files from fs.

Author: Dilip Biswal <dbiswal@us.ibm.com>

Closes #9830 from dilipbiswal/spark-11544.
---
 .../apache/spark/sql/sources/interfaces.scala | 25 ++++++++---
 .../datasources/json/JsonSuite.scala          | 41 ++++++++++++++++++-
 2 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index b3d3bdf50df63..f9465157c936d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -21,7 +21,8 @@ import scala.collection.mutable
 import scala.util.Try
 
 import org.apache.hadoop.conf.Configuration
-import org.apache.hadoop.fs.{FileStatus, FileSystem, Path}
+import org.apache.hadoop.fs.{PathFilter, FileStatus, FileSystem, Path}
+import org.apache.hadoop.mapred.{JobConf, FileInputFormat}
 import org.apache.hadoop.mapreduce.{Job, TaskAttemptContext}
 
 import org.apache.spark.{Logging, SparkContext}
@@ -447,9 +448,15 @@ abstract class HadoopFsRelation private[sql](
           val hdfsPath = new Path(path)
           val fs = hdfsPath.getFileSystem(hadoopConf)
           val qualified = hdfsPath.makeQualified(fs.getUri, fs.getWorkingDirectory)
-
           logInfo(s"Listing $qualified on driver")
-          Try(fs.listStatus(qualified)).getOrElse(Array.empty)
+          // Dummy jobconf to get to the pathFilter defined in configuration
+          val jobConf = new JobConf(hadoopConf, this.getClass())
+          val pathFilter = FileInputFormat.getInputPathFilter(jobConf)
+          if (pathFilter != null) {
+            Try(fs.listStatus(qualified, pathFilter)).getOrElse(Array.empty)
+          } else {
+            Try(fs.listStatus(qualified)).getOrElse(Array.empty)
+          }
         }.filterNot { status =>
           val name = status.getPath.getName
           name.toLowerCase == "_temporary" || name.startsWith(".")
@@ -847,8 +854,16 @@ private[sql] object HadoopFsRelation extends Logging {
     if (name == "_temporary" || name.startsWith(".")) {
       Array.empty
     } else {
-      val (dirs, files) = fs.listStatus(status.getPath).partition(_.isDir)
-      files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
+      // Dummy jobconf to get to the pathFilter defined in configuration
+      val jobConf = new JobConf(fs.getConf, this.getClass())
+      val pathFilter = FileInputFormat.getInputPathFilter(jobConf)
+      if (pathFilter != null) {
+        val (dirs, files) = fs.listStatus(status.getPath, pathFilter).partition(_.isDir)
+        files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
+      } else {
+        val (dirs, files) = fs.listStatus(status.getPath).partition(_.isDir)
+        files ++ dirs.flatMap(dir => listLeafFiles(fs, dir))
+      }
     }
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
index 6042b1178affe..ba7718c864637 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
@@ -19,19 +19,27 @@ package org.apache.spark.sql.execution.datasources.json
 
 import java.io.{File, StringWriter}
 import java.sql.{Date, Timestamp}
+import scala.collection.JavaConverters._
 
 import com.fasterxml.jackson.core.JsonFactory
-import org.apache.spark.rdd.RDD
+import org.apache.commons.io.FileUtils
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{Path, PathFilter}
 import org.scalactic.Tolerance._
 
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
-import org.apache.spark.sql.execution.datasources.{ResolvedDataSource, LogicalRelation}
+import org.apache.spark.sql.execution.datasources.{LogicalRelation, ResolvedDataSource}
 import org.apache.spark.sql.execution.datasources.json.InferSchema.compatibleType
 import org.apache.spark.sql.test.SharedSQLContext
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
 
+class TestFileFilter extends PathFilter {
+  override def accept(path: Path): Boolean = path.getParent.getName != "p=2"
+}
+
 class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
   import testImplicits._
 
@@ -1390,4 +1398,33 @@ class JsonSuite extends QueryTest with SharedSQLContext with TestJsonData {
       )
     }
   }
+
+  test("SPARK-11544 test pathfilter") {
+    withTempPath { dir =>
+      val path = dir.getCanonicalPath
+
+      val df = sqlContext.range(2)
+      df.write.json(path + "/p=1")
+      df.write.json(path + "/p=2")
+      assert(sqlContext.read.json(path).count() === 4)
+
+      val clonedConf = new Configuration(hadoopConfiguration)
+      try {
+        // Setting it twice as the name of the propery has changed between hadoop versions.
+        hadoopConfiguration.setClass(
+          "mapred.input.pathFilter.class",
+          classOf[TestFileFilter],
+          classOf[PathFilter])
+        hadoopConfiguration.setClass(
+          "mapreduce.input.pathFilter.class",
+          classOf[TestFileFilter],
+          classOf[PathFilter])
+        assert(sqlContext.read.json(path).count() === 2)
+      } finally {
+        // Hadoop 1 doesn't have `Configuration.unset`
+        hadoopConfiguration.clear()
+        clonedConf.asScala.foreach(entry => hadoopConfiguration.set(entry.getKey, entry.getValue))
+      }
+    }
+  }
 }

From 4114ce20fbe820f111e55e891ae3889b0e6e0006 Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Thu, 19 Nov 2015 22:01:02 -0800
Subject: [PATCH 375/518] [SPARK-11846] Add save/load for AFTSurvivalRegression
 and IsotonicRegression

https://issues.apache.org/jira/browse/SPARK-11846

mengxr

Author: Xusen Yin <yinxusen@gmail.com>

Closes #9836 from yinxusen/SPARK-11846.
---
 .../ml/regression/AFTSurvivalRegression.scala | 78 +++++++++++++++--
 .../ml/regression/IsotonicRegression.scala    | 83 +++++++++++++++++--
 .../AFTSurvivalRegressionSuite.scala          | 37 ++++++++-
 .../regression/IsotonicRegressionSuite.scala  | 34 +++++++-
 4 files changed, 210 insertions(+), 22 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
index b7d095872ffa5..aedfb48058dc5 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
@@ -21,20 +21,20 @@ import scala.collection.mutable
 
 import breeze.linalg.{DenseVector => BDV}
 import breeze.optimize.{CachedDiffFunction, DiffFunction, LBFGS => BreezeLBFGS}
+import org.apache.hadoop.fs.Path
 
-import org.apache.spark.{SparkException, Logging}
-import org.apache.spark.annotation.{Since, Experimental}
-import org.apache.spark.ml.{Model, Estimator}
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.{SchemaUtils, Identifiable}
-import org.apache.spark.mllib.linalg.{Vector, Vectors, VectorUDT}
-import org.apache.spark.mllib.linalg.BLAS
+import org.apache.spark.ml.util._
+import org.apache.spark.ml.{Estimator, Model}
+import org.apache.spark.mllib.linalg.{BLAS, Vector, VectorUDT, Vectors}
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{Row, DataFrame}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.{DoubleType, StructType}
+import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.storage.StorageLevel
+import org.apache.spark.{Logging, SparkException}
 
 /**
  * Params for accelerated failure time (AFT) regression.
@@ -120,7 +120,8 @@ private[regression] trait AFTSurvivalRegressionParams extends Params
 @Experimental
 @Since("1.6.0")
 class AFTSurvivalRegression @Since("1.6.0") (@Since("1.6.0") override val uid: String)
-  extends Estimator[AFTSurvivalRegressionModel] with AFTSurvivalRegressionParams with Logging {
+  extends Estimator[AFTSurvivalRegressionModel] with AFTSurvivalRegressionParams
+  with DefaultParamsWritable with Logging {
 
   @Since("1.6.0")
   def this() = this(Identifiable.randomUID("aftSurvReg"))
@@ -243,6 +244,13 @@ class AFTSurvivalRegression @Since("1.6.0") (@Since("1.6.0") override val uid: S
   override def copy(extra: ParamMap): AFTSurvivalRegression = defaultCopy(extra)
 }
 
+@Since("1.6.0")
+object AFTSurvivalRegression extends DefaultParamsReadable[AFTSurvivalRegression] {
+
+  @Since("1.6.0")
+  override def load(path: String): AFTSurvivalRegression = super.load(path)
+}
+
 /**
  * :: Experimental ::
  * Model produced by [[AFTSurvivalRegression]].
@@ -254,7 +262,7 @@ class AFTSurvivalRegressionModel private[ml] (
     @Since("1.6.0") val coefficients: Vector,
     @Since("1.6.0") val intercept: Double,
     @Since("1.6.0") val scale: Double)
-  extends Model[AFTSurvivalRegressionModel] with AFTSurvivalRegressionParams {
+  extends Model[AFTSurvivalRegressionModel] with AFTSurvivalRegressionParams with MLWritable {
 
   /** @group setParam */
   @Since("1.6.0")
@@ -312,6 +320,58 @@ class AFTSurvivalRegressionModel private[ml] (
     copyValues(new AFTSurvivalRegressionModel(uid, coefficients, intercept, scale), extra)
       .setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: MLWriter =
+    new AFTSurvivalRegressionModel.AFTSurvivalRegressionModelWriter(this)
+}
+
+@Since("1.6.0")
+object AFTSurvivalRegressionModel extends MLReadable[AFTSurvivalRegressionModel] {
+
+  @Since("1.6.0")
+  override def read: MLReader[AFTSurvivalRegressionModel] = new AFTSurvivalRegressionModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): AFTSurvivalRegressionModel = super.load(path)
+
+  /** [[MLWriter]] instance for [[AFTSurvivalRegressionModel]] */
+  private[AFTSurvivalRegressionModel] class AFTSurvivalRegressionModelWriter (
+      instance: AFTSurvivalRegressionModel
+    ) extends MLWriter with Logging {
+
+    private case class Data(coefficients: Vector, intercept: Double, scale: Double)
+
+    override protected def saveImpl(path: String): Unit = {
+      // Save metadata and Params
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      // Save model data: coefficients, intercept, scale
+      val data = Data(instance.coefficients, instance.intercept, instance.scale)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class AFTSurvivalRegressionModelReader extends MLReader[AFTSurvivalRegressionModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = classOf[AFTSurvivalRegressionModel].getName
+
+    override def load(path: String): AFTSurvivalRegressionModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath)
+        .select("coefficients", "intercept", "scale").head()
+      val coefficients = data.getAs[Vector](0)
+      val intercept = data.getDouble(1)
+      val scale = data.getDouble(2)
+      val model = new AFTSurvivalRegressionModel(metadata.uid, coefficients, intercept, scale)
+
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
 }
 
 /**
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
index a1fe01b047108..bbb1c7ac0a51e 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
@@ -17,18 +17,22 @@
 
 package org.apache.spark.ml.regression
 
+import org.apache.hadoop.fs.Path
+
 import org.apache.spark.Logging
 import org.apache.spark.annotation.{Experimental, Since}
-import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.param._
-import org.apache.spark.ml.param.shared.{HasFeaturesCol, HasLabelCol, HasPredictionCol, HasWeightCol}
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.param.shared._
+import org.apache.spark.ml.regression.IsotonicRegressionModel.IsotonicRegressionModelWriter
+import org.apache.spark.ml.util._
+import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT, Vectors}
-import org.apache.spark.mllib.regression.{IsotonicRegression => MLlibIsotonicRegression, IsotonicRegressionModel => MLlibIsotonicRegressionModel}
+import org.apache.spark.mllib.regression.{IsotonicRegression => MLlibIsotonicRegression}
+import org.apache.spark.mllib.regression.{IsotonicRegressionModel => MLlibIsotonicRegressionModel}
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.sql.functions.{col, lit, udf}
 import org.apache.spark.sql.types.{DoubleType, StructType}
+import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.storage.StorageLevel
 
 /**
@@ -127,7 +131,8 @@ private[regression] trait IsotonicRegressionBase extends Params with HasFeatures
 @Since("1.5.0")
 @Experimental
 class IsotonicRegression @Since("1.5.0") (@Since("1.5.0") override val uid: String)
-  extends Estimator[IsotonicRegressionModel] with IsotonicRegressionBase {
+  extends Estimator[IsotonicRegressionModel]
+  with IsotonicRegressionBase with DefaultParamsWritable {
 
   @Since("1.5.0")
   def this() = this(Identifiable.randomUID("isoReg"))
@@ -179,6 +184,13 @@ class IsotonicRegression @Since("1.5.0") (@Since("1.5.0") override val uid: Stri
   }
 }
 
+@Since("1.6.0")
+object IsotonicRegression extends DefaultParamsReadable[IsotonicRegression] {
+
+  @Since("1.6.0")
+  override def load(path: String): IsotonicRegression = super.load(path)
+}
+
 /**
  * :: Experimental ::
  * Model fitted by IsotonicRegression.
@@ -194,7 +206,7 @@ class IsotonicRegression @Since("1.5.0") (@Since("1.5.0") override val uid: Stri
 class IsotonicRegressionModel private[ml] (
     override val uid: String,
     private val oldModel: MLlibIsotonicRegressionModel)
-  extends Model[IsotonicRegressionModel] with IsotonicRegressionBase {
+  extends Model[IsotonicRegressionModel] with IsotonicRegressionBase with MLWritable {
 
   /** @group setParam */
   @Since("1.5.0")
@@ -240,4 +252,61 @@ class IsotonicRegressionModel private[ml] (
   override def transformSchema(schema: StructType): StructType = {
     validateAndTransformSchema(schema, fitting = false)
   }
+
+  @Since("1.6.0")
+  override def write: MLWriter =
+    new IsotonicRegressionModelWriter(this)
+}
+
+@Since("1.6.0")
+object IsotonicRegressionModel extends MLReadable[IsotonicRegressionModel] {
+
+  @Since("1.6.0")
+  override def read: MLReader[IsotonicRegressionModel] = new IsotonicRegressionModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): IsotonicRegressionModel = super.load(path)
+
+  /** [[MLWriter]] instance for [[IsotonicRegressionModel]] */
+  private[IsotonicRegressionModel] class IsotonicRegressionModelWriter (
+      instance: IsotonicRegressionModel
+    ) extends MLWriter with Logging {
+
+    private case class Data(
+        boundaries: Array[Double],
+        predictions: Array[Double],
+        isotonic: Boolean)
+
+    override protected def saveImpl(path: String): Unit = {
+      // Save metadata and Params
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      // Save model data: boundaries, predictions, isotonic
+      val data = Data(
+        instance.oldModel.boundaries, instance.oldModel.predictions, instance.oldModel.isotonic)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class IsotonicRegressionModelReader extends MLReader[IsotonicRegressionModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = classOf[IsotonicRegressionModel].getName
+
+    override def load(path: String): IsotonicRegressionModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath)
+        .select("boundaries", "predictions", "isotonic").head()
+      val boundaries = data.getAs[Seq[Double]](0).toArray
+      val predictions = data.getAs[Seq[Double]](1).toArray
+      val isotonic = data.getBoolean(2)
+      val model = new IsotonicRegressionModel(
+        metadata.uid, new MLlibIsotonicRegressionModel(boundaries, predictions, isotonic))
+
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
index 359f31027172b..d718ef63b531a 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/AFTSurvivalRegressionSuite.scala
@@ -21,14 +21,15 @@ import scala.util.Random
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.random.{ExponentialGenerator, WeibullGenerator}
-import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.mllib.util.MLlibTestSparkContext
-import org.apache.spark.sql.{Row, DataFrame}
+import org.apache.spark.mllib.util.TestingUtils._
+import org.apache.spark.sql.{DataFrame, Row}
 
-class AFTSurvivalRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
+class AFTSurvivalRegressionSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   @transient var datasetUnivariate: DataFrame = _
   @transient var datasetMultivariate: DataFrame = _
@@ -332,4 +333,32 @@ class AFTSurvivalRegressionSuite extends SparkFunSuite with MLlibTestSparkContex
           assert(prediction ~== model.predict(features) relTol 1E-5)
     }
   }
+
+  test("read/write") {
+    def checkModelData(
+        model: AFTSurvivalRegressionModel,
+        model2: AFTSurvivalRegressionModel): Unit = {
+      assert(model.intercept === model2.intercept)
+      assert(model.coefficients === model2.coefficients)
+      assert(model.scale === model2.scale)
+    }
+    val aft = new AFTSurvivalRegression()
+    testEstimatorAndModelReadWrite(aft, datasetMultivariate,
+      AFTSurvivalRegressionSuite.allParamSettings, checkModelData)
+  }
+}
+
+object AFTSurvivalRegressionSuite {
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = Map(
+    "predictionCol" -> "myPrediction",
+    "fitIntercept" -> true,
+    "maxIter" -> 2,
+    "tol" -> 0.01
+  )
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/IsotonicRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/IsotonicRegressionSuite.scala
index 59f4193abc8f0..f067c29d27a7d 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/IsotonicRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/IsotonicRegressionSuite.scala
@@ -19,12 +19,14 @@ package org.apache.spark.ml.regression
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.Vectors
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql.{DataFrame, Row}
 
-class IsotonicRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
+class IsotonicRegressionSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
+
   private def generateIsotonicInput(labels: Seq[Double]): DataFrame = {
     sqlContext.createDataFrame(
       labels.zipWithIndex.map { case (label, i) => (label, i.toDouble, 1.0) }
@@ -164,4 +166,32 @@ class IsotonicRegressionSuite extends SparkFunSuite with MLlibTestSparkContext {
 
     assert(predictions === Array(3.5, 5.0, 5.0, 5.0))
   }
+
+  test("read/write") {
+    val dataset = generateIsotonicInput(Seq(1, 2, 3, 1, 6, 17, 16, 17, 18))
+
+    def checkModelData(model: IsotonicRegressionModel, model2: IsotonicRegressionModel): Unit = {
+      assert(model.boundaries === model2.boundaries)
+      assert(model.predictions === model2.predictions)
+      assert(model.isotonic === model2.isotonic)
+    }
+
+    val ir = new IsotonicRegression()
+    testEstimatorAndModelReadWrite(ir, dataset, IsotonicRegressionSuite.allParamSettings,
+      checkModelData)
+  }
+}
+
+object IsotonicRegressionSuite {
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = Map(
+    "predictionCol" -> "myPrediction",
+    "isotonic" -> true,
+    "featureIndex" -> 0
+  )
 }

From 3b7f056da87a23f3a96f0311b3a947a9b698f38b Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 19 Nov 2015 22:02:17 -0800
Subject: [PATCH 376/518] [SPARK-11829][ML] Add read/write to estimators under
 ml.feature (II)

Add read/write support to the following estimators under spark.ml:
* ChiSqSelector
* PCA
* VectorIndexer
* Word2Vec

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9838 from yanboliang/spark-11829.
---
 .../spark/ml/feature/ChiSqSelector.scala      | 65 ++++++++++++++++--
 .../org/apache/spark/ml/feature/PCA.scala     | 67 +++++++++++++++++--
 .../spark/ml/feature/VectorIndexer.scala      | 66 ++++++++++++++++--
 .../apache/spark/ml/feature/Word2Vec.scala    | 67 +++++++++++++++++--
 .../apache/spark/mllib/feature/Word2Vec.scala |  6 +-
 .../spark/ml/feature/ChiSqSelectorSuite.scala | 22 +++++-
 .../apache/spark/ml/feature/PCASuite.scala    | 26 ++++++-
 .../spark/ml/feature/VectorIndexerSuite.scala | 22 +++++-
 .../spark/ml/feature/Word2VecSuite.scala      | 30 ++++++++-
 9 files changed, 338 insertions(+), 33 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/ChiSqSelector.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/ChiSqSelector.scala
index 5e4061fba5494..dfec03828f4b7 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/ChiSqSelector.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/ChiSqSelector.scala
@@ -17,13 +17,14 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml._
 import org.apache.spark.ml.attribute.{AttributeGroup, _}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.Identifiable
-import org.apache.spark.ml.util.SchemaUtils
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.feature
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT}
 import org.apache.spark.mllib.regression.LabeledPoint
@@ -60,7 +61,7 @@ private[feature] trait ChiSqSelectorParams extends Params
  */
 @Experimental
 final class ChiSqSelector(override val uid: String)
-  extends Estimator[ChiSqSelectorModel] with ChiSqSelectorParams {
+  extends Estimator[ChiSqSelectorModel] with ChiSqSelectorParams with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("chiSqSelector"))
 
@@ -95,6 +96,13 @@ final class ChiSqSelector(override val uid: String)
   override def copy(extra: ParamMap): ChiSqSelector = defaultCopy(extra)
 }
 
+@Since("1.6.0")
+object ChiSqSelector extends DefaultParamsReadable[ChiSqSelector] {
+
+  @Since("1.6.0")
+  override def load(path: String): ChiSqSelector = super.load(path)
+}
+
 /**
  * :: Experimental ::
  * Model fitted by [[ChiSqSelector]].
@@ -103,7 +111,12 @@ final class ChiSqSelector(override val uid: String)
 final class ChiSqSelectorModel private[ml] (
     override val uid: String,
     private val chiSqSelector: feature.ChiSqSelectorModel)
-  extends Model[ChiSqSelectorModel] with ChiSqSelectorParams {
+  extends Model[ChiSqSelectorModel] with ChiSqSelectorParams with MLWritable {
+
+  import ChiSqSelectorModel._
+
+  /** list of indices to select (filter). Must be ordered asc */
+  val selectedFeatures: Array[Int] = chiSqSelector.selectedFeatures
 
   /** @group setParam */
   def setFeaturesCol(value: String): this.type = set(featuresCol, value)
@@ -147,4 +160,46 @@ final class ChiSqSelectorModel private[ml] (
     val copied = new ChiSqSelectorModel(uid, chiSqSelector)
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: MLWriter = new ChiSqSelectorModelWriter(this)
+}
+
+@Since("1.6.0")
+object ChiSqSelectorModel extends MLReadable[ChiSqSelectorModel] {
+
+  private[ChiSqSelectorModel]
+  class ChiSqSelectorModelWriter(instance: ChiSqSelectorModel) extends MLWriter {
+
+    private case class Data(selectedFeatures: Seq[Int])
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = Data(instance.selectedFeatures.toSeq)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class ChiSqSelectorModelReader extends MLReader[ChiSqSelectorModel] {
+
+    private val className = classOf[ChiSqSelectorModel].getName
+
+    override def load(path: String): ChiSqSelectorModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath).select("selectedFeatures").head()
+      val selectedFeatures = data.getAs[Seq[Int]](0).toArray
+      val oldModel = new feature.ChiSqSelectorModel(selectedFeatures)
+      val model = new ChiSqSelectorModel(metadata.uid, oldModel)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: MLReader[ChiSqSelectorModel] = new ChiSqSelectorModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): ChiSqSelectorModel = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
index 539084704b653..32d7afee6e73b 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
@@ -17,13 +17,15 @@
 
 package org.apache.spark.ml.feature
 
-import org.apache.spark.annotation.Experimental
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml._
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.feature
-import org.apache.spark.mllib.linalg.{Vector, VectorUDT}
+import org.apache.spark.mllib.linalg._
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.{StructField, StructType}
@@ -49,7 +51,8 @@ private[feature] trait PCAParams extends Params with HasInputCol with HasOutputC
  * PCA trains a model to project vectors to a low-dimensional space using PCA.
  */
 @Experimental
-class PCA (override val uid: String) extends Estimator[PCAModel] with PCAParams {
+class PCA (override val uid: String) extends Estimator[PCAModel] with PCAParams
+  with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("pca"))
 
@@ -86,6 +89,13 @@ class PCA (override val uid: String) extends Estimator[PCAModel] with PCAParams
   override def copy(extra: ParamMap): PCA = defaultCopy(extra)
 }
 
+@Since("1.6.0")
+object PCA extends DefaultParamsReadable[PCA] {
+
+  @Since("1.6.0")
+  override def load(path: String): PCA = super.load(path)
+}
+
 /**
  * :: Experimental ::
  * Model fitted by [[PCA]].
@@ -94,7 +104,12 @@ class PCA (override val uid: String) extends Estimator[PCAModel] with PCAParams
 class PCAModel private[ml] (
     override val uid: String,
     pcaModel: feature.PCAModel)
-  extends Model[PCAModel] with PCAParams {
+  extends Model[PCAModel] with PCAParams with MLWritable {
+
+  import PCAModel._
+
+  /** a principal components Matrix. Each column is one principal component. */
+  val pc: DenseMatrix = pcaModel.pc
 
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
@@ -127,4 +142,46 @@ class PCAModel private[ml] (
     val copied = new PCAModel(uid, pcaModel)
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: MLWriter = new PCAModelWriter(this)
+}
+
+@Since("1.6.0")
+object PCAModel extends MLReadable[PCAModel] {
+
+  private[PCAModel] class PCAModelWriter(instance: PCAModel) extends MLWriter {
+
+    private case class Data(k: Int, pc: DenseMatrix)
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = Data(instance.getK, instance.pc)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class PCAModelReader extends MLReader[PCAModel] {
+
+    private val className = classOf[PCAModel].getName
+
+    override def load(path: String): PCAModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val Row(k: Int, pc: DenseMatrix) = sqlContext.read.parquet(dataPath)
+        .select("k", "pc")
+        .head()
+      val oldModel = new feature.PCAModel(k, pc)
+      val model = new PCAModel(metadata.uid, oldModel)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: MLReader[PCAModel] = new PCAModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): PCAModel = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala
index 52e0599e38d83..a637a6f2881de 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala
@@ -22,12 +22,14 @@ import java.util.{Map => JMap}
 
 import scala.collection.JavaConverters._
 
-import org.apache.spark.annotation.Experimental
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.attribute._
-import org.apache.spark.ml.param.{IntParam, ParamMap, ParamValidators, Params}
+import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector, VectorUDT}
 import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.sql.functions.udf
@@ -93,7 +95,7 @@ private[ml] trait VectorIndexerParams extends Params with HasInputCol with HasOu
  */
 @Experimental
 class VectorIndexer(override val uid: String) extends Estimator[VectorIndexerModel]
-  with VectorIndexerParams {
+  with VectorIndexerParams with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("vecIdx"))
 
@@ -136,7 +138,11 @@ class VectorIndexer(override val uid: String) extends Estimator[VectorIndexerMod
   override def copy(extra: ParamMap): VectorIndexer = defaultCopy(extra)
 }
 
-private object VectorIndexer {
+@Since("1.6.0")
+object VectorIndexer extends DefaultParamsReadable[VectorIndexer] {
+
+  @Since("1.6.0")
+  override def load(path: String): VectorIndexer = super.load(path)
 
   /**
    * Helper class for tracking unique values for each feature.
@@ -146,7 +152,7 @@ private object VectorIndexer {
    * @param numFeatures  This class fails if it encounters a Vector whose length is not numFeatures.
    * @param maxCategories  This class caps the number of unique values collected at maxCategories.
    */
-  class CategoryStats(private val numFeatures: Int, private val maxCategories: Int)
+  private class CategoryStats(private val numFeatures: Int, private val maxCategories: Int)
     extends Serializable {
 
     /** featureValueSets[feature index] = set of unique values */
@@ -252,7 +258,9 @@ class VectorIndexerModel private[ml] (
     override val uid: String,
     val numFeatures: Int,
     val categoryMaps: Map[Int, Map[Double, Int]])
-  extends Model[VectorIndexerModel] with VectorIndexerParams {
+  extends Model[VectorIndexerModel] with VectorIndexerParams with MLWritable {
+
+  import VectorIndexerModel._
 
   /** Java-friendly version of [[categoryMaps]] */
   def javaCategoryMaps: JMap[JInt, JMap[JDouble, JInt]] = {
@@ -408,4 +416,48 @@ class VectorIndexerModel private[ml] (
     val copied = new VectorIndexerModel(uid, numFeatures, categoryMaps)
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: MLWriter = new VectorIndexerModelWriter(this)
+}
+
+@Since("1.6.0")
+object VectorIndexerModel extends MLReadable[VectorIndexerModel] {
+
+  private[VectorIndexerModel]
+  class VectorIndexerModelWriter(instance: VectorIndexerModel) extends MLWriter {
+
+    private case class Data(numFeatures: Int, categoryMaps: Map[Int, Map[Double, Int]])
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = Data(instance.numFeatures, instance.categoryMaps)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class VectorIndexerModelReader extends MLReader[VectorIndexerModel] {
+
+    private val className = classOf[VectorIndexerModel].getName
+
+    override def load(path: String): VectorIndexerModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath)
+        .select("numFeatures", "categoryMaps")
+        .head()
+      val numFeatures = data.getAs[Int](0)
+      val categoryMaps = data.getAs[Map[Int, Map[Double, Int]]](1)
+      val model = new VectorIndexerModel(metadata.uid, numFeatures, categoryMaps)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: MLReader[VectorIndexerModel] = new VectorIndexerModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): VectorIndexerModel = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
index 708dbeef84db4..a8d61b6dea00b 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala
@@ -17,15 +17,17 @@
 
 package org.apache.spark.ml.feature
 
+import org.apache.hadoop.fs.Path
+
 import org.apache.spark.SparkContext
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.feature
 import org.apache.spark.mllib.linalg.{BLAS, Vector, VectorUDT, Vectors}
-import org.apache.spark.sql.{DataFrame, SQLContext}
+import org.apache.spark.sql.{DataFrame, Row, SQLContext}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
 
@@ -90,7 +92,8 @@ private[feature] trait Word2VecBase extends Params
  * natural language processing or machine learning process.
  */
 @Experimental
-final class Word2Vec(override val uid: String) extends Estimator[Word2VecModel] with Word2VecBase {
+final class Word2Vec(override val uid: String) extends Estimator[Word2VecModel] with Word2VecBase
+  with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("w2v"))
 
@@ -139,6 +142,13 @@ final class Word2Vec(override val uid: String) extends Estimator[Word2VecModel]
   override def copy(extra: ParamMap): Word2Vec = defaultCopy(extra)
 }
 
+@Since("1.6.0")
+object Word2Vec extends DefaultParamsReadable[Word2Vec] {
+
+  @Since("1.6.0")
+  override def load(path: String): Word2Vec = super.load(path)
+}
+
 /**
  * :: Experimental ::
  * Model fitted by [[Word2Vec]].
@@ -147,7 +157,9 @@ final class Word2Vec(override val uid: String) extends Estimator[Word2VecModel]
 class Word2VecModel private[ml] (
     override val uid: String,
     @transient private val wordVectors: feature.Word2VecModel)
-  extends Model[Word2VecModel] with Word2VecBase {
+  extends Model[Word2VecModel] with Word2VecBase with MLWritable {
+
+  import Word2VecModel._
 
   /**
    * Returns a dataframe with two fields, "word" and "vector", with "word" being a String and
@@ -224,4 +236,49 @@ class Word2VecModel private[ml] (
     val copied = new Word2VecModel(uid, wordVectors)
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: MLWriter = new Word2VecModelWriter(this)
+}
+
+@Since("1.6.0")
+object Word2VecModel extends MLReadable[Word2VecModel] {
+
+  private[Word2VecModel]
+  class Word2VecModelWriter(instance: Word2VecModel) extends MLWriter {
+
+    private case class Data(wordIndex: Map[String, Int], wordVectors: Seq[Float])
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val data = Data(instance.wordVectors.wordIndex, instance.wordVectors.wordVectors.toSeq)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class Word2VecModelReader extends MLReader[Word2VecModel] {
+
+    private val className = classOf[Word2VecModel].getName
+
+    override def load(path: String): Word2VecModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath)
+        .select("wordIndex", "wordVectors")
+        .head()
+      val wordIndex = data.getAs[Map[String, Int]](0)
+      val wordVectors = data.getAs[Seq[Float]](1).toArray
+      val oldModel = new feature.Word2VecModel(wordIndex, wordVectors)
+      val model = new Word2VecModel(metadata.uid, oldModel)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: MLReader[Word2VecModel] = new Word2VecModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): Word2VecModel = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
index 7ab0d89d23a3f..a47f27b0afb14 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
@@ -432,9 +432,9 @@ class Word2Vec extends Serializable with Logging {
  *                    (i * vectorSize, i * vectorSize + vectorSize)
  */
 @Since("1.1.0")
-class Word2VecModel private[mllib] (
-    private val wordIndex: Map[String, Int],
-    private val wordVectors: Array[Float]) extends Serializable with Saveable {
+class Word2VecModel private[spark] (
+    private[spark] val wordIndex: Map[String, Int],
+    private[spark] val wordVectors: Array[Float]) extends Serializable with Saveable {
 
   private val numWords = wordIndex.size
   // vectorSize: Dimension of each word's vector.
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala
index e5a42967bd2c8..7827db2794cf3 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala
@@ -18,13 +18,17 @@
 package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
+import org.apache.spark.mllib.feature
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.regression.LabeledPoint
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.{Row, SQLContext}
 
-class ChiSqSelectorSuite extends SparkFunSuite with MLlibTestSparkContext {
+class ChiSqSelectorSuite extends SparkFunSuite with MLlibTestSparkContext
+  with DefaultReadWriteTest {
+
   test("Test Chi-Square selector") {
     val sqlContext = SQLContext.getOrCreate(sc)
     import sqlContext.implicits._
@@ -58,4 +62,20 @@ class ChiSqSelectorSuite extends SparkFunSuite with MLlibTestSparkContext {
         assert(vec1 ~== vec2 absTol 1e-1)
     }
   }
+
+  test("ChiSqSelector read/write") {
+    val t = new ChiSqSelector()
+      .setFeaturesCol("myFeaturesCol")
+      .setLabelCol("myLabelCol")
+      .setOutputCol("myOutputCol")
+      .setNumTopFeatures(2)
+    testDefaultReadWrite(t)
+  }
+
+  test("ChiSqSelectorModel read/write") {
+    val oldModel = new feature.ChiSqSelectorModel(Array(1, 3))
+    val instance = new ChiSqSelectorModel("myChiSqSelectorModel", oldModel)
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.selectedFeatures === instance.selectedFeatures)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
index 30c500f87a769..5a21cd20ceede 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
@@ -19,15 +19,15 @@ package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.distributed.RowMatrix
-import org.apache.spark.mllib.linalg.{Vector, Vectors, DenseMatrix, Matrices}
+import org.apache.spark.mllib.linalg._
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.mllib.feature.{PCAModel => OldPCAModel}
 import org.apache.spark.sql.Row
 
-class PCASuite extends SparkFunSuite with MLlibTestSparkContext {
+class PCASuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new PCA)
@@ -65,4 +65,24 @@ class PCASuite extends SparkFunSuite with MLlibTestSparkContext {
         assert(x ~== y absTol 1e-5, "Transformed vector is different with expected vector.")
     }
   }
+
+  test("read/write") {
+
+    def checkModelData(model1: PCAModel, model2: PCAModel): Unit = {
+      assert(model1.pc === model2.pc)
+    }
+    val allParams: Map[String, Any] = Map(
+      "k" -> 3,
+      "inputCol" -> "features",
+      "outputCol" -> "pca_features"
+    )
+    val data = Seq(
+      (0.0, Vectors.sparse(5, Seq((1, 1.0), (3, 7.0)))),
+      (1.0, Vectors.dense(2.0, 0.0, 3.0, 4.0, 5.0)),
+      (2.0, Vectors.dense(4.0, 0.0, 0.0, 6.0, 7.0))
+    )
+    val df = sqlContext.createDataFrame(data).toDF("id", "features")
+    val pca = new PCA().setK(3)
+    testEstimatorAndModelReadWrite(pca, df, allParams, checkModelData)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala
index 8cb0a2cf14d37..67817fa4baf56 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala
@@ -22,13 +22,14 @@ import scala.beans.{BeanInfo, BeanProperty}
 import org.apache.spark.{Logging, SparkException, SparkFunSuite}
 import org.apache.spark.ml.attribute._
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.{SparseVector, Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.DataFrame
 
-class VectorIndexerSuite extends SparkFunSuite with MLlibTestSparkContext with Logging {
+class VectorIndexerSuite extends SparkFunSuite with MLlibTestSparkContext
+  with DefaultReadWriteTest with Logging {
 
   import VectorIndexerSuite.FeatureData
 
@@ -251,6 +252,23 @@ class VectorIndexerSuite extends SparkFunSuite with MLlibTestSparkContext with L
       }
     }
   }
+
+  test("VectorIndexer read/write") {
+    val t = new VectorIndexer()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setMaxCategories(30)
+    testDefaultReadWrite(t)
+  }
+
+  test("VectorIndexerModel read/write") {
+    val categoryMaps = Map(0 -> Map(0.0 -> 0, 1.0 -> 1), 1 -> Map(0.0 -> 0, 1.0 -> 1,
+      2.0 -> 2, 3.0 -> 3), 2 -> Map(0.0 -> 0, -1.0 -> 1, 2.0 -> 2))
+    val instance = new VectorIndexerModel("myVectorIndexerModel", 3, categoryMaps)
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.numFeatures === instance.numFeatures)
+    assert(newInstance.categoryMaps === instance.categoryMaps)
+  }
 }
 
 private[feature] object VectorIndexerSuite {
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
index 23dfdaa9f8fc6..a773244cd735e 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala
@@ -19,14 +19,14 @@ package org.apache.spark.ml.feature
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.{Row, SQLContext}
 import org.apache.spark.mllib.feature.{Word2VecModel => OldWord2VecModel}
 
-class Word2VecSuite extends SparkFunSuite with MLlibTestSparkContext {
+class Word2VecSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new Word2Vec)
@@ -143,5 +143,31 @@ class Word2VecSuite extends SparkFunSuite with MLlibTestSparkContext {
     }
 
   }
+
+  test("Word2Vec read/write") {
+    val t = new Word2Vec()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setMaxIter(2)
+      .setMinCount(8)
+      .setNumPartitions(1)
+      .setSeed(42L)
+      .setStepSize(0.01)
+      .setVectorSize(100)
+    testDefaultReadWrite(t)
+  }
+
+  test("Word2VecModel read/write") {
+    val word2VecMap = Map(
+      ("china", Array(0.50f, 0.50f, 0.50f, 0.50f)),
+      ("japan", Array(0.40f, 0.50f, 0.50f, 0.50f)),
+      ("taiwan", Array(0.60f, 0.50f, 0.50f, 0.50f)),
+      ("korea", Array(0.45f, 0.60f, 0.60f, 0.60f))
+    )
+    val oldModel = new OldWord2VecModel(word2VecMap)
+    val instance = new Word2VecModel("myWord2VecModel", oldModel)
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.getVectors.collect() === instance.getVectors.collect())
+  }
 }
 

From 7216f405454f6f3557b5b1f72df8f393605faf60 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 19 Nov 2015 22:14:01 -0800
Subject: [PATCH 377/518] [SPARK-11875][ML][PYSPARK] Update doc for PySpark
 HasCheckpointInterval

* Update doc for PySpark ```HasCheckpointInterval``` that users can understand how to disable checkpoint.
* Update doc for PySpark ```cacheNodeIds``` of ```DecisionTreeParams``` to notify the relationship between ```cacheNodeIds``` and ```checkpointInterval```.

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9856 from yanboliang/spark-11875.
---
 python/pyspark/ml/param/_shared_params_code_gen.py |  6 ++++--
 python/pyspark/ml/param/shared.py                  | 14 +++++++-------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/python/pyspark/ml/param/_shared_params_code_gen.py b/python/pyspark/ml/param/_shared_params_code_gen.py
index 070c5db01ae73..0528dc1e3a6b9 100644
--- a/python/pyspark/ml/param/_shared_params_code_gen.py
+++ b/python/pyspark/ml/param/_shared_params_code_gen.py
@@ -118,7 +118,8 @@ def get$Name(self):
         ("inputCols", "input column names.", None),
         ("outputCol", "output column name.", "self.uid + '__output'"),
         ("numFeatures", "number of features.", None),
-        ("checkpointInterval", "checkpoint interval (>= 1).", None),
+        ("checkpointInterval", "set checkpoint interval (>= 1) or disable checkpoint (-1). " +
+         "E.g. 10 means that the cache will get checkpointed every 10 iterations.", None),
         ("seed", "random seed.", "hash(type(self).__name__)"),
         ("tol", "the convergence tolerance for iterative algorithms.", None),
         ("stepSize", "Step size to be used for each iteration of optimization.", None),
@@ -157,7 +158,8 @@ def get$Name(self):
         ("maxMemoryInMB", "Maximum memory in MB allocated to histogram aggregation."),
         ("cacheNodeIds", "If false, the algorithm will pass trees to executors to match " +
          "instances with nodes. If true, the algorithm will cache node IDs for each instance. " +
-         "Caching can speed up training of deeper trees.")]
+         "Caching can speed up training of deeper trees. Users can set how often should the " +
+         "cache be checkpointed or disable it by setting checkpointInterval.")]
 
     decisionTreeCode = '''class DecisionTreeParams(Params):
     """
diff --git a/python/pyspark/ml/param/shared.py b/python/pyspark/ml/param/shared.py
index 4bdf2a8cc563f..4d960801502c2 100644
--- a/python/pyspark/ml/param/shared.py
+++ b/python/pyspark/ml/param/shared.py
@@ -325,16 +325,16 @@ def getNumFeatures(self):
 
 class HasCheckpointInterval(Params):
     """
-    Mixin for param checkpointInterval: checkpoint interval (>= 1).
+    Mixin for param checkpointInterval: set checkpoint interval (>= 1) or disable checkpoint (-1). E.g. 10 means that the cache will get checkpointed every 10 iterations.
     """
 
     # a placeholder to make it appear in the generated doc
-    checkpointInterval = Param(Params._dummy(), "checkpointInterval", "checkpoint interval (>= 1).")
+    checkpointInterval = Param(Params._dummy(), "checkpointInterval", "set checkpoint interval (>= 1) or disable checkpoint (-1). E.g. 10 means that the cache will get checkpointed every 10 iterations.")
 
     def __init__(self):
         super(HasCheckpointInterval, self).__init__()
-        #: param for checkpoint interval (>= 1).
-        self.checkpointInterval = Param(self, "checkpointInterval", "checkpoint interval (>= 1).")
+        #: param for set checkpoint interval (>= 1) or disable checkpoint (-1). E.g. 10 means that the cache will get checkpointed every 10 iterations.
+        self.checkpointInterval = Param(self, "checkpointInterval", "set checkpoint interval (>= 1) or disable checkpoint (-1). E.g. 10 means that the cache will get checkpointed every 10 iterations.")
 
     def setCheckpointInterval(self, value):
         """
@@ -636,7 +636,7 @@ class DecisionTreeParams(Params):
     minInstancesPerNode = Param(Params._dummy(), "minInstancesPerNode", "Minimum number of instances each child must have after split. If a split causes the left or right child to have fewer than minInstancesPerNode, the split will be discarded as invalid. Should be >= 1.")
     minInfoGain = Param(Params._dummy(), "minInfoGain", "Minimum information gain for a split to be considered at a tree node.")
     maxMemoryInMB = Param(Params._dummy(), "maxMemoryInMB", "Maximum memory in MB allocated to histogram aggregation.")
-    cacheNodeIds = Param(Params._dummy(), "cacheNodeIds", "If false, the algorithm will pass trees to executors to match instances with nodes. If true, the algorithm will cache node IDs for each instance. Caching can speed up training of deeper trees.")
+    cacheNodeIds = Param(Params._dummy(), "cacheNodeIds", "If false, the algorithm will pass trees to executors to match instances with nodes. If true, the algorithm will cache node IDs for each instance. Caching can speed up training of deeper trees. Users can set how often should the cache be checkpointed or disable it by setting checkpointInterval.")
     
 
     def __init__(self):
@@ -651,8 +651,8 @@ def __init__(self):
         self.minInfoGain = Param(self, "minInfoGain", "Minimum information gain for a split to be considered at a tree node.")
         #: param for Maximum memory in MB allocated to histogram aggregation.
         self.maxMemoryInMB = Param(self, "maxMemoryInMB", "Maximum memory in MB allocated to histogram aggregation.")
-        #: param for If false, the algorithm will pass trees to executors to match instances with nodes. If true, the algorithm will cache node IDs for each instance. Caching can speed up training of deeper trees.
-        self.cacheNodeIds = Param(self, "cacheNodeIds", "If false, the algorithm will pass trees to executors to match instances with nodes. If true, the algorithm will cache node IDs for each instance. Caching can speed up training of deeper trees.")
+        #: param for If false, the algorithm will pass trees to executors to match instances with nodes. If true, the algorithm will cache node IDs for each instance. Caching can speed up training of deeper trees. Users can set how often should the cache be checkpointed or disable it by setting checkpointInterval.
+        self.cacheNodeIds = Param(self, "cacheNodeIds", "If false, the algorithm will pass trees to executors to match instances with nodes. If true, the algorithm will cache node IDs for each instance. Caching can speed up training of deeper trees. Users can set how often should the cache be checkpointed or disable it by setting checkpointInterval.")
         
     def setMaxDepth(self, value):
         """

From 0fff8eb3e476165461658d4e16682ec64269fdfe Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Thu, 19 Nov 2015 23:42:24 -0800
Subject: [PATCH 378/518] [SPARK-11869][ML] Clean up TempDirectory properly in
 ML tests

Need to remove parent directory (```className```) rather than just tempDir (```className/random_name```)

I tested this with IDFSuite, which has 2 read/write tests, and it fixes the problem.

CC: mengxr  Can you confirm this is fine?  I believe it is since the same ```random_name``` is used for all tests in a suite; we basically have an extra unneeded level of nesting.

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9851 from jkbradley/tempdir-cleanup.
---
 .../src/test/scala/org/apache/spark/ml/util/TempDirectory.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mllib/src/test/scala/org/apache/spark/ml/util/TempDirectory.scala b/mllib/src/test/scala/org/apache/spark/ml/util/TempDirectory.scala
index 2742026a69c2e..c8a0bb16247b4 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/util/TempDirectory.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/util/TempDirectory.scala
@@ -35,7 +35,7 @@ trait TempDirectory extends BeforeAndAfterAll { self: Suite =>
 
   override def beforeAll(): Unit = {
     super.beforeAll()
-    _tempDir = Utils.createTempDir(this.getClass.getName)
+    _tempDir = Utils.createTempDir(namePrefix = this.getClass.getName)
   }
 
   override def afterAll(): Unit = {

From 3e1d120cedb4bd9e1595e95d4d531cf61da6684d Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Thu, 19 Nov 2015 23:43:18 -0800
Subject: [PATCH 379/518] [SPARK-11867] Add save/load for kmeans and naive
 bayes

https://issues.apache.org/jira/browse/SPARK-11867

Author: Xusen Yin <yinxusen@gmail.com>

Closes #9849 from yinxusen/SPARK-11867.
---
 .../spark/ml/classification/NaiveBayes.scala  | 68 +++++++++++++++++--
 .../apache/spark/ml/clustering/KMeans.scala   | 67 ++++++++++++++++--
 .../ml/classification/NaiveBayesSuite.scala   | 47 +++++++++++--
 .../spark/ml/clustering/KMeansSuite.scala     | 41 ++++++++---
 4 files changed, 195 insertions(+), 28 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
index a14dcecbaf5b9..c512a2cb8bf3d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
@@ -17,12 +17,15 @@
 
 package org.apache.spark.ml.classification
 
+import org.apache.hadoop.fs.Path
+
 import org.apache.spark.SparkException
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.PredictorParams
 import org.apache.spark.ml.param.{DoubleParam, Param, ParamMap, ParamValidators}
-import org.apache.spark.ml.util.Identifiable
-import org.apache.spark.mllib.classification.{NaiveBayes => OldNaiveBayes, NaiveBayesModel => OldNaiveBayesModel}
+import org.apache.spark.ml.util._
+import org.apache.spark.mllib.classification.{NaiveBayes => OldNaiveBayes}
+import org.apache.spark.mllib.classification.{NaiveBayesModel => OldNaiveBayesModel}
 import org.apache.spark.mllib.linalg._
 import org.apache.spark.mllib.regression.LabeledPoint
 import org.apache.spark.rdd.RDD
@@ -72,7 +75,7 @@ private[ml] trait NaiveBayesParams extends PredictorParams {
 @Experimental
 class NaiveBayes(override val uid: String)
   extends ProbabilisticClassifier[Vector, NaiveBayes, NaiveBayesModel]
-  with NaiveBayesParams {
+  with NaiveBayesParams with DefaultParamsWritable {
 
   def this() = this(Identifiable.randomUID("nb"))
 
@@ -102,6 +105,13 @@ class NaiveBayes(override val uid: String)
   override def copy(extra: ParamMap): NaiveBayes = defaultCopy(extra)
 }
 
+@Since("1.6.0")
+object NaiveBayes extends DefaultParamsReadable[NaiveBayes] {
+
+  @Since("1.6.0")
+  override def load(path: String): NaiveBayes = super.load(path)
+}
+
 /**
  * :: Experimental ::
  * Model produced by [[NaiveBayes]]
@@ -114,7 +124,8 @@ class NaiveBayesModel private[ml] (
     override val uid: String,
     val pi: Vector,
     val theta: Matrix)
-  extends ProbabilisticClassificationModel[Vector, NaiveBayesModel] with NaiveBayesParams {
+  extends ProbabilisticClassificationModel[Vector, NaiveBayesModel]
+  with NaiveBayesParams with MLWritable {
 
   import OldNaiveBayes.{Bernoulli, Multinomial}
 
@@ -203,12 +214,15 @@ class NaiveBayesModel private[ml] (
     s"NaiveBayesModel (uid=$uid) with ${pi.size} classes"
   }
 
+  @Since("1.6.0")
+  override def write: MLWriter = new NaiveBayesModel.NaiveBayesModelWriter(this)
 }
 
-private[ml] object NaiveBayesModel {
+@Since("1.6.0")
+object NaiveBayesModel extends MLReadable[NaiveBayesModel] {
 
   /** Convert a model from the old API */
-  def fromOld(
+  private[ml] def fromOld(
       oldModel: OldNaiveBayesModel,
       parent: NaiveBayes): NaiveBayesModel = {
     val uid = if (parent != null) parent.uid else Identifiable.randomUID("nb")
@@ -218,4 +232,44 @@ private[ml] object NaiveBayesModel {
       oldModel.theta.flatten, true)
     new NaiveBayesModel(uid, pi, theta)
   }
+
+  @Since("1.6.0")
+  override def read: MLReader[NaiveBayesModel] = new NaiveBayesModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): NaiveBayesModel = super.load(path)
+
+  /** [[MLWriter]] instance for [[NaiveBayesModel]] */
+  private[NaiveBayesModel] class NaiveBayesModelWriter(instance: NaiveBayesModel) extends MLWriter {
+
+    private case class Data(pi: Vector, theta: Matrix)
+
+    override protected def saveImpl(path: String): Unit = {
+      // Save metadata and Params
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      // Save model data: pi, theta
+      val data = Data(instance.pi, instance.theta)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class NaiveBayesModelReader extends MLReader[NaiveBayesModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = classOf[NaiveBayesModel].getName
+
+    override def load(path: String): NaiveBayesModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath).select("pi", "theta").head()
+      val pi = data.getAs[Vector](0)
+      val theta = data.getAs[Matrix](1)
+      val model = new NaiveBayesModel(metadata.uid, pi, theta)
+
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
index 509be63002396..71e968497500f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
@@ -17,10 +17,12 @@
 
 package org.apache.spark.ml.clustering
 
-import org.apache.spark.annotation.{Since, Experimental}
-import org.apache.spark.ml.param.{Param, Params, IntParam, ParamMap}
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.param.{IntParam, Param, ParamMap, Params}
+import org.apache.spark.ml.util._
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.mllib.clustering.{KMeans => MLlibKMeans, KMeansModel => MLlibKMeansModel}
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT}
@@ -28,7 +30,6 @@ import org.apache.spark.sql.functions.{col, udf}
 import org.apache.spark.sql.types.{IntegerType, StructType}
 import org.apache.spark.sql.{DataFrame, Row}
 
-
 /**
  * Common params for KMeans and KMeansModel
  */
@@ -94,7 +95,8 @@ private[clustering] trait KMeansParams extends Params with HasMaxIter with HasFe
 @Experimental
 class KMeansModel private[ml] (
     @Since("1.5.0") override val uid: String,
-    private val parentModel: MLlibKMeansModel) extends Model[KMeansModel] with KMeansParams {
+    private val parentModel: MLlibKMeansModel)
+  extends Model[KMeansModel] with KMeansParams with MLWritable {
 
   @Since("1.5.0")
   override def copy(extra: ParamMap): KMeansModel = {
@@ -129,6 +131,52 @@ class KMeansModel private[ml] (
     val data = dataset.select(col($(featuresCol))).map { case Row(point: Vector) => point }
     parentModel.computeCost(data)
   }
+
+  @Since("1.6.0")
+  override def write: MLWriter = new KMeansModel.KMeansModelWriter(this)
+}
+
+@Since("1.6.0")
+object KMeansModel extends MLReadable[KMeansModel] {
+
+  @Since("1.6.0")
+  override def read: MLReader[KMeansModel] = new KMeansModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): KMeansModel = super.load(path)
+
+  /** [[MLWriter]] instance for [[KMeansModel]] */
+  private[KMeansModel] class KMeansModelWriter(instance: KMeansModel) extends MLWriter {
+
+    private case class Data(clusterCenters: Array[Vector])
+
+    override protected def saveImpl(path: String): Unit = {
+      // Save metadata and Params
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      // Save model data: cluster centers
+      val data = Data(instance.clusterCenters)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class KMeansModelReader extends MLReader[KMeansModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = classOf[KMeansModel].getName
+
+    override def load(path: String): KMeansModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath).select("clusterCenters").head()
+      val clusterCenters = data.getAs[Seq[Vector]](0).toArray
+      val model = new KMeansModel(metadata.uid, new MLlibKMeansModel(clusterCenters))
+
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
 }
 
 /**
@@ -141,7 +189,7 @@ class KMeansModel private[ml] (
 @Experimental
 class KMeans @Since("1.5.0") (
     @Since("1.5.0") override val uid: String)
-  extends Estimator[KMeansModel] with KMeansParams {
+  extends Estimator[KMeansModel] with KMeansParams with DefaultParamsWritable {
 
   setDefault(
     k -> 2,
@@ -210,3 +258,10 @@ class KMeans @Since("1.5.0") (
   }
 }
 
+@Since("1.6.0")
+object KMeans extends DefaultParamsReadable[KMeans] {
+
+  @Since("1.6.0")
+  override def load(path: String): KMeans = super.load(path)
+}
+
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/NaiveBayesSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/NaiveBayesSuite.scala
index 98bc9511163e7..082a6bcd211ab 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/NaiveBayesSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/NaiveBayesSuite.scala
@@ -21,15 +21,30 @@ import breeze.linalg.{Vector => BV}
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
-import org.apache.spark.mllib.classification.NaiveBayes.{Multinomial, Bernoulli}
+import org.apache.spark.ml.util.DefaultReadWriteTest
+import org.apache.spark.mllib.classification.NaiveBayes.{Bernoulli, Multinomial}
+import org.apache.spark.mllib.classification.NaiveBayesSuite._
 import org.apache.spark.mllib.linalg._
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
-import org.apache.spark.mllib.classification.NaiveBayesSuite._
-import org.apache.spark.sql.DataFrame
-import org.apache.spark.sql.Row
+import org.apache.spark.sql.{DataFrame, Row}
+
+class NaiveBayesSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
+
+  @transient var dataset: DataFrame = _
+
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+
+    val pi = Array(0.5, 0.1, 0.4).map(math.log)
+    val theta = Array(
+      Array(0.70, 0.10, 0.10, 0.10), // label 0
+      Array(0.10, 0.70, 0.10, 0.10), // label 1
+      Array(0.10, 0.10, 0.70, 0.10)  // label 2
+    ).map(_.map(math.log))
 
-class NaiveBayesSuite extends SparkFunSuite with MLlibTestSparkContext {
+    dataset = sqlContext.createDataFrame(generateNaiveBayesInput(pi, theta, 100, 42))
+  }
 
   def validatePrediction(predictionAndLabels: DataFrame): Unit = {
     val numOfErrorPredictions = predictionAndLabels.collect().count {
@@ -161,4 +176,26 @@ class NaiveBayesSuite extends SparkFunSuite with MLlibTestSparkContext {
       .select("features", "probability")
     validateProbabilities(featureAndProbabilities, model, "bernoulli")
   }
+
+  test("read/write") {
+    def checkModelData(model: NaiveBayesModel, model2: NaiveBayesModel): Unit = {
+      assert(model.pi === model2.pi)
+      assert(model.theta === model2.theta)
+    }
+    val nb = new NaiveBayes()
+    testEstimatorAndModelReadWrite(nb, dataset, NaiveBayesSuite.allParamSettings, checkModelData)
+  }
+}
+
+object NaiveBayesSuite {
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = Map(
+    "predictionCol" -> "myPrediction",
+    "smoothing" -> 0.1
+  )
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
index c05f90550d161..2724e51f31aa4 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.ml.clustering
 
 import org.apache.spark.SparkFunSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.clustering.{KMeans => MLlibKMeans}
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
@@ -25,16 +26,7 @@ import org.apache.spark.sql.{DataFrame, SQLContext}
 
 private[clustering] case class TestRow(features: Vector)
 
-object KMeansSuite {
-  def generateKMeansData(sql: SQLContext, rows: Int, dim: Int, k: Int): DataFrame = {
-    val sc = sql.sparkContext
-    val rdd = sc.parallelize(1 to rows).map(i => Vectors.dense(Array.fill(dim)((i % k).toDouble)))
-      .map(v => new TestRow(v))
-    sql.createDataFrame(rdd)
-  }
-}
-
-class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
+class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   final val k = 5
   @transient var dataset: DataFrame = _
@@ -106,4 +98,33 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext {
     assert(clusters === Set(0, 1, 2, 3, 4))
     assert(model.computeCost(dataset) < 0.1)
   }
+
+  test("read/write") {
+    def checkModelData(model: KMeansModel, model2: KMeansModel): Unit = {
+      assert(model.clusterCenters === model2.clusterCenters)
+    }
+    val kmeans = new KMeans()
+    testEstimatorAndModelReadWrite(kmeans, dataset, KMeansSuite.allParamSettings, checkModelData)
+  }
+}
+
+object KMeansSuite {
+  def generateKMeansData(sql: SQLContext, rows: Int, dim: Int, k: Int): DataFrame = {
+    val sc = sql.sparkContext
+    val rdd = sc.parallelize(1 to rows).map(i => Vectors.dense(Array.fill(dim)((i % k).toDouble)))
+      .map(v => new TestRow(v))
+    sql.createDataFrame(rdd)
+  }
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = Map(
+    "predictionCol" -> "myPrediction",
+    "k" -> 3,
+    "maxIter" -> 2,
+    "tol" -> 0.01
+  )
 }

From a66142decee48bf5689fb7f4f33646d7bb1ac08d Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Fri, 20 Nov 2015 00:46:29 -0800
Subject: [PATCH 380/518] [SPARK-11877] Prevent agg. fallback conf. from
 leaking across test suites

This patch fixes an issue where the `spark.sql.TungstenAggregate.testFallbackStartsAt` SQLConf setting was not properly reset / cleared at the end of `TungstenAggregationQueryWithControlledFallbackSuite`. This ended up causing test failures in HiveCompatibilitySuite in Maven builds by causing spilling to occur way too frequently.

This configuration leak was inadvertently introduced during test cleanup in #9618.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9857 from JoshRosen/clear-fallback-prop-in-test-teardown.
---
 .../execution/AggregationQuerySuite.scala     | 44 +++++++++----------
 1 file changed, 21 insertions(+), 23 deletions(-)

diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
index 6dde79f74d3d8..39c0a2a0de045 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/AggregationQuerySuite.scala
@@ -868,29 +868,27 @@ class TungstenAggregationQueryWithControlledFallbackSuite extends AggregationQue
 
   override protected def checkAnswer(actual: => DataFrame, expectedAnswer: Seq[Row]): Unit = {
     (0 to 2).foreach { fallbackStartsAt =>
-      sqlContext.setConf(
-        "spark.sql.TungstenAggregate.testFallbackStartsAt",
-        fallbackStartsAt.toString)
-
-      // Create a new df to make sure its physical operator picks up
-      // spark.sql.TungstenAggregate.testFallbackStartsAt.
-      // todo: remove it?
-      val newActual = DataFrame(sqlContext, actual.logicalPlan)
-
-      QueryTest.checkAnswer(newActual, expectedAnswer) match {
-        case Some(errorMessage) =>
-          val newErrorMessage =
-            s"""
-              |The following aggregation query failed when using TungstenAggregate with
-              |controlled fallback (it falls back to sort-based aggregation once it has processed
-              |$fallbackStartsAt input rows). The query is
-              |${actual.queryExecution}
-              |
-              |$errorMessage
-            """.stripMargin
-
-          fail(newErrorMessage)
-        case None =>
+      withSQLConf("spark.sql.TungstenAggregate.testFallbackStartsAt" -> fallbackStartsAt.toString) {
+        // Create a new df to make sure its physical operator picks up
+        // spark.sql.TungstenAggregate.testFallbackStartsAt.
+        // todo: remove it?
+        val newActual = DataFrame(sqlContext, actual.logicalPlan)
+
+        QueryTest.checkAnswer(newActual, expectedAnswer) match {
+          case Some(errorMessage) =>
+            val newErrorMessage =
+              s"""
+                |The following aggregation query failed when using TungstenAggregate with
+                |controlled fallback (it falls back to sort-based aggregation once it has processed
+                |$fallbackStartsAt input rows). The query is
+                |${actual.queryExecution}
+                |
+                |$errorMessage
+              """.stripMargin
+
+            fail(newErrorMessage)
+          case None =>
+        }
       }
     }
   }

From 9ace2e5c8d7fbd360a93bc5fc4eace64a697b44f Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Fri, 20 Nov 2015 09:55:53 -0800
Subject: [PATCH 381/518] [SPARK-11852][ML] StandardScaler minor refactor

```withStd``` and ```withMean``` should be params of ```StandardScaler``` and ```StandardScalerModel```.

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9839 from yanboliang/standardScaler-refactor.
---
 .../spark/ml/feature/StandardScaler.scala     | 60 +++++++++----------
 .../ml/feature/StandardScalerSuite.scala      | 11 ++--
 2 files changed, 32 insertions(+), 39 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
index 6d545219ebf49..d76a9c6275e6b 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala
@@ -36,20 +36,30 @@ import org.apache.spark.sql.types.{StructField, StructType}
 private[feature] trait StandardScalerParams extends Params with HasInputCol with HasOutputCol {
 
   /**
-   * Centers the data with mean before scaling.
+   * Whether to center the data with mean before scaling.
    * It will build a dense output, so this does not work on sparse input
    * and will raise an exception.
    * Default: false
    * @group param
    */
-  val withMean: BooleanParam = new BooleanParam(this, "withMean", "Center data with mean")
+  val withMean: BooleanParam = new BooleanParam(this, "withMean",
+    "Whether to center data with mean")
+
+  /** @group getParam */
+  def getWithMean: Boolean = $(withMean)
 
   /**
-   * Scales the data to unit standard deviation.
+   * Whether to scale the data to unit standard deviation.
    * Default: true
    * @group param
    */
-  val withStd: BooleanParam = new BooleanParam(this, "withStd", "Scale to unit standard deviation")
+  val withStd: BooleanParam = new BooleanParam(this, "withStd",
+    "Whether to scale the data to unit standard deviation")
+
+  /** @group getParam */
+  def getWithStd: Boolean = $(withStd)
+
+  setDefault(withMean -> false, withStd -> true)
 }
 
 /**
@@ -63,8 +73,6 @@ class StandardScaler(override val uid: String) extends Estimator[StandardScalerM
 
   def this() = this(Identifiable.randomUID("stdScal"))
 
-  setDefault(withMean -> false, withStd -> true)
-
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
 
@@ -82,7 +90,7 @@ class StandardScaler(override val uid: String) extends Estimator[StandardScalerM
     val input = dataset.select($(inputCol)).map { case Row(v: Vector) => v }
     val scaler = new feature.StandardScaler(withMean = $(withMean), withStd = $(withStd))
     val scalerModel = scaler.fit(input)
-    copyValues(new StandardScalerModel(uid, scalerModel).setParent(this))
+    copyValues(new StandardScalerModel(uid, scalerModel.std, scalerModel.mean).setParent(this))
   }
 
   override def transformSchema(schema: StructType): StructType = {
@@ -108,29 +116,19 @@ object StandardScaler extends DefaultParamsReadable[StandardScaler] {
 /**
  * :: Experimental ::
  * Model fitted by [[StandardScaler]].
+ *
+ * @param std Standard deviation of the StandardScalerModel
+ * @param mean Mean of the StandardScalerModel
  */
 @Experimental
 class StandardScalerModel private[ml] (
     override val uid: String,
-    scaler: feature.StandardScalerModel)
+    val std: Vector,
+    val mean: Vector)
   extends Model[StandardScalerModel] with StandardScalerParams with MLWritable {
 
   import StandardScalerModel._
 
-  /** Standard deviation of the StandardScalerModel */
-  val std: Vector = scaler.std
-
-  /** Mean of the StandardScalerModel */
-  val mean: Vector = scaler.mean
-
-  /** Whether to scale to unit standard deviation. */
-  @Since("1.6.0")
-  def getWithStd: Boolean = scaler.withStd
-
-  /** Whether to center data with mean. */
-  @Since("1.6.0")
-  def getWithMean: Boolean = scaler.withMean
-
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
 
@@ -139,6 +137,7 @@ class StandardScalerModel private[ml] (
 
   override def transform(dataset: DataFrame): DataFrame = {
     transformSchema(dataset.schema, logging = true)
+    val scaler = new feature.StandardScalerModel(std, mean, $(withStd), $(withMean))
     val scale = udf { scaler.transform _ }
     dataset.withColumn($(outputCol), scale(col($(inputCol))))
   }
@@ -154,7 +153,7 @@ class StandardScalerModel private[ml] (
   }
 
   override def copy(extra: ParamMap): StandardScalerModel = {
-    val copied = new StandardScalerModel(uid, scaler)
+    val copied = new StandardScalerModel(uid, std, mean)
     copyValues(copied, extra).setParent(parent)
   }
 
@@ -168,11 +167,11 @@ object StandardScalerModel extends MLReadable[StandardScalerModel] {
   private[StandardScalerModel]
   class StandardScalerModelWriter(instance: StandardScalerModel) extends MLWriter {
 
-    private case class Data(std: Vector, mean: Vector, withStd: Boolean, withMean: Boolean)
+    private case class Data(std: Vector, mean: Vector)
 
     override protected def saveImpl(path: String): Unit = {
       DefaultParamsWriter.saveMetadata(instance, path, sc)
-      val data = Data(instance.std, instance.mean, instance.getWithStd, instance.getWithMean)
+      val data = Data(instance.std, instance.mean)
       val dataPath = new Path(path, "data").toString
       sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
     }
@@ -185,13 +184,10 @@ object StandardScalerModel extends MLReadable[StandardScalerModel] {
     override def load(path: String): StandardScalerModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
       val dataPath = new Path(path, "data").toString
-      val Row(std: Vector, mean: Vector, withStd: Boolean, withMean: Boolean) =
-        sqlContext.read.parquet(dataPath)
-          .select("std", "mean", "withStd", "withMean")
-          .head()
-      // This is very likely to change in the future because withStd and withMean should be params.
-      val oldModel = new feature.StandardScalerModel(std, mean, withStd, withMean)
-      val model = new StandardScalerModel(metadata.uid, oldModel)
+      val Row(std: Vector, mean: Vector) = sqlContext.read.parquet(dataPath)
+        .select("std", "mean")
+        .head()
+      val model = new StandardScalerModel(metadata.uid, std, mean)
       DefaultParamsReader.getAndSetParams(model, metadata)
       model
     }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala
index 49a4b2efe0c29..1eae125a524ef 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/StandardScalerSuite.scala
@@ -70,8 +70,8 @@ class StandardScalerSuite extends SparkFunSuite with MLlibTestSparkContext
 
   test("params") {
     ParamsSuite.checkParams(new StandardScaler)
-    val oldModel = new feature.StandardScalerModel(Vectors.dense(1.0), Vectors.dense(2.0))
-    ParamsSuite.checkParams(new StandardScalerModel("empty", oldModel))
+    ParamsSuite.checkParams(new StandardScalerModel("empty",
+      Vectors.dense(1.0), Vectors.dense(2.0)))
   }
 
   test("Standardization with default parameter") {
@@ -126,13 +126,10 @@ class StandardScalerSuite extends SparkFunSuite with MLlibTestSparkContext
   }
 
   test("StandardScalerModel read/write") {
-    val oldModel = new feature.StandardScalerModel(
-      Vectors.dense(1.0, 2.0), Vectors.dense(3.0, 4.0), false, true)
-    val instance = new StandardScalerModel("myStandardScalerModel", oldModel)
+    val instance = new StandardScalerModel("myStandardScalerModel",
+      Vectors.dense(1.0, 2.0), Vectors.dense(3.0, 4.0))
     val newInstance = testDefaultReadWrite(instance)
     assert(newInstance.std === instance.std)
     assert(newInstance.mean === instance.mean)
-    assert(newInstance.getWithStd === instance.getWithStd)
-    assert(newInstance.getWithMean === instance.getWithMean)
   }
 }

From e359d5dcf5bd300213054ebeae9fe75c4f7eb9e7 Mon Sep 17 00:00:00 2001
From: Yuhao Yang <hhbyyh@gmail.com>
Date: Fri, 20 Nov 2015 09:57:09 -0800
Subject: [PATCH 382/518] [SPARK-11689][ML] Add user guide and example code for
 LDA under spark.ml

jira: https://issues.apache.org/jira/browse/SPARK-11689

Add simple user guide for LDA under spark.ml and example code under examples/. Use include_example to include example code in the user guide markdown. Check SPARK-11606 for instructions.

Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #9722 from hhbyyh/ldaMLExample.
---
 docs/ml-clustering.md                         | 30 ++++++
 docs/ml-guide.md                              |  3 +-
 docs/mllib-guide.md                           |  1 +
 .../spark/examples/ml/JavaLDAExample.java     | 94 +++++++++++++++++++
 .../apache/spark/examples/ml/LDAExample.scala | 77 +++++++++++++++
 5 files changed, 204 insertions(+), 1 deletion(-)
 create mode 100644 docs/ml-clustering.md
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala

diff --git a/docs/ml-clustering.md b/docs/ml-clustering.md
new file mode 100644
index 0000000000000..1743ef43a6ddf
--- /dev/null
+++ b/docs/ml-clustering.md
@@ -0,0 +1,30 @@
+---
+layout: global
+title: Clustering - ML
+displayTitle: <a href="ml-guide.html">ML</a> - Clustering
+---
+
+In this section, we introduce the pipeline API for [clustering in mllib](mllib-clustering.html).
+
+## Latent Dirichlet allocation (LDA)
+
+`LDA` is implemented as an `Estimator` that supports both `EMLDAOptimizer` and `OnlineLDAOptimizer`,
+and generates a `LDAModel` as the base models. Expert users may cast a `LDAModel` generated by
+`EMLDAOptimizer` to a `DistributedLDAModel` if needed.
+
+<div class="codetabs">
+
+Refer to the [Scala API docs](api/scala/index.html#org.apache.spark.ml.clustering.LDA) for more details.
+
+<div data-lang="scala" markdown="1">
+{% include_example scala/org/apache/spark/examples/ml/LDAExample.scala %}
+</div>
+
+<div data-lang="java" markdown="1">
+
+Refer to the [Java API docs](api/java/org/apache/spark/ml/clustering/LDA.html) for more details.
+
+{% include_example java/org/apache/spark/examples/ml/JavaLDAExample.java %}
+</div>
+
+</div>
\ No newline at end of file
diff --git a/docs/ml-guide.md b/docs/ml-guide.md
index be18a05361a17..6f35b30c3d4df 100644
--- a/docs/ml-guide.md
+++ b/docs/ml-guide.md
@@ -40,6 +40,7 @@ Also, some algorithms have additional capabilities in the `spark.ml` API; e.g.,
 provide class probabilities, and linear models provide model summaries.
 
 * [Feature extraction, transformation, and selection](ml-features.html)
+* [Clustering](ml-clustering.html)
 * [Decision Trees for classification and regression](ml-decision-tree.html)
 * [Ensembles](ml-ensembles.html)
 * [Linear methods with elastic net regularization](ml-linear-methods.html)
@@ -950,4 +951,4 @@ model.transform(test)
 {% endhighlight %}
 </div>
 
-</div>
+</div>
\ No newline at end of file
diff --git a/docs/mllib-guide.md b/docs/mllib-guide.md
index 91e50ccfecec4..54e35fcbb15af 100644
--- a/docs/mllib-guide.md
+++ b/docs/mllib-guide.md
@@ -69,6 +69,7 @@ We list major functionality from both below, with links to detailed guides.
 concepts. It also contains sections on using algorithms within the Pipelines API, for example:
 
 * [Feature extraction, transformation, and selection](ml-features.html)
+* [Clustering](ml-clustering.html)
 * [Decision trees for classification and regression](ml-decision-tree.html)
 * [Ensembles](ml-ensembles.html)
 * [Linear methods with elastic net regularization](ml-linear-methods.html)
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
new file mode 100644
index 0000000000000..b3a7d2eb29780
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+
+import java.util.regex.Pattern;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.ml.clustering.LDA;
+import org.apache.spark.ml.clustering.LDAModel;
+import org.apache.spark.mllib.linalg.Vector;
+import org.apache.spark.mllib.linalg.VectorUDT;
+import org.apache.spark.mllib.linalg.Vectors;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.sql.catalyst.expressions.GenericRow;
+import org.apache.spark.sql.types.Metadata;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+
+/**
+ * An example demonstrating LDA
+ * Run with
+ * <pre>
+ * bin/run-example ml.JavaLDAExample
+ * </pre>
+ */
+public class JavaLDAExample {
+
+  private static class ParseVector implements Function<String, Row> {
+    private static final Pattern separator = Pattern.compile(" ");
+
+    @Override
+    public Row call(String line) {
+      String[] tok = separator.split(line);
+      double[] point = new double[tok.length];
+      for (int i = 0; i < tok.length; ++i) {
+        point[i] = Double.parseDouble(tok[i]);
+      }
+      Vector[] points = {Vectors.dense(point)};
+      return new GenericRow(points);
+    }
+  }
+
+  public static void main(String[] args) {
+
+    String inputFile = "data/mllib/sample_lda_data.txt";
+
+    // Parses the arguments
+    SparkConf conf = new SparkConf().setAppName("JavaLDAExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // Loads data
+    JavaRDD<Row> points = jsc.textFile(inputFile).map(new ParseVector());
+    StructField[] fields = {new StructField("features", new VectorUDT(), false, Metadata.empty())};
+    StructType schema = new StructType(fields);
+    DataFrame dataset = sqlContext.createDataFrame(points, schema);
+
+    // Trains a LDA model
+    LDA lda = new LDA()
+      .setK(10)
+      .setMaxIter(10);
+    LDAModel model = lda.fit(dataset);
+
+    System.out.println(model.logLikelihood(dataset));
+    System.out.println(model.logPerplexity(dataset));
+
+    // Shows the result
+    DataFrame topics = model.describeTopics(3);
+    topics.show(false);
+    model.transform(dataset).show(false);
+
+    jsc.stop();
+  }
+}
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala
new file mode 100644
index 0000000000000..419ce3d87a6ac
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml
+
+// scalastyle:off println
+import org.apache.spark.{SparkContext, SparkConf}
+import org.apache.spark.mllib.linalg.{VectorUDT, Vectors}
+// $example on$
+import org.apache.spark.ml.clustering.LDA
+import org.apache.spark.sql.{Row, SQLContext}
+import org.apache.spark.sql.types.{StructField, StructType}
+// $example off$
+
+/**
+ * An example demonstrating a LDA of ML pipeline.
+ * Run with
+ * {{{
+ * bin/run-example ml.LDAExample
+ * }}}
+ */
+object LDAExample {
+
+  final val FEATURES_COL = "features"
+
+  def main(args: Array[String]): Unit = {
+
+    val input = "data/mllib/sample_lda_data.txt"
+    // Creates a Spark context and a SQL context
+    val conf = new SparkConf().setAppName(s"${this.getClass.getSimpleName}")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+
+    // $example on$
+    // Loads data
+    val rowRDD = sc.textFile(input).filter(_.nonEmpty)
+      .map(_.split(" ").map(_.toDouble)).map(Vectors.dense).map(Row(_))
+    val schema = StructType(Array(StructField(FEATURES_COL, new VectorUDT, false)))
+    val dataset = sqlContext.createDataFrame(rowRDD, schema)
+
+    // Trains a LDA model
+    val lda = new LDA()
+      .setK(10)
+      .setMaxIter(10)
+      .setFeaturesCol(FEATURES_COL)
+    val model = lda.fit(dataset)
+    val transformed = model.transform(dataset)
+
+    val ll = model.logLikelihood(dataset)
+    val lp = model.logPerplexity(dataset)
+
+    // describeTopics
+    val topics = model.describeTopics(3)
+
+    // Shows the result
+    topics.show(false)
+    transformed.show(false)
+
+    // $example off$
+    sc.stop()
+  }
+}
+// scalastyle:on println

From bef361c589c0a38740232fd8d0a45841e4fc969a Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Fri, 20 Nov 2015 11:20:47 -0800
Subject: [PATCH 383/518] [SPARK-11876][SQL] Support printSchema in DataSet API

DataSet APIs look great! However, I am lost when doing multiple level joins.  For example,
```
val ds1 = Seq(("a", 1), ("b", 2)).toDS().as("a")
val ds2 = Seq(("a", 1), ("b", 2)).toDS().as("b")
val ds3 = Seq(("a", 1), ("b", 2)).toDS().as("c")

ds1.joinWith(ds2, $"a._2" === $"b._2").as("ab").joinWith(ds3, $"ab._1._2" === $"c._2").printSchema()
```

The printed schema is like
```
root
 |-- _1: struct (nullable = true)
 |    |-- _1: struct (nullable = true)
 |    |    |-- _1: string (nullable = true)
 |    |    |-- _2: integer (nullable = true)
 |    |-- _2: struct (nullable = true)
 |    |    |-- _1: string (nullable = true)
 |    |    |-- _2: integer (nullable = true)
 |-- _2: struct (nullable = true)
 |    |-- _1: string (nullable = true)
 |    |-- _2: integer (nullable = true)
```

Personally, I think we need the printSchema function. Sometimes, I do not know how to specify the column, especially when their data types are mixed. For example, if I want to write the following select for the above multi-level join, I have to know the schema:
```
newDS.select(expr("_1._2._2 + 1").as[Int]).collect()
```

marmbrus rxin cloud-fan  Do you have the same feeling?

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9855 from gatorsmile/printSchemaDataSet.
---
 .../src/main/scala/org/apache/spark/sql/DataFrame.scala  | 9 ---------
 .../scala/org/apache/spark/sql/execution/Queryable.scala | 9 +++++++++
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 98358127e2709..7abcecaa2880e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -299,15 +299,6 @@ class DataFrame private[sql](
    */
   def columns: Array[String] = schema.fields.map(_.name)
 
-  /**
-   * Prints the schema to the console in a nice tree format.
-   * @group basic
-   * @since 1.3.0
-   */
-  // scalastyle:off println
-  def printSchema(): Unit = println(schema.treeString)
-  // scalastyle:on println
-
   /**
    * Returns true if the `collect` and `take` methods can be run locally
    * (without any Spark executors).
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
index e86a52c149a2f..321e2c783537f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
@@ -37,6 +37,15 @@ private[sql] trait Queryable {
     }
   }
 
+  /**
+   * Prints the schema to the console in a nice tree format.
+   * @group basic
+   * @since 1.3.0
+   */
+  // scalastyle:off println
+  def printSchema(): Unit = println(schema.treeString)
+  // scalastyle:on println
+
   /**
    * Prints the plans (logical and physical) to the console for debugging purposes.
    * @since 1.3.0

From 60bfb113325c71491f8dcf98b6036b0caa2144fe Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Fri, 20 Nov 2015 11:43:45 -0800
Subject: [PATCH 384/518] [SPARK-11817][SQL] Truncating the fractional seconds
 to prevent inserting a NULL

JIRA: https://issues.apache.org/jira/browse/SPARK-11817

Instead of return None, we should truncate the fractional seconds to prevent inserting NULL.

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #9834 from viirya/truncate-fractional-sec.
---
 .../apache/spark/sql/catalyst/util/DateTimeUtils.scala    | 5 +++++
 .../spark/sql/catalyst/util/DateTimeUtilsSuite.scala      | 8 ++++++++
 2 files changed, 13 insertions(+)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
index 17a5527f3fb29..2b93882919487 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala
@@ -327,6 +327,11 @@ object DateTimeUtils {
       return None
     }
 
+    // Instead of return None, we truncate the fractional seconds to prevent inserting NULL
+    if (segments(6) > 999999) {
+      segments(6) = segments(6).toString.take(6).toInt
+    }
+
     if (segments(3) < 0 || segments(3) > 23 || segments(4) < 0 || segments(4) > 59 ||
         segments(5) < 0 || segments(5) > 59 || segments(6) < 0 || segments(6) > 999999 ||
         segments(7) < 0 || segments(7) > 23 || segments(8) < 0 || segments(8) > 59) {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
index faca128badfd6..0ce5a2fb69505 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala
@@ -343,6 +343,14 @@ class DateTimeUtilsSuite extends SparkFunSuite {
       UTF8String.fromString("2015-03-18T12:03.17-0:70")).isEmpty)
     assert(stringToTimestamp(
       UTF8String.fromString("2015-03-18T12:03.17-1:0:0")).isEmpty)
+
+    // Truncating the fractional seconds
+    c = Calendar.getInstance(TimeZone.getTimeZone("GMT+00:00"))
+    c.set(2015, 2, 18, 12, 3, 17)
+    c.set(Calendar.MILLISECOND, 0)
+    assert(stringToTimestamp(
+      UTF8String.fromString("2015-03-18T12:03:17.123456789+0:00")).get ===
+        c.getTimeInMillis * 1000 + 123456)
   }
 
   test("hours") {

From 3b9d2a347f9c796b90852173d84189834e499e25 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Fri, 20 Nov 2015 12:04:42 -0800
Subject: [PATCH 385/518] [SPARK-11819][SQL] nice error message for missing
 encoder

before this PR, when users try to get an encoder for an un-supported class, they will only get a very simple error message like `Encoder for type xxx is not supported`.

After this PR, the error message become more friendly, for example:
```
No Encoder found for abc.xyz.NonEncodable
- array element class: "abc.xyz.NonEncodable"
- field (class: "scala.Array", name: "arrayField")
- root class: "abc.xyz.AnotherClass"
```

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9810 from cloud-fan/error-message.
---
 .../spark/sql/catalyst/ScalaReflection.scala  | 90 ++++++++++++++-----
 .../encoders/EncoderErrorMessageSuite.scala   | 62 +++++++++++++
 2 files changed, 129 insertions(+), 23 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 33ae700706dae..918050b531c02 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -63,7 +63,7 @@ object ScalaReflection extends ScalaReflection {
       case t if t <:< definitions.BooleanTpe => BooleanType
       case t if t <:< localTypeOf[Array[Byte]] => BinaryType
       case _ =>
-        val className: String = tpe.erasure.typeSymbol.asClass.fullName
+        val className = getClassNameFromType(tpe)
         className match {
           case "scala.Array" =>
             val TypeRef(_, _, Seq(elementType)) = tpe
@@ -320,9 +320,23 @@ object ScalaReflection extends ScalaReflection {
     }
   }
 
-  /** Returns expressions for extracting all the fields from the given type. */
+  /**
+   * Returns expressions for extracting all the fields from the given type.
+   *
+   * If the given type is not supported, i.e. there is no encoder can be built for this type,
+   * an [[UnsupportedOperationException]] will be thrown with detailed error message to explain
+   * the type path walked so far and which class we are not supporting.
+   * There are 4 kinds of type path:
+   *  * the root type: `root class: "abc.xyz.MyClass"`
+   *  * the value type of [[Option]]: `option value class: "abc.xyz.MyClass"`
+   *  * the element type of [[Array]] or [[Seq]]: `array element class: "abc.xyz.MyClass"`
+   *  * the field of [[Product]]: `field (class: "abc.xyz.MyClass", name: "myField")`
+   */
   def extractorsFor[T : TypeTag](inputObject: Expression): CreateNamedStruct = {
-    extractorFor(inputObject, localTypeOf[T]) match {
+    val tpe = localTypeOf[T]
+    val clsName = getClassNameFromType(tpe)
+    val walkedTypePath = s"""- root class: "${clsName}"""" :: Nil
+    extractorFor(inputObject, tpe, walkedTypePath) match {
       case s: CreateNamedStruct => s
       case other => CreateNamedStruct(expressions.Literal("value") :: other :: Nil)
     }
@@ -331,7 +345,28 @@ object ScalaReflection extends ScalaReflection {
   /** Helper for extracting internal fields from a case class. */
   private def extractorFor(
       inputObject: Expression,
-      tpe: `Type`): Expression = ScalaReflectionLock.synchronized {
+      tpe: `Type`,
+      walkedTypePath: Seq[String]): Expression = ScalaReflectionLock.synchronized {
+
+    def toCatalystArray(input: Expression, elementType: `Type`): Expression = {
+      val externalDataType = dataTypeFor(elementType)
+      val Schema(catalystType, nullable) = silentSchemaFor(elementType)
+      if (isNativeType(catalystType)) {
+        NewInstance(
+          classOf[GenericArrayData],
+          input :: Nil,
+          dataType = ArrayType(catalystType, nullable))
+      } else {
+        val clsName = getClassNameFromType(elementType)
+        val newPath = s"""- array element class: "$clsName"""" +: walkedTypePath
+        // `MapObjects` will run `extractorFor` lazily, we need to eagerly call `extractorFor` here
+        // to trigger the type check.
+        extractorFor(inputObject, elementType, newPath)
+
+        MapObjects(extractorFor(_, elementType, newPath), input, externalDataType)
+      }
+    }
+
     if (!inputObject.dataType.isInstanceOf[ObjectType]) {
       inputObject
     } else {
@@ -378,15 +413,16 @@ object ScalaReflection extends ScalaReflection {
 
             // For non-primitives, we can just extract the object from the Option and then recurse.
             case other =>
-              val className: String = optType.erasure.typeSymbol.asClass.fullName
+              val className = getClassNameFromType(optType)
               val classObj = Utils.classForName(className)
               val optionObjectType = ObjectType(classObj)
+              val newPath = s"""- option value class: "$className"""" +: walkedTypePath
 
               val unwrapped = UnwrapOption(optionObjectType, inputObject)
               expressions.If(
                 IsNull(unwrapped),
-                expressions.Literal.create(null, schemaFor(optType).dataType),
-                extractorFor(unwrapped, optType))
+                expressions.Literal.create(null, silentSchemaFor(optType).dataType),
+                extractorFor(unwrapped, optType, newPath))
           }
 
         case t if t <:< localTypeOf[Product] =>
@@ -412,7 +448,10 @@ object ScalaReflection extends ScalaReflection {
             val fieldName = p.name.toString
             val fieldType = p.typeSignature.substituteTypes(formalTypeArgs, actualTypeArgs)
             val fieldValue = Invoke(inputObject, fieldName, dataTypeFor(fieldType))
-            expressions.Literal(fieldName) :: extractorFor(fieldValue, fieldType) :: Nil
+            val clsName = getClassNameFromType(fieldType)
+            val newPath = s"""- field (class: "$clsName", name: "$fieldName")""" +: walkedTypePath
+
+            expressions.Literal(fieldName) :: extractorFor(fieldValue, fieldType, newPath) :: Nil
           })
 
         case t if t <:< localTypeOf[Array[_]] =>
@@ -500,23 +539,11 @@ object ScalaReflection extends ScalaReflection {
           Invoke(inputObject, "booleanValue", BooleanType)
 
         case other =>
-          throw new UnsupportedOperationException(s"Extractor for type $other is not supported")
+          throw new UnsupportedOperationException(
+            s"No Encoder found for $tpe\n" + walkedTypePath.mkString("\n"))
       }
     }
   }
-
-  private def toCatalystArray(input: Expression, elementType: `Type`): Expression = {
-    val externalDataType = dataTypeFor(elementType)
-    val Schema(catalystType, nullable) = schemaFor(elementType)
-    if (isNativeType(catalystType)) {
-      NewInstance(
-        classOf[GenericArrayData],
-        input :: Nil,
-        dataType = ArrayType(catalystType, nullable))
-    } else {
-      MapObjects(extractorFor(_, elementType), input, externalDataType)
-    }
-  }
 }
 
 /**
@@ -561,7 +588,7 @@ trait ScalaReflection {
 
   /** Returns a catalyst DataType and its nullability for the given Scala Type using reflection. */
   def schemaFor(tpe: `Type`): Schema = ScalaReflectionLock.synchronized {
-    val className: String = tpe.erasure.typeSymbol.asClass.fullName
+    val className = getClassNameFromType(tpe)
     tpe match {
       case t if Utils.classIsLoadable(className) &&
         Utils.classForName(className).isAnnotationPresent(classOf[SQLUserDefinedType]) =>
@@ -637,6 +664,23 @@ trait ScalaReflection {
     }
   }
 
+  /**
+   * Returns a catalyst DataType and its nullability for the given Scala Type using reflection.
+   *
+   * Unlike `schemaFor`, this method won't throw exception for un-supported type, it will return
+   * `NullType` silently instead.
+   */
+  private def silentSchemaFor(tpe: `Type`): Schema = try {
+    schemaFor(tpe)
+  } catch {
+    case _: UnsupportedOperationException => Schema(NullType, nullable = true)
+  }
+
+  /** Returns the full class name for a type. */
+  private def getClassNameFromType(tpe: `Type`): String = {
+    tpe.erasure.typeSymbol.asClass.fullName
+  }
+
   /**
    * Returns classes of input parameters of scala function object.
    */
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderErrorMessageSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderErrorMessageSuite.scala
index 0b2a10bb04c10..8c766ef829923 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderErrorMessageSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/EncoderErrorMessageSuite.scala
@@ -17,9 +17,22 @@
 
 package org.apache.spark.sql.catalyst.encoders
 
+import scala.reflect.ClassTag
+
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.Encoders
 
+class NonEncodable(i: Int)
+
+case class ComplexNonEncodable1(name1: NonEncodable)
+
+case class ComplexNonEncodable2(name2: ComplexNonEncodable1)
+
+case class ComplexNonEncodable3(name3: Option[NonEncodable])
+
+case class ComplexNonEncodable4(name4: Array[NonEncodable])
+
+case class ComplexNonEncodable5(name5: Option[Array[NonEncodable]])
 
 class EncoderErrorMessageSuite extends SparkFunSuite {
 
@@ -37,4 +50,53 @@ class EncoderErrorMessageSuite extends SparkFunSuite {
     intercept[UnsupportedOperationException] { Encoders.javaSerialization[Long] }
     intercept[UnsupportedOperationException] { Encoders.javaSerialization[Char] }
   }
+
+  test("nice error message for missing encoder") {
+    val errorMsg1 =
+      intercept[UnsupportedOperationException](ExpressionEncoder[ComplexNonEncodable1]).getMessage
+    assert(errorMsg1.contains(
+      s"""root class: "${clsName[ComplexNonEncodable1]}""""))
+    assert(errorMsg1.contains(
+      s"""field (class: "${clsName[NonEncodable]}", name: "name1")"""))
+
+    val errorMsg2 =
+      intercept[UnsupportedOperationException](ExpressionEncoder[ComplexNonEncodable2]).getMessage
+    assert(errorMsg2.contains(
+      s"""root class: "${clsName[ComplexNonEncodable2]}""""))
+    assert(errorMsg2.contains(
+      s"""field (class: "${clsName[ComplexNonEncodable1]}", name: "name2")"""))
+    assert(errorMsg1.contains(
+      s"""field (class: "${clsName[NonEncodable]}", name: "name1")"""))
+
+    val errorMsg3 =
+      intercept[UnsupportedOperationException](ExpressionEncoder[ComplexNonEncodable3]).getMessage
+    assert(errorMsg3.contains(
+      s"""root class: "${clsName[ComplexNonEncodable3]}""""))
+    assert(errorMsg3.contains(
+      s"""field (class: "scala.Option", name: "name3")"""))
+    assert(errorMsg3.contains(
+      s"""option value class: "${clsName[NonEncodable]}""""))
+
+    val errorMsg4 =
+      intercept[UnsupportedOperationException](ExpressionEncoder[ComplexNonEncodable4]).getMessage
+    assert(errorMsg4.contains(
+      s"""root class: "${clsName[ComplexNonEncodable4]}""""))
+    assert(errorMsg4.contains(
+      s"""field (class: "scala.Array", name: "name4")"""))
+    assert(errorMsg4.contains(
+      s"""array element class: "${clsName[NonEncodable]}""""))
+
+    val errorMsg5 =
+      intercept[UnsupportedOperationException](ExpressionEncoder[ComplexNonEncodable5]).getMessage
+    assert(errorMsg5.contains(
+      s"""root class: "${clsName[ComplexNonEncodable5]}""""))
+    assert(errorMsg5.contains(
+      s"""field (class: "scala.Option", name: "name5")"""))
+    assert(errorMsg5.contains(
+      s"""option value class: "scala.Array""""))
+    assert(errorMsg5.contains(
+      s"""array element class: "${clsName[NonEncodable]}""""))
+  }
+
+  private def clsName[T : ClassTag]: String = implicitly[ClassTag[T]].runtimeClass.getName
 }

From 652def318e47890bd0a0977dc982cc07f99fb06a Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Fri, 20 Nov 2015 13:17:35 -0800
Subject: [PATCH 386/518] [SPARK-11650] Reduce RPC timeouts to speed up slow
 AkkaUtilsSuite test

This patch reduces some RPC timeouts in order to speed up the slow "AkkaUtilsSuite.remote fetch ssl on - untrusted server", which used to take two minutes to run.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9869 from JoshRosen/SPARK-11650.
---
 core/src/test/scala/org/apache/spark/util/AkkaUtilsSuite.scala | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/core/src/test/scala/org/apache/spark/util/AkkaUtilsSuite.scala b/core/src/test/scala/org/apache/spark/util/AkkaUtilsSuite.scala
index 61601016e005e..0af4b6098bb0a 100644
--- a/core/src/test/scala/org/apache/spark/util/AkkaUtilsSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/AkkaUtilsSuite.scala
@@ -340,10 +340,11 @@ class AkkaUtilsSuite extends SparkFunSuite with LocalSparkContext with ResetSyst
       new MapOutputTrackerMasterEndpoint(rpcEnv, masterTracker, conf))
 
     val slaveConf = sparkSSLConfig()
+      .set("spark.rpc.askTimeout", "5s")
+      .set("spark.rpc.lookupTimeout", "5s")
     val securityManagerBad = new SecurityManager(slaveConf)
 
     val slaveRpcEnv = RpcEnv.create("spark-slave", hostname, 0, slaveConf, securityManagerBad)
-    val slaveTracker = new MapOutputTrackerWorker(conf)
     try {
       slaveRpcEnv.setupEndpointRef("spark", rpcEnv.address, MapOutputTracker.ENDPOINT_NAME)
       fail("should receive either ActorNotFound or TimeoutException")

From 9ed4ad4265cf9d3135307eb62dae6de0b220fc21 Mon Sep 17 00:00:00 2001
From: Nong Li <nong@databricks.com>
Date: Fri, 20 Nov 2015 14:19:34 -0800
Subject: [PATCH 387/518] [SPARK-11724][SQL] Change casting between int and
 timestamp to consistently treat int in seconds.

Hive has since changed this behavior as well. https://issues.apache.org/jira/browse/HIVE-3454

Author: Nong Li <nong@databricks.com>
Author: Nong Li <nongli@gmail.com>
Author: Yin Huai <yhuai@databricks.com>

Closes #9685 from nongli/spark-11724.
---
 .../spark/sql/catalyst/expressions/Cast.scala |  6 ++--
 .../sql/catalyst/expressions/CastSuite.scala  | 16 +++++----
 .../apache/spark/sql/DateFunctionsSuite.scala |  3 ++
 ...esting-0-237a6af90a857da1efcbe98f6bbbf9d6} |  2 +-
 ... cast #3-0-76ee270337f664b36cacfc6528ac109 |  1 -
 ...cast #5-0-dbd7bcd167d322d6617b884c02c7f247 |  1 -
 ...cast #7-0-1d70654217035f8ce5f64344f4c5a80f |  1 -
 .../sql/hive/execution/HiveQuerySuite.scala   | 34 +++++++++++++------
 8 files changed, 39 insertions(+), 25 deletions(-)
 rename sql/hive/src/test/resources/golden/{constant null testing-0-9a02bc7de09bcabcbd4c91f54a814c20 => constant null testing-0-237a6af90a857da1efcbe98f6bbbf9d6} (52%)
 delete mode 100644 sql/hive/src/test/resources/golden/timestamp cast #3-0-76ee270337f664b36cacfc6528ac109
 delete mode 100644 sql/hive/src/test/resources/golden/timestamp cast #5-0-dbd7bcd167d322d6617b884c02c7f247
 delete mode 100644 sql/hive/src/test/resources/golden/timestamp cast #7-0-1d70654217035f8ce5f64344f4c5a80f

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
index 5564e242b0472..533d17ea5c172 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
@@ -204,8 +204,8 @@ case class Cast(child: Expression, dataType: DataType)
     if (d.isNaN || d.isInfinite) null else (d * 1000000L).toLong
   }
 
-  // converting milliseconds to us
-  private[this] def longToTimestamp(t: Long): Long = t * 1000L
+  // converting seconds to us
+  private[this] def longToTimestamp(t: Long): Long = t * 1000000L
   // converting us to seconds
   private[this] def timestampToLong(ts: Long): Long = math.floor(ts.toDouble / 1000000L).toLong
   // converting us to seconds in double
@@ -647,7 +647,7 @@ case class Cast(child: Expression, dataType: DataType)
 
   private[this] def decimalToTimestampCode(d: String): String =
     s"($d.toBigDecimal().bigDecimal().multiply(new java.math.BigDecimal(1000000L))).longValue()"
-  private[this] def longToTimeStampCode(l: String): String = s"$l * 1000L"
+  private[this] def longToTimeStampCode(l: String): String = s"$l * 1000000L"
   private[this] def timestampToIntegerCode(ts: String): String =
     s"java.lang.Math.floor((double) $ts / 1000000L)"
   private[this] def timestampToDoubleCode(ts: String): String = s"$ts / 1000000.0"
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuite.scala
index f4db4da7646f8..ab77a764483e8 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CastSuite.scala
@@ -258,8 +258,8 @@ class CastSuite extends SparkFunSuite with ExpressionEvalHelper {
 
   test("cast from int 2") {
     checkEvaluation(cast(1, LongType), 1.toLong)
-    checkEvaluation(cast(cast(1000, TimestampType), LongType), 1.toLong)
-    checkEvaluation(cast(cast(-1200, TimestampType), LongType), -2.toLong)
+    checkEvaluation(cast(cast(1000, TimestampType), LongType), 1000.toLong)
+    checkEvaluation(cast(cast(-1200, TimestampType), LongType), -1200.toLong)
 
     checkEvaluation(cast(123, DecimalType.USER_DEFAULT), Decimal(123))
     checkEvaluation(cast(123, DecimalType(3, 0)), Decimal(123))
@@ -348,14 +348,14 @@ class CastSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(
       cast(cast(cast(cast(cast(cast("5", ByteType), TimestampType),
         DecimalType.SYSTEM_DEFAULT), LongType), StringType), ShortType),
-      0.toShort)
+      5.toShort)
     checkEvaluation(
       cast(cast(cast(cast(cast(cast("5", TimestampType), ByteType),
         DecimalType.SYSTEM_DEFAULT), LongType), StringType), ShortType),
       null)
     checkEvaluation(cast(cast(cast(cast(cast(cast("5", DecimalType.SYSTEM_DEFAULT),
       ByteType), TimestampType), LongType), StringType), ShortType),
-      0.toShort)
+      5.toShort)
 
     checkEvaluation(cast("23", DoubleType), 23d)
     checkEvaluation(cast("23", IntegerType), 23)
@@ -479,10 +479,12 @@ class CastSuite extends SparkFunSuite with ExpressionEvalHelper {
     checkEvaluation(cast(ts, LongType), 15.toLong)
     checkEvaluation(cast(ts, FloatType), 15.003f)
     checkEvaluation(cast(ts, DoubleType), 15.003)
-    checkEvaluation(cast(cast(tss, ShortType), TimestampType), DateTimeUtils.fromJavaTimestamp(ts))
+    checkEvaluation(cast(cast(tss, ShortType), TimestampType),
+      DateTimeUtils.fromJavaTimestamp(ts) * 1000)
     checkEvaluation(cast(cast(tss, IntegerType), TimestampType),
-      DateTimeUtils.fromJavaTimestamp(ts))
-    checkEvaluation(cast(cast(tss, LongType), TimestampType), DateTimeUtils.fromJavaTimestamp(ts))
+      DateTimeUtils.fromJavaTimestamp(ts) * 1000)
+    checkEvaluation(cast(cast(tss, LongType), TimestampType),
+      DateTimeUtils.fromJavaTimestamp(ts) * 1000)
     checkEvaluation(
       cast(cast(millis.toFloat / 1000, TimestampType), FloatType),
       millis.toFloat / 1000)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
index 241cbd0115070..a61c3aa48a73f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DateFunctionsSuite.scala
@@ -448,6 +448,9 @@ class DateFunctionsSuite extends QueryTest with SharedSQLContext {
       Row(date1.getTime / 1000L), Row(date2.getTime / 1000L)))
     checkAnswer(df.selectExpr(s"unix_timestamp(s, '$fmt')"), Seq(
       Row(ts1.getTime / 1000L), Row(ts2.getTime / 1000L)))
+
+    val now = sql("select unix_timestamp()").collect().head.getLong(0)
+    checkAnswer(sql(s"select cast ($now as timestamp)"), Row(new java.util.Date(now * 1000)))
   }
 
   test("to_unix_timestamp") {
diff --git a/sql/hive/src/test/resources/golden/constant null testing-0-9a02bc7de09bcabcbd4c91f54a814c20 b/sql/hive/src/test/resources/golden/constant null testing-0-237a6af90a857da1efcbe98f6bbbf9d6
similarity index 52%
rename from sql/hive/src/test/resources/golden/constant null testing-0-9a02bc7de09bcabcbd4c91f54a814c20
rename to sql/hive/src/test/resources/golden/constant null testing-0-237a6af90a857da1efcbe98f6bbbf9d6
index 7c41615f8c184..a01c2622c68e2 100644
--- a/sql/hive/src/test/resources/golden/constant null testing-0-9a02bc7de09bcabcbd4c91f54a814c20	
+++ b/sql/hive/src/test/resources/golden/constant null testing-0-237a6af90a857da1efcbe98f6bbbf9d6	
@@ -1 +1 @@
-1	NULL	1	NULL	1.0	NULL	true	NULL	1	NULL	1.0	NULL	1	NULL	1	NULL	1	NULL	1970-01-01	NULL	1969-12-31 16:00:00.001	NULL	1	NULL
+1	NULL	1	NULL	1.0	NULL	true	NULL	1	NULL	1.0	NULL	1	NULL	1	NULL	1	NULL	1970-01-01	NULL	NULL	1	NULL
diff --git a/sql/hive/src/test/resources/golden/timestamp cast #3-0-76ee270337f664b36cacfc6528ac109 b/sql/hive/src/test/resources/golden/timestamp cast #3-0-76ee270337f664b36cacfc6528ac109
deleted file mode 100644
index d00491fd7e5bb..0000000000000
--- a/sql/hive/src/test/resources/golden/timestamp cast #3-0-76ee270337f664b36cacfc6528ac109	
+++ /dev/null
@@ -1 +0,0 @@
-1
diff --git a/sql/hive/src/test/resources/golden/timestamp cast #5-0-dbd7bcd167d322d6617b884c02c7f247 b/sql/hive/src/test/resources/golden/timestamp cast #5-0-dbd7bcd167d322d6617b884c02c7f247
deleted file mode 100644
index 84a31a5a6970b..0000000000000
--- a/sql/hive/src/test/resources/golden/timestamp cast #5-0-dbd7bcd167d322d6617b884c02c7f247	
+++ /dev/null
@@ -1 +0,0 @@
--0.001
diff --git a/sql/hive/src/test/resources/golden/timestamp cast #7-0-1d70654217035f8ce5f64344f4c5a80f b/sql/hive/src/test/resources/golden/timestamp cast #7-0-1d70654217035f8ce5f64344f4c5a80f
deleted file mode 100644
index 3fbedf693b51d..0000000000000
--- a/sql/hive/src/test/resources/golden/timestamp cast #7-0-1d70654217035f8ce5f64344f4c5a80f	
+++ /dev/null
@@ -1 +0,0 @@
--2
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
index f0a7a6cc7a1e3..8a5acaf3e10bc 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveQuerySuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.hive.execution
 
 import java.io.File
+import java.sql.Timestamp
 import java.util.{Locale, TimeZone}
 
 import scala.util.Try
@@ -248,12 +249,17 @@ class HiveQuerySuite extends HiveComparisonTest with BeforeAndAfter {
       |IF(TRUE, CAST(NULL AS BINARY), CAST("1" AS BINARY)) AS COL18,
       |IF(FALSE, CAST(NULL AS DATE), CAST("1970-01-01" AS DATE)) AS COL19,
       |IF(TRUE, CAST(NULL AS DATE), CAST("1970-01-01" AS DATE)) AS COL20,
-      |IF(FALSE, CAST(NULL AS TIMESTAMP), CAST(1 AS TIMESTAMP)) AS COL21,
-      |IF(TRUE, CAST(NULL AS TIMESTAMP), CAST(1 AS TIMESTAMP)) AS COL22,
-      |IF(FALSE, CAST(NULL AS DECIMAL), CAST(1 AS DECIMAL)) AS COL23,
-      |IF(TRUE, CAST(NULL AS DECIMAL), CAST(1 AS DECIMAL)) AS COL24
+      |IF(TRUE, CAST(NULL AS TIMESTAMP), CAST(1 AS TIMESTAMP)) AS COL21,
+      |IF(FALSE, CAST(NULL AS DECIMAL), CAST(1 AS DECIMAL)) AS COL22,
+      |IF(TRUE, CAST(NULL AS DECIMAL), CAST(1 AS DECIMAL)) AS COL23
       |FROM src LIMIT 1""".stripMargin)
 
+  test("constant null testing timestamp") {
+    val r1 = sql("SELECT IF(FALSE, CAST(NULL AS TIMESTAMP), CAST(1 AS TIMESTAMP)) AS COL20")
+      .collect().head
+    assert(new Timestamp(1000) == r1.getTimestamp(0))
+  }
+
   createQueryTest("constant array",
   """
     |SELECT sort_array(
@@ -603,26 +609,32 @@ class HiveQuerySuite extends HiveComparisonTest with BeforeAndAfter {
   // Jdk version leads to different query output for double, so not use createQueryTest here
   test("timestamp cast #1") {
     val res = sql("SELECT CAST(CAST(1 AS TIMESTAMP) AS DOUBLE) FROM src LIMIT 1").collect().head
-    assert(0.001 == res.getDouble(0))
+    assert(1 == res.getDouble(0))
   }
 
   createQueryTest("timestamp cast #2",
     "SELECT CAST(CAST(1.2 AS TIMESTAMP) AS DOUBLE) FROM src LIMIT 1")
 
-  createQueryTest("timestamp cast #3",
-    "SELECT CAST(CAST(1200 AS TIMESTAMP) AS INT) FROM src LIMIT 1")
+  test("timestamp cast #3") {
+    val res = sql("SELECT CAST(CAST(1200 AS TIMESTAMP) AS INT) FROM src LIMIT 1").collect().head
+    assert(1200 == res.getInt(0))
+  }
 
   createQueryTest("timestamp cast #4",
     "SELECT CAST(CAST(1.2 AS TIMESTAMP) AS DOUBLE) FROM src LIMIT 1")
 
-  createQueryTest("timestamp cast #5",
-    "SELECT CAST(CAST(-1 AS TIMESTAMP) AS DOUBLE) FROM src LIMIT 1")
+  test("timestamp cast #5") {
+    val res = sql("SELECT CAST(CAST(-1 AS TIMESTAMP) AS DOUBLE) FROM src LIMIT 1").collect().head
+    assert(-1 == res.get(0))
+  }
 
   createQueryTest("timestamp cast #6",
     "SELECT CAST(CAST(-1.2 AS TIMESTAMP) AS DOUBLE) FROM src LIMIT 1")
 
-  createQueryTest("timestamp cast #7",
-    "SELECT CAST(CAST(-1200 AS TIMESTAMP) AS INT) FROM src LIMIT 1")
+  test("timestamp cast #7") {
+    val res = sql("SELECT CAST(CAST(-1200 AS TIMESTAMP) AS INT) FROM src LIMIT 1").collect().head
+    assert(-1200 == res.getInt(0))
+  }
 
   createQueryTest("timestamp cast #8",
     "SELECT CAST(CAST(-1.2 AS TIMESTAMP) AS DOUBLE) FROM src LIMIT 1")

From be7a2cfd978143f6f265eca63e9e24f755bc9f22 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Fri, 20 Nov 2015 14:23:01 -0800
Subject: [PATCH 388/518] [SPARK-11870][STREAMING][PYSPARK] Rethrow the
 exceptions in TransformFunction and TransformFunctionSerializer

TransformFunction and TransformFunctionSerializer don't rethrow the exception, so when any exception happens, it just return None. This will cause some weird NPE and confuse people.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9847 from zsxwing/pyspark-streaming-exception.
---
 python/pyspark/streaming/tests.py | 16 ++++++++++++++++
 python/pyspark/streaming/util.py  |  3 +++
 2 files changed, 19 insertions(+)

diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index 3403f6d20d789..a0e0267cafa58 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -403,6 +403,22 @@ def func(dstream):
         expected = [[('k', v)] for v in expected]
         self._test_func(input, func, expected)
 
+    def test_failed_func(self):
+        input = [self.sc.parallelize([d], 1) for d in range(4)]
+        input_stream = self.ssc.queueStream(input)
+
+        def failed_func(i):
+            raise ValueError("failed")
+
+        input_stream.map(failed_func).pprint()
+        self.ssc.start()
+        try:
+            self.ssc.awaitTerminationOrTimeout(10)
+        except:
+            return
+
+        self.fail("a failed func should throw an error")
+
 
 class StreamingListenerTests(PySparkStreamingTestCase):
 
diff --git a/python/pyspark/streaming/util.py b/python/pyspark/streaming/util.py
index b20613b1283bd..767c732eb90b4 100644
--- a/python/pyspark/streaming/util.py
+++ b/python/pyspark/streaming/util.py
@@ -64,6 +64,7 @@ def call(self, milliseconds, jrdds):
                 return r._jrdd
         except Exception:
             traceback.print_exc()
+            raise
 
     def __repr__(self):
         return "TransformFunction(%s)" % self.func
@@ -95,6 +96,7 @@ def dumps(self, id):
             return bytearray(self.serializer.dumps((func.func, func.deserializers)))
         except Exception:
             traceback.print_exc()
+            raise
 
     def loads(self, data):
         try:
@@ -102,6 +104,7 @@ def loads(self, data):
             return TransformFunction(self.ctx, f, *deserializers)
         except Exception:
             traceback.print_exc()
+            raise
 
     def __repr__(self):
         return "TransformFunctionSerializer(%s)" % self.serializer

From 89fd9bd06160fa89dedbf685bfe159ffe4a06ec6 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Fri, 20 Nov 2015 14:31:26 -0800
Subject: [PATCH 389/518] [SPARK-11887] Close PersistenceEngine at the end of
 PersistenceEngineSuite tests

In PersistenceEngineSuite, we do not call `close()` on the PersistenceEngine at the end of the test. For the ZooKeeperPersistenceEngine, this causes us to leak a ZooKeeper client, causing the logs of unrelated tests to be periodically spammed with connection error messages from that client:

```
15/11/20 05:13:35.789 pool-1-thread-1-ScalaTest-running-PersistenceEngineSuite-SendThread(localhost:15741) INFO ClientCnxn: Opening socket connection to server localhost/127.0.0.1:15741. Will not attempt to authenticate using SASL (unknown error)
15/11/20 05:13:35.790 pool-1-thread-1-ScalaTest-running-PersistenceEngineSuite-SendThread(localhost:15741) WARN ClientCnxn: Session 0x15124ff48dd0000 for server null, unexpected error, closing socket connection and attempting reconnect
java.net.ConnectException: Connection refused
	at sun.nio.ch.SocketChannelImpl.checkConnect(Native Method)
	at sun.nio.ch.SocketChannelImpl.finishConnect(SocketChannelImpl.java:739)
	at org.apache.zookeeper.ClientCnxnSocketNIO.doTransport(ClientCnxnSocketNIO.java:350)
	at org.apache.zookeeper.ClientCnxn$SendThread.run(ClientCnxn.java:1068)
```

This patch fixes this by using a `finally` block.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9864 from JoshRosen/close-zookeeper-client-in-tests.
---
 .../master/PersistenceEngineSuite.scala       | 100 +++++++++---------
 1 file changed, 52 insertions(+), 48 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/deploy/master/PersistenceEngineSuite.scala b/core/src/test/scala/org/apache/spark/deploy/master/PersistenceEngineSuite.scala
index 34775577de8a3..7a44728675680 100644
--- a/core/src/test/scala/org/apache/spark/deploy/master/PersistenceEngineSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/master/PersistenceEngineSuite.scala
@@ -63,56 +63,60 @@ class PersistenceEngineSuite extends SparkFunSuite {
       conf: SparkConf, persistenceEngineCreator: Serializer => PersistenceEngine): Unit = {
     val serializer = new JavaSerializer(conf)
     val persistenceEngine = persistenceEngineCreator(serializer)
-    persistenceEngine.persist("test_1", "test_1_value")
-    assert(Seq("test_1_value") === persistenceEngine.read[String]("test_"))
-    persistenceEngine.persist("test_2", "test_2_value")
-    assert(Set("test_1_value", "test_2_value") === persistenceEngine.read[String]("test_").toSet)
-    persistenceEngine.unpersist("test_1")
-    assert(Seq("test_2_value") === persistenceEngine.read[String]("test_"))
-    persistenceEngine.unpersist("test_2")
-    assert(persistenceEngine.read[String]("test_").isEmpty)
-
-    // Test deserializing objects that contain RpcEndpointRef
-    val testRpcEnv = RpcEnv.create("test", "localhost", 12345, conf, new SecurityManager(conf))
     try {
-      // Create a real endpoint so that we can test RpcEndpointRef deserialization
-      val workerEndpoint = testRpcEnv.setupEndpoint("worker", new RpcEndpoint {
-        override val rpcEnv: RpcEnv = testRpcEnv
-      })
-
-      val workerToPersist = new WorkerInfo(
-        id = "test_worker",
-        host = "127.0.0.1",
-        port = 10000,
-        cores = 0,
-        memory = 0,
-        endpoint = workerEndpoint,
-        webUiPort = 0,
-        publicAddress = ""
-      )
-
-      persistenceEngine.addWorker(workerToPersist)
-
-      val (storedApps, storedDrivers, storedWorkers) =
-        persistenceEngine.readPersistedData(testRpcEnv)
-
-      assert(storedApps.isEmpty)
-      assert(storedDrivers.isEmpty)
-
-      // Check deserializing WorkerInfo
-      assert(storedWorkers.size == 1)
-      val recoveryWorkerInfo = storedWorkers.head
-      assert(workerToPersist.id === recoveryWorkerInfo.id)
-      assert(workerToPersist.host === recoveryWorkerInfo.host)
-      assert(workerToPersist.port === recoveryWorkerInfo.port)
-      assert(workerToPersist.cores === recoveryWorkerInfo.cores)
-      assert(workerToPersist.memory === recoveryWorkerInfo.memory)
-      assert(workerToPersist.endpoint === recoveryWorkerInfo.endpoint)
-      assert(workerToPersist.webUiPort === recoveryWorkerInfo.webUiPort)
-      assert(workerToPersist.publicAddress === recoveryWorkerInfo.publicAddress)
+      persistenceEngine.persist("test_1", "test_1_value")
+      assert(Seq("test_1_value") === persistenceEngine.read[String]("test_"))
+      persistenceEngine.persist("test_2", "test_2_value")
+      assert(Set("test_1_value", "test_2_value") === persistenceEngine.read[String]("test_").toSet)
+      persistenceEngine.unpersist("test_1")
+      assert(Seq("test_2_value") === persistenceEngine.read[String]("test_"))
+      persistenceEngine.unpersist("test_2")
+      assert(persistenceEngine.read[String]("test_").isEmpty)
+
+      // Test deserializing objects that contain RpcEndpointRef
+      val testRpcEnv = RpcEnv.create("test", "localhost", 12345, conf, new SecurityManager(conf))
+      try {
+        // Create a real endpoint so that we can test RpcEndpointRef deserialization
+        val workerEndpoint = testRpcEnv.setupEndpoint("worker", new RpcEndpoint {
+          override val rpcEnv: RpcEnv = testRpcEnv
+        })
+
+        val workerToPersist = new WorkerInfo(
+          id = "test_worker",
+          host = "127.0.0.1",
+          port = 10000,
+          cores = 0,
+          memory = 0,
+          endpoint = workerEndpoint,
+          webUiPort = 0,
+          publicAddress = ""
+        )
+
+        persistenceEngine.addWorker(workerToPersist)
+
+        val (storedApps, storedDrivers, storedWorkers) =
+          persistenceEngine.readPersistedData(testRpcEnv)
+
+        assert(storedApps.isEmpty)
+        assert(storedDrivers.isEmpty)
+
+        // Check deserializing WorkerInfo
+        assert(storedWorkers.size == 1)
+        val recoveryWorkerInfo = storedWorkers.head
+        assert(workerToPersist.id === recoveryWorkerInfo.id)
+        assert(workerToPersist.host === recoveryWorkerInfo.host)
+        assert(workerToPersist.port === recoveryWorkerInfo.port)
+        assert(workerToPersist.cores === recoveryWorkerInfo.cores)
+        assert(workerToPersist.memory === recoveryWorkerInfo.memory)
+        assert(workerToPersist.endpoint === recoveryWorkerInfo.endpoint)
+        assert(workerToPersist.webUiPort === recoveryWorkerInfo.webUiPort)
+        assert(workerToPersist.publicAddress === recoveryWorkerInfo.publicAddress)
+      } finally {
+        testRpcEnv.shutdown()
+        testRpcEnv.awaitTermination()
+      }
     } finally {
-      testRpcEnv.shutdown()
-      testRpcEnv.awaitTermination()
+      persistenceEngine.close()
     }
   }
 

From 03ba56d78f50747710d01c27d409ba2be42ae557 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-Baptiste=20Onofr=C3=A9?= <jbonofre@apache.org>
Date: Fri, 20 Nov 2015 14:45:40 -0800
Subject: [PATCH 390/518] [SPARK-11716][SQL] UDFRegistration just drops the
 input type when re-creating the UserDefinedFunction
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

https://issues.apache.org/jira/browse/SPARK-11716

This is one is #9739 and a regression test. When commit it, please make sure the author is jbonofre.

You can find the original PR at https://github.com/apache/spark/pull/9739

closes #9739

Author: Jean-Baptiste Onofré <jbonofre@apache.org>
Author: Yin Huai <yhuai@databricks.com>

Closes #9868 from yhuai/SPARK-11716.
---
 .../apache/spark/sql/UDFRegistration.scala    | 48 +++++++++----------
 .../scala/org/apache/spark/sql/UDFSuite.scala | 15 ++++++
 2 files changed, 39 insertions(+), 24 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala b/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
index fc4d0938c533a..051694c0d43a6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/UDFRegistration.scala
@@ -88,7 +88,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
           val inputTypes = Try($inputTypes).getOrElse(Nil)
           def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
           functionRegistry.registerFunction(name, builder)
-          UserDefinedFunction(func, dataType)
+          UserDefinedFunction(func, dataType, inputTypes)
         }""")
     }
 
@@ -120,7 +120,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -133,7 +133,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -146,7 +146,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -159,7 +159,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -172,7 +172,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -185,7 +185,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -198,7 +198,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -211,7 +211,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -224,7 +224,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -237,7 +237,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -250,7 +250,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -263,7 +263,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -276,7 +276,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -289,7 +289,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -302,7 +302,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -315,7 +315,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -328,7 +328,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -341,7 +341,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -354,7 +354,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -367,7 +367,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -380,7 +380,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -393,7 +393,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: ScalaReflection.schemaFor[A21].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   /**
@@ -406,7 +406,7 @@ class UDFRegistration private[sql] (sqlContext: SQLContext) extends Logging {
     val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: ScalaReflection.schemaFor[A2].dataType :: ScalaReflection.schemaFor[A3].dataType :: ScalaReflection.schemaFor[A4].dataType :: ScalaReflection.schemaFor[A5].dataType :: ScalaReflection.schemaFor[A6].dataType :: ScalaReflection.schemaFor[A7].dataType :: ScalaReflection.schemaFor[A8].dataType :: ScalaReflection.schemaFor[A9].dataType :: ScalaReflection.schemaFor[A10].dataType :: ScalaReflection.schemaFor[A11].dataType :: ScalaReflection.schemaFor[A12].dataType :: ScalaReflection.schemaFor[A13].dataType :: ScalaReflection.schemaFor[A14].dataType :: ScalaReflection.schemaFor[A15].dataType :: ScalaReflection.schemaFor[A16].dataType :: ScalaReflection.schemaFor[A17].dataType :: ScalaReflection.schemaFor[A18].dataType :: ScalaReflection.schemaFor[A19].dataType :: ScalaReflection.schemaFor[A20].dataType :: ScalaReflection.schemaFor[A21].dataType :: ScalaReflection.schemaFor[A22].dataType :: Nil).getOrElse(Nil)
     def builder(e: Seq[Expression]) = ScalaUDF(func, dataType, e, inputTypes)
     functionRegistry.registerFunction(name, builder)
-    UserDefinedFunction(func, dataType)
+    UserDefinedFunction(func, dataType, inputTypes)
   }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
index 9837fa6bdb357..fd736718af12c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala
@@ -232,4 +232,19 @@ class UDFSuite extends QueryTest with SharedSQLContext {
            | (SELECT complexDataFunc(m, a, b) AS t FROM complexData) tmp
           """.stripMargin).toDF(), complexData.select("m", "a", "b"))
   }
+
+  test("SPARK-11716 UDFRegistration does not include the input data type in returned UDF") {
+    val myUDF = sqlContext.udf.register("testDataFunc", (n: Int, s: String) => { (n, s.toInt) })
+
+    // Without the fix, this will fail because we fail to cast data type of b to string
+    // because myUDF does not know its input data type. With the fix, this query should not
+    // fail.
+    checkAnswer(
+      testData2.select(myUDF($"a", $"b").as("t")),
+      testData2.selectExpr("struct(a, b)"))
+
+    checkAnswer(
+      sql("SELECT tmp.t.* FROM (SELECT testDataFunc(a, b) AS t from testData2) tmp").toDF(),
+      testData2)
+  }
 }

From a6239d587c638691f52eca3eee905c53fbf35a12 Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Fri, 20 Nov 2015 15:10:55 -0800
Subject: [PATCH 391/518] [SPARK-11756][SPARKR] Fix use of aliases - SparkR can
 not output help information for SparkR:::summary correctly

Fix use of aliases and changes uses of rdname and seealso
`aliases` is the hint for `?` - it should not be linked to some other name - those should be seealso
https://cran.r-project.org/web/packages/roxygen2/vignettes/rd.html

Clean up usage on family, as multiple use of family with the same rdname is causing duplicated See Also html blocks (like http://spark.apache.org/docs/latest/api/R/count.html)
Also changing some rdname for dplyr-like variant for better R user visibility in R doc, eg. rbind, summary, mutate, summarize

shivaram yanboliang

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9750 from felixcheung/rdocaliases.
---
 R/pkg/R/DataFrame.R | 96 ++++++++++++---------------------------------
 R/pkg/R/broadcast.R |  1 -
 R/pkg/R/generics.R  | 12 +++---
 R/pkg/R/group.R     | 12 +++---
 4 files changed, 37 insertions(+), 84 deletions(-)

diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index 06b0108b1389e..8a13e7a36766d 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -254,7 +254,6 @@ setMethod("dtypes",
 #' @family DataFrame functions
 #' @rdname columns
 #' @name columns
-#' @aliases names
 #' @export
 #' @examples
 #'\dontrun{
@@ -272,7 +271,6 @@ setMethod("columns",
             })
           })
 
-#' @family DataFrame functions
 #' @rdname columns
 #' @name names
 setMethod("names",
@@ -281,7 +279,6 @@ setMethod("names",
             columns(x)
           })
 
-#' @family DataFrame functions
 #' @rdname columns
 #' @name names<-
 setMethod("names<-",
@@ -533,14 +530,8 @@ setMethod("distinct",
             dataFrame(sdf)
           })
 
-#' @title Distinct rows in a DataFrame
-#
-#' @description Returns a new DataFrame containing distinct rows in this DataFrame
-#'
-#' @family DataFrame functions
-#' @rdname unique
+#' @rdname distinct
 #' @name unique
-#' @aliases distinct
 setMethod("unique",
           signature(x = "DataFrame"),
           function(x) {
@@ -557,7 +548,7 @@ setMethod("unique",
 #'
 #' @family DataFrame functions
 #' @rdname sample
-#' @aliases sample_frac
+#' @name sample
 #' @export
 #' @examples
 #'\dontrun{
@@ -579,7 +570,6 @@ setMethod("sample",
             dataFrame(sdf)
           })
 
-#' @family DataFrame functions
 #' @rdname sample
 #' @name sample_frac
 setMethod("sample_frac",
@@ -589,16 +579,15 @@ setMethod("sample_frac",
             sample(x, withReplacement, fraction)
           })
 
-#' Count
+#' nrow
 #'
 #' Returns the number of rows in a DataFrame
 #'
 #' @param x A SparkSQL DataFrame
 #'
 #' @family DataFrame functions
-#' @rdname count
+#' @rdname nrow
 #' @name count
-#' @aliases nrow
 #' @export
 #' @examples
 #'\dontrun{
@@ -614,14 +603,8 @@ setMethod("count",
             callJMethod(x@sdf, "count")
           })
 
-#' @title Number of rows for a DataFrame
-#' @description Returns number of rows in a DataFrames
-#'
 #' @name nrow
-#'
-#' @family DataFrame functions
 #' @rdname nrow
-#' @aliases count
 setMethod("nrow",
           signature(x = "DataFrame"),
           function(x) {
@@ -870,7 +853,6 @@ setMethod("toRDD",
 #' @param x a DataFrame
 #' @return a GroupedData
 #' @seealso GroupedData
-#' @aliases group_by
 #' @family DataFrame functions
 #' @rdname groupBy
 #' @name groupBy
@@ -896,7 +878,6 @@ setMethod("groupBy",
              groupedData(sgd)
            })
 
-#' @family DataFrame functions
 #' @rdname groupBy
 #' @name group_by
 setMethod("group_by",
@@ -913,7 +894,6 @@ setMethod("group_by",
 #' @family DataFrame functions
 #' @rdname agg
 #' @name agg
-#' @aliases summarize
 #' @export
 setMethod("agg",
           signature(x = "DataFrame"),
@@ -921,7 +901,6 @@ setMethod("agg",
             agg(groupBy(x), ...)
           })
 
-#' @family DataFrame functions
 #' @rdname agg
 #' @name summarize
 setMethod("summarize",
@@ -1092,7 +1071,6 @@ setMethod("[", signature(x = "DataFrame", i = "Column"),
 #' @family DataFrame functions
 #' @rdname subset
 #' @name subset
-#' @aliases [
 #' @family subsetting functions
 #' @examples
 #' \dontrun{
@@ -1216,7 +1194,7 @@ setMethod("selectExpr",
 #' @family DataFrame functions
 #' @rdname withColumn
 #' @name withColumn
-#' @aliases mutate transform
+#' @seealso \link{rename} \link{mutate}
 #' @export
 #' @examples
 #'\dontrun{
@@ -1231,7 +1209,6 @@ setMethod("withColumn",
           function(x, colName, col) {
             select(x, x$"*", alias(col, colName))
           })
-
 #' Mutate
 #'
 #' Return a new DataFrame with the specified columns added.
@@ -1240,9 +1217,9 @@ setMethod("withColumn",
 #' @param col a named argument of the form name = col
 #' @return A new DataFrame with the new columns added.
 #' @family DataFrame functions
-#' @rdname withColumn
+#' @rdname mutate
 #' @name mutate
-#' @aliases withColumn transform
+#' @seealso \link{rename} \link{withColumn}
 #' @export
 #' @examples
 #'\dontrun{
@@ -1273,17 +1250,15 @@ setMethod("mutate",
           })
 
 #' @export
-#' @family DataFrame functions
-#' @rdname withColumn
+#' @rdname mutate
 #' @name transform
-#' @aliases withColumn mutate
 setMethod("transform",
           signature(`_data` = "DataFrame"),
           function(`_data`, ...) {
             mutate(`_data`, ...)
           })
 
-#' WithColumnRenamed
+#' rename
 #'
 #' Rename an existing column in a DataFrame.
 #'
@@ -1292,8 +1267,9 @@ setMethod("transform",
 #' @param newCol The new column name.
 #' @return A DataFrame with the column name changed.
 #' @family DataFrame functions
-#' @rdname withColumnRenamed
+#' @rdname rename
 #' @name withColumnRenamed
+#' @seealso \link{mutate}
 #' @export
 #' @examples
 #'\dontrun{
@@ -1316,17 +1292,9 @@ setMethod("withColumnRenamed",
             select(x, cols)
           })
 
-#' Rename
-#'
-#' Rename an existing column in a DataFrame.
-#'
-#' @param x A DataFrame
-#' @param newCol A named pair of the form new_column_name = existing_column
-#' @return A DataFrame with the column name changed.
-#' @family DataFrame functions
-#' @rdname withColumnRenamed
+#' @param newColPair A named pair of the form new_column_name = existing_column
+#' @rdname rename
 #' @name rename
-#' @aliases withColumnRenamed
 #' @export
 #' @examples
 #'\dontrun{
@@ -1371,7 +1339,6 @@ setClassUnion("characterOrColumn", c("character", "Column"))
 #' @family DataFrame functions
 #' @rdname arrange
 #' @name arrange
-#' @aliases orderby
 #' @export
 #' @examples
 #'\dontrun{
@@ -1395,8 +1362,8 @@ setMethod("arrange",
             dataFrame(sdf)
           })
 
-#' @family DataFrame functions
 #' @rdname arrange
+#' @name arrange
 #' @export
 setMethod("arrange",
           signature(x = "DataFrame", col = "character"),
@@ -1427,9 +1394,9 @@ setMethod("arrange",
             do.call("arrange", c(x, jcols))
           })
 
-#' @family DataFrame functions
 #' @rdname arrange
-#' @name orderby
+#' @name orderBy
+#' @export
 setMethod("orderBy",
           signature(x = "DataFrame", col = "characterOrColumn"),
           function(x, col) {
@@ -1492,6 +1459,7 @@ setMethod("where",
 #' @family DataFrame functions
 #' @rdname join
 #' @name join
+#' @seealso \link{merge}
 #' @export
 #' @examples
 #'\dontrun{
@@ -1528,9 +1496,7 @@ setMethod("join",
             dataFrame(sdf)
           })
 
-#'
 #' @name merge
-#' @aliases join
 #' @title Merges two data frames
 #' @param x the first data frame to be joined
 #' @param y the second data frame to be joined
@@ -1550,6 +1516,7 @@ setMethod("join",
 #'   outer join will be returned.
 #' @family DataFrame functions
 #' @rdname merge
+#' @seealso \link{join}
 #' @export
 #' @examples
 #'\dontrun{
@@ -1671,7 +1638,7 @@ generateAliasesForIntersectedCols <- function (x, intersectedColNames, suffix) {
   cols
 }
 
-#' UnionAll
+#' rbind
 #'
 #' Return a new DataFrame containing the union of rows in this DataFrame
 #' and another DataFrame. This is equivalent to `UNION ALL` in SQL.
@@ -1681,7 +1648,7 @@ generateAliasesForIntersectedCols <- function (x, intersectedColNames, suffix) {
 #' @param y A Spark DataFrame
 #' @return A DataFrame containing the result of the union.
 #' @family DataFrame functions
-#' @rdname unionAll
+#' @rdname rbind
 #' @name unionAll
 #' @export
 #' @examples
@@ -1700,13 +1667,11 @@ setMethod("unionAll",
           })
 
 #' @title Union two or more DataFrames
-#'
 #' @description Returns a new DataFrame containing rows of all parameters.
 #'
-#' @family DataFrame functions
 #' @rdname rbind
 #' @name rbind
-#' @aliases unionAll
+#' @export
 setMethod("rbind",
           signature(... = "DataFrame"),
           function(x, ..., deparse.level = 1) {
@@ -1795,7 +1760,6 @@ setMethod("except",
 #' @family DataFrame functions
 #' @rdname write.df
 #' @name write.df
-#' @aliases saveDF
 #' @export
 #' @examples
 #'\dontrun{
@@ -1828,7 +1792,6 @@ setMethod("write.df",
             callJMethod(df@sdf, "save", source, jmode, options)
           })
 
-#' @family DataFrame functions
 #' @rdname write.df
 #' @name saveDF
 #' @export
@@ -1891,7 +1854,7 @@ setMethod("saveAsTable",
             callJMethod(df@sdf, "saveAsTable", tableName, source, jmode, options)
           })
 
-#' describe
+#' summary
 #'
 #' Computes statistics for numeric columns.
 #' If no columns are given, this function computes statistics for all numerical columns.
@@ -1901,9 +1864,8 @@ setMethod("saveAsTable",
 #' @param ... Additional expressions
 #' @return A DataFrame
 #' @family DataFrame functions
-#' @rdname describe
+#' @rdname summary
 #' @name describe
-#' @aliases summary
 #' @export
 #' @examples
 #'\dontrun{
@@ -1923,8 +1885,7 @@ setMethod("describe",
             dataFrame(sdf)
           })
 
-#' @family DataFrame functions
-#' @rdname describe
+#' @rdname summary
 #' @name describe
 setMethod("describe",
           signature(x = "DataFrame"),
@@ -1934,11 +1895,6 @@ setMethod("describe",
             dataFrame(sdf)
           })
 
-#' @title Summary
-#'
-#' @description Computes statistics for numeric columns of the DataFrame
-#'
-#' @family DataFrame functions
 #' @rdname summary
 #' @name summary
 setMethod("summary",
@@ -1966,7 +1922,6 @@ setMethod("summary",
 #' @family DataFrame functions
 #' @rdname nafunctions
 #' @name dropna
-#' @aliases na.omit
 #' @export
 #' @examples
 #'\dontrun{
@@ -1993,7 +1948,6 @@ setMethod("dropna",
             dataFrame(sdf)
           })
 
-#' @family DataFrame functions
 #' @rdname nafunctions
 #' @name na.omit
 #' @export
@@ -2019,9 +1973,7 @@ setMethod("na.omit",
 #'             type are ignored. For example, if value is a character, and
 #'             subset contains a non-character column, then the non-character
 #'             column is simply ignored.
-#' @return A DataFrame
 #'
-#' @family DataFrame functions
 #' @rdname nafunctions
 #' @name fillna
 #' @export
diff --git a/R/pkg/R/broadcast.R b/R/pkg/R/broadcast.R
index 2403925b267c8..38f0eed95e065 100644
--- a/R/pkg/R/broadcast.R
+++ b/R/pkg/R/broadcast.R
@@ -51,7 +51,6 @@ Broadcast <- function(id, value, jBroadcastRef, objName) {
 #
 # @param bcast The broadcast variable to get
 # @rdname broadcast
-# @aliases value,Broadcast-method
 setMethod("value",
           signature(bcast = "Broadcast"),
           function(bcast) {
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 71004a05ba611..1b3f10ea04643 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -397,7 +397,7 @@ setGeneric("cov", function(x, col1, col2) {standardGeneric("cov") })
 #' @export
 setGeneric("corr", function(x, col1, col2, method = "pearson") {standardGeneric("corr") })
 
-#' @rdname describe
+#' @rdname summary
 #' @export
 setGeneric("describe", function(x, col, ...) { standardGeneric("describe") })
 
@@ -459,11 +459,11 @@ setGeneric("isLocal", function(x) { standardGeneric("isLocal") })
 #' @export
 setGeneric("limit", function(x, num) {standardGeneric("limit") })
 
-#' rdname merge
+#' @rdname merge
 #' @export
 setGeneric("merge")
 
-#' @rdname withColumn
+#' @rdname mutate
 #' @export
 setGeneric("mutate", function(.data, ...) {standardGeneric("mutate") })
 
@@ -475,7 +475,7 @@ setGeneric("orderBy", function(x, col) { standardGeneric("orderBy") })
 #' @export
 setGeneric("printSchema", function(x) { standardGeneric("printSchema") })
 
-#' @rdname withColumnRenamed
+#' @rdname rename
 #' @export
 setGeneric("rename", function(x, ...) { standardGeneric("rename") })
 
@@ -553,7 +553,7 @@ setGeneric("toJSON", function(x) { standardGeneric("toJSON") })
 
 setGeneric("toRDD", function(x) { standardGeneric("toRDD") })
 
-#' @rdname unionAll
+#' @rdname rbind
 #' @export
 setGeneric("unionAll", function(x, y) { standardGeneric("unionAll") })
 
@@ -565,7 +565,7 @@ setGeneric("where", function(x, condition) { standardGeneric("where") })
 #' @export
 setGeneric("withColumn", function(x, colName, col) { standardGeneric("withColumn") })
 
-#' @rdname withColumnRenamed
+#' @rdname rename
 #' @export
 setGeneric("withColumnRenamed",
            function(x, existingCol, newCol) { standardGeneric("withColumnRenamed") })
diff --git a/R/pkg/R/group.R b/R/pkg/R/group.R
index e5f702faee65d..23b49aebda05f 100644
--- a/R/pkg/R/group.R
+++ b/R/pkg/R/group.R
@@ -68,7 +68,7 @@ setMethod("count",
             dataFrame(callJMethod(x@sgd, "count"))
           })
 
-#' Agg
+#' summarize
 #'
 #' Aggregates on the entire DataFrame without groups.
 #' The resulting DataFrame will also contain the grouping columns.
@@ -78,12 +78,14 @@ setMethod("count",
 #'
 #' @param x a GroupedData
 #' @return a DataFrame
-#' @rdname agg
+#' @rdname summarize
+#' @name agg
 #' @family agg_funcs
 #' @examples
 #' \dontrun{
 #'  df2 <- agg(df, age = "sum")  # new column name will be created as 'SUM(age#0)'
-#'  df2 <- agg(df, ageSum = sum(df$age)) # Creates a new column named ageSum
+#'  df3 <- agg(df, ageSum = sum(df$age)) # Creates a new column named ageSum
+#'  df4 <- summarize(df, ageSum = max(df$age))
 #' }
 setMethod("agg",
           signature(x = "GroupedData"),
@@ -110,8 +112,8 @@ setMethod("agg",
             dataFrame(sdf)
           })
 
-#' @rdname agg
-#' @aliases agg
+#' @rdname summarize
+#' @name summarize
 setMethod("summarize",
           signature(x = "GroupedData"),
           function(x, ...) {

From 4b84c72dfbb9ddb415fee35f69305b5d7b280891 Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Fri, 20 Nov 2015 15:17:17 -0800
Subject: [PATCH 392/518] [SPARK-11636][SQL] Support classes defined in the
 REPL with Encoders

#theScaryParts (i.e. changes to the repl, executor classloaders and codegen)...

Author: Michael Armbrust <michael@databricks.com>
Author: Yin Huai <yhuai@databricks.com>

Closes #9825 from marmbrus/dataset-replClasses2.
---
 .../org/apache/spark/repl/SparkIMain.scala    | 14 +++++++----
 .../org/apache/spark/repl/ReplSuite.scala     | 24 +++++++++++++++++++
 .../spark/repl/ExecutorClassLoader.scala      |  8 ++++++-
 .../expressions/codegen/CodeGenerator.scala   |  4 ++--
 4 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/repl/scala-2.10/src/main/scala/org/apache/spark/repl/SparkIMain.scala b/repl/scala-2.10/src/main/scala/org/apache/spark/repl/SparkIMain.scala
index 4ee605fd7f11e..829b12269fd2b 100644
--- a/repl/scala-2.10/src/main/scala/org/apache/spark/repl/SparkIMain.scala
+++ b/repl/scala-2.10/src/main/scala/org/apache/spark/repl/SparkIMain.scala
@@ -1221,10 +1221,16 @@ import org.apache.spark.annotation.DeveloperApi
         )
       }
 
-      val preamble = """
-        |class %s extends Serializable {
-        |  %s%s%s
-      """.stripMargin.format(lineRep.readName, envLines.map("  " + _ + ";\n").mkString, importsPreamble, indentCode(toCompute))
+      val preamble = s"""
+        |class ${lineRep.readName} extends Serializable {
+        |  ${envLines.map("  " + _ + ";\n").mkString}
+        |  $importsPreamble
+        |
+        |  // If we need to construct any objects defined in the REPL on an executor we will need
+        |  // to pass the outer scope to the appropriate encoder.
+        |  org.apache.spark.sql.catalyst.encoders.OuterScopes.addOuterScope(this)
+        |  ${indentCode(toCompute)}
+      """.stripMargin
       val postamble = importsTrailer + "\n}" + "\n" +
         "object " + lineRep.readName + " {\n" +
         "  val INSTANCE = new " + lineRep.readName + "();\n" +
diff --git a/repl/scala-2.10/src/test/scala/org/apache/spark/repl/ReplSuite.scala b/repl/scala-2.10/src/test/scala/org/apache/spark/repl/ReplSuite.scala
index 5674dcd669bee..081aa03002cc6 100644
--- a/repl/scala-2.10/src/test/scala/org/apache/spark/repl/ReplSuite.scala
+++ b/repl/scala-2.10/src/test/scala/org/apache/spark/repl/ReplSuite.scala
@@ -262,6 +262,9 @@ class ReplSuite extends SparkFunSuite {
         |import sqlContext.implicits._
         |case class TestCaseClass(value: Int)
         |sc.parallelize(1 to 10).map(x => TestCaseClass(x)).toDF().collect()
+        |
+        |// Test Dataset Serialization in the REPL
+        |Seq(TestCaseClass(1)).toDS().collect()
       """.stripMargin)
     assertDoesNotContain("error:", output)
     assertDoesNotContain("Exception", output)
@@ -278,6 +281,27 @@ class ReplSuite extends SparkFunSuite {
     assertDoesNotContain("java.lang.ClassNotFoundException", output)
   }
 
+  test("Datasets and encoders") {
+    val output = runInterpreter("local",
+      """
+        |import org.apache.spark.sql.functions._
+        |import org.apache.spark.sql.Encoder
+        |import org.apache.spark.sql.expressions.Aggregator
+        |import org.apache.spark.sql.TypedColumn
+        |val simpleSum = new Aggregator[Int, Int, Int] with Serializable {
+        |  def zero: Int = 0                     // The initial value.
+        |  def reduce(b: Int, a: Int) = b + a    // Add an element to the running total
+        |  def merge(b1: Int, b2: Int) = b1 + b2 // Merge intermediate values.
+        |  def finish(b: Int) = b                // Return the final result.
+        |}.toColumn
+        |
+        |val ds = Seq(1, 2, 3, 4).toDS()
+        |ds.select(simpleSum).collect
+      """.stripMargin)
+    assertDoesNotContain("error:", output)
+    assertDoesNotContain("Exception", output)
+  }
+
   test("SPARK-2632 importing a method from non serializable class and not using it.") {
     val output = runInterpreter("local",
     """
diff --git a/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala b/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
index 3d2d235a00c93..a976e96809cb8 100644
--- a/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
+++ b/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
@@ -65,7 +65,13 @@ class ExecutorClassLoader(conf: SparkConf, classUri: String, parent: ClassLoader
           case e: ClassNotFoundException => {
             val classOption = findClassLocally(name)
             classOption match {
-              case None => throw new ClassNotFoundException(name, e)
+              case None =>
+                // If this class has a cause, it will break the internal assumption of Janino
+                // (the compiler used for Spark SQL code-gen).
+                // See org.codehaus.janino.ClassLoaderIClassLoader's findIClass, you will see
+                // its behavior will be changed if there is a cause and the compilation
+                // of generated class will fail.
+                throw new ClassNotFoundException(name)
               case Some(a) => a
             }
           }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
index 1b7260cdfe515..2f3d6aeb86c5b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
@@ -31,7 +31,7 @@ import org.apache.spark.sql.catalyst.util.{MapData, ArrayData}
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.Platform
 import org.apache.spark.unsafe.types._
-
+import org.apache.spark.util.Utils
 
 /**
  * Java source for evaluating an [[Expression]] given a [[InternalRow]] of input.
@@ -536,7 +536,7 @@ abstract class CodeGenerator[InType <: AnyRef, OutType <: AnyRef] extends Loggin
    */
   private[this] def doCompile(code: String): GeneratedClass = {
     val evaluator = new ClassBodyEvaluator()
-    evaluator.setParentClassLoader(getClass.getClassLoader)
+    evaluator.setParentClassLoader(Utils.getContextOrSparkClassLoader)
     // Cannot be under package codegen, or fail with java.lang.InstantiationException
     evaluator.setClassName("org.apache.spark.sql.catalyst.expressions.GeneratedClass")
     evaluator.setDefaultImports(Array(

From ed47b1e660b830e2d4fac8d6df93f634b260393c Mon Sep 17 00:00:00 2001
From: Vikas Nelamangala <vikasnelamangala@Vikass-MacBook-Pro.local>
Date: Fri, 20 Nov 2015 15:18:41 -0800
Subject: [PATCH 393/518] [SPARK-11549][DOCS] Replace example code in
 mllib-evaluation-metrics.md using include_example

Author: Vikas Nelamangala <vikasnelamangala@Vikass-MacBook-Pro.local>

Closes #9689 from vikasnp/master.
---
 docs/mllib-evaluation-metrics.md              | 940 +-----------------
 ...avaBinaryClassificationMetricsExample.java | 113 +++
 ...ultiLabelClassificationMetricsExample.java |  80 ++
 ...ulticlassClassificationMetricsExample.java |  97 ++
 .../mllib/JavaRankingMetricsExample.java      | 176 ++++
 .../mllib/JavaRegressionMetricsExample.java   |  91 ++
 .../binary_classification_metrics_example.py  |  55 +
 .../mllib/multi_class_metrics_example.py      |  69 ++
 .../mllib/multi_label_metrics_example.py      |  61 ++
 .../python/mllib/ranking_metrics_example.py   |  55 +
 .../mllib/regression_metrics_example.py       |  59 ++
 .../BinaryClassificationMetricsExample.scala  | 103 ++
 .../mllib/MultiLabelMetricsExample.scala      |  69 ++
 .../mllib/MulticlassMetricsExample.scala      |  99 ++
 .../mllib/RankingMetricsExample.scala         | 110 ++
 .../mllib/RegressionMetricsExample.scala      |  67 ++
 16 files changed, 1319 insertions(+), 925 deletions(-)
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaBinaryClassificationMetricsExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaMultiLabelClassificationMetricsExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaMulticlassClassificationMetricsExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaRankingMetricsExample.java
 create mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaRegressionMetricsExample.java
 create mode 100644 examples/src/main/python/mllib/binary_classification_metrics_example.py
 create mode 100644 examples/src/main/python/mllib/multi_class_metrics_example.py
 create mode 100644 examples/src/main/python/mllib/multi_label_metrics_example.py
 create mode 100644 examples/src/main/python/mllib/ranking_metrics_example.py
 create mode 100644 examples/src/main/python/mllib/regression_metrics_example.py
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/BinaryClassificationMetricsExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/MultiLabelMetricsExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/MulticlassMetricsExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/RankingMetricsExample.scala
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/mllib/RegressionMetricsExample.scala

diff --git a/docs/mllib-evaluation-metrics.md b/docs/mllib-evaluation-metrics.md
index f73eff637dc36..6924037b941f3 100644
--- a/docs/mllib-evaluation-metrics.md
+++ b/docs/mllib-evaluation-metrics.md
@@ -104,214 +104,21 @@ data, and evaluate the performance of the algorithm by several binary evaluation
 <div data-lang="scala" markdown="1">
 Refer to the [`LogisticRegressionWithLBFGS` Scala docs](api/scala/index.html#org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS) and [`BinaryClassificationMetrics` Scala docs](api/scala/index.html#org.apache.spark.mllib.evaluation.BinaryClassificationMetrics) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS
-import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics
-import org.apache.spark.mllib.regression.LabeledPoint
-import org.apache.spark.mllib.util.MLUtils
-
-// Load training data in LIBSVM format
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_binary_classification_data.txt")
-
-// Split data into training (60%) and test (40%)
-val Array(training, test) = data.randomSplit(Array(0.6, 0.4), seed = 11L)
-training.cache()
-
-// Run training algorithm to build the model
-val model = new LogisticRegressionWithLBFGS()
-  .setNumClasses(2)
-  .run(training)
-
-// Clear the prediction threshold so the model will return probabilities
-model.clearThreshold
-
-// Compute raw scores on the test set
-val predictionAndLabels = test.map { case LabeledPoint(label, features) =>
-  val prediction = model.predict(features)
-  (prediction, label)
-}
-
-// Instantiate metrics object
-val metrics = new BinaryClassificationMetrics(predictionAndLabels)
-
-// Precision by threshold
-val precision = metrics.precisionByThreshold
-precision.foreach { case (t, p) =>
-    println(s"Threshold: $t, Precision: $p")
-}
-
-// Recall by threshold
-val recall = metrics.recallByThreshold
-recall.foreach { case (t, r) =>
-    println(s"Threshold: $t, Recall: $r")
-}
-
-// Precision-Recall Curve
-val PRC = metrics.pr
-
-// F-measure
-val f1Score = metrics.fMeasureByThreshold
-f1Score.foreach { case (t, f) =>
-    println(s"Threshold: $t, F-score: $f, Beta = 1")
-}
-
-val beta = 0.5
-val fScore = metrics.fMeasureByThreshold(beta)
-f1Score.foreach { case (t, f) =>
-    println(s"Threshold: $t, F-score: $f, Beta = 0.5")
-}
-
-// AUPRC
-val auPRC = metrics.areaUnderPR
-println("Area under precision-recall curve = " + auPRC)
-
-// Compute thresholds used in ROC and PR curves
-val thresholds = precision.map(_._1)
-
-// ROC Curve
-val roc = metrics.roc
-
-// AUROC
-val auROC = metrics.areaUnderROC
-println("Area under ROC = " + auROC)
-
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/BinaryClassificationMetricsExample.scala %}
 
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`LogisticRegressionModel` Java docs](api/java/org/apache/spark/mllib/classification/LogisticRegressionModel.html) and [`LogisticRegressionWithLBFGS` Java docs](api/java/org/apache/spark/mllib/classification/LogisticRegressionWithLBFGS.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-
-import org.apache.spark.api.java.*;
-import org.apache.spark.rdd.RDD;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.mllib.classification.LogisticRegressionModel;
-import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS;
-import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.SparkConf;
-import org.apache.spark.SparkContext;
-
-public class BinaryClassification {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf().setAppName("Binary Classification Metrics");
-    SparkContext sc = new SparkContext(conf);
-    String path = "data/mllib/sample_binary_classification_data.txt";
-    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD();
-
-    // Split initial RDD into two... [60% training data, 40% testing data].
-    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[] {0.6, 0.4}, 11L);
-    JavaRDD<LabeledPoint> training = splits[0].cache();
-    JavaRDD<LabeledPoint> test = splits[1];
-
-    // Run training algorithm to build the model.
-    final LogisticRegressionModel model = new LogisticRegressionWithLBFGS()
-      .setNumClasses(2)
-      .run(training.rdd());
-
-    // Clear the prediction threshold so the model will return probabilities
-    model.clearThreshold();
-
-    // Compute raw scores on the test set.
-    JavaRDD<Tuple2<Object, Object>> predictionAndLabels = test.map(
-      new Function<LabeledPoint, Tuple2<Object, Object>>() {
-        public Tuple2<Object, Object> call(LabeledPoint p) {
-          Double prediction = model.predict(p.features());
-          return new Tuple2<Object, Object>(prediction, p.label());
-        }
-      }
-    );
-
-    // Get evaluation metrics.
-    BinaryClassificationMetrics metrics = new BinaryClassificationMetrics(predictionAndLabels.rdd());
-
-    // Precision by threshold
-    JavaRDD<Tuple2<Object, Object>> precision = metrics.precisionByThreshold().toJavaRDD();
-    System.out.println("Precision by threshold: " + precision.toArray());
-
-    // Recall by threshold
-    JavaRDD<Tuple2<Object, Object>> recall = metrics.recallByThreshold().toJavaRDD();
-    System.out.println("Recall by threshold: " + recall.toArray());
-
-    // F Score by threshold
-    JavaRDD<Tuple2<Object, Object>> f1Score = metrics.fMeasureByThreshold().toJavaRDD();
-    System.out.println("F1 Score by threshold: " + f1Score.toArray());
-
-    JavaRDD<Tuple2<Object, Object>> f2Score = metrics.fMeasureByThreshold(2.0).toJavaRDD();
-    System.out.println("F2 Score by threshold: " + f2Score.toArray());
-
-    // Precision-recall curve
-    JavaRDD<Tuple2<Object, Object>> prc = metrics.pr().toJavaRDD();
-    System.out.println("Precision-recall curve: " + prc.toArray());
-
-    // Thresholds
-    JavaRDD<Double> thresholds = precision.map(
-      new Function<Tuple2<Object, Object>, Double>() {
-        public Double call (Tuple2<Object, Object> t) {
-          return new Double(t._1().toString());
-        }
-      }
-    );
-
-    // ROC Curve
-    JavaRDD<Tuple2<Object, Object>> roc = metrics.roc().toJavaRDD();
-    System.out.println("ROC curve: " + roc.toArray());
-
-    // AUPRC
-    System.out.println("Area under precision-recall curve = " + metrics.areaUnderPR());
-
-    // AUROC
-    System.out.println("Area under ROC = " + metrics.areaUnderROC());
-
-    // Save and load model
-    model.save(sc, "myModelPath");
-    LogisticRegressionModel sameModel = LogisticRegressionModel.load(sc, "myModelPath");
-  }
-}
-
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaBinaryClassificationMetricsExample.java %}
 
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`BinaryClassificationMetrics` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.evaluation.BinaryClassificationMetrics) and [`LogisticRegressionWithLBFGS` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.classification.LogisticRegressionWithLBFGS) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.classification import LogisticRegressionWithLBFGS
-from pyspark.mllib.evaluation import BinaryClassificationMetrics
-from pyspark.mllib.regression import LabeledPoint
-from pyspark.mllib.util import MLUtils
-
-# Several of the methods available in scala are currently missing from pyspark
-
-# Load training data in LIBSVM format
-data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_binary_classification_data.txt")
-
-# Split data into training (60%) and test (40%)
-training, test = data.randomSplit([0.6, 0.4], seed = 11L)
-training.cache()
-
-# Run training algorithm to build the model
-model = LogisticRegressionWithLBFGS.train(training)
-
-# Compute raw scores on the test set
-predictionAndLabels = test.map(lambda lp: (float(model.predict(lp.features)), lp.label))
-
-# Instantiate metrics object
-metrics = BinaryClassificationMetrics(predictionAndLabels)
-
-# Area under precision-recall curve
-print("Area under PR = %s" % metrics.areaUnderPR)
-
-# Area under ROC curve
-print("Area under ROC = %s" % metrics.areaUnderROC)
-
-{% endhighlight %}
-
+{% include_example python/mllib/binary_classification_metrics_example.py %}
 </div>
 </div>
 
@@ -433,204 +240,21 @@ the data, and evaluate the performance of the algorithm by several multiclass cl
 <div data-lang="scala" markdown="1">
 Refer to the [`MulticlassMetrics` Scala docs](api/scala/index.html#org.apache.spark.mllib.evaluation.MulticlassMetrics) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS
-import org.apache.spark.mllib.evaluation.MulticlassMetrics
-import org.apache.spark.mllib.regression.LabeledPoint
-import org.apache.spark.mllib.util.MLUtils
-
-// Load training data in LIBSVM format
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt")
-
-// Split data into training (60%) and test (40%)
-val Array(training, test) = data.randomSplit(Array(0.6, 0.4), seed = 11L)
-training.cache()
-
-// Run training algorithm to build the model
-val model = new LogisticRegressionWithLBFGS()
-  .setNumClasses(3)
-  .run(training)
-
-// Compute raw scores on the test set
-val predictionAndLabels = test.map { case LabeledPoint(label, features) =>
-  val prediction = model.predict(features)
-  (prediction, label)
-}
-
-// Instantiate metrics object
-val metrics = new MulticlassMetrics(predictionAndLabels)
-
-// Confusion matrix
-println("Confusion matrix:")
-println(metrics.confusionMatrix)
-
-// Overall Statistics
-val precision = metrics.precision
-val recall = metrics.recall // same as true positive rate
-val f1Score = metrics.fMeasure
-println("Summary Statistics")
-println(s"Precision = $precision")
-println(s"Recall = $recall")
-println(s"F1 Score = $f1Score")
-
-// Precision by label
-val labels = metrics.labels
-labels.foreach { l =>
-    println(s"Precision($l) = " + metrics.precision(l))
-}
-
-// Recall by label
-labels.foreach { l =>
-    println(s"Recall($l) = " + metrics.recall(l))
-}
-
-// False positive rate by label
-labels.foreach { l =>
-    println(s"FPR($l) = " + metrics.falsePositiveRate(l))
-}
-
-// F-measure by label
-labels.foreach { l =>
-    println(s"F1-Score($l) = " + metrics.fMeasure(l))
-}
-
-// Weighted stats
-println(s"Weighted precision: ${metrics.weightedPrecision}")
-println(s"Weighted recall: ${metrics.weightedRecall}")
-println(s"Weighted F1 score: ${metrics.weightedFMeasure}")
-println(s"Weighted false positive rate: ${metrics.weightedFalsePositiveRate}")
-
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/MulticlassMetricsExample.scala %}
 
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`MulticlassMetrics` Java docs](api/java/org/apache/spark/mllib/evaluation/MulticlassMetrics.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-
-import org.apache.spark.api.java.*;
-import org.apache.spark.rdd.RDD;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.mllib.classification.LogisticRegressionModel;
-import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS;
-import org.apache.spark.mllib.evaluation.MulticlassMetrics;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.mllib.linalg.Matrix;
-import org.apache.spark.SparkConf;
-import org.apache.spark.SparkContext;
-
-public class MulticlassClassification {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf().setAppName("Multiclass Classification Metrics");
-    SparkContext sc = new SparkContext(conf);
-    String path = "data/mllib/sample_multiclass_classification_data.txt";
-    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD();
-
-    // Split initial RDD into two... [60% training data, 40% testing data].
-    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[] {0.6, 0.4}, 11L);
-    JavaRDD<LabeledPoint> training = splits[0].cache();
-    JavaRDD<LabeledPoint> test = splits[1];
-
-    // Run training algorithm to build the model.
-    final LogisticRegressionModel model = new LogisticRegressionWithLBFGS()
-      .setNumClasses(3)
-      .run(training.rdd());
-
-    // Compute raw scores on the test set.
-    JavaRDD<Tuple2<Object, Object>> predictionAndLabels = test.map(
-      new Function<LabeledPoint, Tuple2<Object, Object>>() {
-        public Tuple2<Object, Object> call(LabeledPoint p) {
-          Double prediction = model.predict(p.features());
-          return new Tuple2<Object, Object>(prediction, p.label());
-        }
-      }
-    );
-
-    // Get evaluation metrics.
-    MulticlassMetrics metrics = new MulticlassMetrics(predictionAndLabels.rdd());
-
-    // Confusion matrix
-    Matrix confusion = metrics.confusionMatrix();
-    System.out.println("Confusion matrix: \n" + confusion);
-
-    // Overall statistics
-    System.out.println("Precision = " + metrics.precision());
-    System.out.println("Recall = " + metrics.recall());
-    System.out.println("F1 Score = " + metrics.fMeasure());
-
-    // Stats by labels
-    for (int i = 0; i < metrics.labels().length; i++) {
-        System.out.format("Class %f precision = %f\n", metrics.labels()[i], metrics.precision(metrics.labels()[i]));
-        System.out.format("Class %f recall = %f\n", metrics.labels()[i], metrics.recall(metrics.labels()[i]));
-        System.out.format("Class %f F1 score = %f\n", metrics.labels()[i], metrics.fMeasure(metrics.labels()[i]));
-    }
-
-    //Weighted stats
-    System.out.format("Weighted precision = %f\n", metrics.weightedPrecision());
-    System.out.format("Weighted recall = %f\n", metrics.weightedRecall());
-    System.out.format("Weighted F1 score = %f\n", metrics.weightedFMeasure());
-    System.out.format("Weighted false positive rate = %f\n", metrics.weightedFalsePositiveRate());
-
-    // Save and load model
-    model.save(sc, "myModelPath");
-    LogisticRegressionModel sameModel = LogisticRegressionModel.load(sc, "myModelPath");
-  }
-}
-
-{% endhighlight %}
+ {% include_example java/org/apache/spark/examples/mllib/JavaMulticlassClassificationMetricsExample.java %}
 
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`MulticlassMetrics` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.evaluation.MulticlassMetrics) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.classification import LogisticRegressionWithLBFGS
-from pyspark.mllib.util import MLUtils
-from pyspark.mllib.evaluation import MulticlassMetrics
-
-# Load training data in LIBSVM format
-data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt")
-
-# Split data into training (60%) and test (40%)
-training, test = data.randomSplit([0.6, 0.4], seed = 11L)
-training.cache()
-
-# Run training algorithm to build the model
-model = LogisticRegressionWithLBFGS.train(training, numClasses=3)
-
-# Compute raw scores on the test set
-predictionAndLabels = test.map(lambda lp: (float(model.predict(lp.features)), lp.label))
-
-# Instantiate metrics object
-metrics = MulticlassMetrics(predictionAndLabels)
-
-# Overall statistics
-precision = metrics.precision()
-recall = metrics.recall()
-f1Score = metrics.fMeasure()
-print("Summary Stats")
-print("Precision = %s" % precision)
-print("Recall = %s" % recall)
-print("F1 Score = %s" % f1Score)
-
-# Statistics by class
-labels = data.map(lambda lp: lp.label).distinct().collect()
-for label in sorted(labels):
-    print("Class %s precision = %s" % (label, metrics.precision(label)))
-    print("Class %s recall = %s" % (label, metrics.recall(label)))
-    print("Class %s F1 Measure = %s" % (label, metrics.fMeasure(label, beta=1.0)))
-
-# Weighted stats
-print("Weighted recall = %s" % metrics.weightedRecall)
-print("Weighted precision = %s" % metrics.weightedPrecision)
-print("Weighted F(1) Score = %s" % metrics.weightedFMeasure())
-print("Weighted F(0.5) Score = %s" % metrics.weightedFMeasure(beta=0.5))
-print("Weighted false positive rate = %s" % metrics.weightedFalsePositiveRate)
-{% endhighlight %}
+{% include_example python/mllib/multi_class_metrics_example.py %}
 
 </div>
 </div>
@@ -766,154 +390,21 @@ True classes:
 <div data-lang="scala" markdown="1">
 Refer to the [`MultilabelMetrics` Scala docs](api/scala/index.html#org.apache.spark.mllib.evaluation.MultilabelMetrics) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.evaluation.MultilabelMetrics
-import org.apache.spark.rdd.RDD;
-
-val scoreAndLabels: RDD[(Array[Double], Array[Double])] = sc.parallelize(
-  Seq((Array(0.0, 1.0), Array(0.0, 2.0)),
-    (Array(0.0, 2.0), Array(0.0, 1.0)),
-    (Array(), Array(0.0)),
-    (Array(2.0), Array(2.0)),
-    (Array(2.0, 0.0), Array(2.0, 0.0)),
-    (Array(0.0, 1.0, 2.0), Array(0.0, 1.0)),
-    (Array(1.0), Array(1.0, 2.0))), 2)
-
-// Instantiate metrics object
-val metrics = new MultilabelMetrics(scoreAndLabels)
-
-// Summary stats
-println(s"Recall = ${metrics.recall}")
-println(s"Precision = ${metrics.precision}")
-println(s"F1 measure = ${metrics.f1Measure}")
-println(s"Accuracy = ${metrics.accuracy}")
-
-// Individual label stats
-metrics.labels.foreach(label => println(s"Class $label precision = ${metrics.precision(label)}"))
-metrics.labels.foreach(label => println(s"Class $label recall = ${metrics.recall(label)}"))
-metrics.labels.foreach(label => println(s"Class $label F1-score = ${metrics.f1Measure(label)}"))
-
-// Micro stats
-println(s"Micro recall = ${metrics.microRecall}")
-println(s"Micro precision = ${metrics.microPrecision}")
-println(s"Micro F1 measure = ${metrics.microF1Measure}")
-
-// Hamming loss
-println(s"Hamming loss = ${metrics.hammingLoss}")
-
-// Subset accuracy
-println(s"Subset accuracy = ${metrics.subsetAccuracy}")
-
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/MultiLabelMetricsExample.scala %}
 
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`MultilabelMetrics` Java docs](api/java/org/apache/spark/mllib/evaluation/MultilabelMetrics.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-
-import org.apache.spark.api.java.*;
-import org.apache.spark.rdd.RDD;
-import org.apache.spark.mllib.evaluation.MultilabelMetrics;
-import org.apache.spark.SparkConf;
-import java.util.Arrays;
-import java.util.List;
-
-public class MultilabelClassification {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf().setAppName("Multilabel Classification Metrics");
-    JavaSparkContext sc = new JavaSparkContext(conf);
-
-    List<Tuple2<double[], double[]>> data = Arrays.asList(
-        new Tuple2<double[], double[]>(new double[]{0.0, 1.0}, new double[]{0.0, 2.0}),
-        new Tuple2<double[], double[]>(new double[]{0.0, 2.0}, new double[]{0.0, 1.0}),
-        new Tuple2<double[], double[]>(new double[]{}, new double[]{0.0}),
-        new Tuple2<double[], double[]>(new double[]{2.0}, new double[]{2.0}),
-        new Tuple2<double[], double[]>(new double[]{2.0, 0.0}, new double[]{2.0, 0.0}),
-        new Tuple2<double[], double[]>(new double[]{0.0, 1.0, 2.0}, new double[]{0.0, 1.0}),
-        new Tuple2<double[], double[]>(new double[]{1.0}, new double[]{1.0, 2.0})
-        );
-    JavaRDD<Tuple2<double[], double[]>> scoreAndLabels = sc.parallelize(data);
-
-    // Instantiate metrics object
-    MultilabelMetrics metrics = new MultilabelMetrics(scoreAndLabels.rdd());
-
-    // Summary stats
-    System.out.format("Recall = %f\n", metrics.recall());
-    System.out.format("Precision = %f\n", metrics.precision());
-    System.out.format("F1 measure = %f\n", metrics.f1Measure());
-    System.out.format("Accuracy = %f\n", metrics.accuracy());
-
-    // Stats by labels
-    for (int i = 0; i < metrics.labels().length - 1; i++) {
-        System.out.format("Class %1.1f precision = %f\n", metrics.labels()[i], metrics.precision(metrics.labels()[i]));
-        System.out.format("Class %1.1f recall = %f\n", metrics.labels()[i], metrics.recall(metrics.labels()[i]));
-        System.out.format("Class %1.1f F1 score = %f\n", metrics.labels()[i], metrics.f1Measure(metrics.labels()[i]));
-    }
-
-    // Micro stats
-    System.out.format("Micro recall = %f\n", metrics.microRecall());
-    System.out.format("Micro precision = %f\n", metrics.microPrecision());
-    System.out.format("Micro F1 measure = %f\n", metrics.microF1Measure());
-
-    // Hamming loss
-    System.out.format("Hamming loss = %f\n", metrics.hammingLoss());
-
-    // Subset accuracy
-    System.out.format("Subset accuracy = %f\n", metrics.subsetAccuracy());
-
-  }
-}
-
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaMultiLabelClassificationMetricsExample.java %}
 
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`MultilabelMetrics` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.evaluation.MultilabelMetrics) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.evaluation import MultilabelMetrics
-
-scoreAndLabels = sc.parallelize([
-    ([0.0, 1.0], [0.0, 2.0]),
-    ([0.0, 2.0], [0.0, 1.0]),
-    ([], [0.0]),
-    ([2.0], [2.0]),
-    ([2.0, 0.0], [2.0, 0.0]),
-    ([0.0, 1.0, 2.0], [0.0, 1.0]),
-    ([1.0], [1.0, 2.0])])
-
-# Instantiate metrics object
-metrics = MultilabelMetrics(scoreAndLabels)
-
-# Summary stats
-print("Recall = %s" % metrics.recall())
-print("Precision = %s" % metrics.precision())
-print("F1 measure = %s" % metrics.f1Measure())
-print("Accuracy = %s" % metrics.accuracy)
-
-# Individual label stats
-labels = scoreAndLabels.flatMap(lambda x: x[1]).distinct().collect()
-for label in labels:
-    print("Class %s precision = %s" % (label, metrics.precision(label)))
-    print("Class %s recall = %s" % (label, metrics.recall(label)))
-    print("Class %s F1 Measure = %s" % (label, metrics.f1Measure(label)))
-
-# Micro stats
-print("Micro precision = %s" % metrics.microPrecision)
-print("Micro recall = %s" % metrics.microRecall)
-print("Micro F1 measure = %s" % metrics.microF1Measure)
-
-# Hamming loss
-print("Hamming loss = %s" % metrics.hammingLoss)
-
-# Subset accuracy
-print("Subset accuracy = %s" % metrics.subsetAccuracy)
-
-{% endhighlight %}
+{% include_example python/mllib/multi_label_metrics_example.py %}
 
 </div>
 </div>
@@ -1027,280 +518,21 @@ expanded world of non-positive weights are "the same as never having interacted
 <div data-lang="scala" markdown="1">
 Refer to the [`RegressionMetrics` Scala docs](api/scala/index.html#org.apache.spark.mllib.evaluation.RegressionMetrics) and [`RankingMetrics` Scala docs](api/scala/index.html#org.apache.spark.mllib.evaluation.RankingMetrics) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.evaluation.{RegressionMetrics, RankingMetrics}
-import org.apache.spark.mllib.recommendation.{ALS, Rating}
-
-// Read in the ratings data
-val ratings = sc.textFile("data/mllib/sample_movielens_data.txt").map { line =>
-  val fields = line.split("::")
-  Rating(fields(0).toInt, fields(1).toInt, fields(2).toDouble - 2.5)
-}.cache()
-
-// Map ratings to 1 or 0, 1 indicating a movie that should be recommended
-val binarizedRatings = ratings.map(r => Rating(r.user, r.product, if (r.rating > 0) 1.0 else 0.0)).cache()
-
-// Summarize ratings
-val numRatings = ratings.count()
-val numUsers = ratings.map(_.user).distinct().count()
-val numMovies = ratings.map(_.product).distinct().count()
-println(s"Got $numRatings ratings from $numUsers users on $numMovies movies.")
-
-// Build the model
-val numIterations = 10
-val rank = 10
-val lambda = 0.01
-val model = ALS.train(ratings, rank, numIterations, lambda)
-
-// Define a function to scale ratings from 0 to 1
-def scaledRating(r: Rating): Rating = {
-  val scaledRating = math.max(math.min(r.rating, 1.0), 0.0)
-  Rating(r.user, r.product, scaledRating)
-}
-
-// Get sorted top ten predictions for each user and then scale from [0, 1]
-val userRecommended = model.recommendProductsForUsers(10).map{ case (user, recs) =>
-  (user, recs.map(scaledRating))
-}
-
-// Assume that any movie a user rated 3 or higher (which maps to a 1) is a relevant document
-// Compare with top ten most relevant documents
-val userMovies = binarizedRatings.groupBy(_.user)
-val relevantDocuments = userMovies.join(userRecommended).map{ case (user, (actual, predictions)) =>
-  (predictions.map(_.product), actual.filter(_.rating > 0.0).map(_.product).toArray)
-}
-
-// Instantiate metrics object
-val metrics = new RankingMetrics(relevantDocuments)
-
-// Precision at K
-Array(1, 3, 5).foreach{ k =>
-  println(s"Precision at $k = ${metrics.precisionAt(k)}")
-}
-
-// Mean average precision
-println(s"Mean average precision = ${metrics.meanAveragePrecision}")
-
-// Normalized discounted cumulative gain
-Array(1, 3, 5).foreach{ k =>
-  println(s"NDCG at $k = ${metrics.ndcgAt(k)}")
-}
-
-// Get predictions for each data point
-val allPredictions = model.predict(ratings.map(r => (r.user, r.product))).map(r => ((r.user, r.product), r.rating))
-val allRatings = ratings.map(r => ((r.user, r.product), r.rating))
-val predictionsAndLabels = allPredictions.join(allRatings).map{ case ((user, product), (predicted, actual)) =>
-  (predicted, actual)
-}
-
-// Get the RMSE using regression metrics
-val regressionMetrics = new RegressionMetrics(predictionsAndLabels)
-println(s"RMSE = ${regressionMetrics.rootMeanSquaredError}")
-
-// R-squared
-println(s"R-squared = ${regressionMetrics.r2}")
-
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/RankingMetricsExample.scala %}
 
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`RegressionMetrics` Java docs](api/java/org/apache/spark/mllib/evaluation/RegressionMetrics.html) and [`RankingMetrics` Java docs](api/java/org/apache/spark/mllib/evaluation/RankingMetrics.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-
-import org.apache.spark.api.java.*;
-import org.apache.spark.rdd.RDD;
-import org.apache.spark.mllib.recommendation.MatrixFactorizationModel;
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.function.Function;
-import java.util.*;
-import org.apache.spark.mllib.evaluation.RegressionMetrics;
-import org.apache.spark.mllib.evaluation.RankingMetrics;
-import org.apache.spark.mllib.recommendation.ALS;
-import org.apache.spark.mllib.recommendation.Rating;
-
-// Read in the ratings data
-public class Ranking {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf().setAppName("Ranking Metrics");
-    JavaSparkContext sc = new JavaSparkContext(conf);
-    String path = "data/mllib/sample_movielens_data.txt";
-    JavaRDD<String> data = sc.textFile(path);
-    JavaRDD<Rating> ratings = data.map(
-      new Function<String, Rating>() {
-        public Rating call(String line) {
-          String[] parts = line.split("::");
-          return new Rating(Integer.parseInt(parts[0]), Integer.parseInt(parts[1]), Double.parseDouble(parts[2]) - 2.5);
-        }
-      }
-    );
-    ratings.cache();
-
-    // Train an ALS model
-    final MatrixFactorizationModel model = ALS.train(JavaRDD.toRDD(ratings), 10, 10, 0.01);
-
-    // Get top 10 recommendations for every user and scale ratings from 0 to 1
-    JavaRDD<Tuple2<Object, Rating[]>> userRecs = model.recommendProductsForUsers(10).toJavaRDD();
-    JavaRDD<Tuple2<Object, Rating[]>> userRecsScaled = userRecs.map(
-      new Function<Tuple2<Object, Rating[]>, Tuple2<Object, Rating[]>>() {
-        public Tuple2<Object, Rating[]> call(Tuple2<Object, Rating[]> t) {
-          Rating[] scaledRatings = new Rating[t._2().length];
-          for (int i = 0; i < scaledRatings.length; i++) {
-            double newRating = Math.max(Math.min(t._2()[i].rating(), 1.0), 0.0);
-            scaledRatings[i] = new Rating(t._2()[i].user(), t._2()[i].product(), newRating);
-          }
-          return new Tuple2<Object, Rating[]>(t._1(), scaledRatings);
-        }
-      }
-    );
-    JavaPairRDD<Object, Rating[]> userRecommended = JavaPairRDD.fromJavaRDD(userRecsScaled);
-
-    // Map ratings to 1 or 0, 1 indicating a movie that should be recommended
-    JavaRDD<Rating> binarizedRatings = ratings.map(
-      new Function<Rating, Rating>() {
-        public Rating call(Rating r) {
-          double binaryRating;
-          if (r.rating() > 0.0) {
-            binaryRating = 1.0;
-          }
-          else {
-            binaryRating = 0.0;
-          }
-          return new Rating(r.user(), r.product(), binaryRating);
-        }
-      }
-    );
-
-    // Group ratings by common user
-    JavaPairRDD<Object, Iterable<Rating>> userMovies = binarizedRatings.groupBy(
-      new Function<Rating, Object>() {
-        public Object call(Rating r) {
-          return r.user();
-        }
-      }
-    );
-
-    // Get true relevant documents from all user ratings
-    JavaPairRDD<Object, List<Integer>> userMoviesList = userMovies.mapValues(
-      new Function<Iterable<Rating>, List<Integer>>() {
-        public List<Integer> call(Iterable<Rating> docs) {
-          List<Integer> products = new ArrayList<Integer>();
-          for (Rating r : docs) {
-            if (r.rating() > 0.0) {
-              products.add(r.product());
-            }
-          }
-          return products;
-        }
-      }
-    );
-
-    // Extract the product id from each recommendation
-    JavaPairRDD<Object, List<Integer>> userRecommendedList = userRecommended.mapValues(
-      new Function<Rating[], List<Integer>>() {
-        public List<Integer> call(Rating[] docs) {
-          List<Integer> products = new ArrayList<Integer>();
-          for (Rating r : docs) {
-            products.add(r.product());
-          }
-          return products;
-        }
-      }
-    );
-    JavaRDD<Tuple2<List<Integer>, List<Integer>>> relevantDocs = userMoviesList.join(userRecommendedList).values();
-
-    // Instantiate the metrics object
-    RankingMetrics metrics = RankingMetrics.of(relevantDocs);
-
-    // Precision and NDCG at k
-    Integer[] kVector = {1, 3, 5};
-    for (Integer k : kVector) {
-      System.out.format("Precision at %d = %f\n", k, metrics.precisionAt(k));
-      System.out.format("NDCG at %d = %f\n", k, metrics.ndcgAt(k));
-    }
-
-    // Mean average precision
-    System.out.format("Mean average precision = %f\n", metrics.meanAveragePrecision());
-
-    // Evaluate the model using numerical ratings and regression metrics
-    JavaRDD<Tuple2<Object, Object>> userProducts = ratings.map(
-      new Function<Rating, Tuple2<Object, Object>>() {
-        public Tuple2<Object, Object> call(Rating r) {
-          return new Tuple2<Object, Object>(r.user(), r.product());
-        }
-      }
-    );
-    JavaPairRDD<Tuple2<Integer, Integer>, Object> predictions = JavaPairRDD.fromJavaRDD(
-      model.predict(JavaRDD.toRDD(userProducts)).toJavaRDD().map(
-        new Function<Rating, Tuple2<Tuple2<Integer, Integer>, Object>>() {
-          public Tuple2<Tuple2<Integer, Integer>, Object> call(Rating r){
-            return new Tuple2<Tuple2<Integer, Integer>, Object>(
-              new Tuple2<Integer, Integer>(r.user(), r.product()), r.rating());
-          }
-        }
-    ));
-    JavaRDD<Tuple2<Object, Object>> ratesAndPreds =
-      JavaPairRDD.fromJavaRDD(ratings.map(
-        new Function<Rating, Tuple2<Tuple2<Integer, Integer>, Object>>() {
-          public Tuple2<Tuple2<Integer, Integer>, Object> call(Rating r){
-            return new Tuple2<Tuple2<Integer, Integer>, Object>(
-              new Tuple2<Integer, Integer>(r.user(), r.product()), r.rating());
-          }
-        }
-    )).join(predictions).values();
-
-    // Create regression metrics object
-    RegressionMetrics regressionMetrics = new RegressionMetrics(ratesAndPreds.rdd());
-
-    // Root mean squared error
-    System.out.format("RMSE = %f\n", regressionMetrics.rootMeanSquaredError());
-
-    // R-squared
-    System.out.format("R-squared = %f\n", regressionMetrics.r2());
-  }
-}
-
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaRankingMetricsExample.java %}
 
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`RegressionMetrics` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.evaluation.RegressionMetrics) and [`RankingMetrics` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.evaluation.RankingMetrics) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.recommendation import ALS, Rating
-from pyspark.mllib.evaluation import RegressionMetrics, RankingMetrics
-
-#  Read in the ratings data
-lines = sc.textFile("data/mllib/sample_movielens_data.txt")
-
-def parseLine(line):
-    fields = line.split("::")
-    return Rating(int(fields[0]), int(fields[1]), float(fields[2]) - 2.5)
-ratings = lines.map(lambda r: parseLine(r))
-
-# Train a model on to predict user-product ratings
-model = ALS.train(ratings, 10, 10, 0.01)
-
-# Get predicted ratings on all existing user-product pairs
-testData = ratings.map(lambda p: (p.user, p.product))
-predictions = model.predictAll(testData).map(lambda r: ((r.user, r.product), r.rating))
-
-ratingsTuple = ratings.map(lambda r: ((r.user, r.product), r.rating))
-scoreAndLabels = predictions.join(ratingsTuple).map(lambda tup: tup[1])
-
-# Instantiate regression metrics to compare predicted and actual ratings
-metrics = RegressionMetrics(scoreAndLabels)
-
-# Root mean sqaured error
-print("RMSE = %s" % metrics.rootMeanSquaredError)
-
-# R-squared
-print("R-squared = %s" % metrics.r2)
-
-{% endhighlight %}
+{% include_example python/mllib/ranking_metrics_example.py %}
 
 </div>
 </div>
@@ -1350,163 +582,21 @@ and evaluate the performance of the algorithm by several regression metrics.
 <div data-lang="scala" markdown="1">
 Refer to the [`RegressionMetrics` Scala docs](api/scala/index.html#org.apache.spark.mllib.evaluation.RegressionMetrics) for details on the API.
 
-{% highlight scala %}
-import org.apache.spark.mllib.regression.LabeledPoint
-import org.apache.spark.mllib.regression.LinearRegressionModel
-import org.apache.spark.mllib.regression.LinearRegressionWithSGD
-import org.apache.spark.mllib.linalg.Vectors
-import org.apache.spark.mllib.evaluation.RegressionMetrics
-import org.apache.spark.mllib.util.MLUtils
-
-// Load the data
-val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_linear_regression_data.txt").cache()
-
-// Build the model
-val numIterations = 100
-val model = LinearRegressionWithSGD.train(data, numIterations)
-
-// Get predictions
-val valuesAndPreds = data.map{ point =>
-  val prediction = model.predict(point.features)
-  (prediction, point.label)
-}
-
-// Instantiate metrics object
-val metrics = new RegressionMetrics(valuesAndPreds)
-
-// Squared error
-println(s"MSE = ${metrics.meanSquaredError}")
-println(s"RMSE = ${metrics.rootMeanSquaredError}")
-
-// R-squared
-println(s"R-squared = ${metrics.r2}")
-
-// Mean absolute error
-println(s"MAE = ${metrics.meanAbsoluteError}")
-
-// Explained variance
-println(s"Explained variance = ${metrics.explainedVariance}")
-
-{% endhighlight %}
+{% include_example scala/org/apache/spark/examples/mllib/RegressionMetricsExample.scala %}
 
 </div>
 
 <div data-lang="java" markdown="1">
 Refer to the [`RegressionMetrics` Java docs](api/java/org/apache/spark/mllib/evaluation/RegressionMetrics.html) for details on the API.
 
-{% highlight java %}
-import scala.Tuple2;
-
-import org.apache.spark.api.java.*;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.mllib.linalg.Vectors;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.regression.LinearRegressionModel;
-import org.apache.spark.mllib.regression.LinearRegressionWithSGD;
-import org.apache.spark.mllib.evaluation.RegressionMetrics;
-import org.apache.spark.SparkConf;
-
-public class LinearRegression {
-  public static void main(String[] args) {
-    SparkConf conf = new SparkConf().setAppName("Linear Regression Example");
-    JavaSparkContext sc = new JavaSparkContext(conf);
-
-    // Load and parse the data
-    String path = "data/mllib/sample_linear_regression_data.txt";
-    JavaRDD<String> data = sc.textFile(path);
-    JavaRDD<LabeledPoint> parsedData = data.map(
-      new Function<String, LabeledPoint>() {
-        public LabeledPoint call(String line) {
-          String[] parts = line.split(" ");
-          double[] v = new double[parts.length - 1];
-          for (int i = 1; i < parts.length - 1; i++)
-            v[i - 1] = Double.parseDouble(parts[i].split(":")[1]);
-          return new LabeledPoint(Double.parseDouble(parts[0]), Vectors.dense(v));
-        }
-      }
-    );
-    parsedData.cache();
-
-    // Building the model
-    int numIterations = 100;
-    final LinearRegressionModel model =
-      LinearRegressionWithSGD.train(JavaRDD.toRDD(parsedData), numIterations);
-
-    // Evaluate model on training examples and compute training error
-    JavaRDD<Tuple2<Object, Object>> valuesAndPreds = parsedData.map(
-      new Function<LabeledPoint, Tuple2<Object, Object>>() {
-        public Tuple2<Object, Object> call(LabeledPoint point) {
-          double prediction = model.predict(point.features());
-          return new Tuple2<Object, Object>(prediction, point.label());
-        }
-      }
-    );
-
-    // Instantiate metrics object
-    RegressionMetrics metrics = new RegressionMetrics(valuesAndPreds.rdd());
-
-    // Squared error
-    System.out.format("MSE = %f\n", metrics.meanSquaredError());
-    System.out.format("RMSE = %f\n", metrics.rootMeanSquaredError());
-
-    // R-squared
-    System.out.format("R Squared = %f\n", metrics.r2());
-
-    // Mean absolute error
-    System.out.format("MAE = %f\n", metrics.meanAbsoluteError());
-
-    // Explained variance
-    System.out.format("Explained Variance = %f\n", metrics.explainedVariance());
-
-    // Save and load model
-    model.save(sc.sc(), "myModelPath");
-    LinearRegressionModel sameModel = LinearRegressionModel.load(sc.sc(), "myModelPath");
-  }
-}
-
-{% endhighlight %}
+{% include_example java/org/apache/spark/examples/mllib/JavaRegressionMetricsExample.java %}
 
 </div>
 
 <div data-lang="python" markdown="1">
 Refer to the [`RegressionMetrics` Python docs](api/python/pyspark.mllib.html#pyspark.mllib.evaluation.RegressionMetrics) for more details on the API.
 
-{% highlight python %}
-from pyspark.mllib.regression import LabeledPoint, LinearRegressionWithSGD
-from pyspark.mllib.evaluation import RegressionMetrics
-from pyspark.mllib.linalg import DenseVector
-
-# Load and parse the data
-def parsePoint(line):
-    values = line.split()
-    return LabeledPoint(float(values[0]), DenseVector([float(x.split(':')[1]) for x in values[1:]]))
-
-data = sc.textFile("data/mllib/sample_linear_regression_data.txt")
-parsedData = data.map(parsePoint)
-
-# Build the model
-model = LinearRegressionWithSGD.train(parsedData)
-
-# Get predictions
-valuesAndPreds = parsedData.map(lambda p: (float(model.predict(p.features)), p.label))
-
-# Instantiate metrics object
-metrics = RegressionMetrics(valuesAndPreds)
-
-# Squared Error
-print("MSE = %s" % metrics.meanSquaredError)
-print("RMSE = %s" % metrics.rootMeanSquaredError)
-
-# R-squared
-print("R-squared = %s" % metrics.r2)
-
-# Mean absolute error
-print("MAE = %s" % metrics.meanAbsoluteError)
-
-# Explained variance
-print("Explained variance = %s" % metrics.explainedVariance)
-
-{% endhighlight %}
+{% include_example python/mllib/regression_metrics_example.py %}
 
 </div>
 </div>
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaBinaryClassificationMetricsExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaBinaryClassificationMetricsExample.java
new file mode 100644
index 0000000000000..980a9108af53f
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaBinaryClassificationMetricsExample.java
@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import scala.Tuple2;
+
+import org.apache.spark.api.java.*;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.mllib.classification.LogisticRegressionModel;
+import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS;
+import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.util.MLUtils;
+// $example off$
+import org.apache.spark.SparkConf;
+import org.apache.spark.SparkContext;
+
+public class JavaBinaryClassificationMetricsExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("Java Binary Classification Metrics Example");
+    SparkContext sc = new SparkContext(conf);
+    // $example on$
+    String path = "data/mllib/sample_binary_classification_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD();
+
+    // Split initial RDD into two... [60% training data, 40% testing data].
+    JavaRDD<LabeledPoint>[] splits =
+      data.randomSplit(new double[]{0.6, 0.4}, 11L);
+    JavaRDD<LabeledPoint> training = splits[0].cache();
+    JavaRDD<LabeledPoint> test = splits[1];
+
+    // Run training algorithm to build the model.
+    final LogisticRegressionModel model = new LogisticRegressionWithLBFGS()
+      .setNumClasses(2)
+      .run(training.rdd());
+
+    // Clear the prediction threshold so the model will return probabilities
+    model.clearThreshold();
+
+    // Compute raw scores on the test set.
+    JavaRDD<Tuple2<Object, Object>> predictionAndLabels = test.map(
+      new Function<LabeledPoint, Tuple2<Object, Object>>() {
+        public Tuple2<Object, Object> call(LabeledPoint p) {
+          Double prediction = model.predict(p.features());
+          return new Tuple2<Object, Object>(prediction, p.label());
+        }
+      }
+    );
+
+    // Get evaluation metrics.
+    BinaryClassificationMetrics metrics = new BinaryClassificationMetrics(predictionAndLabels.rdd());
+
+    // Precision by threshold
+    JavaRDD<Tuple2<Object, Object>> precision = metrics.precisionByThreshold().toJavaRDD();
+    System.out.println("Precision by threshold: " + precision.toArray());
+
+    // Recall by threshold
+    JavaRDD<Tuple2<Object, Object>> recall = metrics.recallByThreshold().toJavaRDD();
+    System.out.println("Recall by threshold: " + recall.toArray());
+
+    // F Score by threshold
+    JavaRDD<Tuple2<Object, Object>> f1Score = metrics.fMeasureByThreshold().toJavaRDD();
+    System.out.println("F1 Score by threshold: " + f1Score.toArray());
+
+    JavaRDD<Tuple2<Object, Object>> f2Score = metrics.fMeasureByThreshold(2.0).toJavaRDD();
+    System.out.println("F2 Score by threshold: " + f2Score.toArray());
+
+    // Precision-recall curve
+    JavaRDD<Tuple2<Object, Object>> prc = metrics.pr().toJavaRDD();
+    System.out.println("Precision-recall curve: " + prc.toArray());
+
+    // Thresholds
+    JavaRDD<Double> thresholds = precision.map(
+      new Function<Tuple2<Object, Object>, Double>() {
+        public Double call(Tuple2<Object, Object> t) {
+          return new Double(t._1().toString());
+        }
+      }
+    );
+
+    // ROC Curve
+    JavaRDD<Tuple2<Object, Object>> roc = metrics.roc().toJavaRDD();
+    System.out.println("ROC curve: " + roc.toArray());
+
+    // AUPRC
+    System.out.println("Area under precision-recall curve = " + metrics.areaUnderPR());
+
+    // AUROC
+    System.out.println("Area under ROC = " + metrics.areaUnderROC());
+
+    // Save and load model
+    model.save(sc, "target/tmp/LogisticRegressionModel");
+    LogisticRegressionModel sameModel = LogisticRegressionModel.load(sc,
+      "target/tmp/LogisticRegressionModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaMultiLabelClassificationMetricsExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaMultiLabelClassificationMetricsExample.java
new file mode 100644
index 0000000000000..b54e1ea3f2bcf
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaMultiLabelClassificationMetricsExample.java
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.Arrays;
+import java.util.List;
+
+import scala.Tuple2;
+
+import org.apache.spark.api.java.*;
+import org.apache.spark.mllib.evaluation.MultilabelMetrics;
+import org.apache.spark.rdd.RDD;
+import org.apache.spark.SparkConf;
+// $example off$
+import org.apache.spark.SparkContext;
+
+public class JavaMultiLabelClassificationMetricsExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("Multilabel Classification Metrics Example");
+    JavaSparkContext sc = new JavaSparkContext(conf);
+    // $example on$
+    List<Tuple2<double[], double[]>> data = Arrays.asList(
+      new Tuple2<double[], double[]>(new double[]{0.0, 1.0}, new double[]{0.0, 2.0}),
+      new Tuple2<double[], double[]>(new double[]{0.0, 2.0}, new double[]{0.0, 1.0}),
+      new Tuple2<double[], double[]>(new double[]{}, new double[]{0.0}),
+      new Tuple2<double[], double[]>(new double[]{2.0}, new double[]{2.0}),
+      new Tuple2<double[], double[]>(new double[]{2.0, 0.0}, new double[]{2.0, 0.0}),
+      new Tuple2<double[], double[]>(new double[]{0.0, 1.0, 2.0}, new double[]{0.0, 1.0}),
+      new Tuple2<double[], double[]>(new double[]{1.0}, new double[]{1.0, 2.0})
+    );
+    JavaRDD<Tuple2<double[], double[]>> scoreAndLabels = sc.parallelize(data);
+
+    // Instantiate metrics object
+    MultilabelMetrics metrics = new MultilabelMetrics(scoreAndLabels.rdd());
+
+    // Summary stats
+    System.out.format("Recall = %f\n", metrics.recall());
+    System.out.format("Precision = %f\n", metrics.precision());
+    System.out.format("F1 measure = %f\n", metrics.f1Measure());
+    System.out.format("Accuracy = %f\n", metrics.accuracy());
+
+    // Stats by labels
+    for (int i = 0; i < metrics.labels().length - 1; i++) {
+      System.out.format("Class %1.1f precision = %f\n", metrics.labels()[i], metrics.precision
+        (metrics.labels()[i]));
+      System.out.format("Class %1.1f recall = %f\n", metrics.labels()[i], metrics.recall(metrics
+        .labels()[i]));
+      System.out.format("Class %1.1f F1 score = %f\n", metrics.labels()[i], metrics.f1Measure
+        (metrics.labels()[i]));
+    }
+
+    // Micro stats
+    System.out.format("Micro recall = %f\n", metrics.microRecall());
+    System.out.format("Micro precision = %f\n", metrics.microPrecision());
+    System.out.format("Micro F1 measure = %f\n", metrics.microF1Measure());
+
+    // Hamming loss
+    System.out.format("Hamming loss = %f\n", metrics.hammingLoss());
+
+    // Subset accuracy
+    System.out.format("Subset accuracy = %f\n", metrics.subsetAccuracy());
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaMulticlassClassificationMetricsExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaMulticlassClassificationMetricsExample.java
new file mode 100644
index 0000000000000..21f628fb51b6e
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaMulticlassClassificationMetricsExample.java
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import scala.Tuple2;
+
+import org.apache.spark.api.java.*;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.mllib.classification.LogisticRegressionModel;
+import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS;
+import org.apache.spark.mllib.evaluation.MulticlassMetrics;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.util.MLUtils;
+import org.apache.spark.mllib.linalg.Matrix;
+// $example off$
+import org.apache.spark.SparkConf;
+import org.apache.spark.SparkContext;
+
+public class JavaMulticlassClassificationMetricsExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("Multi class Classification Metrics Example");
+    SparkContext sc = new SparkContext(conf);
+    // $example on$
+    String path = "data/mllib/sample_multiclass_classification_data.txt";
+    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc, path).toJavaRDD();
+
+    // Split initial RDD into two... [60% training data, 40% testing data].
+    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.6, 0.4}, 11L);
+    JavaRDD<LabeledPoint> training = splits[0].cache();
+    JavaRDD<LabeledPoint> test = splits[1];
+
+    // Run training algorithm to build the model.
+    final LogisticRegressionModel model = new LogisticRegressionWithLBFGS()
+      .setNumClasses(3)
+      .run(training.rdd());
+
+    // Compute raw scores on the test set.
+    JavaRDD<Tuple2<Object, Object>> predictionAndLabels = test.map(
+      new Function<LabeledPoint, Tuple2<Object, Object>>() {
+        public Tuple2<Object, Object> call(LabeledPoint p) {
+          Double prediction = model.predict(p.features());
+          return new Tuple2<Object, Object>(prediction, p.label());
+        }
+      }
+    );
+
+    // Get evaluation metrics.
+    MulticlassMetrics metrics = new MulticlassMetrics(predictionAndLabels.rdd());
+
+    // Confusion matrix
+    Matrix confusion = metrics.confusionMatrix();
+    System.out.println("Confusion matrix: \n" + confusion);
+
+    // Overall statistics
+    System.out.println("Precision = " + metrics.precision());
+    System.out.println("Recall = " + metrics.recall());
+    System.out.println("F1 Score = " + metrics.fMeasure());
+
+    // Stats by labels
+    for (int i = 0; i < metrics.labels().length; i++) {
+      System.out.format("Class %f precision = %f\n", metrics.labels()[i],metrics.precision
+        (metrics.labels()[i]));
+      System.out.format("Class %f recall = %f\n", metrics.labels()[i], metrics.recall(metrics
+        .labels()[i]));
+      System.out.format("Class %f F1 score = %f\n", metrics.labels()[i], metrics.fMeasure
+        (metrics.labels()[i]));
+    }
+
+    //Weighted stats
+    System.out.format("Weighted precision = %f\n", metrics.weightedPrecision());
+    System.out.format("Weighted recall = %f\n", metrics.weightedRecall());
+    System.out.format("Weighted F1 score = %f\n", metrics.weightedFMeasure());
+    System.out.format("Weighted false positive rate = %f\n", metrics.weightedFalsePositiveRate());
+
+    // Save and load model
+    model.save(sc, "target/tmp/LogisticRegressionModel");
+    LogisticRegressionModel sameModel = LogisticRegressionModel.load(sc,
+      "target/tmp/LogisticRegressionModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaRankingMetricsExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRankingMetricsExample.java
new file mode 100644
index 0000000000000..7c4c97e74681f
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRankingMetricsExample.java
@@ -0,0 +1,176 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import java.util.*;
+
+import scala.Tuple2;
+
+import org.apache.spark.api.java.*;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.mllib.evaluation.RegressionMetrics;
+import org.apache.spark.mllib.evaluation.RankingMetrics;
+import org.apache.spark.mllib.recommendation.ALS;
+import org.apache.spark.mllib.recommendation.MatrixFactorizationModel;
+import org.apache.spark.mllib.recommendation.Rating;
+// $example off$
+import org.apache.spark.SparkConf;
+
+public class JavaRankingMetricsExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("Java Ranking Metrics Example");
+    JavaSparkContext sc = new JavaSparkContext(conf);
+    // $example on$
+    String path = "data/mllib/sample_movielens_data.txt";
+    JavaRDD<String> data = sc.textFile(path);
+    JavaRDD<Rating> ratings = data.map(
+      new Function<String, Rating>() {
+        public Rating call(String line) {
+          String[] parts = line.split("::");
+            return new Rating(Integer.parseInt(parts[0]), Integer.parseInt(parts[1]), Double
+              .parseDouble(parts[2]) - 2.5);
+        }
+      }
+    );
+    ratings.cache();
+
+    // Train an ALS model
+    final MatrixFactorizationModel model = ALS.train(JavaRDD.toRDD(ratings), 10, 10, 0.01);
+
+    // Get top 10 recommendations for every user and scale ratings from 0 to 1
+    JavaRDD<Tuple2<Object, Rating[]>> userRecs = model.recommendProductsForUsers(10).toJavaRDD();
+    JavaRDD<Tuple2<Object, Rating[]>> userRecsScaled = userRecs.map(
+      new Function<Tuple2<Object, Rating[]>, Tuple2<Object, Rating[]>>() {
+        public Tuple2<Object, Rating[]> call(Tuple2<Object, Rating[]> t) {
+          Rating[] scaledRatings = new Rating[t._2().length];
+          for (int i = 0; i < scaledRatings.length; i++) {
+            double newRating = Math.max(Math.min(t._2()[i].rating(), 1.0), 0.0);
+            scaledRatings[i] = new Rating(t._2()[i].user(), t._2()[i].product(), newRating);
+          }
+          return new Tuple2<Object, Rating[]>(t._1(), scaledRatings);
+        }
+      }
+    );
+    JavaPairRDD<Object, Rating[]> userRecommended = JavaPairRDD.fromJavaRDD(userRecsScaled);
+
+    // Map ratings to 1 or 0, 1 indicating a movie that should be recommended
+    JavaRDD<Rating> binarizedRatings = ratings.map(
+      new Function<Rating, Rating>() {
+        public Rating call(Rating r) {
+          double binaryRating;
+          if (r.rating() > 0.0) {
+            binaryRating = 1.0;
+          } else {
+            binaryRating = 0.0;
+          }
+          return new Rating(r.user(), r.product(), binaryRating);
+        }
+      }
+    );
+
+    // Group ratings by common user
+    JavaPairRDD<Object, Iterable<Rating>> userMovies = binarizedRatings.groupBy(
+      new Function<Rating, Object>() {
+        public Object call(Rating r) {
+          return r.user();
+        }
+      }
+    );
+
+    // Get true relevant documents from all user ratings
+    JavaPairRDD<Object, List<Integer>> userMoviesList = userMovies.mapValues(
+      new Function<Iterable<Rating>, List<Integer>>() {
+        public List<Integer> call(Iterable<Rating> docs) {
+          List<Integer> products = new ArrayList<Integer>();
+          for (Rating r : docs) {
+            if (r.rating() > 0.0) {
+              products.add(r.product());
+            }
+          }
+          return products;
+        }
+      }
+    );
+
+    // Extract the product id from each recommendation
+    JavaPairRDD<Object, List<Integer>> userRecommendedList = userRecommended.mapValues(
+      new Function<Rating[], List<Integer>>() {
+        public List<Integer> call(Rating[] docs) {
+          List<Integer> products = new ArrayList<Integer>();
+          for (Rating r : docs) {
+            products.add(r.product());
+          }
+          return products;
+        }
+      }
+    );
+    JavaRDD<Tuple2<List<Integer>, List<Integer>>> relevantDocs = userMoviesList.join
+      (userRecommendedList).values();
+
+    // Instantiate the metrics object
+    RankingMetrics metrics = RankingMetrics.of(relevantDocs);
+
+    // Precision and NDCG at k
+    Integer[] kVector = {1, 3, 5};
+    for (Integer k : kVector) {
+      System.out.format("Precision at %d = %f\n", k, metrics.precisionAt(k));
+      System.out.format("NDCG at %d = %f\n", k, metrics.ndcgAt(k));
+    }
+
+    // Mean average precision
+    System.out.format("Mean average precision = %f\n", metrics.meanAveragePrecision());
+
+    // Evaluate the model using numerical ratings and regression metrics
+    JavaRDD<Tuple2<Object, Object>> userProducts = ratings.map(
+      new Function<Rating, Tuple2<Object, Object>>() {
+        public Tuple2<Object, Object> call(Rating r) {
+          return new Tuple2<Object, Object>(r.user(), r.product());
+        }
+      }
+    );
+    JavaPairRDD<Tuple2<Integer, Integer>, Object> predictions = JavaPairRDD.fromJavaRDD(
+      model.predict(JavaRDD.toRDD(userProducts)).toJavaRDD().map(
+        new Function<Rating, Tuple2<Tuple2<Integer, Integer>, Object>>() {
+          public Tuple2<Tuple2<Integer, Integer>, Object> call(Rating r) {
+            return new Tuple2<Tuple2<Integer, Integer>, Object>(
+              new Tuple2<Integer, Integer>(r.user(), r.product()), r.rating());
+          }
+        }
+      ));
+    JavaRDD<Tuple2<Object, Object>> ratesAndPreds =
+      JavaPairRDD.fromJavaRDD(ratings.map(
+        new Function<Rating, Tuple2<Tuple2<Integer, Integer>, Object>>() {
+          public Tuple2<Tuple2<Integer, Integer>, Object> call(Rating r) {
+            return new Tuple2<Tuple2<Integer, Integer>, Object>(
+              new Tuple2<Integer, Integer>(r.user(), r.product()), r.rating());
+          }
+        }
+      )).join(predictions).values();
+
+    // Create regression metrics object
+    RegressionMetrics regressionMetrics = new RegressionMetrics(ratesAndPreds.rdd());
+
+    // Root mean squared error
+    System.out.format("RMSE = %f\n", regressionMetrics.rootMeanSquaredError());
+
+    // R-squared
+    System.out.format("R-squared = %f\n", regressionMetrics.r2());
+    // $example off$
+  }
+}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaRegressionMetricsExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRegressionMetricsExample.java
new file mode 100644
index 0000000000000..d2efc6bf97776
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRegressionMetricsExample.java
@@ -0,0 +1,91 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.mllib;
+
+// $example on$
+import scala.Tuple2;
+
+import org.apache.spark.api.java.*;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.mllib.linalg.Vectors;
+import org.apache.spark.mllib.regression.LabeledPoint;
+import org.apache.spark.mllib.regression.LinearRegressionModel;
+import org.apache.spark.mllib.regression.LinearRegressionWithSGD;
+import org.apache.spark.mllib.evaluation.RegressionMetrics;
+import org.apache.spark.SparkConf;
+// $example off$
+
+public class JavaRegressionMetricsExample {
+  public static void main(String[] args) {
+    SparkConf conf = new SparkConf().setAppName("Java Regression Metrics Example");
+    JavaSparkContext sc = new JavaSparkContext(conf);
+    // $example on$
+    // Load and parse the data
+    String path = "data/mllib/sample_linear_regression_data.txt";
+    JavaRDD<String> data = sc.textFile(path);
+    JavaRDD<LabeledPoint> parsedData = data.map(
+      new Function<String, LabeledPoint>() {
+        public LabeledPoint call(String line) {
+          String[] parts = line.split(" ");
+          double[] v = new double[parts.length - 1];
+          for (int i = 1; i < parts.length - 1; i++)
+            v[i - 1] = Double.parseDouble(parts[i].split(":")[1]);
+          return new LabeledPoint(Double.parseDouble(parts[0]), Vectors.dense(v));
+        }
+      }
+    );
+    parsedData.cache();
+
+    // Building the model
+    int numIterations = 100;
+    final LinearRegressionModel model = LinearRegressionWithSGD.train(JavaRDD.toRDD(parsedData),
+      numIterations);
+
+    // Evaluate model on training examples and compute training error
+    JavaRDD<Tuple2<Object, Object>> valuesAndPreds = parsedData.map(
+      new Function<LabeledPoint, Tuple2<Object, Object>>() {
+        public Tuple2<Object, Object> call(LabeledPoint point) {
+          double prediction = model.predict(point.features());
+          return new Tuple2<Object, Object>(prediction, point.label());
+        }
+      }
+    );
+
+    // Instantiate metrics object
+    RegressionMetrics metrics = new RegressionMetrics(valuesAndPreds.rdd());
+
+    // Squared error
+    System.out.format("MSE = %f\n", metrics.meanSquaredError());
+    System.out.format("RMSE = %f\n", metrics.rootMeanSquaredError());
+
+    // R-squared
+    System.out.format("R Squared = %f\n", metrics.r2());
+
+    // Mean absolute error
+    System.out.format("MAE = %f\n", metrics.meanAbsoluteError());
+
+    // Explained variance
+    System.out.format("Explained Variance = %f\n", metrics.explainedVariance());
+
+    // Save and load model
+    model.save(sc.sc(), "target/tmp/LogisticRegressionModel");
+    LinearRegressionModel sameModel = LinearRegressionModel.load(sc.sc(),
+      "target/tmp/LogisticRegressionModel");
+    // $example off$
+  }
+}
diff --git a/examples/src/main/python/mllib/binary_classification_metrics_example.py b/examples/src/main/python/mllib/binary_classification_metrics_example.py
new file mode 100644
index 0000000000000..437acb998acc3
--- /dev/null
+++ b/examples/src/main/python/mllib/binary_classification_metrics_example.py
@@ -0,0 +1,55 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+Binary Classification Metrics Example.
+"""
+from __future__ import print_function
+import sys
+from pyspark import SparkContext, SQLContext
+# $example on$
+from pyspark.mllib.classification import LogisticRegressionWithLBFGS
+from pyspark.mllib.evaluation import BinaryClassificationMetrics
+from pyspark.mllib.util import MLUtils
+# $example off$
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="BinaryClassificationMetricsExample")
+    sqlContext = SQLContext(sc)
+    # $example on$
+    # Several of the methods available in scala are currently missing from pyspark
+    # Load training data in LIBSVM format
+    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_binary_classification_data.txt")
+
+    # Split data into training (60%) and test (40%)
+    training, test = data.randomSplit([0.6, 0.4], seed=11L)
+    training.cache()
+
+    # Run training algorithm to build the model
+    model = LogisticRegressionWithLBFGS.train(training)
+
+    # Compute raw scores on the test set
+    predictionAndLabels = test.map(lambda lp: (float(model.predict(lp.features)), lp.label))
+
+    # Instantiate metrics object
+    metrics = BinaryClassificationMetrics(predictionAndLabels)
+
+    # Area under precision-recall curve
+    print("Area under PR = %s" % metrics.areaUnderPR)
+
+    # Area under ROC curve
+    print("Area under ROC = %s" % metrics.areaUnderROC)
+    # $example off$
diff --git a/examples/src/main/python/mllib/multi_class_metrics_example.py b/examples/src/main/python/mllib/multi_class_metrics_example.py
new file mode 100644
index 0000000000000..cd56b3c97c778
--- /dev/null
+++ b/examples/src/main/python/mllib/multi_class_metrics_example.py
@@ -0,0 +1,69 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# $example on$
+from pyspark.mllib.classification import LogisticRegressionWithLBFGS
+from pyspark.mllib.util import MLUtils
+from pyspark.mllib.evaluation import MulticlassMetrics
+# $example off$
+
+from pyspark import SparkContext
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="MultiClassMetricsExample")
+
+    # Several of the methods available in scala are currently missing from pyspark
+    # $example on$
+    # Load training data in LIBSVM format
+    data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt")
+
+    # Split data into training (60%) and test (40%)
+    training, test = data.randomSplit([0.6, 0.4], seed=11L)
+    training.cache()
+
+    # Run training algorithm to build the model
+    model = LogisticRegressionWithLBFGS.train(training, numClasses=3)
+
+    # Compute raw scores on the test set
+    predictionAndLabels = test.map(lambda lp: (float(model.predict(lp.features)), lp.label))
+
+    # Instantiate metrics object
+    metrics = MulticlassMetrics(predictionAndLabels)
+
+    # Overall statistics
+    precision = metrics.precision()
+    recall = metrics.recall()
+    f1Score = metrics.fMeasure()
+    print("Summary Stats")
+    print("Precision = %s" % precision)
+    print("Recall = %s" % recall)
+    print("F1 Score = %s" % f1Score)
+
+    # Statistics by class
+    labels = data.map(lambda lp: lp.label).distinct().collect()
+    for label in sorted(labels):
+        print("Class %s precision = %s" % (label, metrics.precision(label)))
+        print("Class %s recall = %s" % (label, metrics.recall(label)))
+        print("Class %s F1 Measure = %s" % (label, metrics.fMeasure(label, beta=1.0)))
+
+    # Weighted stats
+    print("Weighted recall = %s" % metrics.weightedRecall)
+    print("Weighted precision = %s" % metrics.weightedPrecision)
+    print("Weighted F(1) Score = %s" % metrics.weightedFMeasure())
+    print("Weighted F(0.5) Score = %s" % metrics.weightedFMeasure(beta=0.5))
+    print("Weighted false positive rate = %s" % metrics.weightedFalsePositiveRate)
+    # $example off$
diff --git a/examples/src/main/python/mllib/multi_label_metrics_example.py b/examples/src/main/python/mllib/multi_label_metrics_example.py
new file mode 100644
index 0000000000000..960ade6597379
--- /dev/null
+++ b/examples/src/main/python/mllib/multi_label_metrics_example.py
@@ -0,0 +1,61 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# $example on$
+from pyspark.mllib.evaluation import MultilabelMetrics
+# $example off$
+from pyspark import SparkContext
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="MultiLabelMetricsExample")
+    # $example on$
+    scoreAndLabels = sc.parallelize([
+        ([0.0, 1.0], [0.0, 2.0]),
+        ([0.0, 2.0], [0.0, 1.0]),
+        ([], [0.0]),
+        ([2.0], [2.0]),
+        ([2.0, 0.0], [2.0, 0.0]),
+        ([0.0, 1.0, 2.0], [0.0, 1.0]),
+        ([1.0], [1.0, 2.0])])
+
+    # Instantiate metrics object
+    metrics = MultilabelMetrics(scoreAndLabels)
+
+    # Summary stats
+    print("Recall = %s" % metrics.recall())
+    print("Precision = %s" % metrics.precision())
+    print("F1 measure = %s" % metrics.f1Measure())
+    print("Accuracy = %s" % metrics.accuracy)
+
+    # Individual label stats
+    labels = scoreAndLabels.flatMap(lambda x: x[1]).distinct().collect()
+    for label in labels:
+        print("Class %s precision = %s" % (label, metrics.precision(label)))
+        print("Class %s recall = %s" % (label, metrics.recall(label)))
+        print("Class %s F1 Measure = %s" % (label, metrics.f1Measure(label)))
+
+    # Micro stats
+    print("Micro precision = %s" % metrics.microPrecision)
+    print("Micro recall = %s" % metrics.microRecall)
+    print("Micro F1 measure = %s" % metrics.microF1Measure)
+
+    # Hamming loss
+    print("Hamming loss = %s" % metrics.hammingLoss)
+
+    # Subset accuracy
+    print("Subset accuracy = %s" % metrics.subsetAccuracy)
+    # $example off$
diff --git a/examples/src/main/python/mllib/ranking_metrics_example.py b/examples/src/main/python/mllib/ranking_metrics_example.py
new file mode 100644
index 0000000000000..327791966c901
--- /dev/null
+++ b/examples/src/main/python/mllib/ranking_metrics_example.py
@@ -0,0 +1,55 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# $example on$
+from pyspark.mllib.recommendation import ALS, Rating
+from pyspark.mllib.evaluation import RegressionMetrics, RankingMetrics
+# $example off$
+from pyspark import SparkContext
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="Ranking Metrics Example")
+
+    # Several of the methods available in scala are currently missing from pyspark
+    # $example on$
+    # Read in the ratings data
+    lines = sc.textFile("data/mllib/sample_movielens_data.txt")
+
+    def parseLine(line):
+        fields = line.split("::")
+        return Rating(int(fields[0]), int(fields[1]), float(fields[2]) - 2.5)
+    ratings = lines.map(lambda r: parseLine(r))
+
+    # Train a model on to predict user-product ratings
+    model = ALS.train(ratings, 10, 10, 0.01)
+
+    # Get predicted ratings on all existing user-product pairs
+    testData = ratings.map(lambda p: (p.user, p.product))
+    predictions = model.predictAll(testData).map(lambda r: ((r.user, r.product), r.rating))
+
+    ratingsTuple = ratings.map(lambda r: ((r.user, r.product), r.rating))
+    scoreAndLabels = predictions.join(ratingsTuple).map(lambda tup: tup[1])
+
+    # Instantiate regression metrics to compare predicted and actual ratings
+    metrics = RegressionMetrics(scoreAndLabels)
+
+    # Root mean sqaured error
+    print("RMSE = %s" % metrics.rootMeanSquaredError)
+
+    # R-squared
+    print("R-squared = %s" % metrics.r2)
+    # $example off$
diff --git a/examples/src/main/python/mllib/regression_metrics_example.py b/examples/src/main/python/mllib/regression_metrics_example.py
new file mode 100644
index 0000000000000..a3a83aafd7a1f
--- /dev/null
+++ b/examples/src/main/python/mllib/regression_metrics_example.py
@@ -0,0 +1,59 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# $example on$
+from pyspark.mllib.regression import LabeledPoint, LinearRegressionWithSGD
+from pyspark.mllib.evaluation import RegressionMetrics
+from pyspark.mllib.linalg import DenseVector
+# $example off$
+
+from pyspark import SparkContext
+
+if __name__ == "__main__":
+    sc = SparkContext(appName="Regression Metrics Example")
+
+    # $example on$
+    # Load and parse the data
+    def parsePoint(line):
+        values = line.split()
+        return LabeledPoint(float(values[0]),
+                            DenseVector([float(x.split(':')[1]) for x in values[1:]]))
+
+    data = sc.textFile("data/mllib/sample_linear_regression_data.txt")
+    parsedData = data.map(parsePoint)
+
+    # Build the model
+    model = LinearRegressionWithSGD.train(parsedData)
+
+    # Get predictions
+    valuesAndPreds = parsedData.map(lambda p: (float(model.predict(p.features)), p.label))
+
+    # Instantiate metrics object
+    metrics = RegressionMetrics(valuesAndPreds)
+
+    # Squared Error
+    print("MSE = %s" % metrics.meanSquaredError)
+    print("RMSE = %s" % metrics.rootMeanSquaredError)
+
+    # R-squared
+    print("R-squared = %s" % metrics.r2)
+
+    # Mean absolute error
+    print("MAE = %s" % metrics.meanAbsoluteError)
+
+    # Explained variance
+    print("Explained variance = %s" % metrics.explainedVariance)
+    # $example off$
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/BinaryClassificationMetricsExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/BinaryClassificationMetricsExample.scala
new file mode 100644
index 0000000000000..13a37827ab935
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/BinaryClassificationMetricsExample.scala
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS
+import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics
+import org.apache.spark.mllib.regression.LabeledPoint
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+import org.apache.spark.{SparkContext, SparkConf}
+
+object BinaryClassificationMetricsExample {
+
+  def main(args: Array[String]): Unit = {
+
+    val conf = new SparkConf().setAppName("BinaryClassificationMetricsExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    // Load training data in LIBSVM format
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_binary_classification_data.txt")
+
+    // Split data into training (60%) and test (40%)
+    val Array(training, test) = data.randomSplit(Array(0.6, 0.4), seed = 11L)
+    training.cache()
+
+    // Run training algorithm to build the model
+    val model = new LogisticRegressionWithLBFGS()
+      .setNumClasses(2)
+      .run(training)
+
+    // Clear the prediction threshold so the model will return probabilities
+    model.clearThreshold
+
+    // Compute raw scores on the test set
+    val predictionAndLabels = test.map { case LabeledPoint(label, features) =>
+      val prediction = model.predict(features)
+      (prediction, label)
+    }
+
+    // Instantiate metrics object
+    val metrics = new BinaryClassificationMetrics(predictionAndLabels)
+
+    // Precision by threshold
+    val precision = metrics.precisionByThreshold
+    precision.foreach { case (t, p) =>
+      println(s"Threshold: $t, Precision: $p")
+    }
+
+    // Recall by threshold
+    val recall = metrics.recallByThreshold
+    recall.foreach { case (t, r) =>
+      println(s"Threshold: $t, Recall: $r")
+    }
+
+    // Precision-Recall Curve
+    val PRC = metrics.pr
+
+    // F-measure
+    val f1Score = metrics.fMeasureByThreshold
+    f1Score.foreach { case (t, f) =>
+      println(s"Threshold: $t, F-score: $f, Beta = 1")
+    }
+
+    val beta = 0.5
+    val fScore = metrics.fMeasureByThreshold(beta)
+    f1Score.foreach { case (t, f) =>
+      println(s"Threshold: $t, F-score: $f, Beta = 0.5")
+    }
+
+    // AUPRC
+    val auPRC = metrics.areaUnderPR
+    println("Area under precision-recall curve = " + auPRC)
+
+    // Compute thresholds used in ROC and PR curves
+    val thresholds = precision.map(_._1)
+
+    // ROC Curve
+    val roc = metrics.roc
+
+    // AUROC
+    val auROC = metrics.areaUnderROC
+    println("Area under ROC = " + auROC)
+    // $example off$
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/MultiLabelMetricsExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/MultiLabelMetricsExample.scala
new file mode 100644
index 0000000000000..4503c15360adc
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/MultiLabelMetricsExample.scala
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.evaluation.MultilabelMetrics
+import org.apache.spark.rdd.RDD
+// $example off$
+import org.apache.spark.{SparkContext, SparkConf}
+
+object MultiLabelMetricsExample {
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("MultiLabelMetricsExample")
+    val sc = new SparkContext(conf)
+    // $example on$
+    val scoreAndLabels: RDD[(Array[Double], Array[Double])] = sc.parallelize(
+      Seq((Array(0.0, 1.0), Array(0.0, 2.0)),
+        (Array(0.0, 2.0), Array(0.0, 1.0)),
+        (Array.empty[Double], Array(0.0)),
+        (Array(2.0), Array(2.0)),
+        (Array(2.0, 0.0), Array(2.0, 0.0)),
+        (Array(0.0, 1.0, 2.0), Array(0.0, 1.0)),
+        (Array(1.0), Array(1.0, 2.0))), 2)
+
+    // Instantiate metrics object
+    val metrics = new MultilabelMetrics(scoreAndLabels)
+
+    // Summary stats
+    println(s"Recall = ${metrics.recall}")
+    println(s"Precision = ${metrics.precision}")
+    println(s"F1 measure = ${metrics.f1Measure}")
+    println(s"Accuracy = ${metrics.accuracy}")
+
+    // Individual label stats
+    metrics.labels.foreach(label =>
+      println(s"Class $label precision = ${metrics.precision(label)}"))
+    metrics.labels.foreach(label => println(s"Class $label recall = ${metrics.recall(label)}"))
+    metrics.labels.foreach(label => println(s"Class $label F1-score = ${metrics.f1Measure(label)}"))
+
+    // Micro stats
+    println(s"Micro recall = ${metrics.microRecall}")
+    println(s"Micro precision = ${metrics.microPrecision}")
+    println(s"Micro F1 measure = ${metrics.microF1Measure}")
+
+    // Hamming loss
+    println(s"Hamming loss = ${metrics.hammingLoss}")
+
+    // Subset accuracy
+    println(s"Subset accuracy = ${metrics.subsetAccuracy}")
+    // $example off$
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/MulticlassMetricsExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/MulticlassMetricsExample.scala
new file mode 100644
index 0000000000000..0904449245989
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/MulticlassMetricsExample.scala
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS
+import org.apache.spark.mllib.evaluation.MulticlassMetrics
+import org.apache.spark.mllib.regression.LabeledPoint
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+import org.apache.spark.{SparkContext, SparkConf}
+
+object MulticlassMetricsExample {
+
+  def main(args: Array[String]): Unit = {
+    val conf = new SparkConf().setAppName("MulticlassMetricsExample")
+    val sc = new SparkContext(conf)
+
+    // $example on$
+    // Load training data in LIBSVM format
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_multiclass_classification_data.txt")
+
+    // Split data into training (60%) and test (40%)
+    val Array(training, test) = data.randomSplit(Array(0.6, 0.4), seed = 11L)
+    training.cache()
+
+    // Run training algorithm to build the model
+    val model = new LogisticRegressionWithLBFGS()
+      .setNumClasses(3)
+      .run(training)
+
+    // Compute raw scores on the test set
+    val predictionAndLabels = test.map { case LabeledPoint(label, features) =>
+      val prediction = model.predict(features)
+      (prediction, label)
+    }
+
+    // Instantiate metrics object
+    val metrics = new MulticlassMetrics(predictionAndLabels)
+
+    // Confusion matrix
+    println("Confusion matrix:")
+    println(metrics.confusionMatrix)
+
+    // Overall Statistics
+    val precision = metrics.precision
+    val recall = metrics.recall // same as true positive rate
+    val f1Score = metrics.fMeasure
+    println("Summary Statistics")
+    println(s"Precision = $precision")
+    println(s"Recall = $recall")
+    println(s"F1 Score = $f1Score")
+
+    // Precision by label
+    val labels = metrics.labels
+    labels.foreach { l =>
+      println(s"Precision($l) = " + metrics.precision(l))
+    }
+
+    // Recall by label
+    labels.foreach { l =>
+      println(s"Recall($l) = " + metrics.recall(l))
+    }
+
+    // False positive rate by label
+    labels.foreach { l =>
+      println(s"FPR($l) = " + metrics.falsePositiveRate(l))
+    }
+
+    // F-measure by label
+    labels.foreach { l =>
+      println(s"F1-Score($l) = " + metrics.fMeasure(l))
+    }
+
+    // Weighted stats
+    println(s"Weighted precision: ${metrics.weightedPrecision}")
+    println(s"Weighted recall: ${metrics.weightedRecall}")
+    println(s"Weighted F1 score: ${metrics.weightedFMeasure}")
+    println(s"Weighted false positive rate: ${metrics.weightedFalsePositiveRate}")
+    // $example off$
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/RankingMetricsExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/RankingMetricsExample.scala
new file mode 100644
index 0000000000000..cffa03d5cc9f4
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/RankingMetricsExample.scala
@@ -0,0 +1,110 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// scalastyle:off println
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.evaluation.{RegressionMetrics, RankingMetrics}
+import org.apache.spark.mllib.recommendation.{ALS, Rating}
+// $example off$
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkContext, SparkConf}
+
+object RankingMetricsExample {
+  def main(args: Array[String]) {
+    val conf = new SparkConf().setAppName("RankingMetricsExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+    import sqlContext.implicits._
+    // $example on$
+    // Read in the ratings data
+    val ratings = sc.textFile("data/mllib/sample_movielens_data.txt").map { line =>
+      val fields = line.split("::")
+      Rating(fields(0).toInt, fields(1).toInt, fields(2).toDouble - 2.5)
+    }.cache()
+
+    // Map ratings to 1 or 0, 1 indicating a movie that should be recommended
+    val binarizedRatings = ratings.map(r => Rating(r.user, r.product,
+      if (r.rating > 0) 1.0 else 0.0)).cache()
+
+    // Summarize ratings
+    val numRatings = ratings.count()
+    val numUsers = ratings.map(_.user).distinct().count()
+    val numMovies = ratings.map(_.product).distinct().count()
+    println(s"Got $numRatings ratings from $numUsers users on $numMovies movies.")
+
+    // Build the model
+    val numIterations = 10
+    val rank = 10
+    val lambda = 0.01
+    val model = ALS.train(ratings, rank, numIterations, lambda)
+
+    // Define a function to scale ratings from 0 to 1
+    def scaledRating(r: Rating): Rating = {
+      val scaledRating = math.max(math.min(r.rating, 1.0), 0.0)
+      Rating(r.user, r.product, scaledRating)
+    }
+
+    // Get sorted top ten predictions for each user and then scale from [0, 1]
+    val userRecommended = model.recommendProductsForUsers(10).map { case (user, recs) =>
+      (user, recs.map(scaledRating))
+    }
+
+    // Assume that any movie a user rated 3 or higher (which maps to a 1) is a relevant document
+    // Compare with top ten most relevant documents
+    val userMovies = binarizedRatings.groupBy(_.user)
+    val relevantDocuments = userMovies.join(userRecommended).map { case (user, (actual,
+    predictions)) =>
+      (predictions.map(_.product), actual.filter(_.rating > 0.0).map(_.product).toArray)
+    }
+
+    // Instantiate metrics object
+    val metrics = new RankingMetrics(relevantDocuments)
+
+    // Precision at K
+    Array(1, 3, 5).foreach { k =>
+      println(s"Precision at $k = ${metrics.precisionAt(k)}")
+    }
+
+    // Mean average precision
+    println(s"Mean average precision = ${metrics.meanAveragePrecision}")
+
+    // Normalized discounted cumulative gain
+    Array(1, 3, 5).foreach { k =>
+      println(s"NDCG at $k = ${metrics.ndcgAt(k)}")
+    }
+
+    // Get predictions for each data point
+    val allPredictions = model.predict(ratings.map(r => (r.user, r.product))).map(r => ((r.user,
+      r.product), r.rating))
+    val allRatings = ratings.map(r => ((r.user, r.product), r.rating))
+    val predictionsAndLabels = allPredictions.join(allRatings).map { case ((user, product),
+    (predicted, actual)) =>
+      (predicted, actual)
+    }
+
+    // Get the RMSE using regression metrics
+    val regressionMetrics = new RegressionMetrics(predictionsAndLabels)
+    println(s"RMSE = ${regressionMetrics.rootMeanSquaredError}")
+
+    // R-squared
+    println(s"R-squared = ${regressionMetrics.r2}")
+    // $example off$
+  }
+}
+// scalastyle:on println
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/RegressionMetricsExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/RegressionMetricsExample.scala
new file mode 100644
index 0000000000000..47d44532521ca
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/RegressionMetricsExample.scala
@@ -0,0 +1,67 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+// scalastyle:off println
+
+package org.apache.spark.examples.mllib
+
+// $example on$
+import org.apache.spark.mllib.regression.LinearRegressionWithSGD
+import org.apache.spark.mllib.evaluation.RegressionMetrics
+import org.apache.spark.mllib.util.MLUtils
+// $example off$
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.{SparkConf, SparkContext}
+
+object RegressionMetricsExample {
+  def main(args: Array[String]) : Unit = {
+    val conf = new SparkConf().setAppName("RegressionMetricsExample")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+    // $example on$
+    // Load the data
+    val data = MLUtils.loadLibSVMFile(sc, "data/mllib/sample_linear_regression_data.txt").cache()
+
+    // Build the model
+    val numIterations = 100
+    val model = LinearRegressionWithSGD.train(data, numIterations)
+
+    // Get predictions
+    val valuesAndPreds = data.map{ point =>
+      val prediction = model.predict(point.features)
+      (prediction, point.label)
+    }
+
+    // Instantiate metrics object
+    val metrics = new RegressionMetrics(valuesAndPreds)
+
+    // Squared error
+    println(s"MSE = ${metrics.meanSquaredError}")
+    println(s"RMSE = ${metrics.rootMeanSquaredError}")
+
+    // R-squared
+    println(s"R-squared = ${metrics.r2}")
+
+    // Mean absolute error
+    println(s"MAE = ${metrics.meanAbsoluteError}")
+
+    // Explained variance
+    println(s"Explained variance = ${metrics.explainedVariance}")
+    // $example off$
+  }
+}
+// scalastyle:on println
+

From 58b4e4f88a330135c4cec04a30d24ef91bc61d91 Mon Sep 17 00:00:00 2001
From: Nong Li <nong@databricks.com>
Date: Fri, 20 Nov 2015 15:30:53 -0800
Subject: [PATCH 394/518] [SPARK-11787][SPARK-11883][SQL][FOLLOW-UP] Cleanup
 for this patch.

This mainly moves SqlNewHadoopRDD to the sql package. There is some state that is
shared between core and I've left that in core. This allows some other associated
minor cleanup.

Author: Nong Li <nong@databricks.com>

Closes #9845 from nongli/spark-11787.
---
 .../org/apache/spark/rdd/HadoopRDD.scala      |  6 +-
 .../spark/rdd/SqlNewHadoopRDDState.scala      | 41 +++++++++++++
 .../sql/catalyst/expressions/UnsafeRow.java   | 59 ++++++++++++++----
 .../catalyst/expressions/InputFileName.scala  |  6 +-
 .../parquet/UnsafeRowParquetRecordReader.java | 14 +++++
 .../scala/org/apache/spark/sql/SQLConf.scala  |  5 ++
 .../datasources}/SqlNewHadoopRDD.scala        | 60 +++++++------------
 .../datasources/parquet/ParquetRelation.scala |  2 +-
 .../parquet/ParquetFilterSuite.scala          | 43 ++++++-------
 .../datasources/parquet/ParquetIOSuite.scala  | 19 ++++++
 10 files changed, 175 insertions(+), 80 deletions(-)
 create mode 100644 core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDDState.scala
 rename {core/src/main/scala/org/apache/spark/rdd => sql/core/src/main/scala/org/apache/spark/sql/execution/datasources}/SqlNewHadoopRDD.scala (86%)

diff --git a/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala b/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
index 7db583468792e..f37c95bedc0a5 100644
--- a/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/HadoopRDD.scala
@@ -215,8 +215,8 @@ class HadoopRDD[K, V](
 
       // Sets the thread local variable for the file's name
       split.inputSplit.value match {
-        case fs: FileSplit => SqlNewHadoopRDD.setInputFileName(fs.getPath.toString)
-        case _ => SqlNewHadoopRDD.unsetInputFileName()
+        case fs: FileSplit => SqlNewHadoopRDDState.setInputFileName(fs.getPath.toString)
+        case _ => SqlNewHadoopRDDState.unsetInputFileName()
       }
 
       // Find a function that will return the FileSystem bytes read by this thread. Do this before
@@ -256,7 +256,7 @@ class HadoopRDD[K, V](
 
       override def close() {
         if (reader != null) {
-          SqlNewHadoopRDD.unsetInputFileName()
+          SqlNewHadoopRDDState.unsetInputFileName()
           // Close the reader and release it. Note: it's very important that we don't close the
           // reader more than once, since that exposes us to MAPREDUCE-5918 when running against
           // Hadoop 1.x and older Hadoop 2.x releases. That bug can lead to non-deterministic
diff --git a/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDDState.scala b/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDDState.scala
new file mode 100644
index 0000000000000..3f15fff793661
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDDState.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.rdd
+
+import org.apache.spark.unsafe.types.UTF8String
+
+/**
+ * State for SqlNewHadoopRDD objects. This is split this way because of the package splits.
+ * TODO: Move/Combine this with org.apache.spark.sql.datasources.SqlNewHadoopRDD
+ */
+private[spark] object SqlNewHadoopRDDState {
+  /**
+   * The thread variable for the name of the current file being read. This is used by
+   * the InputFileName function in Spark SQL.
+   */
+  private[this] val inputFileName: ThreadLocal[UTF8String] = new ThreadLocal[UTF8String] {
+    override protected def initialValue(): UTF8String = UTF8String.fromString("")
+  }
+
+  def getInputFileName(): UTF8String = inputFileName.get()
+
+  private[spark] def setInputFileName(file: String) = inputFileName.set(UTF8String.fromString(file))
+
+  private[spark] def unsetInputFileName(): Unit = inputFileName.remove()
+
+}
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
index 33769363a0ed5..b6979d0c82977 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java
@@ -17,7 +17,11 @@
 
 package org.apache.spark.sql.catalyst.expressions;
 
-import java.io.*;
+import java.io.Externalizable;
+import java.io.IOException;
+import java.io.ObjectInput;
+import java.io.ObjectOutput;
+import java.io.OutputStream;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
@@ -26,12 +30,26 @@
 import java.util.HashSet;
 import java.util.Set;
 
-import com.esotericsoftware.kryo.Kryo;
-import com.esotericsoftware.kryo.KryoSerializable;
-import com.esotericsoftware.kryo.io.Input;
-import com.esotericsoftware.kryo.io.Output;
-
-import org.apache.spark.sql.types.*;
+import org.apache.spark.sql.types.ArrayType;
+import org.apache.spark.sql.types.BinaryType;
+import org.apache.spark.sql.types.BooleanType;
+import org.apache.spark.sql.types.ByteType;
+import org.apache.spark.sql.types.CalendarIntervalType;
+import org.apache.spark.sql.types.DataType;
+import org.apache.spark.sql.types.DateType;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.types.DecimalType;
+import org.apache.spark.sql.types.DoubleType;
+import org.apache.spark.sql.types.FloatType;
+import org.apache.spark.sql.types.IntegerType;
+import org.apache.spark.sql.types.LongType;
+import org.apache.spark.sql.types.MapType;
+import org.apache.spark.sql.types.NullType;
+import org.apache.spark.sql.types.ShortType;
+import org.apache.spark.sql.types.StringType;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.types.TimestampType;
+import org.apache.spark.sql.types.UserDefinedType;
 import org.apache.spark.unsafe.Platform;
 import org.apache.spark.unsafe.array.ByteArrayMethods;
 import org.apache.spark.unsafe.bitset.BitSetMethods;
@@ -39,9 +57,23 @@
 import org.apache.spark.unsafe.types.CalendarInterval;
 import org.apache.spark.unsafe.types.UTF8String;
 
-import static org.apache.spark.sql.types.DataTypes.*;
+import static org.apache.spark.sql.types.DataTypes.BooleanType;
+import static org.apache.spark.sql.types.DataTypes.ByteType;
+import static org.apache.spark.sql.types.DataTypes.DateType;
+import static org.apache.spark.sql.types.DataTypes.DoubleType;
+import static org.apache.spark.sql.types.DataTypes.FloatType;
+import static org.apache.spark.sql.types.DataTypes.IntegerType;
+import static org.apache.spark.sql.types.DataTypes.LongType;
+import static org.apache.spark.sql.types.DataTypes.NullType;
+import static org.apache.spark.sql.types.DataTypes.ShortType;
+import static org.apache.spark.sql.types.DataTypes.TimestampType;
 import static org.apache.spark.unsafe.Platform.BYTE_ARRAY_OFFSET;
 
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.KryoSerializable;
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.Output;
+
 /**
  * An Unsafe implementation of Row which is backed by raw memory instead of Java objects.
  *
@@ -116,11 +148,6 @@ public static boolean isMutable(DataType dt) {
   /** The size of this row's backing data, in bytes) */
   private int sizeInBytes;
 
-  private void setNotNullAt(int i) {
-    assertIndexIsValid(i);
-    BitSetMethods.unset(baseObject, baseOffset, i);
-  }
-
   /** The width of the null tracking bit set, in bytes */
   private int bitSetWidthInBytes;
 
@@ -187,6 +214,12 @@ public void pointTo(byte[] buf, int sizeInBytes) {
     pointTo(buf, numFields, sizeInBytes);
   }
 
+
+  public void setNotNullAt(int i) {
+    assertIndexIsValid(i);
+    BitSetMethods.unset(baseObject, baseOffset, i);
+  }
+
   @Override
   public void setNullAt(int i) {
     assertIndexIsValid(i);
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/InputFileName.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/InputFileName.scala
index d809877817a5b..bf215783fc27d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/InputFileName.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/InputFileName.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
-import org.apache.spark.rdd.SqlNewHadoopRDD
+import org.apache.spark.rdd.SqlNewHadoopRDDState
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.codegen.{GeneratedExpressionCode, CodeGenContext}
 import org.apache.spark.sql.types.{DataType, StringType}
@@ -37,13 +37,13 @@ case class InputFileName() extends LeafExpression with Nondeterministic {
   override protected def initInternal(): Unit = {}
 
   override protected def evalInternal(input: InternalRow): UTF8String = {
-    SqlNewHadoopRDD.getInputFileName()
+    SqlNewHadoopRDDState.getInputFileName()
   }
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
     ev.isNull = "false"
     s"final ${ctx.javaType(dataType)} ${ev.value} = " +
-      "org.apache.spark.rdd.SqlNewHadoopRDD.getInputFileName();"
+      "org.apache.spark.rdd.SqlNewHadoopRDDState.getInputFileName();"
   }
 
 }
diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/UnsafeRowParquetRecordReader.java b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/UnsafeRowParquetRecordReader.java
index 8a92e489ccb7c..dade488ca281b 100644
--- a/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/UnsafeRowParquetRecordReader.java
+++ b/sql/core/src/main/java/org/apache/spark/sql/execution/datasources/parquet/UnsafeRowParquetRecordReader.java
@@ -108,6 +108,19 @@ public class UnsafeRowParquetRecordReader extends SpecificParquetRecordReaderBas
    */
   private static final int DEFAULT_VAR_LEN_SIZE = 32;
 
+  /**
+   * Tries to initialize the reader for this split. Returns true if this reader supports reading
+   * this split and false otherwise.
+   */
+  public boolean tryInitialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext) {
+    try {
+      initialize(inputSplit, taskAttemptContext);
+      return true;
+    } catch (Exception e) {
+      return false;
+    }
+  }
+
   /**
    * Implementation of RecordReader API.
    */
@@ -326,6 +339,7 @@ private void decodeBinaryBatch(int col, int num) throws IOException {
         } else {
           rowWriters[n].write(col, bytes.array(), bytes.position(), len);
         }
+        rows[n].setNotNullAt(col);
       } else {
         rows[n].setNullAt(col);
       }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
index f40e603cd1939..5ef3a48c56a87 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala
@@ -323,6 +323,11 @@ private[spark] object SQLConf {
       "option must be set in Hadoop Configuration.  2. This option overrides " +
       "\"spark.sql.sources.outputCommitterClass\".")
 
+  val PARQUET_UNSAFE_ROW_RECORD_READER_ENABLED = booleanConf(
+    key = "spark.sql.parquet.enableUnsafeRowRecordReader",
+    defaultValue = Some(true),
+    doc = "Enables using the custom ParquetUnsafeRowRecordReader.")
+
   val ORC_FILTER_PUSHDOWN_ENABLED = booleanConf("spark.sql.orc.filterPushdown",
     defaultValue = Some(false),
     doc = "When true, enable filter pushdown for ORC files.")
diff --git a/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SqlNewHadoopRDD.scala
similarity index 86%
rename from core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDD.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SqlNewHadoopRDD.scala
index 4d176332b69ce..56cb63d9eff2a 100644
--- a/core/src/main/scala/org/apache/spark/rdd/SqlNewHadoopRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SqlNewHadoopRDD.scala
@@ -20,6 +20,8 @@ package org.apache.spark.rdd
 import java.text.SimpleDateFormat
 import java.util.Date
 
+import scala.reflect.ClassTag
+
 import org.apache.hadoop.conf.{Configurable, Configuration}
 import org.apache.hadoop.io.Writable
 import org.apache.hadoop.mapreduce._
@@ -28,13 +30,12 @@ import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.executor.DataReadMethod
 import org.apache.spark.mapreduce.SparkHadoopMapReduceUtil
+import org.apache.spark.sql.{SQLConf, SQLContext}
+import org.apache.spark.sql.execution.datasources.parquet.UnsafeRowParquetRecordReader
 import org.apache.spark.storage.StorageLevel
-import org.apache.spark.unsafe.types.UTF8String
-import org.apache.spark.util.{Utils, SerializableConfiguration, ShutdownHookManager}
+import org.apache.spark.util.{SerializableConfiguration, ShutdownHookManager}
 import org.apache.spark.{Partition => SparkPartition, _}
 
-import scala.reflect.ClassTag
-
 
 private[spark] class SqlNewHadoopPartition(
     rddId: Int,
@@ -61,13 +62,13 @@ private[spark] class SqlNewHadoopPartition(
  * changes based on [[org.apache.spark.rdd.HadoopRDD]].
  */
 private[spark] class SqlNewHadoopRDD[V: ClassTag](
-    sc : SparkContext,
+    sqlContext: SQLContext,
     broadcastedConf: Broadcast[SerializableConfiguration],
     @transient private val initDriverSideJobFuncOpt: Option[Job => Unit],
     initLocalJobFuncOpt: Option[Job => Unit],
     inputFormatClass: Class[_ <: InputFormat[Void, V]],
     valueClass: Class[V])
-  extends RDD[V](sc, Nil)
+    extends RDD[V](sqlContext.sparkContext, Nil)
   with SparkHadoopMapReduceUtil
   with Logging {
 
@@ -99,7 +100,7 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
   // If true, enable using the custom RecordReader for parquet. This only works for
   // a subset of the types (no complex types).
   protected val enableUnsafeRowParquetReader: Boolean =
-      sc.conf.getBoolean("spark.parquet.enableUnsafeRowRecordReader", true)
+    sqlContext.getConf(SQLConf.PARQUET_UNSAFE_ROW_RECORD_READER_ENABLED.key).toBoolean
 
   override def getPartitions: Array[SparkPartition] = {
     val conf = getConf(isDriverSide = true)
@@ -120,8 +121,8 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
   }
 
   override def compute(
-      theSplit: SparkPartition,
-      context: TaskContext): Iterator[V] = {
+    theSplit: SparkPartition,
+    context: TaskContext): Iterator[V] = {
     val iter = new Iterator[V] {
       val split = theSplit.asInstanceOf[SqlNewHadoopPartition]
       logInfo("Input split: " + split.serializableHadoopSplit)
@@ -132,8 +133,8 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
 
       // Sets the thread local variable for the file's name
       split.serializableHadoopSplit.value match {
-        case fs: FileSplit => SqlNewHadoopRDD.setInputFileName(fs.getPath.toString)
-        case _ => SqlNewHadoopRDD.unsetInputFileName()
+        case fs: FileSplit => SqlNewHadoopRDDState.setInputFileName(fs.getPath.toString)
+        case _ => SqlNewHadoopRDDState.unsetInputFileName()
       }
 
       // Find a function that will return the FileSystem bytes read by this thread. Do this before
@@ -163,15 +164,13 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
         * TODO: plumb this through a different way?
         */
       if (enableUnsafeRowParquetReader &&
-          format.getClass.getName == "org.apache.parquet.hadoop.ParquetInputFormat") {
-        // TODO: move this class to sql.execution and remove this.
-        reader = Utils.classForName(
-          "org.apache.spark.sql.execution.datasources.parquet.UnsafeRowParquetRecordReader")
-          .newInstance().asInstanceOf[RecordReader[Void, V]]
-        try {
-          reader.initialize(split.serializableHadoopSplit.value, hadoopAttemptContext)
-        } catch {
-          case e: Exception => reader = null
+        format.getClass.getName == "org.apache.parquet.hadoop.ParquetInputFormat") {
+        val parquetReader: UnsafeRowParquetRecordReader = new UnsafeRowParquetRecordReader()
+        if (!parquetReader.tryInitialize(
+            split.serializableHadoopSplit.value, hadoopAttemptContext)) {
+          parquetReader.close()
+        } else {
+          reader = parquetReader.asInstanceOf[RecordReader[Void, V]]
         }
       }
 
@@ -217,7 +216,7 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
 
       private def close() {
         if (reader != null) {
-          SqlNewHadoopRDD.unsetInputFileName()
+          SqlNewHadoopRDDState.unsetInputFileName()
           // Close the reader and release it. Note: it's very important that we don't close the
           // reader more than once, since that exposes us to MAPREDUCE-5918 when running against
           // Hadoop 1.x and older Hadoop 2.x releases. That bug can lead to non-deterministic
@@ -235,7 +234,7 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
           if (bytesReadCallback.isDefined) {
             inputMetrics.updateBytesRead()
           } else if (split.serializableHadoopSplit.value.isInstanceOf[FileSplit] ||
-                     split.serializableHadoopSplit.value.isInstanceOf[CombineFileSplit]) {
+            split.serializableHadoopSplit.value.isInstanceOf[CombineFileSplit]) {
             // If we can't get the bytes read from the FS stats, fall back to the split size,
             // which may be inaccurate.
             try {
@@ -276,23 +275,6 @@ private[spark] class SqlNewHadoopRDD[V: ClassTag](
     }
     super.persist(storageLevel)
   }
-}
-
-private[spark] object SqlNewHadoopRDD {
-
-  /**
-   * The thread variable for the name of the current file being read. This is used by
-   * the InputFileName function in Spark SQL.
-   */
-  private[this] val inputFileName: ThreadLocal[UTF8String] = new ThreadLocal[UTF8String] {
-    override protected def initialValue(): UTF8String = UTF8String.fromString("")
-  }
-
-  def getInputFileName(): UTF8String = inputFileName.get()
-
-  private[spark] def setInputFileName(file: String) = inputFileName.set(UTF8String.fromString(file))
-
-  private[spark] def unsetInputFileName(): Unit = inputFileName.remove()
 
   /**
    * Analogous to [[org.apache.spark.rdd.MapPartitionsRDD]], but passes in an InputSplit to
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
index cb0aab8cc0d09..fdd745f48e973 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetRelation.scala
@@ -319,7 +319,7 @@ private[sql] class ParquetRelation(
 
     Utils.withDummyCallSite(sqlContext.sparkContext) {
       new SqlNewHadoopRDD(
-        sc = sqlContext.sparkContext,
+        sqlContext = sqlContext,
         broadcastedConf = broadcastedConf,
         initDriverSideJobFuncOpt = Some(setInputPaths),
         initLocalJobFuncOpt = Some(initLocalJobFuncOpt),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
index c8028a5ef5528..cc5aae03d5516 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala
@@ -337,29 +337,30 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
     }
   }
 
-  // Renable when we can toggle custom ParquetRecordReader on/off. The custom reader does
-  // not do row by row filtering (and we probably don't want to push that).
-  ignore("SPARK-11661 Still pushdown filters returned by unhandledFilters") {
+  // The unsafe row RecordReader does not support row by row filtering so run it with it disabled.
+  test("SPARK-11661 Still pushdown filters returned by unhandledFilters") {
     import testImplicits._
     withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true") {
-      withTempPath { dir =>
-        val path = s"${dir.getCanonicalPath}/part=1"
-        (1 to 3).map(i => (i, i.toString)).toDF("a", "b").write.parquet(path)
-        val df = sqlContext.read.parquet(path).filter("a = 2")
-
-        // This is the source RDD without Spark-side filtering.
-        val childRDD =
-          df
-            .queryExecution
-            .executedPlan.asInstanceOf[org.apache.spark.sql.execution.Filter]
-            .child
-            .execute()
-
-        // The result should be single row.
-        // When a filter is pushed to Parquet, Parquet can apply it to every row.
-        // So, we can check the number of rows returned from the Parquet
-        // to make sure our filter pushdown work.
-        assert(childRDD.count == 1)
+      withSQLConf(SQLConf.PARQUET_UNSAFE_ROW_RECORD_READER_ENABLED.key -> "false") {
+        withTempPath { dir =>
+          val path = s"${dir.getCanonicalPath}/part=1"
+          (1 to 3).map(i => (i, i.toString)).toDF("a", "b").write.parquet(path)
+          val df = sqlContext.read.parquet(path).filter("a = 2")
+
+          // This is the source RDD without Spark-side filtering.
+          val childRDD =
+            df
+              .queryExecution
+              .executedPlan.asInstanceOf[org.apache.spark.sql.execution.Filter]
+              .child
+              .execute()
+
+          // The result should be single row.
+          // When a filter is pushed to Parquet, Parquet can apply it to every row.
+          // So, we can check the number of rows returned from the Parquet
+          // to make sure our filter pushdown work.
+          assert(childRDD.count == 1)
+        }
       }
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
index 177ab42f7767c..0c5d4887ed799 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetIOSuite.scala
@@ -579,6 +579,25 @@ class ParquetIOSuite extends QueryTest with ParquetTest with SharedSQLContext {
     }
   }
 
+  test("null and non-null strings") {
+    // Create a dataset where the first values are NULL and then some non-null values. The
+    // number of non-nulls needs to be bigger than the ParquetReader batch size.
+    val data = sqlContext.range(200).map { i =>
+      if (i.getLong(0) < 150) Row(None)
+      else Row("a")
+    }
+    val df = sqlContext.createDataFrame(data, StructType(StructField("col", StringType) :: Nil))
+    assert(df.agg("col" -> "count").collect().head.getLong(0) == 50)
+
+    withTempPath { dir =>
+      val path = s"${dir.getCanonicalPath}/data"
+      df.write.parquet(path)
+
+      val df2 = sqlContext.read.parquet(path)
+      assert(df2.agg("col" -> "count").collect().head.getLong(0) == 50)
+    }
+  }
+
   test("read dictionary encoded decimals written as INT32") {
     checkAnswer(
       // Decimal column in this file is encoded using plain dictionary

From 968acf3bd9a502fcad15df3e53e359695ae702cc Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Fri, 20 Nov 2015 15:36:30 -0800
Subject: [PATCH 395/518] [SPARK-11889][SQL] Fix type inference for
 GroupedDataset.agg in REPL

In this PR I delete a method that breaks type inference for aggregators (only in the REPL)

The error when this method is present is:
```
<console>:38: error: missing parameter type for expanded function ((x$2) => x$2._2)
              ds.groupBy(_._1).agg(sum(_._2), sum(_._3)).collect()
```

Author: Michael Armbrust <michael@databricks.com>

Closes #9870 from marmbrus/dataset-repl-agg.
---
 .../org/apache/spark/repl/ReplSuite.scala     | 24 +++++++++++++++++
 .../org/apache/spark/sql/GroupedDataset.scala | 27 +++----------------
 .../apache/spark/sql/JavaDatasetSuite.java    |  8 +++---
 3 files changed, 30 insertions(+), 29 deletions(-)

diff --git a/repl/scala-2.10/src/test/scala/org/apache/spark/repl/ReplSuite.scala b/repl/scala-2.10/src/test/scala/org/apache/spark/repl/ReplSuite.scala
index 081aa03002cc6..cbcccb11f14ae 100644
--- a/repl/scala-2.10/src/test/scala/org/apache/spark/repl/ReplSuite.scala
+++ b/repl/scala-2.10/src/test/scala/org/apache/spark/repl/ReplSuite.scala
@@ -339,6 +339,30 @@ class ReplSuite extends SparkFunSuite {
     }
   }
 
+  test("Datasets agg type-inference") {
+    val output = runInterpreter("local",
+      """
+        |import org.apache.spark.sql.functions._
+        |import org.apache.spark.sql.Encoder
+        |import org.apache.spark.sql.expressions.Aggregator
+        |import org.apache.spark.sql.TypedColumn
+        |/** An `Aggregator` that adds up any numeric type returned by the given function. */
+        |class SumOf[I, N : Numeric](f: I => N) extends Aggregator[I, N, N] with Serializable {
+        |  val numeric = implicitly[Numeric[N]]
+        |  override def zero: N = numeric.zero
+        |  override def reduce(b: N, a: I): N = numeric.plus(b, f(a))
+        |  override def merge(b1: N,b2: N): N = numeric.plus(b1, b2)
+        |  override def finish(reduction: N): N = reduction
+        |}
+        |
+        |def sum[I, N : Numeric : Encoder](f: I => N): TypedColumn[I, N] = new SumOf(f).toColumn
+        |val ds = Seq((1, 1, 2L), (1, 2, 3L), (1, 3, 4L), (2, 1, 5L)).toDS()
+        |ds.groupBy(_._1).agg(sum(_._2), sum(_._3)).collect()
+      """.stripMargin)
+    assertDoesNotContain("error:", output)
+    assertDoesNotContain("Exception", output)
+  }
+
   test("collecting objects of class defined in repl") {
     val output = runInterpreter("local[2]",
       """
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 6de3dd626576a..263f049104762 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -146,31 +146,10 @@ class GroupedDataset[K, T] private[sql](
     reduce(f.call _)
   }
 
-  /**
-   * Compute aggregates by specifying a series of aggregate columns, and return a [[DataFrame]].
-   * We can call `as[T : Encoder]` to turn the returned [[DataFrame]] to [[Dataset]] again.
-   *
-   * The available aggregate methods are defined in [[org.apache.spark.sql.functions]].
-   *
-   * {{{
-   *   // Selects the age of the oldest employee and the aggregate expense for each department
-   *
-   *   // Scala:
-   *   import org.apache.spark.sql.functions._
-   *   df.groupBy("department").agg(max("age"), sum("expense"))
-   *
-   *   // Java:
-   *   import static org.apache.spark.sql.functions.*;
-   *   df.groupBy("department").agg(max("age"), sum("expense"));
-   * }}}
-   *
-   * We can also use `Aggregator.toColumn` to pass in typed aggregate functions.
-   *
-   * @since 1.6.0
-   */
+  // This is here to prevent us from adding overloads that would be ambiguous.
   @scala.annotation.varargs
-  def agg(expr: Column, exprs: Column*): DataFrame =
-    groupedData.agg(withEncoder(expr), exprs.map(withEncoder): _*)
+  private def agg(exprs: Column*): DataFrame =
+    groupedData.agg(withEncoder(exprs.head), exprs.tail.map(withEncoder): _*)
 
   private def withEncoder(c: Column): Column = c match {
     case tc: TypedColumn[_, _] =>
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index ce40dd856f679..f7249b8945c49 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -404,11 +404,9 @@ public String call(Tuple2<String, Integer> value) throws Exception {
       grouped.agg(new IntSumOf().toColumn(Encoders.INT(), Encoders.INT()));
     Assert.assertEquals(Arrays.asList(tuple2("a", 3), tuple2("b", 3)), agged.collectAsList());
 
-    Dataset<Tuple4<String, Integer, Long, Long>> agged2 = grouped.agg(
-      new IntSumOf().toColumn(Encoders.INT(), Encoders.INT()),
-      expr("sum(_2)"),
-      count("*"))
-      .as(Encoders.tuple(Encoders.STRING(), Encoders.INT(), Encoders.LONG(), Encoders.LONG()));
+    Dataset<Tuple2<String, Integer>> agged2 = grouped.agg(
+      new IntSumOf().toColumn(Encoders.INT(), Encoders.INT()))
+      .as(Encoders.tuple(Encoders.STRING(), Encoders.INT()));
     Assert.assertEquals(
       Arrays.asList(
         new Tuple4<>("a", 3, 3L, 2L),

From 68ed046836975b492b594967256d3c7951b568a5 Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Fri, 20 Nov 2015 15:38:04 -0800
Subject: [PATCH 396/518] [SPARK-11890][SQL] Fix compilation for Scala 2.11

Author: Michael Armbrust <michael@databricks.com>

Closes #9871 from marmbrus/scala211-break.
---
 .../scala/org/apache/spark/sql/catalyst/ScalaReflection.scala | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 918050b531c02..4a4a62ed1a468 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -670,14 +670,14 @@ trait ScalaReflection {
    * Unlike `schemaFor`, this method won't throw exception for un-supported type, it will return
    * `NullType` silently instead.
    */
-  private def silentSchemaFor(tpe: `Type`): Schema = try {
+  protected def silentSchemaFor(tpe: `Type`): Schema = try {
     schemaFor(tpe)
   } catch {
     case _: UnsupportedOperationException => Schema(NullType, nullable = true)
   }
 
   /** Returns the full class name for a type. */
-  private def getClassNameFromType(tpe: `Type`): String = {
+  protected def getClassNameFromType(tpe: `Type`): String = {
     tpe.erasure.typeSymbol.asClass.fullName
   }
 

From 47815878ad5e47e89bfbd57acb848be2ce67a4a5 Mon Sep 17 00:00:00 2001
From: Michael Armbrust <michael@databricks.com>
Date: Fri, 20 Nov 2015 16:02:03 -0800
Subject: [PATCH 397/518] [HOTFIX] Fix Java Dataset Tests

---
 .../test/java/test/org/apache/spark/sql/JavaDatasetSuite.java | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index f7249b8945c49..f32374b4c04df 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -409,8 +409,8 @@ public String call(Tuple2<String, Integer> value) throws Exception {
       .as(Encoders.tuple(Encoders.STRING(), Encoders.INT()));
     Assert.assertEquals(
       Arrays.asList(
-        new Tuple4<>("a", 3, 3L, 2L),
-        new Tuple4<>("b", 3, 3L, 1L)),
+        new Tuple2<>("a", 3),
+        new Tuple2<>("b", 3)),
       agged2.collectAsList());
   }
 

From a2dce22e0a25922e2052318d32f32877b7c27ec2 Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Fri, 20 Nov 2015 16:51:47 -0800
Subject: [PATCH 398/518] Revert "[SPARK-11689][ML] Add user guide and example
 code for LDA under spark.ml"

This reverts commit e359d5dcf5bd300213054ebeae9fe75c4f7eb9e7.
---
 docs/ml-clustering.md                         | 30 ------
 docs/ml-guide.md                              |  3 +-
 docs/mllib-guide.md                           |  1 -
 .../spark/examples/ml/JavaLDAExample.java     | 94 -------------------
 .../apache/spark/examples/ml/LDAExample.scala | 77 ---------------
 5 files changed, 1 insertion(+), 204 deletions(-)
 delete mode 100644 docs/ml-clustering.md
 delete mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
 delete mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala

diff --git a/docs/ml-clustering.md b/docs/ml-clustering.md
deleted file mode 100644
index 1743ef43a6ddf..0000000000000
--- a/docs/ml-clustering.md
+++ /dev/null
@@ -1,30 +0,0 @@
----
-layout: global
-title: Clustering - ML
-displayTitle: <a href="ml-guide.html">ML</a> - Clustering
----
-
-In this section, we introduce the pipeline API for [clustering in mllib](mllib-clustering.html).
-
-## Latent Dirichlet allocation (LDA)
-
-`LDA` is implemented as an `Estimator` that supports both `EMLDAOptimizer` and `OnlineLDAOptimizer`,
-and generates a `LDAModel` as the base models. Expert users may cast a `LDAModel` generated by
-`EMLDAOptimizer` to a `DistributedLDAModel` if needed.
-
-<div class="codetabs">
-
-Refer to the [Scala API docs](api/scala/index.html#org.apache.spark.ml.clustering.LDA) for more details.
-
-<div data-lang="scala" markdown="1">
-{% include_example scala/org/apache/spark/examples/ml/LDAExample.scala %}
-</div>
-
-<div data-lang="java" markdown="1">
-
-Refer to the [Java API docs](api/java/org/apache/spark/ml/clustering/LDA.html) for more details.
-
-{% include_example java/org/apache/spark/examples/ml/JavaLDAExample.java %}
-</div>
-
-</div>
\ No newline at end of file
diff --git a/docs/ml-guide.md b/docs/ml-guide.md
index 6f35b30c3d4df..be18a05361a17 100644
--- a/docs/ml-guide.md
+++ b/docs/ml-guide.md
@@ -40,7 +40,6 @@ Also, some algorithms have additional capabilities in the `spark.ml` API; e.g.,
 provide class probabilities, and linear models provide model summaries.
 
 * [Feature extraction, transformation, and selection](ml-features.html)
-* [Clustering](ml-clustering.html)
 * [Decision Trees for classification and regression](ml-decision-tree.html)
 * [Ensembles](ml-ensembles.html)
 * [Linear methods with elastic net regularization](ml-linear-methods.html)
@@ -951,4 +950,4 @@ model.transform(test)
 {% endhighlight %}
 </div>
 
-</div>
\ No newline at end of file
+</div>
diff --git a/docs/mllib-guide.md b/docs/mllib-guide.md
index 54e35fcbb15af..91e50ccfecec4 100644
--- a/docs/mllib-guide.md
+++ b/docs/mllib-guide.md
@@ -69,7 +69,6 @@ We list major functionality from both below, with links to detailed guides.
 concepts. It also contains sections on using algorithms within the Pipelines API, for example:
 
 * [Feature extraction, transformation, and selection](ml-features.html)
-* [Clustering](ml-clustering.html)
 * [Decision trees for classification and regression](ml-decision-tree.html)
 * [Ensembles](ml-ensembles.html)
 * [Linear methods with elastic net regularization](ml-linear-methods.html)
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
deleted file mode 100644
index b3a7d2eb29780..0000000000000
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
+++ /dev/null
@@ -1,94 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.examples.ml;
-
-import java.util.regex.Pattern;
-
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.ml.clustering.LDA;
-import org.apache.spark.ml.clustering.LDAModel;
-import org.apache.spark.mllib.linalg.Vector;
-import org.apache.spark.mllib.linalg.VectorUDT;
-import org.apache.spark.mllib.linalg.Vectors;
-import org.apache.spark.sql.DataFrame;
-import org.apache.spark.sql.Row;
-import org.apache.spark.sql.SQLContext;
-import org.apache.spark.sql.catalyst.expressions.GenericRow;
-import org.apache.spark.sql.types.Metadata;
-import org.apache.spark.sql.types.StructField;
-import org.apache.spark.sql.types.StructType;
-
-/**
- * An example demonstrating LDA
- * Run with
- * <pre>
- * bin/run-example ml.JavaLDAExample
- * </pre>
- */
-public class JavaLDAExample {
-
-  private static class ParseVector implements Function<String, Row> {
-    private static final Pattern separator = Pattern.compile(" ");
-
-    @Override
-    public Row call(String line) {
-      String[] tok = separator.split(line);
-      double[] point = new double[tok.length];
-      for (int i = 0; i < tok.length; ++i) {
-        point[i] = Double.parseDouble(tok[i]);
-      }
-      Vector[] points = {Vectors.dense(point)};
-      return new GenericRow(points);
-    }
-  }
-
-  public static void main(String[] args) {
-
-    String inputFile = "data/mllib/sample_lda_data.txt";
-
-    // Parses the arguments
-    SparkConf conf = new SparkConf().setAppName("JavaLDAExample");
-    JavaSparkContext jsc = new JavaSparkContext(conf);
-    SQLContext sqlContext = new SQLContext(jsc);
-
-    // Loads data
-    JavaRDD<Row> points = jsc.textFile(inputFile).map(new ParseVector());
-    StructField[] fields = {new StructField("features", new VectorUDT(), false, Metadata.empty())};
-    StructType schema = new StructType(fields);
-    DataFrame dataset = sqlContext.createDataFrame(points, schema);
-
-    // Trains a LDA model
-    LDA lda = new LDA()
-      .setK(10)
-      .setMaxIter(10);
-    LDAModel model = lda.fit(dataset);
-
-    System.out.println(model.logLikelihood(dataset));
-    System.out.println(model.logPerplexity(dataset));
-
-    // Shows the result
-    DataFrame topics = model.describeTopics(3);
-    topics.show(false);
-    model.transform(dataset).show(false);
-
-    jsc.stop();
-  }
-}
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala
deleted file mode 100644
index 419ce3d87a6ac..0000000000000
--- a/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.examples.ml
-
-// scalastyle:off println
-import org.apache.spark.{SparkContext, SparkConf}
-import org.apache.spark.mllib.linalg.{VectorUDT, Vectors}
-// $example on$
-import org.apache.spark.ml.clustering.LDA
-import org.apache.spark.sql.{Row, SQLContext}
-import org.apache.spark.sql.types.{StructField, StructType}
-// $example off$
-
-/**
- * An example demonstrating a LDA of ML pipeline.
- * Run with
- * {{{
- * bin/run-example ml.LDAExample
- * }}}
- */
-object LDAExample {
-
-  final val FEATURES_COL = "features"
-
-  def main(args: Array[String]): Unit = {
-
-    val input = "data/mllib/sample_lda_data.txt"
-    // Creates a Spark context and a SQL context
-    val conf = new SparkConf().setAppName(s"${this.getClass.getSimpleName}")
-    val sc = new SparkContext(conf)
-    val sqlContext = new SQLContext(sc)
-
-    // $example on$
-    // Loads data
-    val rowRDD = sc.textFile(input).filter(_.nonEmpty)
-      .map(_.split(" ").map(_.toDouble)).map(Vectors.dense).map(Row(_))
-    val schema = StructType(Array(StructField(FEATURES_COL, new VectorUDT, false)))
-    val dataset = sqlContext.createDataFrame(rowRDD, schema)
-
-    // Trains a LDA model
-    val lda = new LDA()
-      .setK(10)
-      .setMaxIter(10)
-      .setFeaturesCol(FEATURES_COL)
-    val model = lda.fit(dataset)
-    val transformed = model.transform(dataset)
-
-    val ll = model.logLikelihood(dataset)
-    val lp = model.logPerplexity(dataset)
-
-    // describeTopics
-    val topics = model.describeTopics(3)
-
-    // Shows the result
-    topics.show(false)
-    transformed.show(false)
-
-    // $example off$
-    sc.stop()
-  }
-}
-// scalastyle:on println

From 7d3f922c4ba76c4193f98234ae662065c39cdfb1 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Fri, 20 Nov 2015 23:31:19 -0800
Subject: [PATCH 399/518] [SPARK-11819][SQL][FOLLOW-UP] fix scala 2.11 build

seems scala 2.11 doesn't support: define private methods in `trait xxx` and use it in `object xxx extend xxx`.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9879 from cloud-fan/follow.
---
 .../scala/org/apache/spark/sql/catalyst/ScalaReflection.scala | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 4a4a62ed1a468..476becec4dd52 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -670,14 +670,14 @@ trait ScalaReflection {
    * Unlike `schemaFor`, this method won't throw exception for un-supported type, it will return
    * `NullType` silently instead.
    */
-  protected def silentSchemaFor(tpe: `Type`): Schema = try {
+  def silentSchemaFor(tpe: `Type`): Schema = try {
     schemaFor(tpe)
   } catch {
     case _: UnsupportedOperationException => Schema(NullType, nullable = true)
   }
 
   /** Returns the full class name for a type. */
-  protected def getClassNameFromType(tpe: `Type`): String = {
+  def getClassNameFromType(tpe: `Type`): String = {
     tpe.erasure.typeSymbol.asClass.fullName
   }
 

From 54328b6d862fe62ae01bdd87df4798ceb9d506d6 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Sat, 21 Nov 2015 00:10:13 -0800
Subject: [PATCH 400/518] [SPARK-11900][SQL] Add since version for all encoders

Author: Reynold Xin <rxin@databricks.com>

Closes #9881 from rxin/SPARK-11900.
---
 .../scala/org/apache/spark/sql/Encoder.scala  | 63 +++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index 86bb536459035..5cb8edf64e87c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -45,13 +45,52 @@ trait Encoder[T] extends Serializable {
  */
 object Encoders {
 
+  /**
+   * An encoder for nullable boolean type.
+   * @since 1.6.0
+   */
   def BOOLEAN: Encoder[java.lang.Boolean] = ExpressionEncoder()
+
+  /**
+   * An encoder for nullable byte type.
+   * @since 1.6.0
+   */
   def BYTE: Encoder[java.lang.Byte] = ExpressionEncoder()
+
+  /**
+   * An encoder for nullable short type.
+   * @since 1.6.0
+   */
   def SHORT: Encoder[java.lang.Short] = ExpressionEncoder()
+
+  /**
+   * An encoder for nullable int type.
+   * @since 1.6.0
+   */
   def INT: Encoder[java.lang.Integer] = ExpressionEncoder()
+
+  /**
+   * An encoder for nullable long type.
+   * @since 1.6.0
+   */
   def LONG: Encoder[java.lang.Long] = ExpressionEncoder()
+
+  /**
+   * An encoder for nullable float type.
+   * @since 1.6.0
+   */
   def FLOAT: Encoder[java.lang.Float] = ExpressionEncoder()
+
+  /**
+   * An encoder for nullable double type.
+   * @since 1.6.0
+   */
   def DOUBLE: Encoder[java.lang.Double] = ExpressionEncoder()
+
+  /**
+   * An encoder for nullable string type.
+   * @since 1.6.0
+   */
   def STRING: Encoder[java.lang.String] = ExpressionEncoder()
 
   /**
@@ -59,6 +98,8 @@ object Encoders {
    * This encoder maps T into a single byte array (binary) field.
    *
    * T must be publicly accessible.
+   *
+   * @since 1.6.0
    */
   def kryo[T: ClassTag]: Encoder[T] = genericSerializer(useKryo = true)
 
@@ -67,6 +108,8 @@ object Encoders {
    * This encoder maps T into a single byte array (binary) field.
    *
    * T must be publicly accessible.
+   *
+   * @since 1.6.0
    */
   def kryo[T](clazz: Class[T]): Encoder[T] = kryo(ClassTag[T](clazz))
 
@@ -77,6 +120,8 @@ object Encoders {
    * Note that this is extremely inefficient and should only be used as the last resort.
    *
    * T must be publicly accessible.
+   *
+   * @since 1.6.0
    */
   def javaSerialization[T: ClassTag]: Encoder[T] = genericSerializer(useKryo = false)
 
@@ -87,6 +132,8 @@ object Encoders {
    * Note that this is extremely inefficient and should only be used as the last resort.
    *
    * T must be publicly accessible.
+   *
+   * @since 1.6.0
    */
   def javaSerialization[T](clazz: Class[T]): Encoder[T] = javaSerialization(ClassTag[T](clazz))
 
@@ -120,12 +167,20 @@ object Encoders {
     )
   }
 
+  /**
+   * An encoder for 2-ary tuples.
+   * @since 1.6.0
+   */
   def tuple[T1, T2](
       e1: Encoder[T1],
       e2: Encoder[T2]): Encoder[(T1, T2)] = {
     ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2))
   }
 
+  /**
+   * An encoder for 3-ary tuples.
+   * @since 1.6.0
+   */
   def tuple[T1, T2, T3](
       e1: Encoder[T1],
       e2: Encoder[T2],
@@ -133,6 +188,10 @@ object Encoders {
     ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2), encoderFor(e3))
   }
 
+  /**
+   * An encoder for 4-ary tuples.
+   * @since 1.6.0
+   */
   def tuple[T1, T2, T3, T4](
       e1: Encoder[T1],
       e2: Encoder[T2],
@@ -141,6 +200,10 @@ object Encoders {
     ExpressionEncoder.tuple(encoderFor(e1), encoderFor(e2), encoderFor(e3), encoderFor(e4))
   }
 
+  /**
+   * An encoder for 5-ary tuples.
+   * @since 1.6.0
+   */
   def tuple[T1, T2, T3, T4, T5](
       e1: Encoder[T1],
       e2: Encoder[T2],

From 596710268e29e8f624c3ba2fade08b66ec7084eb Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Sat, 21 Nov 2015 00:54:18 -0800
Subject: [PATCH 401/518] [SPARK-11901][SQL] API audit for Aggregator.

Author: Reynold Xin <rxin@databricks.com>

Closes #9882 from rxin/SPARK-11901.
---
 .../scala/org/apache/spark/sql/Dataset.scala  |  1 -
 .../spark/sql/expressions/Aggregator.scala    | 39 ++++++++++++-------
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index bdcdc5d47cbae..07647508421a4 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -22,7 +22,6 @@ import scala.collection.JavaConverters._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.rdd.RDD
 import org.apache.spark.api.java.function._
-import org.apache.spark.sql.catalyst.InternalRow
 
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
index 72610e735f782..b0cd32b5f73e6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
@@ -17,11 +17,10 @@
 
 package org.apache.spark.sql.expressions
 
-import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.encoderFor
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete}
 import org.apache.spark.sql.execution.aggregate.TypedAggregateExpression
-import org.apache.spark.sql.{Dataset, DataFrame, TypedColumn}
+import org.apache.spark.sql.{DataFrame, Dataset, Encoder, TypedColumn}
 
 /**
  * A base class for user-defined aggregations, which can be used in [[DataFrame]] and [[Dataset]]
@@ -32,55 +31,65 @@ import org.apache.spark.sql.{Dataset, DataFrame, TypedColumn}
  *   case class Data(i: Int)
  *
  *   val customSummer =  new Aggregator[Data, Int, Int] {
- *     def zero = 0
- *     def reduce(b: Int, a: Data) = b + a.i
- *     def present(r: Int) = r
+ *     def zero: Int = 0
+ *     def reduce(b: Int, a: Data): Int = b + a.i
+ *     def merge(b1: Int, b2: Int): Int = b1 + b2
+ *     def present(r: Int): Int = r
  *   }.toColumn()
  *
- *   val ds: Dataset[Data]
+ *   val ds: Dataset[Data] = ...
  *   val aggregated = ds.select(customSummer)
  * }}}
  *
  * Based loosely on Aggregator from Algebird: https://github.com/twitter/algebird
  *
- * @tparam A The input type for the aggregation.
+ * @tparam I The input type for the aggregation.
  * @tparam B The type of the intermediate value of the reduction.
- * @tparam C The type of the final result.
+ * @tparam O The type of the final output result.
+ *
+ * @since 1.6.0
  */
-abstract class Aggregator[-A, B, C] extends Serializable {
+abstract class Aggregator[-I, B, O] extends Serializable {
 
-  /** A zero value for this aggregation. Should satisfy the property that any b + zero = b */
+  /**
+   * A zero value for this aggregation. Should satisfy the property that any b + zero = b.
+   * @since 1.6.0
+   */
   def zero: B
 
   /**
    * Combine two values to produce a new value.  For performance, the function may modify `b` and
    * return it instead of constructing new object for b.
+   * @since 1.6.0
    */
-  def reduce(b: B, a: A): B
+  def reduce(b: B, a: I): B
 
   /**
-   * Merge two intermediate values
+   * Merge two intermediate values.
+   * @since 1.6.0
    */
   def merge(b1: B, b2: B): B
 
   /**
    * Transform the output of the reduction.
+   * @since 1.6.0
    */
-  def finish(reduction: B): C
+  def finish(reduction: B): O
 
   /**
    * Returns this `Aggregator` as a [[TypedColumn]] that can be used in [[Dataset]] or [[DataFrame]]
    * operations.
+   * @since 1.6.0
    */
   def toColumn(
       implicit bEncoder: Encoder[B],
-      cEncoder: Encoder[C]): TypedColumn[A, C] = {
+      cEncoder: Encoder[O]): TypedColumn[I, O] = {
     val expr =
       new AggregateExpression(
         TypedAggregateExpression(this),
         Complete,
         false)
 
-    new TypedColumn[A, C](expr, encoderFor[C])
+    new TypedColumn[I, O](expr, encoderFor[O])
   }
 }

From ff442bbcffd4f93cfcc2f76d160011e725d2fb3f Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Sat, 21 Nov 2015 15:00:37 -0800
Subject: [PATCH 402/518] [SPARK-11899][SQL] API audit for GroupedDataset.

1. Renamed map to mapGroup, flatMap to flatMapGroup.
2. Renamed asKey -> keyAs.
3. Added more documentation.
4. Changed type parameter T to V on GroupedDataset.
5. Added since versions for all functions.

Author: Reynold Xin <rxin@databricks.com>

Closes #9880 from rxin/SPARK-11899.
---
 .../api/java/function/MapGroupFunction.java   |   2 +-
 .../scala/org/apache/spark/sql/Encoder.scala  |   4 +
 .../sql/catalyst/JavaTypeInference.scala      |   3 +-
 .../scala/org/apache/spark/sql/Column.scala   |   2 +
 .../org/apache/spark/sql/DataFrame.scala      |   1 -
 .../org/apache/spark/sql/GroupedDataset.scala | 132 ++++++++++++++----
 .../apache/spark/sql/JavaDatasetSuite.java    |   8 +-
 .../spark/sql/DatasetPrimitiveSuite.scala     |   4 +-
 .../org/apache/spark/sql/DatasetSuite.scala   |  20 +--
 9 files changed, 131 insertions(+), 45 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java b/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java
index 2935f9986a560..4f3f222e064bb 100644
--- a/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java
+++ b/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java
@@ -21,7 +21,7 @@
 import java.util.Iterator;
 
 /**
- * Base interface for a map function used in GroupedDataset's map function.
+ * Base interface for a map function used in GroupedDataset's mapGroup function.
  */
 public interface MapGroupFunction<K, V, R> extends Serializable {
   R call(K key, Iterator<V> values) throws Exception;
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
index 5cb8edf64e87c..03aa25eda807f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/Encoder.scala
@@ -30,6 +30,8 @@ import org.apache.spark.sql.types._
  *
  * Encoders are not intended to be thread-safe and thus they are allow to avoid internal locking
  * and reuse internal buffers to improve performance.
+ *
+ * @since 1.6.0
  */
 trait Encoder[T] extends Serializable {
 
@@ -42,6 +44,8 @@ trait Encoder[T] extends Serializable {
 
 /**
  * Methods for creating encoders.
+ *
+ * @since 1.6.0
  */
 object Encoders {
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/JavaTypeInference.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/JavaTypeInference.scala
index 88a457f87ce4e..7d4cfbe6faecb 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/JavaTypeInference.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/JavaTypeInference.scala
@@ -29,7 +29,7 @@ import org.apache.spark.sql.types._
 /**
  * Type-inference utilities for POJOs and Java collections.
  */
-private [sql] object JavaTypeInference {
+object JavaTypeInference {
 
   private val iterableType = TypeToken.of(classOf[JIterable[_]])
   private val mapType = TypeToken.of(classOf[JMap[_, _]])
@@ -53,7 +53,6 @@ private [sql] object JavaTypeInference {
    * @return (SQL data type, nullable)
    */
   private def inferDataType(typeToken: TypeToken[_]): (DataType, Boolean) = {
-    // TODO: All of this could probably be moved to Catalyst as it is mostly not Spark specific.
     typeToken.getRawType match {
       case c: Class[_] if c.isAnnotationPresent(classOf[SQLUserDefinedType]) =>
         (c.getAnnotation(classOf[SQLUserDefinedType]).udt().newInstance(), true)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index 82e9cd7f50a31..30c554a85e693 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -46,6 +46,8 @@ private[sql] object Column {
  * @tparam T The input type expected for this expression.  Can be `Any` if the expression is type
  *           checked by the analyzer instead of the compiler (i.e. `expr("sum(...)")`).
  * @tparam U The output type of this column.
+ *
+ * @since 1.6.0
  */
 class TypedColumn[-T, U](
     expr: Expression,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 7abcecaa2880e..5586fc994b98a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -110,7 +110,6 @@ private[sql] object DataFrame {
  * @groupname action Actions
  * @since 1.3.0
  */
-// TODO: Improve documentation.
 @Experimental
 class DataFrame private[sql](
     @transient val sqlContext: SQLContext,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 263f049104762..7f43ce16901b9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, encoderFor, Ou
 import org.apache.spark.sql.catalyst.expressions.{Alias, CreateStruct, Attribute}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.QueryExecution
+import org.apache.spark.sql.expressions.Aggregator
 
 /**
  * :: Experimental ::
@@ -36,11 +37,13 @@ import org.apache.spark.sql.execution.QueryExecution
  * making this change to the class hierarchy would break some function signatures. As such, this
  * class should be considered a preview of the final API.  Changes will be made to the interface
  * after Spark 1.6.
+ *
+ * @since 1.6.0
  */
 @Experimental
-class GroupedDataset[K, T] private[sql](
+class GroupedDataset[K, V] private[sql](
     kEncoder: Encoder[K],
-    tEncoder: Encoder[T],
+    tEncoder: Encoder[V],
     val queryExecution: QueryExecution,
     private val dataAttributes: Seq[Attribute],
     private val groupingAttributes: Seq[Attribute]) extends Serializable {
@@ -67,8 +70,10 @@ class GroupedDataset[K, T] private[sql](
   /**
    * Returns a new [[GroupedDataset]] where the type of the key has been mapped to the specified
    * type. The mapping of key columns to the type follows the same rules as `as` on [[Dataset]].
+   *
+   * @since 1.6.0
    */
-  def asKey[L : Encoder]: GroupedDataset[L, T] =
+  def keyAs[L : Encoder]: GroupedDataset[L, V] =
     new GroupedDataset(
       encoderFor[L],
       unresolvedTEncoder,
@@ -78,6 +83,8 @@ class GroupedDataset[K, T] private[sql](
 
   /**
    * Returns a [[Dataset]] that contains each unique key.
+   *
+   * @since 1.6.0
    */
   def keys: Dataset[K] = {
     new Dataset[K](
@@ -92,12 +99,18 @@ class GroupedDataset[K, T] private[sql](
    * function can return an iterator containing elements of an arbitrary type which will be returned
    * as a new [[Dataset]].
    *
+   * This function does not support partial aggregation, and as a result requires shuffling all
+   * the data in the [[Dataset]]. If an application intends to perform an aggregation over each
+   * key, it is best to use the reduce function or an [[Aggregator]].
+   *
    * Internally, the implementation will spill to disk if any given group is too large to fit into
    * memory.  However, users must take care to avoid materializing the whole iterator for a group
    * (for example, by calling `toList`) unless they are sure that this is possible given the memory
    * constraints of their cluster.
+   *
+   * @since 1.6.0
    */
-  def flatMap[U : Encoder](f: (K, Iterator[T]) => TraversableOnce[U]): Dataset[U] = {
+  def flatMapGroup[U : Encoder](f: (K, Iterator[V]) => TraversableOnce[U]): Dataset[U] = {
     new Dataset[U](
       sqlContext,
       MapGroups(
@@ -108,8 +121,25 @@ class GroupedDataset[K, T] private[sql](
         logicalPlan))
   }
 
-  def flatMap[U](f: FlatMapGroupFunction[K, T, U], encoder: Encoder[U]): Dataset[U] = {
-    flatMap((key, data) => f.call(key, data.asJava).asScala)(encoder)
+  /**
+   * Applies the given function to each group of data.  For each unique group, the function will
+   * be passed the group key and an iterator that contains all of the elements in the group. The
+   * function can return an iterator containing elements of an arbitrary type which will be returned
+   * as a new [[Dataset]].
+   *
+   * This function does not support partial aggregation, and as a result requires shuffling all
+   * the data in the [[Dataset]]. If an application intends to perform an aggregation over each
+   * key, it is best to use the reduce function or an [[Aggregator]].
+   *
+   * Internally, the implementation will spill to disk if any given group is too large to fit into
+   * memory.  However, users must take care to avoid materializing the whole iterator for a group
+   * (for example, by calling `toList`) unless they are sure that this is possible given the memory
+   * constraints of their cluster.
+   *
+   * @since 1.6.0
+   */
+  def flatMapGroup[U](f: FlatMapGroupFunction[K, V, U], encoder: Encoder[U]): Dataset[U] = {
+    flatMapGroup((key, data) => f.call(key, data.asJava).asScala)(encoder)
   }
 
   /**
@@ -117,32 +147,62 @@ class GroupedDataset[K, T] private[sql](
    * be passed the group key and an iterator that contains all of the elements in the group. The
    * function can return an element of arbitrary type which will be returned as a new [[Dataset]].
    *
+   * This function does not support partial aggregation, and as a result requires shuffling all
+   * the data in the [[Dataset]]. If an application intends to perform an aggregation over each
+   * key, it is best to use the reduce function or an [[Aggregator]].
+   *
    * Internally, the implementation will spill to disk if any given group is too large to fit into
    * memory.  However, users must take care to avoid materializing the whole iterator for a group
    * (for example, by calling `toList`) unless they are sure that this is possible given the memory
    * constraints of their cluster.
+   *
+   * @since 1.6.0
    */
-  def map[U : Encoder](f: (K, Iterator[T]) => U): Dataset[U] = {
-    val func = (key: K, it: Iterator[T]) => Iterator(f(key, it))
-    flatMap(func)
+  def mapGroup[U : Encoder](f: (K, Iterator[V]) => U): Dataset[U] = {
+    val func = (key: K, it: Iterator[V]) => Iterator(f(key, it))
+    flatMapGroup(func)
   }
 
-  def map[U](f: MapGroupFunction[K, T, U], encoder: Encoder[U]): Dataset[U] = {
-    map((key, data) => f.call(key, data.asJava))(encoder)
+  /**
+   * Applies the given function to each group of data.  For each unique group, the function will
+   * be passed the group key and an iterator that contains all of the elements in the group. The
+   * function can return an element of arbitrary type which will be returned as a new [[Dataset]].
+   *
+   * This function does not support partial aggregation, and as a result requires shuffling all
+   * the data in the [[Dataset]]. If an application intends to perform an aggregation over each
+   * key, it is best to use the reduce function or an [[Aggregator]].
+   *
+   * Internally, the implementation will spill to disk if any given group is too large to fit into
+   * memory.  However, users must take care to avoid materializing the whole iterator for a group
+   * (for example, by calling `toList`) unless they are sure that this is possible given the memory
+   * constraints of their cluster.
+   *
+   * @since 1.6.0
+   */
+  def mapGroup[U](f: MapGroupFunction[K, V, U], encoder: Encoder[U]): Dataset[U] = {
+    mapGroup((key, data) => f.call(key, data.asJava))(encoder)
   }
 
   /**
    * Reduces the elements of each group of data using the specified binary function.
    * The given function must be commutative and associative or the result may be non-deterministic.
+   *
+   * @since 1.6.0
    */
-  def reduce(f: (T, T) => T): Dataset[(K, T)] = {
-    val func = (key: K, it: Iterator[T]) => Iterator(key -> it.reduce(f))
+  def reduce(f: (V, V) => V): Dataset[(K, V)] = {
+    val func = (key: K, it: Iterator[V]) => Iterator((key, it.reduce(f)))
 
     implicit val resultEncoder = ExpressionEncoder.tuple(unresolvedKEncoder, unresolvedTEncoder)
-    flatMap(func)
+    flatMapGroup(func)
   }
 
-  def reduce(f: ReduceFunction[T]): Dataset[(K, T)] = {
+  /**
+   * Reduces the elements of each group of data using the specified binary function.
+   * The given function must be commutative and associative or the result may be non-deterministic.
+   *
+   * @since 1.6.0
+   */
+  def reduce(f: ReduceFunction[V]): Dataset[(K, V)] = {
     reduce(f.call _)
   }
 
@@ -185,41 +245,51 @@ class GroupedDataset[K, T] private[sql](
   /**
    * Computes the given aggregation, returning a [[Dataset]] of tuples for each unique key
    * and the result of computing this aggregation over all elements in the group.
+   *
+   * @since 1.6.0
    */
-  def agg[U1](col1: TypedColumn[T, U1]): Dataset[(K, U1)] =
+  def agg[U1](col1: TypedColumn[V, U1]): Dataset[(K, U1)] =
     aggUntyped(col1).asInstanceOf[Dataset[(K, U1)]]
 
   /**
    * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
    * and the result of computing these aggregations over all elements in the group.
+   *
+   * @since 1.6.0
    */
-  def agg[U1, U2](col1: TypedColumn[T, U1], col2: TypedColumn[T, U2]): Dataset[(K, U1, U2)] =
+  def agg[U1, U2](col1: TypedColumn[V, U1], col2: TypedColumn[V, U2]): Dataset[(K, U1, U2)] =
     aggUntyped(col1, col2).asInstanceOf[Dataset[(K, U1, U2)]]
 
   /**
    * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
    * and the result of computing these aggregations over all elements in the group.
+   *
+   * @since 1.6.0
    */
   def agg[U1, U2, U3](
-      col1: TypedColumn[T, U1],
-      col2: TypedColumn[T, U2],
-      col3: TypedColumn[T, U3]): Dataset[(K, U1, U2, U3)] =
+      col1: TypedColumn[V, U1],
+      col2: TypedColumn[V, U2],
+      col3: TypedColumn[V, U3]): Dataset[(K, U1, U2, U3)] =
     aggUntyped(col1, col2, col3).asInstanceOf[Dataset[(K, U1, U2, U3)]]
 
   /**
    * Computes the given aggregations, returning a [[Dataset]] of tuples for each unique key
    * and the result of computing these aggregations over all elements in the group.
+   *
+   * @since 1.6.0
    */
   def agg[U1, U2, U3, U4](
-      col1: TypedColumn[T, U1],
-      col2: TypedColumn[T, U2],
-      col3: TypedColumn[T, U3],
-      col4: TypedColumn[T, U4]): Dataset[(K, U1, U2, U3, U4)] =
+      col1: TypedColumn[V, U1],
+      col2: TypedColumn[V, U2],
+      col3: TypedColumn[V, U3],
+      col4: TypedColumn[V, U4]): Dataset[(K, U1, U2, U3, U4)] =
     aggUntyped(col1, col2, col3, col4).asInstanceOf[Dataset[(K, U1, U2, U3, U4)]]
 
   /**
    * Returns a [[Dataset]] that contains a tuple with each key and the number of items present
    * for that key.
+   *
+   * @since 1.6.0
    */
   def count(): Dataset[(K, Long)] = agg(functions.count("*").as(ExpressionEncoder[Long]))
 
@@ -228,10 +298,12 @@ class GroupedDataset[K, T] private[sql](
    * be passed the grouping key and 2 iterators containing all elements in the group from
    * [[Dataset]] `this` and `other`.  The function can return an iterator containing elements of an
    * arbitrary type which will be returned as a new [[Dataset]].
+   *
+   * @since 1.6.0
    */
   def cogroup[U, R : Encoder](
       other: GroupedDataset[K, U])(
-      f: (K, Iterator[T], Iterator[U]) => TraversableOnce[R]): Dataset[R] = {
+      f: (K, Iterator[V], Iterator[U]) => TraversableOnce[R]): Dataset[R] = {
     implicit def uEnc: Encoder[U] = other.unresolvedTEncoder
     new Dataset[R](
       sqlContext,
@@ -243,9 +315,17 @@ class GroupedDataset[K, T] private[sql](
         other.logicalPlan))
   }
 
+  /**
+   * Applies the given function to each cogrouped data.  For each unique group, the function will
+   * be passed the grouping key and 2 iterators containing all elements in the group from
+   * [[Dataset]] `this` and `other`.  The function can return an iterator containing elements of an
+   * arbitrary type which will be returned as a new [[Dataset]].
+   *
+   * @since 1.6.0
+   */
   def cogroup[U, R](
       other: GroupedDataset[K, U],
-      f: CoGroupFunction[K, T, U, R],
+      f: CoGroupFunction[K, V, U, R],
       encoder: Encoder[R]): Dataset[R] = {
     cogroup(other)((key, left, right) => f.call(key, left.asJava, right.asJava).asScala)(encoder)
   }
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index f32374b4c04df..cf335efdd23b8 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -170,7 +170,7 @@ public Integer call(String v) throws Exception {
       }
     }, Encoders.INT());
 
-    Dataset<String> mapped = grouped.map(new MapGroupFunction<Integer, String, String>() {
+    Dataset<String> mapped = grouped.mapGroup(new MapGroupFunction<Integer, String, String>() {
       @Override
       public String call(Integer key, Iterator<String> values) throws Exception {
         StringBuilder sb = new StringBuilder(key.toString());
@@ -183,7 +183,7 @@ public String call(Integer key, Iterator<String> values) throws Exception {
 
     Assert.assertEquals(Arrays.asList("1a", "3foobar"), mapped.collectAsList());
 
-    Dataset<String> flatMapped = grouped.flatMap(
+    Dataset<String> flatMapped = grouped.flatMapGroup(
       new FlatMapGroupFunction<Integer, String, String>() {
         @Override
         public Iterable<String> call(Integer key, Iterator<String> values) throws Exception {
@@ -247,9 +247,9 @@ public void testGroupByColumn() {
     List<String> data = Arrays.asList("a", "foo", "bar");
     Dataset<String> ds = context.createDataset(data, Encoders.STRING());
     GroupedDataset<Integer, String> grouped =
-      ds.groupBy(length(col("value"))).asKey(Encoders.INT());
+      ds.groupBy(length(col("value"))).keyAs(Encoders.INT());
 
-    Dataset<String> mapped = grouped.map(
+    Dataset<String> mapped = grouped.mapGroup(
       new MapGroupFunction<Integer, String, String>() {
         @Override
         public String call(Integer key, Iterator<String> data) throws Exception {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
index 63b00975e4eb1..d387710357be0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
@@ -86,7 +86,7 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
   test("groupBy function, map") {
     val ds = Seq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11).toDS()
     val grouped = ds.groupBy(_ % 2)
-    val agged = grouped.map { case (g, iter) =>
+    val agged = grouped.mapGroup { case (g, iter) =>
       val name = if (g == 0) "even" else "odd"
       (name, iter.size)
     }
@@ -99,7 +99,7 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
   test("groupBy function, flatMap") {
     val ds = Seq("a", "b", "c", "xyz", "hello").toDS()
     val grouped = ds.groupBy(_.length)
-    val agged = grouped.flatMap { case (g, iter) => Iterator(g.toString, iter.mkString) }
+    val agged = grouped.flatMapGroup { case (g, iter) => Iterator(g.toString, iter.mkString) }
 
     checkAnswer(
       agged,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 89d964aa3e469..9da02550b39ce 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -224,7 +224,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy function, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy(v => (v._1, "word"))
-    val agged = grouped.map { case (g, iter) => (g._1, iter.map(_._2).sum) }
+    val agged = grouped.mapGroup { case (g, iter) => (g._1, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
@@ -234,7 +234,9 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy function, flatMap") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy(v => (v._1, "word"))
-    val agged = grouped.flatMap { case (g, iter) => Iterator(g._1, iter.map(_._2).sum.toString) }
+    val agged = grouped.flatMapGroup { case (g, iter) =>
+      Iterator(g._1, iter.map(_._2).sum.toString)
+    }
 
     checkAnswer(
       agged,
@@ -253,7 +255,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy columns, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1")
-    val agged = grouped.map { case (g, iter) => (g.getString(0), iter.map(_._2).sum) }
+    val agged = grouped.mapGroup { case (g, iter) => (g.getString(0), iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
@@ -262,8 +264,8 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
 
   test("groupBy columns asKey, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
-    val grouped = ds.groupBy($"_1").asKey[String]
-    val agged = grouped.map { case (g, iter) => (g, iter.map(_._2).sum) }
+    val grouped = ds.groupBy($"_1").keyAs[String]
+    val agged = grouped.mapGroup { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
@@ -272,8 +274,8 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
 
   test("groupBy columns asKey tuple, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
-    val grouped = ds.groupBy($"_1", lit(1)).asKey[(String, Int)]
-    val agged = grouped.map { case (g, iter) => (g, iter.map(_._2).sum) }
+    val grouped = ds.groupBy($"_1", lit(1)).keyAs[(String, Int)]
+    val agged = grouped.mapGroup { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
@@ -282,8 +284,8 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
 
   test("groupBy columns asKey class, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
-    val grouped = ds.groupBy($"_1".as("a"), lit(1).as("b")).asKey[ClassData]
-    val agged = grouped.map { case (g, iter) => (g, iter.map(_._2).sum) }
+    val grouped = ds.groupBy($"_1".as("a"), lit(1).as("b")).keyAs[ClassData]
+    val agged = grouped.mapGroup { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,

From 426004a9c9a864f90494d08601e6974709091a56 Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Sun, 22 Nov 2015 10:36:47 -0800
Subject: [PATCH 403/518] [SPARK-11908][SQL] Add NullType support to RowEncoder

JIRA: https://issues.apache.org/jira/browse/SPARK-11908

We should add NullType support to RowEncoder.

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #9891 from viirya/rowencoder-nulltype.
---
 .../org/apache/spark/sql/catalyst/encoders/RowEncoder.scala  | 5 +++--
 .../org/apache/spark/sql/catalyst/expressions/objects.scala  | 3 +++
 .../apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala | 3 +++
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
index 4cda4824acdc3..fa553e7c5324c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
@@ -48,7 +48,7 @@ object RowEncoder {
   private def extractorsFor(
       inputObject: Expression,
       inputType: DataType): Expression = inputType match {
-    case BooleanType | ByteType | ShortType | IntegerType | LongType |
+    case NullType | BooleanType | ByteType | ShortType | IntegerType | LongType |
          FloatType | DoubleType | BinaryType => inputObject
 
     case udt: UserDefinedType[_] =>
@@ -143,6 +143,7 @@ object RowEncoder {
     case _: MapType => ObjectType(classOf[scala.collection.Map[_, _]])
     case _: StructType => ObjectType(classOf[Row])
     case udt: UserDefinedType[_] => ObjectType(udt.userClass)
+    case _: NullType => ObjectType(classOf[java.lang.Object])
   }
 
   private def constructorFor(schema: StructType): Expression = {
@@ -158,7 +159,7 @@ object RowEncoder {
   }
 
   private def constructorFor(input: Expression): Expression = input.dataType match {
-    case BooleanType | ByteType | ShortType | IntegerType | LongType |
+    case NullType | BooleanType | ByteType | ShortType | IntegerType | LongType |
          FloatType | DoubleType | BinaryType => input
 
     case udt: UserDefinedType[_] =>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index ef7399e0196ab..82317d3385167 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -369,6 +369,9 @@ case class MapObjects(
   private lazy val completeFunction = function(loopAttribute)
 
   private def itemAccessorMethod(dataType: DataType): String => String = dataType match {
+    case NullType =>
+      val nullTypeClassName = NullType.getClass.getName + ".MODULE$"
+      (i: String) => s".get($i, $nullTypeClassName)"
     case IntegerType => (i: String) => s".getInt($i)"
     case LongType => (i: String) => s".getLong($i)"
     case FloatType => (i: String) => s".getFloat($i)"
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
index 46c6e0d98d349..0ea51ece4bc5e 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala
@@ -80,11 +80,13 @@ class RowEncoderSuite extends SparkFunSuite {
   private val structOfString = new StructType().add("str", StringType)
   private val structOfUDT = new StructType().add("udt", new ExamplePointUDT, false)
   private val arrayOfString = ArrayType(StringType)
+  private val arrayOfNull = ArrayType(NullType)
   private val mapOfString = MapType(StringType, StringType)
   private val arrayOfUDT = ArrayType(new ExamplePointUDT, false)
 
   encodeDecodeTest(
     new StructType()
+      .add("null", NullType)
       .add("boolean", BooleanType)
       .add("byte", ByteType)
       .add("short", ShortType)
@@ -101,6 +103,7 @@ class RowEncoderSuite extends SparkFunSuite {
 
   encodeDecodeTest(
     new StructType()
+      .add("arrayOfNull", arrayOfNull)
       .add("arrayOfString", arrayOfString)
       .add("arrayOfArrayOfString", ArrayType(arrayOfString))
       .add("arrayOfArrayOfInt", ArrayType(ArrayType(IntegerType)))

From fe89c1817d668e46adf70d0896c42c22a547c76a Mon Sep 17 00:00:00 2001
From: Xiangrui Meng <meng@databricks.com>
Date: Sun, 22 Nov 2015 21:45:46 -0800
Subject: [PATCH 404/518] [SPARK-11895][ML] rename and refactor DatasetExample
 under mllib/examples

We used the name `Dataset` to refer to `SchemaRDD` in 1.2 in ML pipelines and created this example file. Since `Dataset` has a new meaning in Spark 1.6, we should rename it to avoid confusion. This PR also removes support for dense format to simplify the example code.

cc: yinxusen

Author: Xiangrui Meng <meng@databricks.com>

Closes #9873 from mengxr/SPARK-11895.
---
 .../DataFrameExample.scala}                   | 71 +++++++------------
 1 file changed, 26 insertions(+), 45 deletions(-)
 rename examples/src/main/scala/org/apache/spark/examples/{mllib/DatasetExample.scala => ml/DataFrameExample.scala} (51%)

diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/DatasetExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/DataFrameExample.scala
similarity index 51%
rename from examples/src/main/scala/org/apache/spark/examples/mllib/DatasetExample.scala
rename to examples/src/main/scala/org/apache/spark/examples/ml/DataFrameExample.scala
index dc13f82488af7..424f00158c2f2 100644
--- a/examples/src/main/scala/org/apache/spark/examples/mllib/DatasetExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/DataFrameExample.scala
@@ -16,7 +16,7 @@
  */
 
 // scalastyle:off println
-package org.apache.spark.examples.mllib
+package org.apache.spark.examples.ml
 
 import java.io.File
 
@@ -24,25 +24,22 @@ import com.google.common.io.Files
 import scopt.OptionParser
 
 import org.apache.spark.{SparkConf, SparkContext}
+import org.apache.spark.examples.mllib.AbstractParams
 import org.apache.spark.mllib.linalg.Vector
-import org.apache.spark.mllib.regression.LabeledPoint
 import org.apache.spark.mllib.stat.MultivariateOnlineSummarizer
-import org.apache.spark.mllib.util.MLUtils
-import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.{Row, SQLContext, DataFrame}
+import org.apache.spark.sql.{DataFrame, Row, SQLContext}
 
 /**
- * An example of how to use [[org.apache.spark.sql.DataFrame]] as a Dataset for ML. Run with
+ * An example of how to use [[org.apache.spark.sql.DataFrame]] for ML. Run with
  * {{{
- * ./bin/run-example org.apache.spark.examples.mllib.DatasetExample [options]
+ * ./bin/run-example ml.DataFrameExample [options]
  * }}}
  * If you use it as a template to create your own app, please use `spark-submit` to submit your app.
  */
-object DatasetExample {
+object DataFrameExample {
 
-  case class Params(
-      input: String = "data/mllib/sample_libsvm_data.txt",
-      dataFormat: String = "libsvm") extends AbstractParams[Params]
+  case class Params(input: String = "data/mllib/sample_libsvm_data.txt")
+    extends AbstractParams[Params]
 
   def main(args: Array[String]) {
     val defaultParams = Params()
@@ -52,9 +49,6 @@ object DatasetExample {
       opt[String]("input")
         .text(s"input path to dataset")
         .action((x, c) => c.copy(input = x))
-      opt[String]("dataFormat")
-        .text("data format: libsvm (default), dense (deprecated in Spark v1.1)")
-        .action((x, c) => c.copy(input = x))
       checkConfig { params =>
         success
       }
@@ -69,55 +63,42 @@ object DatasetExample {
 
   def run(params: Params) {
 
-    val conf = new SparkConf().setAppName(s"DatasetExample with $params")
+    val conf = new SparkConf().setAppName(s"DataFrameExample with $params")
     val sc = new SparkContext(conf)
     val sqlContext = new SQLContext(sc)
-    import sqlContext.implicits._  // for implicit conversions
 
     // Load input data
-    val origData: RDD[LabeledPoint] = params.dataFormat match {
-      case "dense" => MLUtils.loadLabeledPoints(sc, params.input)
-      case "libsvm" => MLUtils.loadLibSVMFile(sc, params.input)
-    }
-    println(s"Loaded ${origData.count()} instances from file: ${params.input}")
-
-    // Convert input data to DataFrame explicitly.
-    val df: DataFrame = origData.toDF()
-    println(s"Inferred schema:\n${df.schema.prettyJson}")
-    println(s"Converted to DataFrame with ${df.count()} records")
-
-    // Select columns
-    val labelsDf: DataFrame = df.select("label")
-    val labels: RDD[Double] = labelsDf.map { case Row(v: Double) => v }
-    val numLabels = labels.count()
-    val meanLabel = labels.fold(0.0)(_ + _) / numLabels
-    println(s"Selected label column with average value $meanLabel")
-
-    val featuresDf: DataFrame = df.select("features")
-    val features: RDD[Vector] = featuresDf.map { case Row(v: Vector) => v }
+    println(s"Loading LIBSVM file with UDT from ${params.input}.")
+    val df: DataFrame = sqlContext.read.format("libsvm").load(params.input).cache()
+    println("Schema from LIBSVM:")
+    df.printSchema()
+    println(s"Loaded training data as a DataFrame with ${df.count()} records.")
+
+    // Show statistical summary of labels.
+    val labelSummary = df.describe("label")
+    labelSummary.show()
+
+    // Convert features column to an RDD of vectors.
+    val features = df.select("features").map { case Row(v: Vector) => v }
     val featureSummary = features.aggregate(new MultivariateOnlineSummarizer())(
       (summary, feat) => summary.add(feat),
       (sum1, sum2) => sum1.merge(sum2))
     println(s"Selected features column with average values:\n ${featureSummary.mean.toString}")
 
+    // Save the records in a parquet file.
     val tmpDir = Files.createTempDir()
     tmpDir.deleteOnExit()
     val outputDir = new File(tmpDir, "dataset").toString
     println(s"Saving to $outputDir as Parquet file.")
     df.write.parquet(outputDir)
 
+    // Load the records back.
     println(s"Loading Parquet file with UDT from $outputDir.")
-    val newDataset = sqlContext.read.parquet(outputDir)
-
-    println(s"Schema from Parquet: ${newDataset.schema.prettyJson}")
-    val newFeatures = newDataset.select("features").map { case Row(v: Vector) => v }
-    val newFeaturesSummary = newFeatures.aggregate(new MultivariateOnlineSummarizer())(
-      (summary, feat) => summary.add(feat),
-      (sum1, sum2) => sum1.merge(sum2))
-    println(s"Selected features column with average values:\n ${newFeaturesSummary.mean.toString}")
+    val newDF = sqlContext.read.parquet(outputDir)
+    println(s"Schema from Parquet:")
+    newDF.printSchema()
 
     sc.stop()
   }
-
 }
 // scalastyle:on println

From a6fda0bfc16a13b28b1cecc96f1ff91363089144 Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Sun, 22 Nov 2015 21:48:48 -0800
Subject: [PATCH 405/518] [SPARK-6791][ML] Add read/write for CrossValidator
 and Evaluators

I believe this works for general estimators within CrossValidator, including compound estimators.  (See the complex unit test.)

Added read/write for all 3 Evaluators as well.

CC: mengxr yanboliang

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9848 from jkbradley/cv-io.
---
 .../scala/org/apache/spark/ml/Pipeline.scala  |  38 +--
 .../BinaryClassificationEvaluator.scala       |  11 +-
 .../MulticlassClassificationEvaluator.scala   |  12 +-
 .../ml/evaluation/RegressionEvaluator.scala   |  11 +-
 .../apache/spark/ml/recommendation/ALS.scala  |  14 +-
 .../spark/ml/tuning/CrossValidator.scala      | 229 +++++++++++++++++-
 .../org/apache/spark/ml/util/ReadWrite.scala  |  48 ++--
 .../org/apache/spark/ml/PipelineSuite.scala   |   4 +-
 .../BinaryClassificationEvaluatorSuite.scala  |  13 +-
 ...lticlassClassificationEvaluatorSuite.scala |  13 +-
 .../evaluation/RegressionEvaluatorSuite.scala |  12 +-
 .../spark/ml/tuning/CrossValidatorSuite.scala | 202 ++++++++++++++-
 12 files changed, 522 insertions(+), 85 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
index 6f15b37abcb30..4b2b3f8489fd0 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
@@ -34,7 +34,6 @@ import org.apache.spark.ml.util.MLWriter
 import org.apache.spark.ml.util._
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.types.StructType
-import org.apache.spark.util.Utils
 
 /**
  * :: DeveloperApi ::
@@ -232,20 +231,9 @@ object Pipeline extends MLReadable[Pipeline] {
         stages: Array[PipelineStage],
         sc: SparkContext,
         path: String): Unit = {
-      // Copied and edited from DefaultParamsWriter.saveMetadata
-      // TODO: modify DefaultParamsWriter.saveMetadata to avoid duplication
-      val uid = instance.uid
-      val cls = instance.getClass.getName
       val stageUids = stages.map(_.uid)
       val jsonParams = List("stageUids" -> parse(compact(render(stageUids.toSeq))))
-      val metadata = ("class" -> cls) ~
-        ("timestamp" -> System.currentTimeMillis()) ~
-        ("sparkVersion" -> sc.version) ~
-        ("uid" -> uid) ~
-        ("paramMap" -> jsonParams)
-      val metadataPath = new Path(path, "metadata").toString
-      val metadataJson = compact(render(metadata))
-      sc.parallelize(Seq(metadataJson), 1).saveAsTextFile(metadataPath)
+      DefaultParamsWriter.saveMetadata(instance, path, sc, paramMap = Some(jsonParams))
 
       // Save stages
       val stagesDir = new Path(path, "stages").toString
@@ -266,30 +254,10 @@ object Pipeline extends MLReadable[Pipeline] {
 
       implicit val format = DefaultFormats
       val stagesDir = new Path(path, "stages").toString
-      val stageUids: Array[String] = metadata.params match {
-        case JObject(pairs) =>
-          if (pairs.length != 1) {
-            // Should not happen unless file is corrupted or we have a bug.
-            throw new RuntimeException(
-              s"Pipeline read expected 1 Param (stageUids), but found ${pairs.length}.")
-          }
-          pairs.head match {
-            case ("stageUids", jsonValue) =>
-              jsonValue.extract[Seq[String]].toArray
-            case (paramName, jsonValue) =>
-              // Should not happen unless file is corrupted or we have a bug.
-              throw new RuntimeException(s"Pipeline read encountered unexpected Param $paramName" +
-                s" in metadata: ${metadata.metadataStr}")
-          }
-        case _ =>
-          throw new IllegalArgumentException(
-            s"Cannot recognize JSON metadata: ${metadata.metadataStr}.")
-      }
+      val stageUids: Array[String] = (metadata.params \ "stageUids").extract[Seq[String]].toArray
       val stages: Array[PipelineStage] = stageUids.zipWithIndex.map { case (stageUid, idx) =>
         val stagePath = SharedReadWrite.getStagePath(stageUid, idx, stageUids.length, stagesDir)
-        val stageMetadata = DefaultParamsReader.loadMetadata(stagePath, sc)
-        val cls = Utils.classForName(stageMetadata.className)
-        cls.getMethod("read").invoke(null).asInstanceOf[MLReader[PipelineStage]].load(stagePath)
+        DefaultParamsReader.loadParamsInstance[PipelineStage](stagePath, sc)
       }
       (metadata.uid, stages)
     }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluator.scala b/mllib/src/main/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluator.scala
index 1fe3abaca81c3..bfb70963b151d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluator.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluator.scala
@@ -20,7 +20,7 @@ package org.apache.spark.ml.evaluation
 import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable, Identifiable, SchemaUtils}
 import org.apache.spark.mllib.evaluation.BinaryClassificationMetrics
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT}
 import org.apache.spark.sql.{DataFrame, Row}
@@ -33,7 +33,7 @@ import org.apache.spark.sql.types.DoubleType
 @Since("1.2.0")
 @Experimental
 class BinaryClassificationEvaluator @Since("1.4.0") (@Since("1.4.0") override val uid: String)
-  extends Evaluator with HasRawPredictionCol with HasLabelCol {
+  extends Evaluator with HasRawPredictionCol with HasLabelCol with DefaultParamsWritable {
 
   @Since("1.2.0")
   def this() = this(Identifiable.randomUID("binEval"))
@@ -105,3 +105,10 @@ class BinaryClassificationEvaluator @Since("1.4.0") (@Since("1.4.0") override va
   @Since("1.4.1")
   override def copy(extra: ParamMap): BinaryClassificationEvaluator = defaultCopy(extra)
 }
+
+@Since("1.6.0")
+object BinaryClassificationEvaluator extends DefaultParamsReadable[BinaryClassificationEvaluator] {
+
+  @Since("1.6.0")
+  override def load(path: String): BinaryClassificationEvaluator = super.load(path)
+}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/evaluation/MulticlassClassificationEvaluator.scala b/mllib/src/main/scala/org/apache/spark/ml/evaluation/MulticlassClassificationEvaluator.scala
index df5f04ca5a8d9..c44db0ec595ea 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/evaluation/MulticlassClassificationEvaluator.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/evaluation/MulticlassClassificationEvaluator.scala
@@ -20,7 +20,7 @@ package org.apache.spark.ml.evaluation
 import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.param.{ParamMap, ParamValidators, Param}
 import org.apache.spark.ml.param.shared.{HasLabelCol, HasPredictionCol}
-import org.apache.spark.ml.util.{SchemaUtils, Identifiable}
+import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable, SchemaUtils, Identifiable}
 import org.apache.spark.mllib.evaluation.MulticlassMetrics
 import org.apache.spark.sql.{Row, DataFrame}
 import org.apache.spark.sql.types.DoubleType
@@ -32,7 +32,7 @@ import org.apache.spark.sql.types.DoubleType
 @Since("1.5.0")
 @Experimental
 class MulticlassClassificationEvaluator @Since("1.5.0") (@Since("1.5.0") override val uid: String)
-  extends Evaluator with HasPredictionCol with HasLabelCol {
+  extends Evaluator with HasPredictionCol with HasLabelCol with DefaultParamsWritable {
 
   @Since("1.5.0")
   def this() = this(Identifiable.randomUID("mcEval"))
@@ -101,3 +101,11 @@ class MulticlassClassificationEvaluator @Since("1.5.0") (@Since("1.5.0") overrid
   @Since("1.5.0")
   override def copy(extra: ParamMap): MulticlassClassificationEvaluator = defaultCopy(extra)
 }
+
+@Since("1.6.0")
+object MulticlassClassificationEvaluator
+  extends DefaultParamsReadable[MulticlassClassificationEvaluator] {
+
+  @Since("1.6.0")
+  override def load(path: String): MulticlassClassificationEvaluator = super.load(path)
+}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/evaluation/RegressionEvaluator.scala b/mllib/src/main/scala/org/apache/spark/ml/evaluation/RegressionEvaluator.scala
index ba012f444d3e0..daaa174a086e0 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/evaluation/RegressionEvaluator.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/evaluation/RegressionEvaluator.scala
@@ -20,7 +20,7 @@ package org.apache.spark.ml.evaluation
 import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml.param.{Param, ParamMap, ParamValidators}
 import org.apache.spark.ml.param.shared.{HasLabelCol, HasPredictionCol}
-import org.apache.spark.ml.util.{Identifiable, SchemaUtils}
+import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable, Identifiable, SchemaUtils}
 import org.apache.spark.mllib.evaluation.RegressionMetrics
 import org.apache.spark.sql.{DataFrame, Row}
 import org.apache.spark.sql.functions._
@@ -33,7 +33,7 @@ import org.apache.spark.sql.types.{DoubleType, FloatType}
 @Since("1.4.0")
 @Experimental
 final class RegressionEvaluator @Since("1.4.0") (@Since("1.4.0") override val uid: String)
-  extends Evaluator with HasPredictionCol with HasLabelCol {
+  extends Evaluator with HasPredictionCol with HasLabelCol with DefaultParamsWritable {
 
   @Since("1.4.0")
   def this() = this(Identifiable.randomUID("regEval"))
@@ -104,3 +104,10 @@ final class RegressionEvaluator @Since("1.4.0") (@Since("1.4.0") override val ui
   @Since("1.5.0")
   override def copy(extra: ParamMap): RegressionEvaluator = defaultCopy(extra)
 }
+
+@Since("1.6.0")
+object RegressionEvaluator extends DefaultParamsReadable[RegressionEvaluator] {
+
+  @Since("1.6.0")
+  override def load(path: String): RegressionEvaluator = super.load(path)
+}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
index 4d35177ad9b0f..b798aa1fab767 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala
@@ -27,9 +27,8 @@ import scala.util.hashing.byteswap64
 
 import com.github.fommil.netlib.BLAS.{getInstance => blas}
 import org.apache.hadoop.fs.{FileSystem, Path}
-import org.json4s.{DefaultFormats, JValue}
+import org.json4s.DefaultFormats
 import org.json4s.JsonDSL._
-import org.json4s.jackson.JsonMethods._
 
 import org.apache.spark.{Logging, Partitioner}
 import org.apache.spark.annotation.{Since, DeveloperApi, Experimental}
@@ -240,7 +239,7 @@ object ALSModel extends MLReadable[ALSModel] {
   private[ALSModel] class ALSModelWriter(instance: ALSModel) extends MLWriter {
 
     override protected def saveImpl(path: String): Unit = {
-      val extraMetadata = render("rank" -> instance.rank)
+      val extraMetadata = "rank" -> instance.rank
       DefaultParamsWriter.saveMetadata(instance, path, sc, Some(extraMetadata))
       val userPath = new Path(path, "userFactors").toString
       instance.userFactors.write.format("parquet").save(userPath)
@@ -257,14 +256,7 @@ object ALSModel extends MLReadable[ALSModel] {
     override def load(path: String): ALSModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
       implicit val format = DefaultFormats
-      val rank: Int = metadata.extraMetadata match {
-        case Some(m: JValue) =>
-          (m \ "rank").extract[Int]
-        case None =>
-          throw new RuntimeException(s"ALSModel loader could not read rank from JSON metadata:" +
-            s" ${metadata.metadataStr}")
-      }
-
+      val rank = (metadata.metadata \ "rank").extract[Int]
       val userPath = new Path(path, "userFactors").toString
       val userFactors = sqlContext.read.format("parquet").load(userPath)
       val itemPath = new Path(path, "itemFactors").toString
diff --git a/mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala b/mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala
index 77d9948ed86b9..83a9048374267 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala
@@ -18,17 +18,24 @@
 package org.apache.spark.ml.tuning
 
 import com.github.fommil.netlib.F2jBLAS
+import org.apache.hadoop.fs.Path
+import org.json4s.{JObject, DefaultFormats}
+import org.json4s.jackson.JsonMethods._
 
-import org.apache.spark.Logging
-import org.apache.spark.annotation.Experimental
+import org.apache.spark.ml.classification.OneVsRestParams
+import org.apache.spark.ml.feature.RFormulaModel
+import org.apache.spark.{SparkContext, Logging}
+import org.apache.spark.annotation.{Experimental, Since}
 import org.apache.spark.ml._
 import org.apache.spark.ml.evaluation.Evaluator
 import org.apache.spark.ml.param._
-import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.ml.util._
+import org.apache.spark.ml.util.DefaultParamsReader.Metadata
 import org.apache.spark.mllib.util.MLUtils
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.types.StructType
 
+
 /**
  * Params for [[CrossValidator]] and [[CrossValidatorModel]].
  */
@@ -53,7 +60,7 @@ private[ml] trait CrossValidatorParams extends ValidatorParams {
  */
 @Experimental
 class CrossValidator(override val uid: String) extends Estimator[CrossValidatorModel]
-  with CrossValidatorParams with Logging {
+  with CrossValidatorParams with MLWritable with Logging {
 
   def this() = this(Identifiable.randomUID("cv"))
 
@@ -131,6 +138,166 @@ class CrossValidator(override val uid: String) extends Estimator[CrossValidatorM
     }
     copied
   }
+
+  // Currently, this only works if all [[Param]]s in [[estimatorParamMaps]] are simple types.
+  // E.g., this may fail if a [[Param]] is an instance of an [[Estimator]].
+  // However, this case should be unusual.
+  @Since("1.6.0")
+  override def write: MLWriter = new CrossValidator.CrossValidatorWriter(this)
+}
+
+@Since("1.6.0")
+object CrossValidator extends MLReadable[CrossValidator] {
+
+  @Since("1.6.0")
+  override def read: MLReader[CrossValidator] = new CrossValidatorReader
+
+  @Since("1.6.0")
+  override def load(path: String): CrossValidator = super.load(path)
+
+  private[CrossValidator] class CrossValidatorWriter(instance: CrossValidator) extends MLWriter {
+
+    SharedReadWrite.validateParams(instance)
+
+    override protected def saveImpl(path: String): Unit =
+      SharedReadWrite.saveImpl(path, instance, sc)
+  }
+
+  private class CrossValidatorReader extends MLReader[CrossValidator] {
+
+    /** Checked against metadata when loading model */
+    private val className = classOf[CrossValidator].getName
+
+    override def load(path: String): CrossValidator = {
+      val (metadata, estimator, evaluator, estimatorParamMaps, numFolds) =
+        SharedReadWrite.load(path, sc, className)
+      new CrossValidator(metadata.uid)
+        .setEstimator(estimator)
+        .setEvaluator(evaluator)
+        .setEstimatorParamMaps(estimatorParamMaps)
+        .setNumFolds(numFolds)
+    }
+  }
+
+  private object CrossValidatorReader {
+    /**
+     * Examine the given estimator (which may be a compound estimator) and extract a mapping
+     * from UIDs to corresponding [[Params]] instances.
+     */
+    def getUidMap(instance: Params): Map[String, Params] = {
+      val uidList = getUidMapImpl(instance)
+      val uidMap = uidList.toMap
+      if (uidList.size != uidMap.size) {
+        throw new RuntimeException("CrossValidator.load found a compound estimator with stages" +
+          s" with duplicate UIDs.  List of UIDs: ${uidList.map(_._1).mkString(", ")}")
+      }
+      uidMap
+    }
+
+    def getUidMapImpl(instance: Params): List[(String, Params)] = {
+      val subStages: Array[Params] = instance match {
+        case p: Pipeline => p.getStages.asInstanceOf[Array[Params]]
+        case pm: PipelineModel => pm.stages.asInstanceOf[Array[Params]]
+        case v: ValidatorParams => Array(v.getEstimator, v.getEvaluator)
+        case ovr: OneVsRestParams =>
+          // TODO: SPARK-11892: This case may require special handling.
+          throw new UnsupportedOperationException("CrossValidator write will fail because it" +
+            " cannot yet handle an estimator containing type: ${ovr.getClass.getName}")
+        case rform: RFormulaModel =>
+          // TODO: SPARK-11891: This case may require special handling.
+          throw new UnsupportedOperationException("CrossValidator write will fail because it" +
+            " cannot yet handle an estimator containing an RFormulaModel")
+        case _: Params => Array()
+      }
+      val subStageMaps = subStages.map(getUidMapImpl).foldLeft(List.empty[(String, Params)])(_ ++ _)
+      List((instance.uid, instance)) ++ subStageMaps
+    }
+  }
+
+  private[tuning] object SharedReadWrite {
+
+    /**
+     * Check that [[CrossValidator.evaluator]] and [[CrossValidator.estimator]] are Writable.
+     * This does not check [[CrossValidator.estimatorParamMaps]].
+     */
+    def validateParams(instance: ValidatorParams): Unit = {
+      def checkElement(elem: Params, name: String): Unit = elem match {
+        case stage: MLWritable => // good
+        case other =>
+          throw new UnsupportedOperationException("CrossValidator write will fail " +
+            s" because it contains $name which does not implement Writable." +
+            s" Non-Writable $name: ${other.uid} of type ${other.getClass}")
+      }
+      checkElement(instance.getEvaluator, "evaluator")
+      checkElement(instance.getEstimator, "estimator")
+      // Check to make sure all Params apply to this estimator.  Throw an error if any do not.
+      // Extraneous Params would cause problems when loading the estimatorParamMaps.
+      val uidToInstance: Map[String, Params] = CrossValidatorReader.getUidMap(instance)
+      instance.getEstimatorParamMaps.foreach { case pMap: ParamMap =>
+        pMap.toSeq.foreach { case ParamPair(p, v) =>
+          require(uidToInstance.contains(p.parent), s"CrossValidator save requires all Params in" +
+            s" estimatorParamMaps to apply to this CrossValidator, its Estimator, or its" +
+            s" Evaluator.  An extraneous Param was found: $p")
+        }
+      }
+    }
+
+    private[tuning] def saveImpl(
+        path: String,
+        instance: CrossValidatorParams,
+        sc: SparkContext,
+        extraMetadata: Option[JObject] = None): Unit = {
+      import org.json4s.JsonDSL._
+
+      val estimatorParamMapsJson = compact(render(
+        instance.getEstimatorParamMaps.map { case paramMap =>
+          paramMap.toSeq.map { case ParamPair(p, v) =>
+            Map("parent" -> p.parent, "name" -> p.name, "value" -> p.jsonEncode(v))
+          }
+        }.toSeq
+      ))
+      val jsonParams = List(
+        "numFolds" -> parse(instance.numFolds.jsonEncode(instance.getNumFolds)),
+        "estimatorParamMaps" -> parse(estimatorParamMapsJson)
+      )
+      DefaultParamsWriter.saveMetadata(instance, path, sc, extraMetadata, Some(jsonParams))
+
+      val evaluatorPath = new Path(path, "evaluator").toString
+      instance.getEvaluator.asInstanceOf[MLWritable].save(evaluatorPath)
+      val estimatorPath = new Path(path, "estimator").toString
+      instance.getEstimator.asInstanceOf[MLWritable].save(estimatorPath)
+    }
+
+    private[tuning] def load[M <: Model[M]](
+        path: String,
+        sc: SparkContext,
+        expectedClassName: String): (Metadata, Estimator[M], Evaluator, Array[ParamMap], Int) = {
+
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, expectedClassName)
+
+      implicit val format = DefaultFormats
+      val evaluatorPath = new Path(path, "evaluator").toString
+      val evaluator = DefaultParamsReader.loadParamsInstance[Evaluator](evaluatorPath, sc)
+      val estimatorPath = new Path(path, "estimator").toString
+      val estimator = DefaultParamsReader.loadParamsInstance[Estimator[M]](estimatorPath, sc)
+
+      val uidToParams = Map(evaluator.uid -> evaluator) ++ CrossValidatorReader.getUidMap(estimator)
+
+      val numFolds = (metadata.params \ "numFolds").extract[Int]
+      val estimatorParamMaps: Array[ParamMap] =
+        (metadata.params \ "estimatorParamMaps").extract[Seq[Seq[Map[String, String]]]].map {
+          pMap =>
+            val paramPairs = pMap.map { case pInfo: Map[String, String] =>
+              val est = uidToParams(pInfo("parent"))
+              val param = est.getParam(pInfo("name"))
+              val value = param.jsonDecode(pInfo("value"))
+              param -> value
+            }
+            ParamMap(paramPairs: _*)
+        }.toArray
+      (metadata, estimator, evaluator, estimatorParamMaps, numFolds)
+    }
+  }
 }
 
 /**
@@ -139,14 +306,14 @@ class CrossValidator(override val uid: String) extends Estimator[CrossValidatorM
  *
  * @param bestModel The best model selected from k-fold cross validation.
  * @param avgMetrics Average cross-validation metrics for each paramMap in
- *                   [[estimatorParamMaps]], in the corresponding order.
+ *                   [[CrossValidator.estimatorParamMaps]], in the corresponding order.
  */
 @Experimental
 class CrossValidatorModel private[ml] (
     override val uid: String,
     val bestModel: Model[_],
     val avgMetrics: Array[Double])
-  extends Model[CrossValidatorModel] with CrossValidatorParams {
+  extends Model[CrossValidatorModel] with CrossValidatorParams with MLWritable {
 
   override def validateParams(): Unit = {
     bestModel.validateParams()
@@ -168,4 +335,54 @@ class CrossValidatorModel private[ml] (
       avgMetrics.clone())
     copyValues(copied, extra).setParent(parent)
   }
+
+  @Since("1.6.0")
+  override def write: MLWriter = new CrossValidatorModel.CrossValidatorModelWriter(this)
+}
+
+@Since("1.6.0")
+object CrossValidatorModel extends MLReadable[CrossValidatorModel] {
+
+  import CrossValidator.SharedReadWrite
+
+  @Since("1.6.0")
+  override def read: MLReader[CrossValidatorModel] = new CrossValidatorModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): CrossValidatorModel = super.load(path)
+
+  private[CrossValidatorModel]
+  class CrossValidatorModelWriter(instance: CrossValidatorModel) extends MLWriter {
+
+    SharedReadWrite.validateParams(instance)
+
+    override protected def saveImpl(path: String): Unit = {
+      import org.json4s.JsonDSL._
+      val extraMetadata = "avgMetrics" -> instance.avgMetrics.toSeq
+      SharedReadWrite.saveImpl(path, instance, sc, Some(extraMetadata))
+      val bestModelPath = new Path(path, "bestModel").toString
+      instance.bestModel.asInstanceOf[MLWritable].save(bestModelPath)
+    }
+  }
+
+  private class CrossValidatorModelReader extends MLReader[CrossValidatorModel] {
+
+    /** Checked against metadata when loading model */
+    private val className = classOf[CrossValidatorModel].getName
+
+    override def load(path: String): CrossValidatorModel = {
+      implicit val format = DefaultFormats
+
+      val (metadata, estimator, evaluator, estimatorParamMaps, numFolds) =
+        SharedReadWrite.load(path, sc, className)
+      val bestModelPath = new Path(path, "bestModel").toString
+      val bestModel = DefaultParamsReader.loadParamsInstance[Model[_]](bestModelPath, sc)
+      val avgMetrics = (metadata.metadata \ "avgMetrics").extract[Seq[Double]].toArray
+      val cv = new CrossValidatorModel(metadata.uid, bestModel, avgMetrics)
+      cv.set(cv.estimator, estimator)
+        .set(cv.evaluator, evaluator)
+        .set(cv.estimatorParamMaps, estimatorParamMaps)
+        .set(cv.numFolds, numFolds)
+    }
+  }
 }
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index ff9322dba122a..8484b1f801066 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -202,25 +202,36 @@ private[ml] object DefaultParamsWriter {
    *  - timestamp
    *  - sparkVersion
    *  - uid
-   *  - paramMap: These must be encodable using [[org.apache.spark.ml.param.Param.jsonEncode()]].
+   *  - paramMap
+   *  - (optionally, extra metadata)
+   * @param extraMetadata  Extra metadata to be saved at same level as uid, paramMap, etc.
+   * @param paramMap  If given, this is saved in the "paramMap" field.
+   *                  Otherwise, all [[org.apache.spark.ml.param.Param]]s are encoded using
+   *                  [[org.apache.spark.ml.param.Param.jsonEncode()]].
    */
   def saveMetadata(
       instance: Params,
       path: String,
       sc: SparkContext,
-      extraMetadata: Option[JValue] = None): Unit = {
+      extraMetadata: Option[JObject] = None,
+      paramMap: Option[JValue] = None): Unit = {
     val uid = instance.uid
     val cls = instance.getClass.getName
     val params = instance.extractParamMap().toSeq.asInstanceOf[Seq[ParamPair[Any]]]
-    val jsonParams = params.map { case ParamPair(p, v) =>
+    val jsonParams = paramMap.getOrElse(render(params.map { case ParamPair(p, v) =>
       p.name -> parse(p.jsonEncode(v))
-    }.toList
-    val metadata = ("class" -> cls) ~
+    }.toList))
+    val basicMetadata = ("class" -> cls) ~
       ("timestamp" -> System.currentTimeMillis()) ~
       ("sparkVersion" -> sc.version) ~
       ("uid" -> uid) ~
-      ("paramMap" -> jsonParams) ~
-      ("extraMetadata" -> extraMetadata)
+      ("paramMap" -> jsonParams)
+    val metadata = extraMetadata match {
+      case Some(jObject) =>
+        basicMetadata ~ jObject
+      case None =>
+        basicMetadata
+    }
     val metadataPath = new Path(path, "metadata").toString
     val metadataJson = compact(render(metadata))
     sc.parallelize(Seq(metadataJson), 1).saveAsTextFile(metadataPath)
@@ -251,8 +262,8 @@ private[ml] object DefaultParamsReader {
   /**
    * All info from metadata file.
    * @param params  paramMap, as a [[JValue]]
-   * @param extraMetadata  Extra metadata saved by [[DefaultParamsWriter.saveMetadata()]]
-   * @param metadataStr  Full metadata file String (for debugging)
+   * @param metadata  All metadata, including the other fields
+   * @param metadataJson  Full metadata file String (for debugging)
    */
   case class Metadata(
       className: String,
@@ -260,8 +271,8 @@ private[ml] object DefaultParamsReader {
       timestamp: Long,
       sparkVersion: String,
       params: JValue,
-      extraMetadata: Option[JValue],
-      metadataStr: String)
+      metadata: JValue,
+      metadataJson: String)
 
   /**
    * Load metadata from file.
@@ -279,13 +290,12 @@ private[ml] object DefaultParamsReader {
     val timestamp = (metadata \ "timestamp").extract[Long]
     val sparkVersion = (metadata \ "sparkVersion").extract[String]
     val params = metadata \ "paramMap"
-    val extraMetadata = (metadata \ "extraMetadata").extract[Option[JValue]]
     if (expectedClassName.nonEmpty) {
       require(className == expectedClassName, s"Error loading metadata: Expected class name" +
         s" $expectedClassName but found class name $className")
     }
 
-    Metadata(className, uid, timestamp, sparkVersion, params, extraMetadata, metadataStr)
+    Metadata(className, uid, timestamp, sparkVersion, params, metadata, metadataStr)
   }
 
   /**
@@ -303,7 +313,17 @@ private[ml] object DefaultParamsReader {
         }
       case _ =>
         throw new IllegalArgumentException(
-          s"Cannot recognize JSON metadata: ${metadata.metadataStr}.")
+          s"Cannot recognize JSON metadata: ${metadata.metadataJson}.")
     }
   }
+
+  /**
+   * Load a [[Params]] instance from the given path, and return it.
+   * This assumes the instance implements [[MLReadable]].
+   */
+  def loadParamsInstance[T](path: String, sc: SparkContext): T = {
+    val metadata = DefaultParamsReader.loadMetadata(path, sc)
+    val cls = Utils.classForName(metadata.className)
+    cls.getMethod("read").invoke(null).asInstanceOf[MLReader[T]].load(path)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
index 12aba6bc6dbeb..8c86767456368 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
@@ -17,11 +17,9 @@
 
 package org.apache.spark.ml
 
-import java.io.File
-
 import scala.collection.JavaConverters._
 
-import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.hadoop.fs.Path
 import org.mockito.Matchers.{any, eq => meq}
 import org.mockito.Mockito.when
 import org.scalatest.mock.MockitoSugar.mock
diff --git a/mllib/src/test/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluatorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluatorSuite.scala
index def869fe66777..a535c1218ecfa 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluatorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluatorSuite.scala
@@ -19,10 +19,21 @@ package org.apache.spark.ml.evaluation
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
+import org.apache.spark.mllib.util.MLlibTestSparkContext
 
-class BinaryClassificationEvaluatorSuite extends SparkFunSuite {
+class BinaryClassificationEvaluatorSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new BinaryClassificationEvaluator)
   }
+
+  test("read/write") {
+    val evaluator = new BinaryClassificationEvaluator()
+      .setRawPredictionCol("myRawPrediction")
+      .setLabelCol("myLabel")
+      .setMetricName("areaUnderPR")
+    testDefaultReadWrite(evaluator)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/evaluation/MulticlassClassificationEvaluatorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/evaluation/MulticlassClassificationEvaluatorSuite.scala
index 6d8412b0b3701..7ee65975d22f7 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/evaluation/MulticlassClassificationEvaluatorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/evaluation/MulticlassClassificationEvaluatorSuite.scala
@@ -19,10 +19,21 @@ package org.apache.spark.ml.evaluation
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
+import org.apache.spark.ml.util.DefaultReadWriteTest
+import org.apache.spark.mllib.util.MLlibTestSparkContext
 
-class MulticlassClassificationEvaluatorSuite extends SparkFunSuite {
+class MulticlassClassificationEvaluatorSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new MulticlassClassificationEvaluator)
   }
+
+  test("read/write") {
+    val evaluator = new MulticlassClassificationEvaluator()
+      .setPredictionCol("myPrediction")
+      .setLabelCol("myLabel")
+      .setMetricName("recall")
+    testDefaultReadWrite(evaluator)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/evaluation/RegressionEvaluatorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/evaluation/RegressionEvaluatorSuite.scala
index aa722da323935..60886bf77d2f0 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/evaluation/RegressionEvaluatorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/evaluation/RegressionEvaluatorSuite.scala
@@ -20,10 +20,12 @@ package org.apache.spark.ml.evaluation
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
 import org.apache.spark.ml.regression.LinearRegression
+import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.{LinearDataGenerator, MLlibTestSparkContext}
 import org.apache.spark.mllib.util.TestingUtils._
 
-class RegressionEvaluatorSuite extends SparkFunSuite with MLlibTestSparkContext {
+class RegressionEvaluatorSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   test("params") {
     ParamsSuite.checkParams(new RegressionEvaluator)
@@ -73,4 +75,12 @@ class RegressionEvaluatorSuite extends SparkFunSuite with MLlibTestSparkContext
     evaluator.setMetricName("mae")
     assert(evaluator.evaluate(predictions) ~== 0.08036075 absTol 0.001)
   }
+
+  test("read/write") {
+    val evaluator = new RegressionEvaluator()
+      .setPredictionCol("myPrediction")
+      .setLabelCol("myLabel")
+      .setMetricName("r2")
+    testDefaultReadWrite(evaluator)
+  }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala
index cbe09292a0337..dd6366050c020 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala
@@ -18,19 +18,22 @@
 package org.apache.spark.ml.tuning
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.ml.util.MLTestingUtils
-import org.apache.spark.ml.{Estimator, Model}
-import org.apache.spark.ml.classification.LogisticRegression
+import org.apache.spark.ml.feature.HashingTF
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
+import org.apache.spark.ml.{Pipeline, Estimator, Model}
+import org.apache.spark.ml.classification.{LogisticRegressionModel, LogisticRegression}
 import org.apache.spark.ml.evaluation.{BinaryClassificationEvaluator, Evaluator, RegressionEvaluator}
-import org.apache.spark.ml.param.ParamMap
+import org.apache.spark.ml.param.{ParamPair, ParamMap}
 import org.apache.spark.ml.param.shared.HasInputCol
 import org.apache.spark.ml.regression.LinearRegression
 import org.apache.spark.mllib.classification.LogisticRegressionSuite.generateLogisticInput
+import org.apache.spark.mllib.linalg.Vectors
 import org.apache.spark.mllib.util.{LinearDataGenerator, MLlibTestSparkContext}
 import org.apache.spark.sql.{DataFrame, SQLContext}
 import org.apache.spark.sql.types.StructType
 
-class CrossValidatorSuite extends SparkFunSuite with MLlibTestSparkContext {
+class CrossValidatorSuite
+  extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   @transient var dataset: DataFrame = _
 
@@ -95,7 +98,7 @@ class CrossValidatorSuite extends SparkFunSuite with MLlibTestSparkContext {
   }
 
   test("validateParams should check estimatorParamMaps") {
-    import CrossValidatorSuite._
+    import CrossValidatorSuite.{MyEstimator, MyEvaluator}
 
     val est = new MyEstimator("est")
     val eval = new MyEvaluator
@@ -116,9 +119,194 @@ class CrossValidatorSuite extends SparkFunSuite with MLlibTestSparkContext {
       cv.validateParams()
     }
   }
+
+  test("read/write: CrossValidator with simple estimator") {
+    val lr = new LogisticRegression().setMaxIter(3)
+    val evaluator = new BinaryClassificationEvaluator()
+      .setMetricName("areaUnderPR")  // not default metric
+    val paramMaps = new ParamGridBuilder()
+        .addGrid(lr.regParam, Array(0.1, 0.2))
+        .build()
+    val cv = new CrossValidator()
+      .setEstimator(lr)
+      .setEvaluator(evaluator)
+      .setNumFolds(20)
+      .setEstimatorParamMaps(paramMaps)
+
+    val cv2 = testDefaultReadWrite(cv, testParams = false)
+
+    assert(cv.uid === cv2.uid)
+    assert(cv.getNumFolds === cv2.getNumFolds)
+
+    assert(cv2.getEvaluator.isInstanceOf[BinaryClassificationEvaluator])
+    val evaluator2 = cv2.getEvaluator.asInstanceOf[BinaryClassificationEvaluator]
+    assert(evaluator.uid === evaluator2.uid)
+    assert(evaluator.getMetricName === evaluator2.getMetricName)
+
+    cv2.getEstimator match {
+      case lr2: LogisticRegression =>
+        assert(lr.uid === lr2.uid)
+        assert(lr.getMaxIter === lr2.getMaxIter)
+      case other =>
+        throw new AssertionError(s"Loaded CrossValidator expected estimator of type" +
+          s" LogisticRegression but found ${other.getClass.getName}")
+    }
+
+    CrossValidatorSuite.compareParamMaps(cv.getEstimatorParamMaps, cv2.getEstimatorParamMaps)
+  }
+
+  test("read/write: CrossValidator with complex estimator") {
+    // workflow: CrossValidator[Pipeline[HashingTF, CrossValidator[LogisticRegression]]]
+    val lrEvaluator = new BinaryClassificationEvaluator()
+      .setMetricName("areaUnderPR")  // not default metric
+
+    val lr = new LogisticRegression().setMaxIter(3)
+    val lrParamMaps = new ParamGridBuilder()
+      .addGrid(lr.regParam, Array(0.1, 0.2))
+      .build()
+    val lrcv = new CrossValidator()
+      .setEstimator(lr)
+      .setEvaluator(lrEvaluator)
+      .setEstimatorParamMaps(lrParamMaps)
+
+    val hashingTF = new HashingTF()
+    val pipeline = new Pipeline().setStages(Array(hashingTF, lrcv))
+    val paramMaps = new ParamGridBuilder()
+      .addGrid(hashingTF.numFeatures, Array(10, 20))
+      .addGrid(lr.elasticNetParam, Array(0.0, 1.0))
+      .build()
+    val evaluator = new BinaryClassificationEvaluator()
+
+    val cv = new CrossValidator()
+      .setEstimator(pipeline)
+      .setEvaluator(evaluator)
+      .setNumFolds(20)
+      .setEstimatorParamMaps(paramMaps)
+
+    val cv2 = testDefaultReadWrite(cv, testParams = false)
+
+    assert(cv.uid === cv2.uid)
+    assert(cv.getNumFolds === cv2.getNumFolds)
+
+    assert(cv2.getEvaluator.isInstanceOf[BinaryClassificationEvaluator])
+    assert(cv.getEvaluator.uid === cv2.getEvaluator.uid)
+
+    CrossValidatorSuite.compareParamMaps(cv.getEstimatorParamMaps, cv2.getEstimatorParamMaps)
+
+    cv2.getEstimator match {
+      case pipeline2: Pipeline =>
+        assert(pipeline.uid === pipeline2.uid)
+        pipeline2.getStages match {
+          case Array(hashingTF2: HashingTF, lrcv2: CrossValidator) =>
+            assert(hashingTF.uid === hashingTF2.uid)
+            lrcv2.getEstimator match {
+              case lr2: LogisticRegression =>
+                assert(lr.uid === lr2.uid)
+                assert(lr.getMaxIter === lr2.getMaxIter)
+              case other =>
+                throw new AssertionError(s"Loaded internal CrossValidator expected to be" +
+                  s" LogisticRegression but found type ${other.getClass.getName}")
+            }
+            assert(lrcv.uid === lrcv2.uid)
+            assert(lrcv2.getEvaluator.isInstanceOf[BinaryClassificationEvaluator])
+            assert(lrEvaluator.uid === lrcv2.getEvaluator.uid)
+            CrossValidatorSuite.compareParamMaps(lrParamMaps, lrcv2.getEstimatorParamMaps)
+          case other =>
+            throw new AssertionError("Loaded Pipeline expected stages (HashingTF, CrossValidator)" +
+              " but found: " + other.map(_.getClass.getName).mkString(", "))
+        }
+      case other =>
+        throw new AssertionError(s"Loaded CrossValidator expected estimator of type" +
+          s" CrossValidator but found ${other.getClass.getName}")
+    }
+  }
+
+  test("read/write: CrossValidator fails for extraneous Param") {
+    val lr = new LogisticRegression()
+    val lr2 = new LogisticRegression()
+    val evaluator = new BinaryClassificationEvaluator()
+    val paramMaps = new ParamGridBuilder()
+      .addGrid(lr.regParam, Array(0.1, 0.2))
+      .addGrid(lr2.regParam, Array(0.1, 0.2))
+      .build()
+    val cv = new CrossValidator()
+      .setEstimator(lr)
+      .setEvaluator(evaluator)
+      .setEstimatorParamMaps(paramMaps)
+    withClue("CrossValidator.write failed to catch extraneous Param error") {
+      intercept[IllegalArgumentException] {
+        cv.write
+      }
+    }
+  }
+
+  test("read/write: CrossValidatorModel") {
+    val lr = new LogisticRegression()
+      .setThreshold(0.6)
+    val lrModel = new LogisticRegressionModel(lr.uid, Vectors.dense(1.0, 2.0), 1.2)
+      .setThreshold(0.6)
+    val evaluator = new BinaryClassificationEvaluator()
+      .setMetricName("areaUnderPR")  // not default metric
+    val paramMaps = new ParamGridBuilder()
+        .addGrid(lr.regParam, Array(0.1, 0.2))
+        .build()
+    val cv = new CrossValidatorModel("cvUid", lrModel, Array(0.3, 0.6))
+    cv.set(cv.estimator, lr)
+      .set(cv.evaluator, evaluator)
+      .set(cv.numFolds, 20)
+      .set(cv.estimatorParamMaps, paramMaps)
+
+    val cv2 = testDefaultReadWrite(cv, testParams = false)
+
+    assert(cv.uid === cv2.uid)
+    assert(cv.getNumFolds === cv2.getNumFolds)
+
+    assert(cv2.getEvaluator.isInstanceOf[BinaryClassificationEvaluator])
+    val evaluator2 = cv2.getEvaluator.asInstanceOf[BinaryClassificationEvaluator]
+    assert(evaluator.uid === evaluator2.uid)
+    assert(evaluator.getMetricName === evaluator2.getMetricName)
+
+    cv2.getEstimator match {
+      case lr2: LogisticRegression =>
+        assert(lr.uid === lr2.uid)
+        assert(lr.getThreshold === lr2.getThreshold)
+      case other =>
+        throw new AssertionError(s"Loaded CrossValidator expected estimator of type" +
+          s" LogisticRegression but found ${other.getClass.getName}")
+    }
+
+    CrossValidatorSuite.compareParamMaps(cv.getEstimatorParamMaps, cv2.getEstimatorParamMaps)
+
+    cv2.bestModel match {
+      case lrModel2: LogisticRegressionModel =>
+        assert(lrModel.uid === lrModel2.uid)
+        assert(lrModel.getThreshold === lrModel2.getThreshold)
+        assert(lrModel.coefficients === lrModel2.coefficients)
+        assert(lrModel.intercept === lrModel2.intercept)
+      case other =>
+        throw new AssertionError(s"Loaded CrossValidator expected bestModel of type" +
+          s" LogisticRegressionModel but found ${other.getClass.getName}")
+    }
+    assert(cv.avgMetrics === cv2.avgMetrics)
+  }
 }
 
-object CrossValidatorSuite {
+object CrossValidatorSuite extends SparkFunSuite {
+
+  /**
+   * Assert sequences of estimatorParamMaps are identical.
+   * Params must be simple types comparable with `===`.
+   */
+  def compareParamMaps(pMaps: Array[ParamMap], pMaps2: Array[ParamMap]): Unit = {
+    assert(pMaps.length === pMaps2.length)
+    pMaps.zip(pMaps2).foreach { case (pMap, pMap2) =>
+      assert(pMap.size === pMap2.size)
+      pMap.toSeq.foreach { case ParamPair(p, v) =>
+        assert(pMap2.contains(p))
+        assert(pMap2(p) === v)
+      }
+    }
+  }
 
   abstract class MyModel extends Model[MyModel]
 

From fc4b792d287095d70379a51f117c225d8d857078 Mon Sep 17 00:00:00 2001
From: Timothy Hunter <timhunter@databricks.com>
Date: Sun, 22 Nov 2015 21:51:42 -0800
Subject: [PATCH 406/518] [SPARK-11835] Adds a sidebar menu to MLlib's
 documentation

This PR adds a sidebar menu when browsing the user guide of MLlib. It uses a YAML file to describe the structure of the documentation. It should be trivial to adapt this to the other projects.

![screen shot 2015-11-18 at 4 46 12 pm](https://cloud.githubusercontent.com/assets/7594753/11259591/a55173f4-8e17-11e5-9340-0aed79d66262.png)

Author: Timothy Hunter <timhunter@databricks.com>

Closes #9826 from thunterdb/spark-11835.
---
 docs/_data/menu-ml.yaml                 | 10 ++++
 docs/_data/menu-mllib.yaml              | 75 +++++++++++++++++++++++++
 docs/_includes/nav-left-wrapper-ml.html |  8 +++
 docs/_includes/nav-left.html            | 17 ++++++
 docs/_layouts/global.html               | 24 +++++---
 docs/css/main.css                       | 37 ++++++++++++
 6 files changed, 163 insertions(+), 8 deletions(-)
 create mode 100644 docs/_data/menu-ml.yaml
 create mode 100644 docs/_data/menu-mllib.yaml
 create mode 100644 docs/_includes/nav-left-wrapper-ml.html
 create mode 100644 docs/_includes/nav-left.html

diff --git a/docs/_data/menu-ml.yaml b/docs/_data/menu-ml.yaml
new file mode 100644
index 0000000000000..dff3d33bf4ed1
--- /dev/null
+++ b/docs/_data/menu-ml.yaml
@@ -0,0 +1,10 @@
+- text: Feature extraction, transformation, and selection
+  url: ml-features.html
+- text: Decision trees for classification and regression
+  url: ml-decision-tree.html
+- text: Ensembles
+  url: ml-ensembles.html
+- text: Linear methods with elastic-net regularization
+  url: ml-linear-methods.html
+- text: Multilayer perceptron classifier
+  url: ml-ann.html
diff --git a/docs/_data/menu-mllib.yaml b/docs/_data/menu-mllib.yaml
new file mode 100644
index 0000000000000..12d22abd52826
--- /dev/null
+++ b/docs/_data/menu-mllib.yaml
@@ -0,0 +1,75 @@
+- text: Data types
+  url: mllib-data-types.html
+- text: Basic statistics
+  url: mllib-statistics.html
+  subitems:
+    - text: Summary statistics
+      url: mllib-statistics.html#summary-statistics
+    - text: Correlations
+      url: mllib-statistics.html#correlations
+    - text: Stratified sampling
+      url: mllib-statistics.html#stratified-sampling
+    - text: Hypothesis testing
+      url: mllib-statistics.html#hypothesis-testing
+    - text: Random data generation
+      url: mllib-statistics.html#random-data-generation
+- text: Classification and regression
+  url: mllib-classification-regression.html
+  subitems:
+    - text: Linear models (SVMs, logistic regression, linear regression)
+      url: mllib-linear-methods.html
+    - text: Naive Bayes
+      url: mllib-naive-bayes.html
+    - text: decision trees
+      url: mllib-decision-tree.html
+    - text: ensembles of trees (Random Forests and Gradient-Boosted Trees)
+      url: mllib-ensembles.html
+    - text: isotonic regression
+      url: mllib-isotonic-regression.html
+- text: Collaborative filtering
+  url: mllib-collaborative-filtering.html
+  subitems:
+    - text: alternating least squares (ALS)
+      url: mllib-collaborative-filtering.html#collaborative-filtering
+- text: Clustering
+  url: mllib-clustering.html
+  subitems:
+    - text: k-means
+      url: mllib-clustering.html#k-means
+    - text: Gaussian mixture
+      url: mllib-clustering.html#gaussian-mixture
+    - text: power iteration clustering (PIC)
+      url: mllib-clustering.html#power-iteration-clustering-pic
+    - text: latent Dirichlet allocation (LDA)
+      url: mllib-clustering.html#latent-dirichlet-allocation-lda
+    - text: streaming k-means
+      url: mllib-clustering.html#streaming-k-means
+- text: Dimensionality reduction
+  url: mllib-dimensionality-reduction.html
+  subitems:
+    - text: singular value decomposition (SVD)
+      url: mllib-dimensionality-reduction.html#singular-value-decomposition-svd
+    - text: principal component analysis (PCA)
+      url: mllib-dimensionality-reduction.html#principal-component-analysis-pca
+- text: Feature extraction and transformation
+  url: mllib-feature-extraction.html
+- text: Frequent pattern mining
+  url: mllib-frequent-pattern-mining.html
+  subitems:
+    - text: FP-growth
+      url: mllib-frequent-pattern-mining.html#fp-growth
+    - text: association rules
+      url: mllib-frequent-pattern-mining.html#association-rules
+    - text: PrefixSpan
+      url: mllib-frequent-pattern-mining.html#prefix-span
+- text: Evaluation metrics
+  url: mllib-evaluation-metrics.html
+- text: PMML model export
+  url: mllib-pmml-model-export.html
+- text: Optimization (developer)
+  url: mllib-optimization.html
+  subitems:
+    - text: stochastic gradient descent
+      url: mllib-optimization.html#stochastic-gradient-descent-sgd
+    - text: limited-memory BFGS (L-BFGS)
+      url: mllib-optimization.html#limited-memory-bfgs-l-bfgs
diff --git a/docs/_includes/nav-left-wrapper-ml.html b/docs/_includes/nav-left-wrapper-ml.html
new file mode 100644
index 0000000000000..0103e890cc21a
--- /dev/null
+++ b/docs/_includes/nav-left-wrapper-ml.html
@@ -0,0 +1,8 @@
+<div class="left-menu-wrapper">
+    <div class="left-menu">
+        <h3>spark.ml package</h3>
+        {% include nav-left.html nav=include.nav-ml %}
+        <h3>spark.mllib package</h3>
+        {% include nav-left.html nav=include.nav-mllib %}
+    </div>
+</div>
\ No newline at end of file
diff --git a/docs/_includes/nav-left.html b/docs/_includes/nav-left.html
new file mode 100644
index 0000000000000..73176f4132554
--- /dev/null
+++ b/docs/_includes/nav-left.html
@@ -0,0 +1,17 @@
+{% assign navurl = page.url | remove: 'index.html' %}
+<ul>
+{% for item in include.nav %}
+    <li>
+        <a href="{{ item.url }}">
+            {% if navurl contains item.url %}
+                <b>{{ item.text }}</b>
+            {% else %}
+                {{ item.text }}
+            {% endif %}
+        </a>
+    </li>
+    {% if item.subitems and navurl contains item.url %}
+        {% include nav-left.html nav=item.subitems %}
+    {% endif %}
+{% endfor %}
+</ul>
diff --git a/docs/_layouts/global.html b/docs/_layouts/global.html
index 467ff7a03fb70..1b09e2221e173 100755
--- a/docs/_layouts/global.html
+++ b/docs/_layouts/global.html
@@ -124,16 +124,24 @@
             </div>
         </div>
 
-        <div class="container" id="content">
-          {% if page.displayTitle %}
-            <h1 class="title">{{ page.displayTitle }}</h1>
-          {% else %}
-            <h1 class="title">{{ page.title }}</h1>
-          {% endif %}
+        <div class="container-wrapper">
 
-          {{ content }}
+            {% if page.url contains "/ml" %}
+              {% include nav-left-wrapper-ml.html nav-mllib=site.data.menu-mllib nav-ml=site.data.menu-ml %}
+            {% endif %}
 
-        </div> <!-- /container -->
+
+            <div class="container" id="content">
+              {% if page.displayTitle %}
+                <h1 class="title">{{ page.displayTitle }}</h1>
+              {% else %}
+                <h1 class="title">{{ page.title }}</h1>
+              {% endif %}
+
+              {{ content }}
+
+            </div> <!-- /container -->
+        </div>
 
         <script src="js/vendor/jquery-1.8.0.min.js"></script>
         <script src="js/vendor/bootstrap.min.js"></script>
diff --git a/docs/css/main.css b/docs/css/main.css
index d770173be1014..356b324d6303b 100755
--- a/docs/css/main.css
+++ b/docs/css/main.css
@@ -39,8 +39,18 @@
   margin-left: 10px;
 }
 
+body .container-wrapper {
+  position: absolute;
+  width: 100%;
+  display: flex;
+}
+
 body #content {
+  position: relative;
+
   line-height: 1.6; /* Inspired by Github's wiki style */
+  background-color: white;
+  padding-left: 15px;
 }
 
 .title {
@@ -155,3 +165,30 @@ ul.nav li.dropdown ul.dropdown-menu li.dropdown-submenu ul.dropdown-menu {
  * AnchorJS (anchor links when hovering over headers)
  */
 a.anchorjs-link:hover { text-decoration: none; }
+
+
+/**
+ * The left navigation bar.
+ */
+.left-menu-wrapper {
+  position: absolute;
+  height: 100%;
+
+  width: 256px;
+  margin-top: -20px;
+  padding-top: 20px;
+  background-color: #F0F8FC;
+}
+
+.left-menu {
+  position: fixed;
+  max-width: 350px;
+
+  padding-right: 10px;
+  width: 256px;
+}
+
+.left-menu h3 {
+  margin-left: 10px;
+  line-height: 30px;
+}
\ No newline at end of file

From d9cf9c21fc6b1aa22e68d66760afd42c4e1c18b8 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Sun, 22 Nov 2015 21:56:07 -0800
Subject: [PATCH 407/518] [SPARK-11912][ML] ml.feature.PCA minor refactor

Like [SPARK-11852](https://issues.apache.org/jira/browse/SPARK-11852), ```k``` is params and we should save it under ```metadata/``` rather than both under ```data/``` and ```metadata/```. Refactor the constructor of ```ml.feature.PCAModel```  to take only ```pc``` but construct ```mllib.feature.PCAModel``` inside ```transform```.

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9897 from yanboliang/spark-11912.
---
 .../org/apache/spark/ml/feature/PCA.scala     | 23 +++++++-------
 .../apache/spark/ml/feature/PCASuite.scala    | 31 ++++++++-----------
 2 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
index 32d7afee6e73b..aa88cb03d23c5 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
@@ -73,7 +73,7 @@ class PCA (override val uid: String) extends Estimator[PCAModel] with PCAParams
     val input = dataset.select($(inputCol)).map { case Row(v: Vector) => v}
     val pca = new feature.PCA(k = $(k))
     val pcaModel = pca.fit(input)
-    copyValues(new PCAModel(uid, pcaModel).setParent(this))
+    copyValues(new PCAModel(uid, pcaModel.pc).setParent(this))
   }
 
   override def transformSchema(schema: StructType): StructType = {
@@ -99,18 +99,17 @@ object PCA extends DefaultParamsReadable[PCA] {
 /**
  * :: Experimental ::
  * Model fitted by [[PCA]].
+ *
+ * @param pc A principal components Matrix. Each column is one principal component.
  */
 @Experimental
 class PCAModel private[ml] (
     override val uid: String,
-    pcaModel: feature.PCAModel)
+    val pc: DenseMatrix)
   extends Model[PCAModel] with PCAParams with MLWritable {
 
   import PCAModel._
 
-  /** a principal components Matrix. Each column is one principal component. */
-  val pc: DenseMatrix = pcaModel.pc
-
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
 
@@ -124,6 +123,7 @@ class PCAModel private[ml] (
    */
   override def transform(dataset: DataFrame): DataFrame = {
     transformSchema(dataset.schema, logging = true)
+    val pcaModel = new feature.PCAModel($(k), pc)
     val pcaOp = udf { pcaModel.transform _ }
     dataset.withColumn($(outputCol), pcaOp(col($(inputCol))))
   }
@@ -139,7 +139,7 @@ class PCAModel private[ml] (
   }
 
   override def copy(extra: ParamMap): PCAModel = {
-    val copied = new PCAModel(uid, pcaModel)
+    val copied = new PCAModel(uid, pc)
     copyValues(copied, extra).setParent(parent)
   }
 
@@ -152,11 +152,11 @@ object PCAModel extends MLReadable[PCAModel] {
 
   private[PCAModel] class PCAModelWriter(instance: PCAModel) extends MLWriter {
 
-    private case class Data(k: Int, pc: DenseMatrix)
+    private case class Data(pc: DenseMatrix)
 
     override protected def saveImpl(path: String): Unit = {
       DefaultParamsWriter.saveMetadata(instance, path, sc)
-      val data = Data(instance.getK, instance.pc)
+      val data = Data(instance.pc)
       val dataPath = new Path(path, "data").toString
       sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
     }
@@ -169,11 +169,10 @@ object PCAModel extends MLReadable[PCAModel] {
     override def load(path: String): PCAModel = {
       val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
       val dataPath = new Path(path, "data").toString
-      val Row(k: Int, pc: DenseMatrix) = sqlContext.read.parquet(dataPath)
-        .select("k", "pc")
+      val Row(pc: DenseMatrix) = sqlContext.read.parquet(dataPath)
+        .select("pc")
         .head()
-      val oldModel = new feature.PCAModel(k, pc)
-      val model = new PCAModel(metadata.uid, oldModel)
+      val model = new PCAModel(metadata.uid, pc)
       DefaultParamsReader.getAndSetParams(model, metadata)
       model
     }
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
index 5a21cd20ceede..edab21e6c3072 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
@@ -32,7 +32,7 @@ class PCASuite extends SparkFunSuite with MLlibTestSparkContext with DefaultRead
   test("params") {
     ParamsSuite.checkParams(new PCA)
     val mat = Matrices.dense(2, 2, Array(0.0, 1.0, 2.0, 3.0)).asInstanceOf[DenseMatrix]
-    val model = new PCAModel("pca", new OldPCAModel(2, mat))
+    val model = new PCAModel("pca", mat)
     ParamsSuite.checkParams(model)
   }
 
@@ -66,23 +66,18 @@ class PCASuite extends SparkFunSuite with MLlibTestSparkContext with DefaultRead
     }
   }
 
-  test("read/write") {
+  test("PCA read/write") {
+    val t = new PCA()
+      .setInputCol("myInputCol")
+      .setOutputCol("myOutputCol")
+      .setK(3)
+    testDefaultReadWrite(t)
+  }
 
-    def checkModelData(model1: PCAModel, model2: PCAModel): Unit = {
-      assert(model1.pc === model2.pc)
-    }
-    val allParams: Map[String, Any] = Map(
-      "k" -> 3,
-      "inputCol" -> "features",
-      "outputCol" -> "pca_features"
-    )
-    val data = Seq(
-      (0.0, Vectors.sparse(5, Seq((1, 1.0), (3, 7.0)))),
-      (1.0, Vectors.dense(2.0, 0.0, 3.0, 4.0, 5.0)),
-      (2.0, Vectors.dense(4.0, 0.0, 0.0, 6.0, 7.0))
-    )
-    val df = sqlContext.createDataFrame(data).toDF("id", "features")
-    val pca = new PCA().setK(3)
-    testEstimatorAndModelReadWrite(pca, df, allParams, checkModelData)
+  test("PCAModel read/write") {
+    val instance = new PCAModel("myPCAModel",
+      Matrices.dense(2, 2, Array(0.0, 1.0, 2.0, 3.0)).asInstanceOf[DenseMatrix])
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.pc === instance.pc)
   }
 }

From 4be360d4ee6cdb4d06306feca38ddef5212608cf Mon Sep 17 00:00:00 2001
From: BenFradet <benjamin.fradet@gmail.com>
Date: Sun, 22 Nov 2015 22:05:01 -0800
Subject: [PATCH 408/518] [SPARK-11902][ML] Unhandled case in
 VectorAssembler#transform

There is an unhandled case in the transform method of VectorAssembler if one of the input columns doesn't have one of the supported type DoubleType, NumericType, BooleanType or VectorUDT.

So, if you try to transform a column of StringType you get a cryptic "scala.MatchError: StringType".

This PR aims to fix this, throwing a SparkException when dealing with an unknown column type.

Author: BenFradet <benjamin.fradet@gmail.com>

Closes #9885 from BenFradet/SPARK-11902.
---
 .../org/apache/spark/ml/feature/VectorAssembler.scala |  2 ++
 .../spark/ml/feature/VectorAssemblerSuite.scala       | 11 +++++++++++
 2 files changed, 13 insertions(+)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
index 0feec0549852b..801096fed27bf 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala
@@ -84,6 +84,8 @@ class VectorAssembler(override val uid: String)
             val numAttrs = group.numAttributes.getOrElse(first.getAs[Vector](index).size)
             Array.fill(numAttrs)(NumericAttribute.defaultAttr)
           }
+        case otherType =>
+          throw new SparkException(s"VectorAssembler does not support the $otherType type")
       }
     }
     val metadata = new AttributeGroup($(outputCol), attrs).toMetadata()
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorAssemblerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorAssemblerSuite.scala
index fb21ab6b9bf2c..9c1c00f41ab1d 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorAssemblerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorAssemblerSuite.scala
@@ -69,6 +69,17 @@ class VectorAssemblerSuite
     }
   }
 
+  test("transform should throw an exception in case of unsupported type") {
+    val df = sqlContext.createDataFrame(Seq(("a", "b", "c"))).toDF("a", "b", "c")
+    val assembler = new VectorAssembler()
+      .setInputCols(Array("a", "b", "c"))
+      .setOutputCol("features")
+    val thrown = intercept[SparkException] {
+      assembler.transform(df)
+    }
+    assert(thrown.getMessage contains "VectorAssembler does not support the StringType type")
+  }
+
   test("ML attributes") {
     val browser = NominalAttribute.defaultAttr.withValues("chrome", "firefox", "safari")
     val hour = NumericAttribute.defaultAttr.withMin(0.0).withMax(24.0)

From 94ce65dfcbba1fe3a1fc9d8002c37d9cd1a11336 Mon Sep 17 00:00:00 2001
From: Xiu Guo <xguo27@gmail.com>
Date: Mon, 23 Nov 2015 08:53:40 -0800
Subject: [PATCH 409/518] [SPARK-11628][SQL] support column datatype of char(x)
 to recognize HiveChar

Can someone review my code to make sure I'm not missing anything? Thanks!

Author: Xiu Guo <xguo27@gmail.com>
Author: Xiu Guo <guoxi@us.ibm.com>

Closes #9612 from xguo27/SPARK-11628.
---
 .../sql/catalyst/util/DataTypeParser.scala    |  6 ++++-
 .../catalyst/util/DataTypeParserSuite.scala   |  8 ++++--
 .../spark/sql/sources/TableScanSuite.scala    |  5 ++++
 .../spark/sql/hive/HiveInspectors.scala       | 25 ++++++++++++++++---
 .../apache/spark/sql/hive/TableReader.scala   |  3 +++
 .../spark/sql/hive/client/HiveShim.scala      |  3 ++-
 6 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DataTypeParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DataTypeParser.scala
index 2b83651f9086d..515c071c283b0 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DataTypeParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DataTypeParser.scala
@@ -52,7 +52,8 @@ private[sql] trait DataTypeParser extends StandardTokenParsers {
     "(?i)decimal".r ^^^ DecimalType.USER_DEFAULT |
     "(?i)date".r ^^^ DateType |
     "(?i)timestamp".r ^^^ TimestampType |
-    varchar
+    varchar |
+    char
 
   protected lazy val fixedDecimalType: Parser[DataType] =
     ("(?i)decimal".r ~> "(" ~> numericLit) ~ ("," ~> numericLit <~ ")") ^^ {
@@ -60,6 +61,9 @@ private[sql] trait DataTypeParser extends StandardTokenParsers {
         DecimalType(precision.toInt, scale.toInt)
     }
 
+  protected lazy val char: Parser[DataType] =
+    "(?i)char".r ~> "(" ~> (numericLit <~ ")") ^^^ StringType
+
   protected lazy val varchar: Parser[DataType] =
     "(?i)varchar".r ~> "(" ~> (numericLit <~ ")") ^^^ StringType
 
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DataTypeParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DataTypeParserSuite.scala
index 1e3409a9db6eb..bebf708965474 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DataTypeParserSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DataTypeParserSuite.scala
@@ -49,7 +49,9 @@ class DataTypeParserSuite extends SparkFunSuite {
   checkDataType("DATE", DateType)
   checkDataType("timestamp", TimestampType)
   checkDataType("string", StringType)
+  checkDataType("ChaR(5)", StringType)
   checkDataType("varchAr(20)", StringType)
+  checkDataType("cHaR(27)", StringType)
   checkDataType("BINARY", BinaryType)
 
   checkDataType("array<doublE>", ArrayType(DoubleType, true))
@@ -83,7 +85,8 @@ class DataTypeParserSuite extends SparkFunSuite {
       |struct<
       |  struct:struct<deciMal:DECimal, anotherDecimal:decimAL(5,2)>,
       |  MAP:Map<timestamp, varchar(10)>,
-      |  arrAy:Array<double>>
+      |  arrAy:Array<double>,
+      |  anotherArray:Array<char(9)>>
     """.stripMargin,
     StructType(
       StructField("struct",
@@ -91,7 +94,8 @@ class DataTypeParserSuite extends SparkFunSuite {
           StructField("deciMal", DecimalType.USER_DEFAULT, true) ::
           StructField("anotherDecimal", DecimalType(5, 2), true) :: Nil), true) ::
       StructField("MAP", MapType(TimestampType, StringType), true) ::
-      StructField("arrAy", ArrayType(DoubleType, true), true) :: Nil)
+      StructField("arrAy", ArrayType(DoubleType, true), true) ::
+      StructField("anotherArray", ArrayType(StringType, true), true) :: Nil)
   )
   // A column name can be a reserved word in our DDL parser and SqlParser.
   checkDataType(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/TableScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/TableScanSuite.scala
index 12af8068c398f..26c1ff520406c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/sources/TableScanSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/TableScanSuite.scala
@@ -85,6 +85,7 @@ case class AllDataTypesScan(
         Date.valueOf("1970-01-01"),
         new Timestamp(20000 + i),
         s"varchar_$i",
+        s"char_$i",
         Seq(i, i + 1),
         Seq(Map(s"str_$i" -> Row(i.toLong))),
         Map(i -> i.toString),
@@ -115,6 +116,7 @@ class TableScanSuite extends DataSourceTest with SharedSQLContext {
       Date.valueOf("1970-01-01"),
       new Timestamp(20000 + i),
       s"varchar_$i",
+      s"char_$i",
       Seq(i, i + 1),
       Seq(Map(s"str_$i" -> Row(i.toLong))),
       Map(i -> i.toString),
@@ -154,6 +156,7 @@ class TableScanSuite extends DataSourceTest with SharedSQLContext {
         |dateField dAte,
         |timestampField tiMestamp,
         |varcharField varchaR(12),
+        |charField ChaR(18),
         |arrayFieldSimple Array<inT>,
         |arrayFieldComplex Array<Map<String, Struct<key:bigInt>>>,
         |mapFieldSimple MAP<iNt, StRing>,
@@ -207,6 +210,7 @@ class TableScanSuite extends DataSourceTest with SharedSQLContext {
       StructField("dateField", DateType, true) ::
       StructField("timestampField", TimestampType, true) ::
       StructField("varcharField", StringType, true) ::
+      StructField("charField", StringType, true) ::
       StructField("arrayFieldSimple", ArrayType(IntegerType), true) ::
       StructField("arrayFieldComplex",
         ArrayType(
@@ -248,6 +252,7 @@ class TableScanSuite extends DataSourceTest with SharedSQLContext {
           | dateField,
           | timestampField,
           | varcharField,
+          | charField,
           | arrayFieldSimple,
           | arrayFieldComplex,
           | mapFieldSimple,
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveInspectors.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveInspectors.scala
index 36f0708f9da3d..95b57d6ad124a 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveInspectors.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveInspectors.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.hive
 
 import scala.collection.JavaConverters._
 
-import org.apache.hadoop.hive.common.`type`.{HiveDecimal, HiveVarchar}
+import org.apache.hadoop.hive.common.`type`.{HiveChar, HiveDecimal, HiveVarchar}
 import org.apache.hadoop.hive.serde2.objectinspector.primitive._
 import org.apache.hadoop.hive.serde2.objectinspector.{StructField => HiveStructField, _}
 import org.apache.hadoop.hive.serde2.typeinfo.{DecimalTypeInfo, TypeInfoFactory}
@@ -61,6 +61,7 @@ import org.apache.spark.unsafe.types.UTF8String
  * Primitive Type
  *   Java Boxed Primitives:
  *       org.apache.hadoop.hive.common.type.HiveVarchar
+ *       org.apache.hadoop.hive.common.type.HiveChar
  *       java.lang.String
  *       java.lang.Integer
  *       java.lang.Boolean
@@ -75,6 +76,7 @@ import org.apache.spark.unsafe.types.UTF8String
  *       java.sql.Timestamp
  *   Writables:
  *       org.apache.hadoop.hive.serde2.io.HiveVarcharWritable
+ *       org.apache.hadoop.hive.serde2.io.HiveCharWritable
  *       org.apache.hadoop.io.Text
  *       org.apache.hadoop.io.IntWritable
  *       org.apache.hadoop.hive.serde2.io.DoubleWritable
@@ -93,7 +95,8 @@ import org.apache.spark.unsafe.types.UTF8String
  *   Struct: Object[] / java.util.List / java POJO
  *   Union: class StandardUnion { byte tag; Object object }
  *
- * NOTICE: HiveVarchar is not supported by catalyst, it will be simply considered as String type.
+ * NOTICE: HiveVarchar/HiveChar is not supported by catalyst, it will be simply considered as
+ *  String type.
  *
  *
  * 2. Hive ObjectInspector is a group of flexible APIs to inspect value in different data
@@ -137,6 +140,7 @@ import org.apache.spark.unsafe.types.UTF8String
  * Primitive Object Inspectors:
  *     WritableConstantStringObjectInspector
  *     WritableConstantHiveVarcharObjectInspector
+ *     WritableConstantHiveCharObjectInspector
  *     WritableConstantHiveDecimalObjectInspector
  *     WritableConstantTimestampObjectInspector
  *     WritableConstantIntObjectInspector
@@ -259,6 +263,8 @@ private[hive] trait HiveInspectors {
       UTF8String.fromString(poi.getWritableConstantValue.toString)
     case poi: WritableConstantHiveVarcharObjectInspector =>
       UTF8String.fromString(poi.getWritableConstantValue.getHiveVarchar.getValue)
+    case poi: WritableConstantHiveCharObjectInspector =>
+      UTF8String.fromString(poi.getWritableConstantValue.getHiveChar.getValue)
     case poi: WritableConstantHiveDecimalObjectInspector =>
       HiveShim.toCatalystDecimal(
         PrimitiveObjectInspectorFactory.javaHiveDecimalObjectInspector,
@@ -303,11 +309,15 @@ private[hive] trait HiveInspectors {
     case _ if data == null => null
     case poi: VoidObjectInspector => null // always be null for void object inspector
     case pi: PrimitiveObjectInspector => pi match {
-      // We think HiveVarchar is also a String
+      // We think HiveVarchar/HiveChar is also a String
       case hvoi: HiveVarcharObjectInspector if hvoi.preferWritable() =>
         UTF8String.fromString(hvoi.getPrimitiveWritableObject(data).getHiveVarchar.getValue)
       case hvoi: HiveVarcharObjectInspector =>
         UTF8String.fromString(hvoi.getPrimitiveJavaObject(data).getValue)
+      case hvoi: HiveCharObjectInspector if hvoi.preferWritable() =>
+        UTF8String.fromString(hvoi.getPrimitiveWritableObject(data).getHiveChar.getValue)
+      case hvoi: HiveCharObjectInspector =>
+        UTF8String.fromString(hvoi.getPrimitiveJavaObject(data).getValue)
       case x: StringObjectInspector if x.preferWritable() =>
         UTF8String.fromString(x.getPrimitiveWritableObject(data).toString)
       case x: StringObjectInspector =>
@@ -377,6 +387,15 @@ private[hive] trait HiveInspectors {
           null
         }
 
+    case _: JavaHiveCharObjectInspector =>
+      (o: Any) =>
+        if (o != null) {
+          val s = o.asInstanceOf[UTF8String].toString
+          new HiveChar(s, s.size)
+        } else {
+          null
+        }
+
     case _: JavaHiveDecimalObjectInspector =>
       (o: Any) =>
         if (o != null) {
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/TableReader.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/TableReader.scala
index 69f481c49a655..70ee02823eeba 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/TableReader.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/TableReader.scala
@@ -382,6 +382,9 @@ private[hive] object HadoopTableReader extends HiveInspectors with Logging {
         case oi: HiveVarcharObjectInspector =>
           (value: Any, row: MutableRow, ordinal: Int) =>
             row.update(ordinal, UTF8String.fromString(oi.getPrimitiveJavaObject(value).getValue))
+        case oi: HiveCharObjectInspector =>
+          (value: Any, row: MutableRow, ordinal: Int) =>
+            row.update(ordinal, UTF8String.fromString(oi.getPrimitiveJavaObject(value).getValue))
         case oi: HiveDecimalObjectInspector =>
           (value: Any, row: MutableRow, ordinal: Int) =>
             row.update(ordinal, HiveShim.toCatalystDecimal(oi, value))
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala
index 48bbb21e6c1de..346840079b853 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala
@@ -321,7 +321,8 @@ private[client] class Shim_v0_13 extends Shim_v0_12 {
   def convertFilters(table: Table, filters: Seq[Expression]): String = {
     // hive varchar is treated as catalyst string, but hive varchar can't be pushed down.
     val varcharKeys = table.getPartitionKeys.asScala
-      .filter(col => col.getType.startsWith(serdeConstants.VARCHAR_TYPE_NAME))
+      .filter(col => col.getType.startsWith(serdeConstants.VARCHAR_TYPE_NAME) ||
+        col.getType.startsWith(serdeConstants.CHAR_TYPE_NAME))
       .map(col => col.getName).toSet
 
     filters.collect {

From 1a5baaa6517872b9a4fd6cd41c4b2cf1e390f6d1 Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Mon, 23 Nov 2015 10:13:59 -0800
Subject: [PATCH 410/518] [SPARK-11894][SQL] fix isNull for GetInternalRowField

We should use `InternalRow.isNullAt` to check if the field is null before calling `InternalRow.getXXX`

Thanks gatorsmile who discovered this bug.

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9904 from cloud-fan/null.
---
 .../sql/catalyst/expressions/objects.scala    | 23 ++++++++-----------
 .../org/apache/spark/sql/DatasetSuite.scala   | 15 +++++++++++-
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index 82317d3385167..4a1f419f0ad8d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -236,11 +236,6 @@ case class NewInstance(
     }
 
     if (propagateNull) {
-      val objNullCheck = if (ctx.defaultValue(dataType) == "null") {
-        s"${ev.isNull} = ${ev.value} == null;"
-      } else {
-        ""
-      }
       val argsNonNull = s"!(${argGen.map(_.isNull).mkString(" || ")})"
 
       s"""
@@ -531,15 +526,15 @@ case class GetInternalRowField(child: Expression, ordinal: Int, dataType: DataTy
     throw new UnsupportedOperationException("Only code-generated evaluation is supported")
 
   override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    val row = child.gen(ctx)
-    s"""
-      ${row.code}
-      final boolean ${ev.isNull} = ${row.isNull};
-      ${ctx.javaType(dataType)} ${ev.value} = ${ctx.defaultValue(dataType)};
-      if (!${ev.isNull}) {
-        ${ev.value} = ${ctx.getValue(row.value, dataType, ordinal.toString)};
-      }
-    """
+    nullSafeCodeGen(ctx, ev, eval => {
+      s"""
+        if ($eval.isNullAt($ordinal)) {
+          ${ev.isNull} = true;
+        } else {
+          ${ev.value} = ${ctx.getValue(eval, dataType, ordinal.toString)};
+        }
+      """
+    })
   }
 }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 9da02550b39ce..cc8e4325fd2f5 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -386,7 +386,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       Seq((JavaData(1), 1L), (JavaData(2), 1L)))
   }
 
-  ignore("Java encoder self join") {
+  test("Java encoder self join") {
     implicit val kryoEncoder = Encoders.javaSerialization[JavaData]
     val ds = Seq(JavaData(1), JavaData(2)).toDS()
     assert(ds.joinWith(ds, lit(true)).collect().toSet ==
@@ -396,6 +396,19 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
         (JavaData(2), JavaData(1)),
         (JavaData(2), JavaData(2))))
   }
+
+  test("SPARK-11894: Incorrect results are returned when using null") {
+    val nullInt = null.asInstanceOf[java.lang.Integer]
+    val ds1 = Seq((nullInt, "1"), (new java.lang.Integer(22), "2")).toDS()
+    val ds2 = Seq((nullInt, "1"), (new java.lang.Integer(22), "2")).toDS()
+
+    checkAnswer(
+      ds1.joinWith(ds2, lit(true)),
+      ((nullInt, "1"), (nullInt, "1")),
+      ((new java.lang.Integer(22), "2"), (nullInt, "1")),
+      ((nullInt, "1"), (new java.lang.Integer(22), "2")),
+      ((new java.lang.Integer(22), "2"), (new java.lang.Integer(22), "2")))
+  }
 }
 
 

From f2996e0d12eeb989b1bfa51a3f6fa54ce1ed4fca Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Mon, 23 Nov 2015 10:15:40 -0800
Subject: [PATCH 411/518] [SPARK-11921][SQL] fix `nullable` of encoder schema

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9906 from cloud-fan/nullable.
---
 .../catalyst/encoders/ExpressionEncoder.scala | 15 +++++++-
 .../encoders/ExpressionEncoderSuite.scala     | 38 ++++++++++++++++++-
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index 6eeba1442c1f3..7bc9aed0b204e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -54,8 +54,13 @@ object ExpressionEncoder {
     val toRowExpression = ScalaReflection.extractorsFor[T](inputObject)
     val fromRowExpression = ScalaReflection.constructorFor[T]
 
+    val schema = ScalaReflection.schemaFor[T] match {
+      case ScalaReflection.Schema(s: StructType, _) => s
+      case ScalaReflection.Schema(dt, nullable) => new StructType().add("value", dt, nullable)
+    }
+
     new ExpressionEncoder[T](
-      toRowExpression.dataType,
+      schema,
       flat,
       toRowExpression.flatten,
       fromRowExpression,
@@ -71,7 +76,13 @@ object ExpressionEncoder {
     encoders.foreach(_.assertUnresolved())
 
     val schema = StructType(encoders.zipWithIndex.map {
-      case (e, i) => StructField(s"_${i + 1}", if (e.flat) e.schema.head.dataType else e.schema)
+      case (e, i) =>
+        val (dataType, nullable) = if (e.flat) {
+          e.schema.head.dataType -> e.schema.head.nullable
+        } else {
+          e.schema -> true
+        }
+        StructField(s"_${i + 1}", dataType, nullable)
     })
 
     val cls = Utils.getContextOrSparkClassLoader.loadClass(s"scala.Tuple${encoders.size}")
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
index 76459b34a484f..d6ca138672ef1 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoderSuite.scala
@@ -30,7 +30,7 @@ import org.apache.spark.sql.Encoders
 import org.apache.spark.sql.catalyst.expressions.AttributeReference
 import org.apache.spark.sql.catalyst.util.ArrayData
 import org.apache.spark.sql.catalyst.{OptionalData, PrimitiveData}
-import org.apache.spark.sql.types.ArrayType
+import org.apache.spark.sql.types.{StructType, ArrayType}
 
 case class RepeatedStruct(s: Seq[PrimitiveData])
 
@@ -238,6 +238,42 @@ class ExpressionEncoderSuite extends SparkFunSuite {
     ExpressionEncoder.tuple(intEnc, ExpressionEncoder.tuple(intEnc, longEnc))
   }
 
+  test("nullable of encoder schema") {
+    def checkNullable[T: ExpressionEncoder](nullable: Boolean*): Unit = {
+      assert(implicitly[ExpressionEncoder[T]].schema.map(_.nullable) === nullable.toSeq)
+    }
+
+    // test for flat encoders
+    checkNullable[Int](false)
+    checkNullable[Option[Int]](true)
+    checkNullable[java.lang.Integer](true)
+    checkNullable[String](true)
+
+    // test for product encoders
+    checkNullable[(String, Int)](true, false)
+    checkNullable[(Int, java.lang.Long)](false, true)
+
+    // test for nested product encoders
+    {
+      val schema = ExpressionEncoder[(Int, (String, Int))].schema
+      assert(schema(0).nullable === false)
+      assert(schema(1).nullable === true)
+      assert(schema(1).dataType.asInstanceOf[StructType](0).nullable === true)
+      assert(schema(1).dataType.asInstanceOf[StructType](1).nullable === false)
+    }
+
+    // test for tupled encoders
+    {
+      val schema = ExpressionEncoder.tuple(
+        ExpressionEncoder[Int],
+        ExpressionEncoder[(String, Int)]).schema
+      assert(schema(0).nullable === false)
+      assert(schema(1).nullable === true)
+      assert(schema(1).dataType.asInstanceOf[StructType](0).nullable === true)
+      assert(schema(1).dataType.asInstanceOf[StructType](1).nullable === false)
+    }
+  }
+
   private val outers: ConcurrentMap[String, AnyRef] = new MapMaker().weakValues().makeMap()
   outers.put(getClass.getName, this)
   private def encodeDecodeTest[T : ExpressionEncoder](

From 946b406519af58c79041217e6f93854b6cf80acd Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Mon, 23 Nov 2015 10:39:33 -0800
Subject: [PATCH 412/518] [SPARK-11913][SQL] support typed aggregate with
 complex buffer schema

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9898 from cloud-fan/agg.
---
 .../aggregate/TypedAggregateExpression.scala  | 25 +++++++----
 .../spark/sql/DatasetAggregatorSuite.scala    | 41 ++++++++++++++++++-
 2 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
index 6ce41aaf01e27..a9719128a626e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala
@@ -23,9 +23,8 @@ import org.apache.spark.Logging
 import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.expressions.Aggregator
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.encoders.encoderFor
+import org.apache.spark.sql.catalyst.encoders.{OuterScopes, encoderFor, ExpressionEncoder}
 import org.apache.spark.sql.catalyst.expressions.aggregate.ImperativeAggregate
-import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.types._
 
@@ -46,14 +45,12 @@ object TypedAggregateExpression {
 /**
  * This class is a rough sketch of how to hook `Aggregator` into the Aggregation system.  It has
  * the following limitations:
- *  - It assumes the aggregator reduces and returns a single column of type `long`.
- *  - It might only work when there is a single aggregator in the first column.
  *  - It assumes the aggregator has a zero, `0`.
  */
 case class TypedAggregateExpression(
     aggregator: Aggregator[Any, Any, Any],
     aEncoder: Option[ExpressionEncoder[Any]], // Should be bound.
-    bEncoder: ExpressionEncoder[Any], // Should be bound.
+    unresolvedBEncoder: ExpressionEncoder[Any],
     cEncoder: ExpressionEncoder[Any],
     children: Seq[Attribute],
     mutableAggBufferOffset: Int,
@@ -80,10 +77,14 @@ case class TypedAggregateExpression(
 
   override lazy val inputTypes: Seq[DataType] = Nil
 
-  override val aggBufferSchema: StructType = bEncoder.schema
+  override val aggBufferSchema: StructType = unresolvedBEncoder.schema
 
   override val aggBufferAttributes: Seq[AttributeReference] = aggBufferSchema.toAttributes
 
+  val bEncoder = unresolvedBEncoder
+    .resolve(aggBufferAttributes, OuterScopes.outerScopes)
+    .bind(aggBufferAttributes)
+
   // Note: although this simply copies aggBufferAttributes, this common code can not be placed
   // in the superclass because that will lead to initialization ordering issues.
   override val inputAggBufferAttributes: Seq[AttributeReference] =
@@ -93,12 +94,18 @@ case class TypedAggregateExpression(
   lazy val boundA = aEncoder.get
 
   private def updateBuffer(buffer: MutableRow, value: InternalRow): Unit = {
-    // todo: need a more neat way to assign the value.
     var i = 0
     while (i < aggBufferAttributes.length) {
+      val offset = mutableAggBufferOffset + i
       aggBufferSchema(i).dataType match {
-        case IntegerType => buffer.setInt(mutableAggBufferOffset + i, value.getInt(i))
-        case LongType => buffer.setLong(mutableAggBufferOffset + i, value.getLong(i))
+        case BooleanType => buffer.setBoolean(offset, value.getBoolean(i))
+        case ByteType => buffer.setByte(offset, value.getByte(i))
+        case ShortType => buffer.setShort(offset, value.getShort(i))
+        case IntegerType => buffer.setInt(offset, value.getInt(i))
+        case LongType => buffer.setLong(offset, value.getLong(i))
+        case FloatType => buffer.setFloat(offset, value.getFloat(i))
+        case DoubleType => buffer.setDouble(offset, value.getDouble(i))
+        case other => buffer.update(offset, value.get(i, other))
       }
       i += 1
     }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
index 9377589790011..19dce5d1e2f37 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetAggregatorSuite.scala
@@ -67,7 +67,7 @@ object ComplexResultAgg extends Aggregator[(String, Int), (Long, Long), (Long, L
 }
 
 case class AggData(a: Int, b: String)
-object ClassInputAgg extends Aggregator[AggData, Int, Int] with Serializable {
+object ClassInputAgg extends Aggregator[AggData, Int, Int] {
   /** A zero value for this aggregation. Should satisfy the property that any b + zero = b */
   override def zero: Int = 0
 
@@ -88,6 +88,28 @@ object ClassInputAgg extends Aggregator[AggData, Int, Int] with Serializable {
   override def merge(b1: Int, b2: Int): Int = b1 + b2
 }
 
+object ComplexBufferAgg extends Aggregator[AggData, (Int, AggData), Int] {
+  /** A zero value for this aggregation. Should satisfy the property that any b + zero = b */
+  override def zero: (Int, AggData) = 0 -> AggData(0, "0")
+
+  /**
+   * Combine two values to produce a new value.  For performance, the function may modify `b` and
+   * return it instead of constructing new object for b.
+   */
+  override def reduce(b: (Int, AggData), a: AggData): (Int, AggData) = (b._1 + 1, a)
+
+  /**
+   * Transform the output of the reduction.
+   */
+  override def finish(reduction: (Int, AggData)): Int = reduction._1
+
+  /**
+   * Merge two intermediate values
+   */
+  override def merge(b1: (Int, AggData), b2: (Int, AggData)): (Int, AggData) =
+    (b1._1 + b2._1, b1._2)
+}
+
 class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
 
   import testImplicits._
@@ -168,4 +190,21 @@ class DatasetAggregatorSuite extends QueryTest with SharedSQLContext {
       ds.groupBy(_.b).agg(ClassInputAgg.toColumn),
       ("one", 1))
   }
+
+  test("typed aggregation: complex input") {
+    val ds = Seq(AggData(1, "one"), AggData(2, "two")).toDS()
+
+    checkAnswer(
+      ds.select(ComplexBufferAgg.toColumn),
+      2
+    )
+
+    checkAnswer(
+      ds.select(expr("avg(a)").as[Double], ComplexBufferAgg.toColumn),
+      (1.5, 2))
+
+    checkAnswer(
+      ds.groupBy(_.b).agg(ComplexBufferAgg.toColumn),
+      ("one", 1), ("two", 1))
+  }
 }

From 5fd86e4fc2e06d2403ca538ae417580c93b69e06 Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Mon, 23 Nov 2015 10:41:17 -0800
Subject: [PATCH 413/518] [SPARK-7173][YARN] Add label expression support for
 application master

Add label expression support for AM to restrict it runs on the specific set of nodes. I tested it locally and works fine.

sryza and vanzin please help to review, thanks a lot.

Author: jerryshao <sshao@hortonworks.com>

Closes #9800 from jerryshao/SPARK-7173.
---
 docs/running-on-yarn.md                       |  9 +++++++
 .../org/apache/spark/deploy/yarn/Client.scala | 26 ++++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/docs/running-on-yarn.md b/docs/running-on-yarn.md
index db6bfa69ee0fe..925a1e0ba6fcf 100644
--- a/docs/running-on-yarn.md
+++ b/docs/running-on-yarn.md
@@ -326,6 +326,15 @@ If you need a reference to the proper location to put log files in the YARN so t
   Otherwise, the client process will exit after submission.
   </td>
 </tr>
+<tr>
+  <td><code>spark.yarn.am.nodeLabelExpression</code></td>
+  <td>(none)</td>
+  <td>
+  A YARN node label expression that restricts the set of nodes AM will be scheduled on.
+  Only versions of YARN greater than or equal to 2.6 support node label expressions, so when
+  running against earlier versions, this property will be ignored.
+  </td>
+</tr>
 <tr>
   <td><code>spark.yarn.executor.nodeLabelExpression</code></td>
   <td>(none)</td>
diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
index ba799884f5689..a77a3e2420e24 100644
--- a/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
+++ b/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
@@ -225,7 +225,31 @@ private[spark] class Client(
     val capability = Records.newRecord(classOf[Resource])
     capability.setMemory(args.amMemory + amMemoryOverhead)
     capability.setVirtualCores(args.amCores)
-    appContext.setResource(capability)
+
+    if (sparkConf.contains("spark.yarn.am.nodeLabelExpression")) {
+      try {
+        val amRequest = Records.newRecord(classOf[ResourceRequest])
+        amRequest.setResourceName(ResourceRequest.ANY)
+        amRequest.setPriority(Priority.newInstance(0))
+        amRequest.setCapability(capability)
+        amRequest.setNumContainers(1)
+        val amLabelExpression = sparkConf.get("spark.yarn.am.nodeLabelExpression")
+        val method = amRequest.getClass.getMethod("setNodeLabelExpression", classOf[String])
+        method.invoke(amRequest, amLabelExpression)
+
+        val setResourceRequestMethod =
+          appContext.getClass.getMethod("setAMContainerResourceRequest", classOf[ResourceRequest])
+        setResourceRequestMethod.invoke(appContext, amRequest)
+      } catch {
+        case e: NoSuchMethodException =>
+          logWarning("Ignoring spark.yarn.am.nodeLabelExpression because the version " +
+            "of YARN does not support it")
+          appContext.setResource(capability)
+      }
+    } else {
+      appContext.setResource(capability)
+    }
+
     appContext
   }
 

From 5231cd5acaae69d735ba3209531705cc222f3cfb Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Mon, 23 Nov 2015 10:45:23 -0800
Subject: [PATCH 414/518] [SPARK-11762][NETWORK] Account for active streams
 when couting outstanding requests.

This way the timeout handling code can correctly close "hung" channels that are
processing streams.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9747 from vanzin/SPARK-11762.
---
 .../network/client/StreamInterceptor.java     | 12 ++++++++-
 .../client/TransportResponseHandler.java      | 15 +++++++++--
 .../TransportResponseHandlerSuite.java        | 27 +++++++++++++++++++
 3 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/network/common/src/main/java/org/apache/spark/network/client/StreamInterceptor.java b/network/common/src/main/java/org/apache/spark/network/client/StreamInterceptor.java
index 02230a00e69fc..88ba3ccebdf20 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/StreamInterceptor.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/StreamInterceptor.java
@@ -30,13 +30,19 @@
  */
 class StreamInterceptor implements TransportFrameDecoder.Interceptor {
 
+  private final TransportResponseHandler handler;
   private final String streamId;
   private final long byteCount;
   private final StreamCallback callback;
 
   private volatile long bytesRead;
 
-  StreamInterceptor(String streamId, long byteCount, StreamCallback callback) {
+  StreamInterceptor(
+      TransportResponseHandler handler,
+      String streamId,
+      long byteCount,
+      StreamCallback callback) {
+    this.handler = handler;
     this.streamId = streamId;
     this.byteCount = byteCount;
     this.callback = callback;
@@ -45,11 +51,13 @@ class StreamInterceptor implements TransportFrameDecoder.Interceptor {
 
   @Override
   public void exceptionCaught(Throwable cause) throws Exception {
+    handler.deactivateStream();
     callback.onFailure(streamId, cause);
   }
 
   @Override
   public void channelInactive() throws Exception {
+    handler.deactivateStream();
     callback.onFailure(streamId, new ClosedChannelException());
   }
 
@@ -65,8 +73,10 @@ public boolean handle(ByteBuf buf) throws Exception {
       RuntimeException re = new IllegalStateException(String.format(
         "Read too many bytes? Expected %d, but read %d.", byteCount, bytesRead));
       callback.onFailure(streamId, re);
+      handler.deactivateStream();
       throw re;
     } else if (bytesRead == byteCount) {
+      handler.deactivateStream();
       callback.onComplete(streamId);
     }
 
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
index ed3f36af58048..cc88991b588c1 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
@@ -24,6 +24,7 @@
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicLong;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.netty.channel.Channel;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -56,6 +57,7 @@ public class TransportResponseHandler extends MessageHandler<ResponseMessage> {
   private final Map<Long, RpcResponseCallback> outstandingRpcs;
 
   private final Queue<StreamCallback> streamCallbacks;
+  private volatile boolean streamActive;
 
   /** Records the time (in system nanoseconds) that the last fetch or RPC request was sent. */
   private final AtomicLong timeOfLastRequestNs;
@@ -87,9 +89,15 @@ public void removeRpcRequest(long requestId) {
   }
 
   public void addStreamCallback(StreamCallback callback) {
+    timeOfLastRequestNs.set(System.nanoTime());
     streamCallbacks.offer(callback);
   }
 
+  @VisibleForTesting
+  public void deactivateStream() {
+    streamActive = false;
+  }
+
   /**
    * Fire the failure callback for all outstanding requests. This is called when we have an
    * uncaught exception or pre-mature connection termination.
@@ -177,14 +185,16 @@ public void handle(ResponseMessage message) {
       StreamResponse resp = (StreamResponse) message;
       StreamCallback callback = streamCallbacks.poll();
       if (callback != null) {
-        StreamInterceptor interceptor = new StreamInterceptor(resp.streamId, resp.byteCount,
+        StreamInterceptor interceptor = new StreamInterceptor(this, resp.streamId, resp.byteCount,
           callback);
         try {
           TransportFrameDecoder frameDecoder = (TransportFrameDecoder)
             channel.pipeline().get(TransportFrameDecoder.HANDLER_NAME);
           frameDecoder.setInterceptor(interceptor);
+          streamActive = true;
         } catch (Exception e) {
           logger.error("Error installing stream handler.", e);
+          deactivateStream();
         }
       } else {
         logger.error("Could not find callback for StreamResponse.");
@@ -208,7 +218,8 @@ public void handle(ResponseMessage message) {
 
   /** Returns total number of outstanding requests (fetch requests + rpcs) */
   public int numOutstandingRequests() {
-    return outstandingFetches.size() + outstandingRpcs.size();
+    return outstandingFetches.size() + outstandingRpcs.size() + streamCallbacks.size() +
+      (streamActive ? 1 : 0);
   }
 
   /** Returns the time in nanoseconds of when the last request was sent out. */
diff --git a/network/common/src/test/java/org/apache/spark/network/TransportResponseHandlerSuite.java b/network/common/src/test/java/org/apache/spark/network/TransportResponseHandlerSuite.java
index 17a03ebe88a93..30144f4a9fc7a 100644
--- a/network/common/src/test/java/org/apache/spark/network/TransportResponseHandlerSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/TransportResponseHandlerSuite.java
@@ -17,6 +17,7 @@
 
 package org.apache.spark.network;
 
+import io.netty.channel.Channel;
 import io.netty.channel.local.LocalChannel;
 import org.junit.Test;
 
@@ -28,12 +29,16 @@
 import org.apache.spark.network.buffer.ManagedBuffer;
 import org.apache.spark.network.client.ChunkReceivedCallback;
 import org.apache.spark.network.client.RpcResponseCallback;
+import org.apache.spark.network.client.StreamCallback;
 import org.apache.spark.network.client.TransportResponseHandler;
 import org.apache.spark.network.protocol.ChunkFetchFailure;
 import org.apache.spark.network.protocol.ChunkFetchSuccess;
 import org.apache.spark.network.protocol.RpcFailure;
 import org.apache.spark.network.protocol.RpcResponse;
 import org.apache.spark.network.protocol.StreamChunkId;
+import org.apache.spark.network.protocol.StreamFailure;
+import org.apache.spark.network.protocol.StreamResponse;
+import org.apache.spark.network.util.TransportFrameDecoder;
 
 public class TransportResponseHandlerSuite {
   @Test
@@ -112,4 +117,26 @@ public void handleFailedRPC() {
     verify(callback, times(1)).onFailure((Throwable) any());
     assertEquals(0, handler.numOutstandingRequests());
   }
+
+  @Test
+  public void testActiveStreams() {
+    Channel c = new LocalChannel();
+    c.pipeline().addLast(TransportFrameDecoder.HANDLER_NAME, new TransportFrameDecoder());
+    TransportResponseHandler handler = new TransportResponseHandler(c);
+
+    StreamResponse response = new StreamResponse("stream", 1234L, null);
+    StreamCallback cb = mock(StreamCallback.class);
+    handler.addStreamCallback(cb);
+    assertEquals(1, handler.numOutstandingRequests());
+    handler.handle(response);
+    assertEquals(1, handler.numOutstandingRequests());
+    handler.deactivateStream();
+    assertEquals(0, handler.numOutstandingRequests());
+
+    StreamFailure failure = new StreamFailure("stream", "uh-oh");
+    handler.addStreamCallback(cb);
+    assertEquals(1, handler.numOutstandingRequests());
+    handler.handle(failure);
+    assertEquals(0, handler.numOutstandingRequests());
+  }
 }

From 98d7ec7df4bb115dbd84cb9acd744b6c8abfebd5 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Mon, 23 Nov 2015 11:51:29 -0800
Subject: [PATCH 415/518] [SPARK-11920][ML][DOC] ML LinearRegression should use
 correct dataset in examples and user guide doc

ML ```LinearRegression``` use ```data/mllib/sample_libsvm_data.txt``` as dataset in examples and user guide doc, but it's actually classification dataset rather than regression dataset. We should use ```data/mllib/sample_linear_regression_data.txt``` instead.
The deeper causes is that ```LinearRegression``` with "normal" solver can not solve this dataset correctly, may be due to the ill condition and unreasonable label. This issue has been reported at [SPARK-11918](https://issues.apache.org/jira/browse/SPARK-11918).
It will confuse users if they run the example code but get exception, so we should make this change which can clearly illustrate the usage of ```LinearRegression``` algorithm.

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9905 from yanboliang/spark-11920.
---
 .../examples/ml/JavaLinearRegressionWithElasticNetExample.java | 2 +-
 .../src/main/python/ml/linear_regression_with_elastic_net.py   | 3 ++-
 .../examples/ml/LinearRegressionWithElasticNetExample.scala    | 3 ++-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java
index 593f8fb3e9fe9..4ad7676c8d32b 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaLinearRegressionWithElasticNetExample.java
@@ -37,7 +37,7 @@ public static void main(String[] args) {
     // $example on$
     // Load training data
     DataFrame training = sqlContext.read().format("libsvm")
-      .load("data/mllib/sample_libsvm_data.txt");
+      .load("data/mllib/sample_linear_regression_data.txt");
 
     LinearRegression lr = new LinearRegression()
       .setMaxIter(10)
diff --git a/examples/src/main/python/ml/linear_regression_with_elastic_net.py b/examples/src/main/python/ml/linear_regression_with_elastic_net.py
index b0278276330c3..a4cd40cf26726 100644
--- a/examples/src/main/python/ml/linear_regression_with_elastic_net.py
+++ b/examples/src/main/python/ml/linear_regression_with_elastic_net.py
@@ -29,7 +29,8 @@
 
     # $example on$
     # Load training data
-    training = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+    training = sqlContext.read.format("libsvm")\
+        .load("data/mllib/sample_linear_regression_data.txt")
 
     lr = LinearRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8)
 
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala
index 5a51ece6f9ba7..22c824cea84d3 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/LinearRegressionWithElasticNetExample.scala
@@ -33,7 +33,8 @@ object LinearRegressionWithElasticNetExample {
 
     // $example on$
     // Load training data
-    val training = sqlCtx.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
+    val training = sqlCtx.read.format("libsvm")
+      .load("data/mllib/sample_linear_regression_data.txt")
 
     val lr = new LinearRegression()
       .setMaxIter(10)

From f6dcc6e96ad3f88563d717d5b6c45394b44db747 Mon Sep 17 00:00:00 2001
From: Mortada Mehyar <mortada.mehyar@gmail.com>
Date: Mon, 23 Nov 2015 12:03:15 -0800
Subject: [PATCH 416/518] [SPARK-11837][EC2] python3 compatibility for
 launching ec2 m3 instances

this currently breaks for python3 because `string` module doesn't have `letters` anymore, instead `ascii_letters` should be used

Author: Mortada Mehyar <mortada.mehyar@gmail.com>

Closes #9797 from mortada/python3_fix.
---
 ec2/spark_ec2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ec2/spark_ec2.py b/ec2/spark_ec2.py
index 9327e21e43db7..9fd652a3df4c4 100755
--- a/ec2/spark_ec2.py
+++ b/ec2/spark_ec2.py
@@ -595,7 +595,7 @@ def launch_cluster(conn, opts, cluster_name):
             dev = BlockDeviceType()
             dev.ephemeral_name = 'ephemeral%d' % i
             # The first ephemeral drive is /dev/sdb.
-            name = '/dev/sd' + string.letters[i + 1]
+            name = '/dev/sd' + string.ascii_letters[i + 1]
             block_map[name] = dev
 
     # Launch slaves

From 1b6e938be836786bac542fa430580248161e5403 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Mon, 23 Nov 2015 13:19:10 -0800
Subject: [PATCH 417/518] [SPARK-4424] Remove
 spark.driver.allowMultipleContexts override in tests

This patch removes `spark.driver.allowMultipleContexts=true` from our test configuration. The multiple SparkContexts check was originally disabled because certain tests suites in SQL needed to create multiple contexts. As far as I know, this configuration change is no longer necessary, so we should remove it in order to make it easier to find test cleanup bugs.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9865 from JoshRosen/SPARK-4424.
---
 pom.xml                  | 2 --
 project/SparkBuild.scala | 1 -
 2 files changed, 3 deletions(-)

diff --git a/pom.xml b/pom.xml
index ad849112ce76c..234fd5dea1a6e 100644
--- a/pom.xml
+++ b/pom.xml
@@ -1958,7 +1958,6 @@
               <spark.master.rest.enabled>false</spark.master.rest.enabled>
               <spark.ui.enabled>false</spark.ui.enabled>
               <spark.ui.showConsoleProgress>false</spark.ui.showConsoleProgress>
-              <spark.driver.allowMultipleContexts>true</spark.driver.allowMultipleContexts>
               <spark.unsafe.exceptionOnMemoryLeak>true</spark.unsafe.exceptionOnMemoryLeak>
               <!-- Needed by sql/hive tests. -->
               <test.src.tables>src</test.src.tables>
@@ -1997,7 +1996,6 @@
               <spark.testing>1</spark.testing>
               <spark.ui.enabled>false</spark.ui.enabled>
               <spark.ui.showConsoleProgress>false</spark.ui.showConsoleProgress>
-              <spark.driver.allowMultipleContexts>true</spark.driver.allowMultipleContexts>
               <spark.unsafe.exceptionOnMemoryLeak>true</spark.unsafe.exceptionOnMemoryLeak>
               <!-- Needed by sql/hive tests. -->
               <test.src.tables>__not_used__</test.src.tables>
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index 67724c4e9e411..f575f0012d59e 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -632,7 +632,6 @@ object TestSettings {
     javaOptions in Test += "-Dspark.master.rest.enabled=false",
     javaOptions in Test += "-Dspark.ui.enabled=false",
     javaOptions in Test += "-Dspark.ui.showConsoleProgress=false",
-    javaOptions in Test += "-Dspark.driver.allowMultipleContexts=true",
     javaOptions in Test += "-Dspark.unsafe.exceptionOnMemoryLeak=true",
     javaOptions in Test += "-Dsun.io.serialization.extendedDebugInfo=true",
     javaOptions in Test += "-Dderby.system.durability=test",

From 1d9120201012213edb1971a09e0849336dbb9415 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Mon, 23 Nov 2015 13:44:30 -0800
Subject: [PATCH 418/518] [SPARK-11836][SQL] udf/cast should not create new
 SQLContext

They should use the existing SQLContext.

Author: Davies Liu <davies@databricks.com>

Closes #9914 from davies/create_udf.
---
 python/pyspark/sql/column.py    | 7 ++++---
 python/pyspark/sql/functions.py | 7 ++++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/python/pyspark/sql/column.py b/python/pyspark/sql/column.py
index 9ca8e1f264cfa..81fd4e782628a 100644
--- a/python/pyspark/sql/column.py
+++ b/python/pyspark/sql/column.py
@@ -346,9 +346,10 @@ def cast(self, dataType):
         if isinstance(dataType, basestring):
             jc = self._jc.cast(dataType)
         elif isinstance(dataType, DataType):
-            sc = SparkContext._active_spark_context
-            ssql_ctx = sc._jvm.SQLContext(sc._jsc.sc())
-            jdt = ssql_ctx.parseDataType(dataType.json())
+            from pyspark.sql import SQLContext
+            sc = SparkContext.getOrCreate()
+            ctx = SQLContext.getOrCreate(sc)
+            jdt = ctx._ssql_ctx.parseDataType(dataType.json())
             jc = self._jc.cast(jdt)
         else:
             raise TypeError("unexpected type: %s" % type(dataType))
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index c3da513c13897..a1ca723bbd7ab 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -1457,14 +1457,15 @@ def __init__(self, func, returnType, name=None):
         self._judf = self._create_judf(name)
 
     def _create_judf(self, name):
+        from pyspark.sql import SQLContext
         f, returnType = self.func, self.returnType  # put them in closure `func`
         func = lambda _, it: map(lambda x: returnType.toInternal(f(*x)), it)
         ser = AutoBatchedSerializer(PickleSerializer())
         command = (func, None, ser, ser)
-        sc = SparkContext._active_spark_context
+        sc = SparkContext.getOrCreate()
         pickled_command, broadcast_vars, env, includes = _prepare_for_python_RDD(sc, command, self)
-        ssql_ctx = sc._jvm.SQLContext(sc._jsc.sc())
-        jdt = ssql_ctx.parseDataType(self.returnType.json())
+        ctx = SQLContext.getOrCreate(sc)
+        jdt = ctx._ssql_ctx.parseDataType(self.returnType.json())
         if name is None:
             name = f.__name__ if hasattr(f, '__name__') else f.__class__.__name__
         judf = sc._jvm.UserDefinedPythonFunction(name, bytearray(pickled_command), env, includes,

From 242be7daed9b01d19794bb2cf1ac421fe5ab7262 Mon Sep 17 00:00:00 2001
From: Luciano Resende <lresende@apache.org>
Date: Mon, 23 Nov 2015 13:46:34 -0800
Subject: [PATCH 419/518] [SPARK-11910][STREAMING][DOCS] Update twitter4j
 dependency version

Author: Luciano Resende <lresende@apache.org>

Closes #9892 from lresende/SPARK-11910.
---
 docs/streaming-programming-guide.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/streaming-programming-guide.md b/docs/streaming-programming-guide.md
index 96b36b7a73209..ed6b28c282135 100644
--- a/docs/streaming-programming-guide.md
+++ b/docs/streaming-programming-guide.md
@@ -723,7 +723,7 @@ Some of these advanced sources are as follows.
 
 - **Kinesis:** Spark Streaming {{site.SPARK_VERSION_SHORT}} is compatible with Kinesis Client Library 1.2.1. See the [Kinesis Integration Guide](streaming-kinesis-integration.html) for more details.
 
-- **Twitter:** Spark Streaming's TwitterUtils uses Twitter4j 3.0.3 to get the public stream of tweets using
+- **Twitter:** Spark Streaming's TwitterUtils uses Twitter4j to get the public stream of tweets using
   [Twitter's Streaming API](https://dev.twitter.com/docs/streaming-apis). Authentication information
   can be provided by any of the [methods](http://twitter4j.org/en/configuration.html) supported by
   Twitter4J library. You can either get the public stream, or get the filtered stream based on a

From 7cfa4c6bc36d97e459d4adee7b03d537d63c337e Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Mon, 23 Nov 2015 13:51:43 -0800
Subject: [PATCH 420/518] [SPARK-11865][NETWORK] Avoid returning inactive
 client in TransportClientFactory.

There's a very narrow race here where it would be possible for the timeout handler
to close a channel after the client factory verified that the channel was still
active. This change makes sure the client is marked as being recently in use so
that the timeout handler does not close it until a new timeout cycle elapses.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9853 from vanzin/SPARK-11865.
---
 .../spark/network/client/TransportClient.java |  9 ++++-
 .../client/TransportClientFactory.java        | 15 ++++++--
 .../client/TransportResponseHandler.java      |  9 +++--
 .../server/TransportChannelHandler.java       | 36 ++++++++++++-------
 4 files changed, 52 insertions(+), 17 deletions(-)

diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java b/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
index a0ba223e340a2..876fcd846791c 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
@@ -73,10 +73,12 @@ public class TransportClient implements Closeable {
   private final Channel channel;
   private final TransportResponseHandler handler;
   @Nullable private String clientId;
+  private volatile boolean timedOut;
 
   public TransportClient(Channel channel, TransportResponseHandler handler) {
     this.channel = Preconditions.checkNotNull(channel);
     this.handler = Preconditions.checkNotNull(handler);
+    this.timedOut = false;
   }
 
   public Channel getChannel() {
@@ -84,7 +86,7 @@ public Channel getChannel() {
   }
 
   public boolean isActive() {
-    return channel.isOpen() || channel.isActive();
+    return !timedOut && (channel.isOpen() || channel.isActive());
   }
 
   public SocketAddress getSocketAddress() {
@@ -263,6 +265,11 @@ public void onFailure(Throwable e) {
     }
   }
 
+  /** Mark this channel as having timed out. */
+  public void timeOut() {
+    this.timedOut = true;
+  }
+
   @Override
   public void close() {
     // close is a local operation and should finish with milliseconds; timeout just to be safe
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java b/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
index 42a4f664e697c..659c47160c7be 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
@@ -136,8 +136,19 @@ public TransportClient createClient(String remoteHost, int remotePort) throws IO
     TransportClient cachedClient = clientPool.clients[clientIndex];
 
     if (cachedClient != null && cachedClient.isActive()) {
-      logger.trace("Returning cached connection to {}: {}", address, cachedClient);
-      return cachedClient;
+      // Make sure that the channel will not timeout by updating the last use time of the
+      // handler. Then check that the client is still alive, in case it timed out before
+      // this code was able to update things.
+      TransportChannelHandler handler = cachedClient.getChannel().pipeline()
+        .get(TransportChannelHandler.class);
+      synchronized (handler) {
+        handler.getResponseHandler().updateTimeOfLastRequest();
+      }
+
+      if (cachedClient.isActive()) {
+        logger.trace("Returning cached connection to {}: {}", address, cachedClient);
+        return cachedClient;
+      }
     }
 
     // If we reach here, we don't have an existing connection open. Let's create a new one.
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
index cc88991b588c1..be181e0660826 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
@@ -71,7 +71,7 @@ public TransportResponseHandler(Channel channel) {
   }
 
   public void addFetchRequest(StreamChunkId streamChunkId, ChunkReceivedCallback callback) {
-    timeOfLastRequestNs.set(System.nanoTime());
+    updateTimeOfLastRequest();
     outstandingFetches.put(streamChunkId, callback);
   }
 
@@ -80,7 +80,7 @@ public void removeFetchRequest(StreamChunkId streamChunkId) {
   }
 
   public void addRpcRequest(long requestId, RpcResponseCallback callback) {
-    timeOfLastRequestNs.set(System.nanoTime());
+    updateTimeOfLastRequest();
     outstandingRpcs.put(requestId, callback);
   }
 
@@ -227,4 +227,9 @@ public long getTimeOfLastRequestNs() {
     return timeOfLastRequestNs.get();
   }
 
+  /** Updates the time of the last request to the current system time. */
+  public void updateTimeOfLastRequest() {
+    timeOfLastRequestNs.set(System.nanoTime());
+  }
+
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java b/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
index f8fcd1c3d7d76..29d688a67578c 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
@@ -116,20 +116,32 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exc
       // there are outstanding requests, we also do a secondary consistency check to ensure
       // there's no race between the idle timeout and incrementing the numOutstandingRequests
       // (see SPARK-7003).
-      boolean isActuallyOverdue =
-        System.nanoTime() - responseHandler.getTimeOfLastRequestNs() > requestTimeoutNs;
-      if (e.state() == IdleState.ALL_IDLE && isActuallyOverdue) {
-        if (responseHandler.numOutstandingRequests() > 0) {
-          String address = NettyUtils.getRemoteAddress(ctx.channel());
-          logger.error("Connection to {} has been quiet for {} ms while there are outstanding " +
-            "requests. Assuming connection is dead; please adjust spark.network.timeout if this " +
-            "is wrong.", address, requestTimeoutNs / 1000 / 1000);
-          ctx.close();
-        } else if (closeIdleConnections) {
-          // While CloseIdleConnections is enable, we also close idle connection
-          ctx.close();
+      //
+      // To avoid a race between TransportClientFactory.createClient() and this code which could
+      // result in an inactive client being returned, this needs to run in a synchronized block.
+      synchronized (this) {
+        boolean isActuallyOverdue =
+          System.nanoTime() - responseHandler.getTimeOfLastRequestNs() > requestTimeoutNs;
+        if (e.state() == IdleState.ALL_IDLE && isActuallyOverdue) {
+          if (responseHandler.numOutstandingRequests() > 0) {
+            String address = NettyUtils.getRemoteAddress(ctx.channel());
+            logger.error("Connection to {} has been quiet for {} ms while there are outstanding " +
+              "requests. Assuming connection is dead; please adjust spark.network.timeout if this " +
+              "is wrong.", address, requestTimeoutNs / 1000 / 1000);
+            client.timeOut();
+            ctx.close();
+          } else if (closeIdleConnections) {
+            // While CloseIdleConnections is enable, we also close idle connection
+            client.timeOut();
+            ctx.close();
+          }
         }
       }
     }
   }
+
+  public TransportResponseHandler getResponseHandler() {
+    return responseHandler;
+  }
+
 }

From c2467dadae8ce44010a912ee91c429310f8add65 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Mon, 23 Nov 2015 13:54:19 -0800
Subject: [PATCH 421/518] [SPARK-11140][CORE] Transfer files using network lib
 when using NettyRpcEnv.

This change abstracts the code that serves jars / files to executors so that
each RpcEnv can have its own implementation; the akka version uses the existing
HTTP-based file serving mechanism, while the netty versions uses the new
stream support added to the network lib, which makes file transfers benefit
from the easier security configuration of the network library, and should also
reduce overhead overall.

The change includes a small fix to TransportChannelHandler so that it propagates
user events to downstream handlers.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9530 from vanzin/SPARK-11140.
---
 .../scala/org/apache/spark/SparkContext.scala |   8 +-
 .../scala/org/apache/spark/SparkEnv.scala     |  14 --
 .../scala/org/apache/spark/rpc/RpcEnv.scala   |  46 ++++++
 .../apache/spark/rpc/akka/AkkaRpcEnv.scala    |  60 +++++++-
 .../apache/spark/rpc/netty/NettyRpcEnv.scala  | 138 ++++++++++++++++--
 .../spark/rpc/netty/NettyStreamManager.scala  |  63 ++++++++
 .../scala/org/apache/spark/util/Utils.scala   |   9 ++
 .../org/apache/spark/rpc/RpcEnvSuite.scala    |  39 ++++-
 .../rpc/netty/NettyRpcHandlerSuite.scala      |  10 +-
 docs/configuration.md                         |   2 +
 docs/security.md                              |   5 +-
 .../launcher/AbstractCommandBuilder.java      |   2 +-
 .../client/TransportClientFactory.java        |   6 +-
 .../server/TransportChannelHandler.java       |   1 +
 14 files changed, 356 insertions(+), 47 deletions(-)
 create mode 100644 core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index af4456c05b0a1..b153a7b08e590 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -1379,7 +1379,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
     }
 
     val key = if (!isLocal && scheme == "file") {
-      env.httpFileServer.addFile(new File(uri.getPath))
+      env.rpcEnv.fileServer.addFile(new File(uri.getPath))
     } else {
       schemeCorrectedPath
     }
@@ -1630,7 +1630,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
       var key = ""
       if (path.contains("\\")) {
         // For local paths with backslashes on Windows, URI throws an exception
-        key = env.httpFileServer.addJar(new File(path))
+        key = env.rpcEnv.fileServer.addJar(new File(path))
       } else {
         val uri = new URI(path)
         key = uri.getScheme match {
@@ -1644,7 +1644,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
               // of the AM to make it show up in the current working directory.
               val fileName = new Path(uri.getPath).getName()
               try {
-                env.httpFileServer.addJar(new File(fileName))
+                env.rpcEnv.fileServer.addJar(new File(fileName))
               } catch {
                 case e: Exception =>
                   // For now just log an error but allow to go through so spark examples work.
@@ -1655,7 +1655,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
               }
             } else {
               try {
-                env.httpFileServer.addJar(new File(uri.getPath))
+                env.rpcEnv.fileServer.addJar(new File(uri.getPath))
               } catch {
                 case exc: FileNotFoundException =>
                   logError(s"Jar not found at $path")
diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala b/core/src/main/scala/org/apache/spark/SparkEnv.scala
index 88df27f733f2a..84230e32a4462 100644
--- a/core/src/main/scala/org/apache/spark/SparkEnv.scala
+++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala
@@ -66,7 +66,6 @@ class SparkEnv (
     val blockTransferService: BlockTransferService,
     val blockManager: BlockManager,
     val securityManager: SecurityManager,
-    val httpFileServer: HttpFileServer,
     val sparkFilesDir: String,
     val metricsSystem: MetricsSystem,
     val memoryManager: MemoryManager,
@@ -91,7 +90,6 @@ class SparkEnv (
     if (!isStopped) {
       isStopped = true
       pythonWorkers.values.foreach(_.stop())
-      Option(httpFileServer).foreach(_.stop())
       mapOutputTracker.stop()
       shuffleManager.stop()
       broadcastManager.stop()
@@ -367,17 +365,6 @@ object SparkEnv extends Logging {
 
     val cacheManager = new CacheManager(blockManager)
 
-    val httpFileServer =
-      if (isDriver) {
-        val fileServerPort = conf.getInt("spark.fileserver.port", 0)
-        val server = new HttpFileServer(conf, securityManager, fileServerPort)
-        server.initialize()
-        conf.set("spark.fileserver.uri", server.serverUri)
-        server
-      } else {
-        null
-      }
-
     val metricsSystem = if (isDriver) {
       // Don't start metrics system right now for Driver.
       // We need to wait for the task scheduler to give us an app ID.
@@ -422,7 +409,6 @@ object SparkEnv extends Logging {
       blockTransferService,
       blockManager,
       securityManager,
-      httpFileServer,
       sparkFilesDir,
       metricsSystem,
       memoryManager,
diff --git a/core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala
index a560fd10cdf76..3d7d281b0dd66 100644
--- a/core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala
@@ -17,6 +17,9 @@
 
 package org.apache.spark.rpc
 
+import java.io.File
+import java.nio.channels.ReadableByteChannel
+
 import scala.concurrent.Future
 
 import org.apache.spark.{SecurityManager, SparkConf}
@@ -132,8 +135,51 @@ private[spark] abstract class RpcEnv(conf: SparkConf) {
    * that contains [[RpcEndpointRef]]s, the deserialization codes should be wrapped by this method.
    */
   def deserialize[T](deserializationAction: () => T): T
+
+  /**
+   * Return the instance of the file server used to serve files. This may be `null` if the
+   * RpcEnv is not operating in server mode.
+   */
+  def fileServer: RpcEnvFileServer
+
+  /**
+   * Open a channel to download a file from the given URI. If the URIs returned by the
+   * RpcEnvFileServer use the "spark" scheme, this method will be called by the Utils class to
+   * retrieve the files.
+   *
+   * @param uri URI with location of the file.
+   */
+  def openChannel(uri: String): ReadableByteChannel
+
 }
 
+/**
+ * A server used by the RpcEnv to server files to other processes owned by the application.
+ *
+ * The file server can return URIs handled by common libraries (such as "http" or "hdfs"), or
+ * it can return "spark" URIs which will be handled by `RpcEnv#fetchFile`.
+ */
+private[spark] trait RpcEnvFileServer {
+
+  /**
+   * Adds a file to be served by this RpcEnv. This is used to serve files from the driver
+   * to executors when they're stored on the driver's local file system.
+   *
+   * @param file Local file to serve.
+   * @return A URI for the location of the file.
+   */
+  def addFile(file: File): String
+
+  /**
+   * Adds a jar to be served by this RpcEnv. Similar to `addFile` but for jars added using
+   * `SparkContext.addJar`.
+   *
+   * @param file Local file to serve.
+   * @return A URI for the location of the file.
+   */
+  def addJar(file: File): String
+
+}
 
 private[spark] case class RpcEnvConfig(
     conf: SparkConf,
diff --git a/core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala
index 059a7e10ec12f..94dbec593c315 100644
--- a/core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/akka/AkkaRpcEnv.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.rpc.akka
 
+import java.io.File
+import java.nio.channels.ReadableByteChannel
 import java.util.concurrent.ConcurrentHashMap
 
 import scala.concurrent.Future
@@ -30,7 +32,7 @@ import akka.pattern.{ask => akkaAsk}
 import akka.remote.{AssociationEvent, AssociatedEvent, DisassociatedEvent, AssociationErrorEvent}
 import akka.serialization.JavaSerializer
 
-import org.apache.spark.{SparkException, Logging, SparkConf}
+import org.apache.spark.{HttpFileServer, Logging, SecurityManager, SparkConf, SparkException}
 import org.apache.spark.rpc._
 import org.apache.spark.util.{ActorLogReceive, AkkaUtils, ThreadUtils}
 
@@ -41,7 +43,10 @@ import org.apache.spark.util.{ActorLogReceive, AkkaUtils, ThreadUtils}
  * remove Akka from the dependencies.
  */
 private[spark] class AkkaRpcEnv private[akka] (
-    val actorSystem: ActorSystem, conf: SparkConf, boundPort: Int)
+    val actorSystem: ActorSystem,
+    val securityManager: SecurityManager,
+    conf: SparkConf,
+    boundPort: Int)
   extends RpcEnv(conf) with Logging {
 
   private val defaultAddress: RpcAddress = {
@@ -64,6 +69,8 @@ private[spark] class AkkaRpcEnv private[akka] (
    */
   private val refToEndpoint = new ConcurrentHashMap[RpcEndpointRef, RpcEndpoint]()
 
+  private val _fileServer = new AkkaFileServer(conf, securityManager)
+
   private def registerEndpoint(endpoint: RpcEndpoint, endpointRef: RpcEndpointRef): Unit = {
     endpointToRef.put(endpoint, endpointRef)
     refToEndpoint.put(endpointRef, endpoint)
@@ -223,6 +230,7 @@ private[spark] class AkkaRpcEnv private[akka] (
 
   override def shutdown(): Unit = {
     actorSystem.shutdown()
+    _fileServer.shutdown()
   }
 
   override def stop(endpoint: RpcEndpointRef): Unit = {
@@ -241,6 +249,52 @@ private[spark] class AkkaRpcEnv private[akka] (
       deserializationAction()
     }
   }
+
+  override def openChannel(uri: String): ReadableByteChannel = {
+    throw new UnsupportedOperationException(
+      "AkkaRpcEnv's files should be retrieved using an HTTP client.")
+  }
+
+  override def fileServer: RpcEnvFileServer = _fileServer
+
+}
+
+private[akka] class AkkaFileServer(
+    conf: SparkConf,
+    securityManager: SecurityManager) extends RpcEnvFileServer {
+
+  @volatile private var httpFileServer: HttpFileServer = _
+
+  override def addFile(file: File): String = {
+    getFileServer().addFile(file)
+  }
+
+  override def addJar(file: File): String = {
+    getFileServer().addJar(file)
+  }
+
+  def shutdown(): Unit = {
+    if (httpFileServer != null) {
+      httpFileServer.stop()
+    }
+  }
+
+  private def getFileServer(): HttpFileServer = {
+    if (httpFileServer == null) synchronized {
+      if (httpFileServer == null) {
+        httpFileServer = startFileServer()
+      }
+    }
+    httpFileServer
+  }
+
+  private def startFileServer(): HttpFileServer = {
+    val fileServerPort = conf.getInt("spark.fileserver.port", 0)
+    val server = new HttpFileServer(conf, securityManager, fileServerPort)
+    server.initialize()
+    server
+  }
+
 }
 
 private[spark] class AkkaRpcEnvFactory extends RpcEnvFactory {
@@ -249,7 +303,7 @@ private[spark] class AkkaRpcEnvFactory extends RpcEnvFactory {
     val (actorSystem, boundPort) = AkkaUtils.createActorSystem(
       config.name, config.host, config.port, config.conf, config.securityManager)
     actorSystem.actorOf(Props(classOf[ErrorMonitor]), "ErrorMonitor")
-    new AkkaRpcEnv(actorSystem, config.conf, boundPort)
+    new AkkaRpcEnv(actorSystem, config.securityManager, config.conf, boundPort)
   }
 }
 
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
index 3ce359868039b..68701f609f77a 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
@@ -20,6 +20,7 @@ import java.io._
 import java.lang.{Boolean => JBoolean}
 import java.net.{InetSocketAddress, URI}
 import java.nio.ByteBuffer
+import java.nio.channels.{Pipe, ReadableByteChannel, WritableByteChannel}
 import java.util.concurrent._
 import java.util.concurrent.atomic.AtomicBoolean
 import javax.annotation.Nullable
@@ -45,27 +46,39 @@ private[netty] class NettyRpcEnv(
     host: String,
     securityManager: SecurityManager) extends RpcEnv(conf) with Logging {
 
-  private val transportConf = SparkTransportConf.fromSparkConf(
+  private[netty] val transportConf = SparkTransportConf.fromSparkConf(
     conf.clone.set("spark.rpc.io.numConnectionsPerPeer", "1"),
     "rpc",
     conf.getInt("spark.rpc.io.threads", 0))
 
   private val dispatcher: Dispatcher = new Dispatcher(this)
 
+  private val streamManager = new NettyStreamManager(this)
+
   private val transportContext = new TransportContext(transportConf,
-    new NettyRpcHandler(dispatcher, this))
+    new NettyRpcHandler(dispatcher, this, streamManager))
 
-  private val clientFactory = {
-    val bootstraps: java.util.List[TransportClientBootstrap] =
-      if (securityManager.isAuthenticationEnabled()) {
-        java.util.Arrays.asList(new SaslClientBootstrap(transportConf, "", securityManager,
-          securityManager.isSaslEncryptionEnabled()))
-      } else {
-        java.util.Collections.emptyList[TransportClientBootstrap]
-      }
-    transportContext.createClientFactory(bootstraps)
+  private def createClientBootstraps(): java.util.List[TransportClientBootstrap] = {
+    if (securityManager.isAuthenticationEnabled()) {
+      java.util.Arrays.asList(new SaslClientBootstrap(transportConf, "", securityManager,
+        securityManager.isSaslEncryptionEnabled()))
+    } else {
+      java.util.Collections.emptyList[TransportClientBootstrap]
+    }
   }
 
+  private val clientFactory = transportContext.createClientFactory(createClientBootstraps())
+
+  /**
+   * A separate client factory for file downloads. This avoids using the same RPC handler as
+   * the main RPC context, so that events caused by these clients are kept isolated from the
+   * main RPC traffic.
+   *
+   * It also allows for different configuration of certain properties, such as the number of
+   * connections per peer.
+   */
+  @volatile private var fileDownloadFactory: TransportClientFactory = _
+
   val timeoutScheduler = ThreadUtils.newDaemonSingleThreadScheduledExecutor("netty-rpc-env-timeout")
 
   // Because TransportClientFactory.createClient is blocking, we need to run it in this thread pool
@@ -292,6 +305,9 @@ private[netty] class NettyRpcEnv(
     if (clientConnectionExecutor != null) {
       clientConnectionExecutor.shutdownNow()
     }
+    if (fileDownloadFactory != null) {
+      fileDownloadFactory.close()
+    }
   }
 
   override def deserialize[T](deserializationAction: () => T): T = {
@@ -300,6 +316,96 @@ private[netty] class NettyRpcEnv(
     }
   }
 
+  override def fileServer: RpcEnvFileServer = streamManager
+
+  override def openChannel(uri: String): ReadableByteChannel = {
+    val parsedUri = new URI(uri)
+    require(parsedUri.getHost() != null, "Host name must be defined.")
+    require(parsedUri.getPort() > 0, "Port must be defined.")
+    require(parsedUri.getPath() != null && parsedUri.getPath().nonEmpty, "Path must be defined.")
+
+    val pipe = Pipe.open()
+    val source = new FileDownloadChannel(pipe.source())
+    try {
+      val client = downloadClient(parsedUri.getHost(), parsedUri.getPort())
+      val callback = new FileDownloadCallback(pipe.sink(), source, client)
+      client.stream(parsedUri.getPath(), callback)
+    } catch {
+      case e: Exception =>
+        pipe.sink().close()
+        source.close()
+        throw e
+    }
+
+    source
+  }
+
+  private def downloadClient(host: String, port: Int): TransportClient = {
+    if (fileDownloadFactory == null) synchronized {
+      if (fileDownloadFactory == null) {
+        val module = "files"
+        val prefix = "spark.rpc.io."
+        val clone = conf.clone()
+
+        // Copy any RPC configuration that is not overridden in the spark.files namespace.
+        conf.getAll.foreach { case (key, value) =>
+          if (key.startsWith(prefix)) {
+            val opt = key.substring(prefix.length())
+            clone.setIfMissing(s"spark.$module.io.$opt", value)
+          }
+        }
+
+        val ioThreads = clone.getInt("spark.files.io.threads", 1)
+        val downloadConf = SparkTransportConf.fromSparkConf(clone, module, ioThreads)
+        val downloadContext = new TransportContext(downloadConf, new NoOpRpcHandler(), true)
+        fileDownloadFactory = downloadContext.createClientFactory(createClientBootstraps())
+      }
+    }
+    fileDownloadFactory.createClient(host, port)
+  }
+
+  private class FileDownloadChannel(source: ReadableByteChannel) extends ReadableByteChannel {
+
+    @volatile private var error: Throwable = _
+
+    def setError(e: Throwable): Unit = error = e
+
+    override def read(dst: ByteBuffer): Int = {
+      if (error != null) {
+        throw error
+      }
+      source.read(dst)
+    }
+
+    override def close(): Unit = source.close()
+
+    override def isOpen(): Boolean = source.isOpen()
+
+  }
+
+  private class FileDownloadCallback(
+      sink: WritableByteChannel,
+      source: FileDownloadChannel,
+      client: TransportClient) extends StreamCallback {
+
+    override def onData(streamId: String, buf: ByteBuffer): Unit = {
+      while (buf.remaining() > 0) {
+        sink.write(buf)
+      }
+    }
+
+    override def onComplete(streamId: String): Unit = {
+      sink.close()
+    }
+
+    override def onFailure(streamId: String, cause: Throwable): Unit = {
+      logError(s"Error downloading stream $streamId.", cause)
+      source.setError(cause)
+      sink.close()
+    }
+
+  }
+
 }
 
 private[netty] object NettyRpcEnv extends Logging {
@@ -420,7 +526,7 @@ private[netty] class NettyRpcEndpointRef(
 
   override def toString: String = s"NettyRpcEndpointRef(${_address})"
 
-  def toURI: URI = new URI(s"spark://${_address}")
+  def toURI: URI = new URI(_address.toString)
 
   final override def equals(that: Any): Boolean = that match {
     case other: NettyRpcEndpointRef => _address == other._address
@@ -471,7 +577,9 @@ private[netty] case class RpcFailure(e: Throwable)
  * with different `RpcAddress` information).
  */
 private[netty] class NettyRpcHandler(
-    dispatcher: Dispatcher, nettyEnv: NettyRpcEnv) extends RpcHandler with Logging {
+    dispatcher: Dispatcher,
+    nettyEnv: NettyRpcEnv,
+    streamManager: StreamManager) extends RpcHandler with Logging {
 
   // TODO: Can we add connection callback (channel registered) to the underlying framework?
   // A variable to track whether we should dispatch the RemoteProcessConnected message.
@@ -498,7 +606,7 @@ private[netty] class NettyRpcHandler(
     dispatcher.postRemoteMessage(messageToDispatch, callback)
   }
 
-  override def getStreamManager: StreamManager = new OneForOneStreamManager
+  override def getStreamManager: StreamManager = streamManager
 
   override def exceptionCaught(cause: Throwable, client: TransportClient): Unit = {
     val addr = client.getChannel.remoteAddress().asInstanceOf[InetSocketAddress]
@@ -516,8 +624,8 @@ private[netty] class NettyRpcHandler(
   override def connectionTerminated(client: TransportClient): Unit = {
     val addr = client.getChannel.remoteAddress().asInstanceOf[InetSocketAddress]
     if (addr != null) {
-      val clientAddr = RpcAddress(addr.getHostName, addr.getPort)
       clients.remove(client)
+      val clientAddr = RpcAddress(addr.getHostName, addr.getPort)
       nettyEnv.removeOutbox(clientAddr)
       dispatcher.postToAll(RemoteProcessDisconnected(clientAddr))
     } else {
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala
new file mode 100644
index 0000000000000..eb1d2604fb235
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.rpc.netty
+
+import java.io.File
+import java.util.concurrent.ConcurrentHashMap
+
+import org.apache.spark.network.buffer.{FileSegmentManagedBuffer, ManagedBuffer}
+import org.apache.spark.network.server.StreamManager
+import org.apache.spark.rpc.RpcEnvFileServer
+
+/**
+ * StreamManager implementation for serving files from a NettyRpcEnv.
+ */
+private[netty] class NettyStreamManager(rpcEnv: NettyRpcEnv)
+  extends StreamManager with RpcEnvFileServer {
+
+  private val files = new ConcurrentHashMap[String, File]()
+  private val jars = new ConcurrentHashMap[String, File]()
+
+  override def getChunk(streamId: Long, chunkIndex: Int): ManagedBuffer = {
+    throw new UnsupportedOperationException()
+  }
+
+  override def openStream(streamId: String): ManagedBuffer = {
+    val Array(ftype, fname) = streamId.stripPrefix("/").split("/", 2)
+    val file = ftype match {
+      case "files" => files.get(fname)
+      case "jars" => jars.get(fname)
+      case _ => throw new IllegalArgumentException(s"Invalid file type: $ftype")
+    }
+
+    require(file != null, s"File not found: $streamId")
+    new FileSegmentManagedBuffer(rpcEnv.transportConf, file, 0, file.length())
+  }
+
+  override def addFile(file: File): String = {
+    require(files.putIfAbsent(file.getName(), file) == null,
+      s"File ${file.getName()} already registered.")
+    s"${rpcEnv.address.toSparkURL}/files/${file.getName()}"
+  }
+
+  override def addJar(file: File): String = {
+    require(jars.putIfAbsent(file.getName(), file) == null,
+      s"JAR ${file.getName()} already registered.")
+    s"${rpcEnv.address.toSparkURL}/jars/${file.getName()}"
+  }
+
+}
diff --git a/core/src/main/scala/org/apache/spark/util/Utils.scala b/core/src/main/scala/org/apache/spark/util/Utils.scala
index 1b3acb8ef7f51..af632349c9cae 100644
--- a/core/src/main/scala/org/apache/spark/util/Utils.scala
+++ b/core/src/main/scala/org/apache/spark/util/Utils.scala
@@ -21,6 +21,7 @@ import java.io._
 import java.lang.management.ManagementFactory
 import java.net._
 import java.nio.ByteBuffer
+import java.nio.channels.Channels
 import java.util.concurrent._
 import java.util.{Locale, Properties, Random, UUID}
 import javax.net.ssl.HttpsURLConnection
@@ -535,6 +536,14 @@ private[spark] object Utils extends Logging {
     val uri = new URI(url)
     val fileOverwrite = conf.getBoolean("spark.files.overwrite", defaultValue = false)
     Option(uri.getScheme).getOrElse("file") match {
+      case "spark" =>
+        if (SparkEnv.get == null) {
+          throw new IllegalStateException(
+            "Cannot retrieve files with 'spark' scheme without an active SparkEnv.")
+        }
+        val source = SparkEnv.get.rpcEnv.openChannel(url)
+        val is = Channels.newInputStream(source)
+        downloadFile(url, is, targetFile, fileOverwrite)
       case "http" | "https" | "ftp" =>
         var uc: URLConnection = null
         if (securityMgr.isAuthenticationEnabled()) {
diff --git a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
index 2f55006420ce1..2b664c6313efa 100644
--- a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
@@ -17,7 +17,9 @@
 
 package org.apache.spark.rpc
 
-import java.io.NotSerializableException
+import java.io.{File, NotSerializableException}
+import java.util.UUID
+import java.nio.charset.StandardCharsets.UTF_8
 import java.util.concurrent.{TimeUnit, CountDownLatch, TimeoutException}
 
 import scala.collection.mutable
@@ -25,10 +27,14 @@ import scala.concurrent.Await
 import scala.concurrent.duration._
 import scala.language.postfixOps
 
+import com.google.common.io.Files
+import org.mockito.Mockito.{mock, when}
 import org.scalatest.BeforeAndAfterAll
 import org.scalatest.concurrent.Eventually._
 
-import org.apache.spark.{SparkConf, SparkException, SparkFunSuite}
+import org.apache.spark.{SecurityManager, SparkConf, SparkEnv, SparkException, SparkFunSuite}
+import org.apache.spark.deploy.SparkHadoopUtil
+import org.apache.spark.util.Utils
 
 /**
  * Common tests for an RpcEnv implementation.
@@ -40,12 +46,17 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
   override def beforeAll(): Unit = {
     val conf = new SparkConf()
     env = createRpcEnv(conf, "local", 0)
+
+    val sparkEnv = mock(classOf[SparkEnv])
+    when(sparkEnv.rpcEnv).thenReturn(env)
+    SparkEnv.set(sparkEnv)
   }
 
   override def afterAll(): Unit = {
     if (env != null) {
       env.shutdown()
     }
+    SparkEnv.set(null)
   }
 
   def createRpcEnv(conf: SparkConf, name: String, port: Int, clientMode: Boolean = false): RpcEnv
@@ -713,6 +724,30 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     assert(shortTimeout.timeoutProp.r.findAllIn(reply4).length === 1)
   }
 
+  test("file server") {
+    val conf = new SparkConf()
+    val tempDir = Utils.createTempDir()
+    val file = new File(tempDir, "file")
+    Files.write(UUID.randomUUID().toString(), file, UTF_8)
+    val jar = new File(tempDir, "jar")
+    Files.write(UUID.randomUUID().toString(), jar, UTF_8)
+
+    val fileUri = env.fileServer.addFile(file)
+    val jarUri = env.fileServer.addJar(jar)
+
+    val destDir = Utils.createTempDir()
+    val destFile = new File(destDir, file.getName())
+    val destJar = new File(destDir, jar.getName())
+
+    val sm = new SecurityManager(conf)
+    val hc = SparkHadoopUtil.get.conf
+    Utils.fetchFile(fileUri, destDir, conf, sm, hc, 0L, false)
+    Utils.fetchFile(jarUri, destDir, conf, sm, hc, 0L, false)
+
+    assert(Files.equal(file, destFile))
+    assert(Files.equal(jar, destJar))
+  }
+
 }
 
 class UnserializableClass
diff --git a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
index f9d8e80c98b66..ccca795683da3 100644
--- a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
@@ -25,17 +25,19 @@ import org.mockito.Matchers._
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.network.client.{TransportResponseHandler, TransportClient}
+import org.apache.spark.network.server.StreamManager
 import org.apache.spark.rpc._
 
 class NettyRpcHandlerSuite extends SparkFunSuite {
 
   val env = mock(classOf[NettyRpcEnv])
-  when(env.deserialize(any(classOf[TransportClient]), any(classOf[Array[Byte]]))(any())).
-    thenReturn(RequestMessage(RpcAddress("localhost", 12345), null, null, false))
+  val sm = mock(classOf[StreamManager])
+  when(env.deserialize(any(classOf[TransportClient]), any(classOf[Array[Byte]]))(any()))
+    .thenReturn(RequestMessage(RpcAddress("localhost", 12345), null, null, false))
 
   test("receive") {
     val dispatcher = mock(classOf[Dispatcher])
-    val nettyRpcHandler = new NettyRpcHandler(dispatcher, env)
+    val nettyRpcHandler = new NettyRpcHandler(dispatcher, env, sm)
 
     val channel = mock(classOf[Channel])
     val client = new TransportClient(channel, mock(classOf[TransportResponseHandler]))
@@ -47,7 +49,7 @@ class NettyRpcHandlerSuite extends SparkFunSuite {
 
   test("connectionTerminated") {
     val dispatcher = mock(classOf[Dispatcher])
-    val nettyRpcHandler = new NettyRpcHandler(dispatcher, env)
+    val nettyRpcHandler = new NettyRpcHandler(dispatcher, env, sm)
 
     val channel = mock(classOf[Channel])
     val client = new TransportClient(channel, mock(classOf[TransportResponseHandler]))
diff --git a/docs/configuration.md b/docs/configuration.md
index c496146e3ed63..4de202d7f7631 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -1020,6 +1020,7 @@ Apart from these, the following properties are also available, and may be useful
   <td>(random)</td>
   <td>
     Port for the executor to listen on. This is used for communicating with the driver.
+    This is only relevant when using the Akka RPC backend.
   </td>
 </tr>
 <tr>
@@ -1027,6 +1028,7 @@ Apart from these, the following properties are also available, and may be useful
   <td>(random)</td>
   <td>
     Port for the driver's HTTP file server to listen on.
+    This is only relevant when using the Akka RPC backend.
   </td>
 </tr>
 <tr>
diff --git a/docs/security.md b/docs/security.md
index 177109415180b..e1af221d446b0 100644
--- a/docs/security.md
+++ b/docs/security.md
@@ -149,7 +149,8 @@ configure those ports.
     <td>(random)</td>
     <td>Schedule tasks</td>
     <td><code>spark.executor.port</code></td>
-    <td>Akka-based. Set to "0" to choose a port randomly.</td>
+    <td>Akka-based. Set to "0" to choose a port randomly. Only used if Akka RPC backend is
+    configured.</td>
   </tr>
   <tr>
     <td>Executor</td>
@@ -157,7 +158,7 @@ configure those ports.
     <td>(random)</td>
     <td>File server for files and jars</td>
     <td><code>spark.fileserver.port</code></td>
-    <td>Jetty-based</td>
+    <td>Jetty-based. Only used if Akka RPC backend is configured.</td>
   </tr>
   <tr>
     <td>Executor</td>
diff --git a/launcher/src/main/java/org/apache/spark/launcher/AbstractCommandBuilder.java b/launcher/src/main/java/org/apache/spark/launcher/AbstractCommandBuilder.java
index 3ee6bd92e47fc..55fe156cf665f 100644
--- a/launcher/src/main/java/org/apache/spark/launcher/AbstractCommandBuilder.java
+++ b/launcher/src/main/java/org/apache/spark/launcher/AbstractCommandBuilder.java
@@ -148,7 +148,7 @@ List<String> buildClassPath(String appClassPath) throws IOException {
       String scala = getScalaVersion();
       List<String> projects = Arrays.asList("core", "repl", "mllib", "bagel", "graphx",
         "streaming", "tools", "sql/catalyst", "sql/core", "sql/hive", "sql/hive-thriftserver",
-        "yarn", "launcher");
+        "yarn", "launcher", "network/common", "network/shuffle", "network/yarn");
       if (prependClasses) {
         if (!isTesting) {
           System.err.println(
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java b/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
index 659c47160c7be..61bafc8380049 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
@@ -170,8 +170,10 @@ public TransportClient createClient(String remoteHost, int remotePort) throws IO
   }
 
   /**
-   * Create a completely new {@link TransportClient} to the given remote host / port
-   * But this connection is not pooled.
+   * Create a completely new {@link TransportClient} to the given remote host / port.
+   * This connection is not pooled.
+   *
+   * As with {@link #createClient(String, int)}, this method is blocking.
    */
   public TransportClient createUnmanagedClient(String remoteHost, int remotePort)
       throws IOException {
diff --git a/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java b/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
index 29d688a67578c..3164e00679035 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
@@ -138,6 +138,7 @@ public void userEventTriggered(ChannelHandlerContext ctx, Object evt) throws Exc
         }
       }
     }
+    ctx.fireUserEventTriggered(evt);
   }
 
   public TransportResponseHandler getResponseHandler() {

From 9db5f601facfdaba6e4333a6b2d2e4a9f009c788 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Mon, 23 Nov 2015 16:33:26 -0800
Subject: [PATCH 422/518] [SPARK-9866][SQL] Speed up VersionsSuite by using
 persistent Ivy cache

This patch attempts to speed up VersionsSuite by storing fetched Hive JARs in an Ivy cache that persists across tests runs. If `SPARK_VERSIONS_SUITE_IVY_PATH` is set, that path will be used for the cache; if it is not set, VersionsSuite will create a temporary Ivy cache which is deleted after the test completes.

Author: Josh Rosen <joshrosen@databricks.com>

Closes #9624 from JoshRosen/SPARK-9866.
---
 .../apache/spark/sql/hive/client/VersionsSuite.scala   | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala
index c6d034a23a1c6..7bc13bc60d30e 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala
@@ -36,10 +36,12 @@ import org.apache.spark.util.Utils
 @ExtendedHiveTest
 class VersionsSuite extends SparkFunSuite with Logging {
 
-  // Do not use a temp path here to speed up subsequent executions of the unit test during
-  // development.
-  private val ivyPath = Some(
-    new File(sys.props("java.io.tmpdir"), "hive-ivy-cache").getAbsolutePath())
+  // In order to speed up test execution during development or in Jenkins, you can specify the path
+  // of an existing Ivy cache:
+  private val ivyPath: Option[String] = {
+    sys.env.get("SPARK_VERSIONS_SUITE_IVY_PATH").orElse(
+      Some(new File(sys.props("java.io.tmpdir"), "hive-ivy-cache").getAbsolutePath))
+  }
 
   private def buildConf() = {
     lazy val warehousePath = Utils.createTempDir()

From 105745645b12afbbc2a350518cb5853a88944183 Mon Sep 17 00:00:00 2001
From: Bryan Cutler <bjcutler@us.ibm.com>
Date: Mon, 23 Nov 2015 17:11:51 -0800
Subject: [PATCH 423/518] [SPARK-10560][PYSPARK][MLLIB][DOCS] Make
 StreamingLogisticRegressionWithSGD Python API equal to Scala one

This is to bring the API documentation of StreamingLogisticReressionWithSGD and StreamingLinearRegressionWithSGC in line with the Scala versions.

-Fixed the algorithm descriptions
-Added default values to parameter descriptions
-Changed StreamingLogisticRegressionWithSGD regParam to default to 0, as in the Scala version

Author: Bryan Cutler <bjcutler@us.ibm.com>

Closes #9141 from BryanCutler/StreamingLogisticRegressionWithSGD-python-api-sync.
---
 python/pyspark/mllib/classification.py | 37 +++++++++++++++++---------
 python/pyspark/mllib/regression.py     | 32 ++++++++++++++--------
 2 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/python/pyspark/mllib/classification.py b/python/pyspark/mllib/classification.py
index aab4015ba80f8..9e6f17ef6e942 100644
--- a/python/pyspark/mllib/classification.py
+++ b/python/pyspark/mllib/classification.py
@@ -652,21 +652,34 @@ def train(cls, data, lambda_=1.0):
 @inherit_doc
 class StreamingLogisticRegressionWithSGD(StreamingLinearAlgorithm):
     """
-    Run LogisticRegression with SGD on a batch of data.
-
-    The weights obtained at the end of training a stream are used as initial
-    weights for the next batch.
-
-    :param stepSize: Step size for each iteration of gradient descent.
-    :param numIterations: Number of iterations run for each batch of data.
-    :param miniBatchFraction: Fraction of data on which SGD is run for each
-                              iteration.
-    :param regParam: L2 Regularization parameter.
-    :param convergenceTol: A condition which decides iteration termination.
+    Train or predict a logistic regression model on streaming data. Training uses
+    Stochastic Gradient Descent to update the model based on each new batch of
+    incoming data from a DStream.
+
+    Each batch of data is assumed to be an RDD of LabeledPoints.
+    The number of data points per batch can vary, but the number
+    of features must be constant. An initial weight
+    vector must be provided.
+
+    :param stepSize:
+      Step size for each iteration of gradient descent.
+      (default: 0.1)
+    :param numIterations:
+      Number of iterations run for each batch of data.
+      (default: 50)
+    :param miniBatchFraction:
+      Fraction of each batch of data to use for updates.
+      (default: 1.0)
+    :param regParam:
+      L2 Regularization parameter.
+      (default: 0.0)
+    :param convergenceTol:
+      Value used to determine when to terminate iterations.
+      (default: 0.001)
 
     .. versionadded:: 1.5.0
     """
-    def __init__(self, stepSize=0.1, numIterations=50, miniBatchFraction=1.0, regParam=0.01,
+    def __init__(self, stepSize=0.1, numIterations=50, miniBatchFraction=1.0, regParam=0.0,
                  convergenceTol=0.001):
         self.stepSize = stepSize
         self.numIterations = numIterations
diff --git a/python/pyspark/mllib/regression.py b/python/pyspark/mllib/regression.py
index 6f00d1df209c0..13b3397501c0b 100644
--- a/python/pyspark/mllib/regression.py
+++ b/python/pyspark/mllib/regression.py
@@ -734,17 +734,27 @@ def predictOnValues(self, dstream):
 @inherit_doc
 class StreamingLinearRegressionWithSGD(StreamingLinearAlgorithm):
     """
-    Run LinearRegression with SGD on a batch of data.
-
-    The problem minimized is (1 / n_samples) * (y - weights'X)**2.
-    After training on a batch of data, the weights obtained at the end of
-    training are used as initial weights for the next batch.
-
-    :param stepSize: Step size for each iteration of gradient descent.
-    :param numIterations: Total number of iterations run.
-    :param miniBatchFraction: Fraction of data on which SGD is run for each
-                              iteration.
-    :param convergenceTol: A condition which decides iteration termination.
+    Train or predict a linear regression model on streaming data. Training uses
+    Stochastic Gradient Descent to update the model based on each new batch of
+    incoming data from a DStream (see `LinearRegressionWithSGD` for model equation).
+
+    Each batch of data is assumed to be an RDD of LabeledPoints.
+    The number of data points per batch can vary, but the number
+    of features must be constant. An initial weight
+    vector must be provided.
+
+    :param stepSize:
+      Step size for each iteration of gradient descent.
+      (default: 0.1)
+    :param numIterations:
+      Number of iterations run for each batch of data.
+      (default: 50)
+    :param miniBatchFraction:
+      Fraction of each batch of data to use for updates.
+      (default: 1.0)
+    :param convergenceTol:
+      Value used to determine when to terminate iterations.
+      (default: 0.001)
 
     .. versionadded:: 1.5.0
     """

From 026ea2eab1f3cde270e8a6391d002915f3e1c6e5 Mon Sep 17 00:00:00 2001
From: Stephen Samuel <sam@sksamuel.com>
Date: Mon, 23 Nov 2015 19:52:12 -0800
Subject: [PATCH 424/518] Updated sql programming guide to include jdbc fetch
 size

Author: Stephen Samuel <sam@sksamuel.com>

Closes #9377 from sksamuel/master.
---
 docs/sql-programming-guide.md | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/docs/sql-programming-guide.md b/docs/sql-programming-guide.md
index e347754055e79..d7b205c2fa0df 100644
--- a/docs/sql-programming-guide.md
+++ b/docs/sql-programming-guide.md
@@ -1820,6 +1820,7 @@ the Data Sources API.  The following options are supported:
       register itself with the JDBC subsystem.
     </td>
   </tr>
+  
   <tr>
     <td><code>partitionColumn, lowerBound, upperBound, numPartitions</code></td>
     <td>
@@ -1831,6 +1832,13 @@ the Data Sources API.  The following options are supported:
       partitioned and returned.
     </td>
   </tr>
+  
+  <tr>
+    <td><code>fetchSize</code></td>
+    <td>
+      The JDBC fetch size, which determines how many rows to fetch per round trip. This can help performance on JDBC drivers which default to low fetch size (eg. Oracle with 10 rows).
+    </td>
+  </tr>
 </table>
 
 <div class="codetabs">

From 8d57524662fad4a0760f3bc924e690c2a110e7f7 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Mon, 23 Nov 2015 22:22:15 -0800
Subject: [PATCH 425/518] [SPARK-11933][SQL] Rename mapGroup -> mapGroups and
 flatMapGroup -> flatMapGroups.

Based on feedback from Matei, this is more consistent with mapPartitions in Spark.

Also addresses some of the cleanups from a previous commit that renames the type variables.

Author: Reynold Xin <rxin@databricks.com>

Closes #9919 from rxin/SPARK-11933.
---
 ...nction.java => FlatMapGroupsFunction.java} |  2 +-
 ...upFunction.java => MapGroupsFunction.java} |  2 +-
 .../org/apache/spark/sql/GroupedDataset.scala | 36 +++++++++----------
 .../apache/spark/sql/JavaDatasetSuite.java    | 10 +++---
 .../spark/sql/DatasetPrimitiveSuite.scala     |  4 +--
 .../org/apache/spark/sql/DatasetSuite.scala   | 12 +++----
 6 files changed, 33 insertions(+), 33 deletions(-)
 rename core/src/main/java/org/apache/spark/api/java/function/{FlatMapGroupFunction.java => FlatMapGroupsFunction.java} (93%)
 rename core/src/main/java/org/apache/spark/api/java/function/{MapGroupFunction.java => MapGroupsFunction.java} (93%)

diff --git a/core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupFunction.java b/core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupsFunction.java
similarity index 93%
rename from core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupFunction.java
rename to core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupsFunction.java
index 18a2d733ca70d..d7a80e7b129b0 100644
--- a/core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupFunction.java
+++ b/core/src/main/java/org/apache/spark/api/java/function/FlatMapGroupsFunction.java
@@ -23,6 +23,6 @@
 /**
  * A function that returns zero or more output records from each grouping key and its values.
  */
-public interface FlatMapGroupFunction<K, V, R> extends Serializable {
+public interface FlatMapGroupsFunction<K, V, R> extends Serializable {
   Iterable<R> call(K key, Iterator<V> values) throws Exception;
 }
diff --git a/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java b/core/src/main/java/org/apache/spark/api/java/function/MapGroupsFunction.java
similarity index 93%
rename from core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java
rename to core/src/main/java/org/apache/spark/api/java/function/MapGroupsFunction.java
index 4f3f222e064bb..faa59eabc8b4f 100644
--- a/core/src/main/java/org/apache/spark/api/java/function/MapGroupFunction.java
+++ b/core/src/main/java/org/apache/spark/api/java/function/MapGroupsFunction.java
@@ -23,6 +23,6 @@
 /**
  * Base interface for a map function used in GroupedDataset's mapGroup function.
  */
-public interface MapGroupFunction<K, V, R> extends Serializable {
+public interface MapGroupsFunction<K, V, R> extends Serializable {
   R call(K key, Iterator<V> values) throws Exception;
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 7f43ce16901b9..793a86b132907 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -43,7 +43,7 @@ import org.apache.spark.sql.expressions.Aggregator
 @Experimental
 class GroupedDataset[K, V] private[sql](
     kEncoder: Encoder[K],
-    tEncoder: Encoder[V],
+    vEncoder: Encoder[V],
     val queryExecution: QueryExecution,
     private val dataAttributes: Seq[Attribute],
     private val groupingAttributes: Seq[Attribute]) extends Serializable {
@@ -53,12 +53,12 @@ class GroupedDataset[K, V] private[sql](
   // queryexecution.
 
   private implicit val unresolvedKEncoder = encoderFor(kEncoder)
-  private implicit val unresolvedTEncoder = encoderFor(tEncoder)
+  private implicit val unresolvedVEncoder = encoderFor(vEncoder)
 
   private val resolvedKEncoder =
     unresolvedKEncoder.resolve(groupingAttributes, OuterScopes.outerScopes)
-  private val resolvedTEncoder =
-    unresolvedTEncoder.resolve(dataAttributes, OuterScopes.outerScopes)
+  private val resolvedVEncoder =
+    unresolvedVEncoder.resolve(dataAttributes, OuterScopes.outerScopes)
 
   private def logicalPlan = queryExecution.analyzed
   private def sqlContext = queryExecution.sqlContext
@@ -76,7 +76,7 @@ class GroupedDataset[K, V] private[sql](
   def keyAs[L : Encoder]: GroupedDataset[L, V] =
     new GroupedDataset(
       encoderFor[L],
-      unresolvedTEncoder,
+      unresolvedVEncoder,
       queryExecution,
       dataAttributes,
       groupingAttributes)
@@ -110,13 +110,13 @@ class GroupedDataset[K, V] private[sql](
    *
    * @since 1.6.0
    */
-  def flatMapGroup[U : Encoder](f: (K, Iterator[V]) => TraversableOnce[U]): Dataset[U] = {
+  def flatMapGroups[U : Encoder](f: (K, Iterator[V]) => TraversableOnce[U]): Dataset[U] = {
     new Dataset[U](
       sqlContext,
       MapGroups(
         f,
         resolvedKEncoder,
-        resolvedTEncoder,
+        resolvedVEncoder,
         groupingAttributes,
         logicalPlan))
   }
@@ -138,8 +138,8 @@ class GroupedDataset[K, V] private[sql](
    *
    * @since 1.6.0
    */
-  def flatMapGroup[U](f: FlatMapGroupFunction[K, V, U], encoder: Encoder[U]): Dataset[U] = {
-    flatMapGroup((key, data) => f.call(key, data.asJava).asScala)(encoder)
+  def flatMapGroups[U](f: FlatMapGroupsFunction[K, V, U], encoder: Encoder[U]): Dataset[U] = {
+    flatMapGroups((key, data) => f.call(key, data.asJava).asScala)(encoder)
   }
 
   /**
@@ -158,9 +158,9 @@ class GroupedDataset[K, V] private[sql](
    *
    * @since 1.6.0
    */
-  def mapGroup[U : Encoder](f: (K, Iterator[V]) => U): Dataset[U] = {
+  def mapGroups[U : Encoder](f: (K, Iterator[V]) => U): Dataset[U] = {
     val func = (key: K, it: Iterator[V]) => Iterator(f(key, it))
-    flatMapGroup(func)
+    flatMapGroups(func)
   }
 
   /**
@@ -179,8 +179,8 @@ class GroupedDataset[K, V] private[sql](
    *
    * @since 1.6.0
    */
-  def mapGroup[U](f: MapGroupFunction[K, V, U], encoder: Encoder[U]): Dataset[U] = {
-    mapGroup((key, data) => f.call(key, data.asJava))(encoder)
+  def mapGroups[U](f: MapGroupsFunction[K, V, U], encoder: Encoder[U]): Dataset[U] = {
+    mapGroups((key, data) => f.call(key, data.asJava))(encoder)
   }
 
   /**
@@ -192,8 +192,8 @@ class GroupedDataset[K, V] private[sql](
   def reduce(f: (V, V) => V): Dataset[(K, V)] = {
     val func = (key: K, it: Iterator[V]) => Iterator((key, it.reduce(f)))
 
-    implicit val resultEncoder = ExpressionEncoder.tuple(unresolvedKEncoder, unresolvedTEncoder)
-    flatMapGroup(func)
+    implicit val resultEncoder = ExpressionEncoder.tuple(unresolvedKEncoder, unresolvedVEncoder)
+    flatMapGroups(func)
   }
 
   /**
@@ -213,7 +213,7 @@ class GroupedDataset[K, V] private[sql](
 
   private def withEncoder(c: Column): Column = c match {
     case tc: TypedColumn[_, _] =>
-      tc.withInputType(resolvedTEncoder.bind(dataAttributes), dataAttributes)
+      tc.withInputType(resolvedVEncoder.bind(dataAttributes), dataAttributes)
     case _ => c
   }
 
@@ -227,7 +227,7 @@ class GroupedDataset[K, V] private[sql](
     val encoders = columns.map(_.encoder)
     val namedColumns =
       columns.map(
-        _.withInputType(resolvedTEncoder, dataAttributes).named)
+        _.withInputType(resolvedVEncoder, dataAttributes).named)
     val keyColumn = if (groupingAttributes.length > 1) {
       Alias(CreateStruct(groupingAttributes), "key")()
     } else {
@@ -304,7 +304,7 @@ class GroupedDataset[K, V] private[sql](
   def cogroup[U, R : Encoder](
       other: GroupedDataset[K, U])(
       f: (K, Iterator[V], Iterator[U]) => TraversableOnce[R]): Dataset[R] = {
-    implicit def uEnc: Encoder[U] = other.unresolvedTEncoder
+    implicit def uEnc: Encoder[U] = other.unresolvedVEncoder
     new Dataset[R](
       sqlContext,
       CoGroup(
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
index cf335efdd23b8..67a3190cb7d4f 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDatasetSuite.java
@@ -170,7 +170,7 @@ public Integer call(String v) throws Exception {
       }
     }, Encoders.INT());
 
-    Dataset<String> mapped = grouped.mapGroup(new MapGroupFunction<Integer, String, String>() {
+    Dataset<String> mapped = grouped.mapGroups(new MapGroupsFunction<Integer, String, String>() {
       @Override
       public String call(Integer key, Iterator<String> values) throws Exception {
         StringBuilder sb = new StringBuilder(key.toString());
@@ -183,8 +183,8 @@ public String call(Integer key, Iterator<String> values) throws Exception {
 
     Assert.assertEquals(Arrays.asList("1a", "3foobar"), mapped.collectAsList());
 
-    Dataset<String> flatMapped = grouped.flatMapGroup(
-      new FlatMapGroupFunction<Integer, String, String>() {
+    Dataset<String> flatMapped = grouped.flatMapGroups(
+      new FlatMapGroupsFunction<Integer, String, String>() {
         @Override
         public Iterable<String> call(Integer key, Iterator<String> values) throws Exception {
           StringBuilder sb = new StringBuilder(key.toString());
@@ -249,8 +249,8 @@ public void testGroupByColumn() {
     GroupedDataset<Integer, String> grouped =
       ds.groupBy(length(col("value"))).keyAs(Encoders.INT());
 
-    Dataset<String> mapped = grouped.mapGroup(
-      new MapGroupFunction<Integer, String, String>() {
+    Dataset<String> mapped = grouped.mapGroups(
+      new MapGroupsFunction<Integer, String, String>() {
         @Override
         public String call(Integer key, Iterator<String> data) throws Exception {
           StringBuilder sb = new StringBuilder(key.toString());
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
index d387710357be0..f75d0961823c4 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetPrimitiveSuite.scala
@@ -86,7 +86,7 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
   test("groupBy function, map") {
     val ds = Seq(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11).toDS()
     val grouped = ds.groupBy(_ % 2)
-    val agged = grouped.mapGroup { case (g, iter) =>
+    val agged = grouped.mapGroups { case (g, iter) =>
       val name = if (g == 0) "even" else "odd"
       (name, iter.size)
     }
@@ -99,7 +99,7 @@ class DatasetPrimitiveSuite extends QueryTest with SharedSQLContext {
   test("groupBy function, flatMap") {
     val ds = Seq("a", "b", "c", "xyz", "hello").toDS()
     val grouped = ds.groupBy(_.length)
-    val agged = grouped.flatMapGroup { case (g, iter) => Iterator(g.toString, iter.mkString) }
+    val agged = grouped.flatMapGroups { case (g, iter) => Iterator(g.toString, iter.mkString) }
 
     checkAnswer(
       agged,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index cc8e4325fd2f5..dbdd7ba14a5b7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -224,7 +224,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy function, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy(v => (v._1, "word"))
-    val agged = grouped.mapGroup { case (g, iter) => (g._1, iter.map(_._2).sum) }
+    val agged = grouped.mapGroups { case (g, iter) => (g._1, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
@@ -234,7 +234,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy function, flatMap") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy(v => (v._1, "word"))
-    val agged = grouped.flatMapGroup { case (g, iter) =>
+    val agged = grouped.flatMapGroups { case (g, iter) =>
       Iterator(g._1, iter.map(_._2).sum.toString)
     }
 
@@ -255,7 +255,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy columns, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1")
-    val agged = grouped.mapGroup { case (g, iter) => (g.getString(0), iter.map(_._2).sum) }
+    val agged = grouped.mapGroups { case (g, iter) => (g.getString(0), iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
@@ -265,7 +265,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy columns asKey, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1").keyAs[String]
-    val agged = grouped.mapGroup { case (g, iter) => (g, iter.map(_._2).sum) }
+    val agged = grouped.mapGroups { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
@@ -275,7 +275,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy columns asKey tuple, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1", lit(1)).keyAs[(String, Int)]
-    val agged = grouped.mapGroup { case (g, iter) => (g, iter.map(_._2).sum) }
+    val agged = grouped.mapGroups { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,
@@ -285,7 +285,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
   test("groupBy columns asKey class, map") {
     val ds = Seq(("a", 10), ("a", 20), ("b", 1), ("b", 2), ("c", 1)).toDS()
     val grouped = ds.groupBy($"_1".as("a"), lit(1).as("b")).keyAs[ClassData]
-    val agged = grouped.mapGroup { case (g, iter) => (g, iter.map(_._2).sum) }
+    val agged = grouped.mapGroups { case (g, iter) => (g, iter.map(_._2).sum) }
 
     checkAnswer(
       agged,

From 6cf51a7007bd72eb93ade149ca9fc53be5b32a17 Mon Sep 17 00:00:00 2001
From: Nicholas Chammas <nicholas.chammas@gmail.com>
Date: Mon, 23 Nov 2015 22:22:50 -0800
Subject: [PATCH 426/518] [SPARK-11903] Remove --skip-java-test

Per [pwendell's comments on SPARK-11903](https://issues.apache.org/jira/browse/SPARK-11903?focusedCommentId=15021511&page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel#comment-15021511) I'm removing this dead code.

If we are concerned about preserving compatibility, I can instead leave the option in and add a warning.

For example:

```sh
echo "Warning: '--skip-java-test' is deprecated and has no effect."
;;
```

cc pwendell, srowen

Author: Nicholas Chammas <nicholas.chammas@gmail.com>

Closes #9924 from nchammas/make-distribution.
---
 make-distribution.sh | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/make-distribution.sh b/make-distribution.sh
index d7d27e253f721..7b417fe7cf619 100755
--- a/make-distribution.sh
+++ b/make-distribution.sh
@@ -69,9 +69,6 @@ while (( "$#" )); do
       echo "Error: '--with-hive' is no longer supported, use Maven options -Phive and -Phive-thriftserver"
       exit_with_usage
       ;;
-    --skip-java-test)
-      SKIP_JAVA_TEST=true
-      ;;
     --with-tachyon)
       SPARK_TACHYON=true
       ;;

From 4021a28ac30b65cb61cf1e041253847253a2d89f Mon Sep 17 00:00:00 2001
From: Mikhail Bautin <mbautin@gmail.com>
Date: Mon, 23 Nov 2015 22:26:08 -0800
Subject: [PATCH 427/518] [SPARK-10707][SQL] Fix nullability computation in
 union output

Author: Mikhail Bautin <mbautin@gmail.com>

Closes #9308 from mbautin/SPARK-10707.
---
 .../plans/logical/basicOperators.scala        | 11 +++++--
 .../spark/sql/execution/basicOperators.scala  |  9 ++++--
 .../org/apache/spark/sql/SQLQuerySuite.scala  | 31 +++++++++++++++++++
 3 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index 0c444482c5e4c..737e62fd59214 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -92,8 +92,10 @@ case class Filter(condition: Expression, child: LogicalPlan) extends UnaryNode {
 }
 
 abstract class SetOperation(left: LogicalPlan, right: LogicalPlan) extends BinaryNode {
-  // TODO: These aren't really the same attributes as nullability etc might change.
-  final override def output: Seq[Attribute] = left.output
+  override def output: Seq[Attribute] =
+    left.output.zip(right.output).map { case (leftAttr, rightAttr) =>
+      leftAttr.withNullability(leftAttr.nullable || rightAttr.nullable)
+    }
 
   final override lazy val resolved: Boolean =
     childrenResolved &&
@@ -115,7 +117,10 @@ case class Union(left: LogicalPlan, right: LogicalPlan) extends SetOperation(lef
 
 case class Intersect(left: LogicalPlan, right: LogicalPlan) extends SetOperation(left, right)
 
-case class Except(left: LogicalPlan, right: LogicalPlan) extends SetOperation(left, right)
+case class Except(left: LogicalPlan, right: LogicalPlan) extends SetOperation(left, right) {
+  /** We don't use right.output because those rows get excluded from the set. */
+  override def output: Seq[Attribute] = left.output
+}
 
 case class Join(
   left: LogicalPlan,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index e79092efdaa3e..d57b8e7a9ed61 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -130,8 +130,13 @@ case class Sample(
  * Union two plans, without a distinct. This is UNION ALL in SQL.
  */
 case class Union(children: Seq[SparkPlan]) extends SparkPlan {
-  // TODO: attributes output by union should be distinct for nullability purposes
-  override def output: Seq[Attribute] = children.head.output
+  override def output: Seq[Attribute] = {
+    children.tail.foldLeft(children.head.output) { case (currentOutput, child) =>
+      currentOutput.zip(child.output).map { case (a1, a2) =>
+        a1.withNullability(a1.nullable || a2.nullable)
+      }
+    }
+  }
   override def outputsUnsafeRows: Boolean = children.forall(_.outputsUnsafeRows)
   override def canProcessUnsafeRows: Boolean = true
   override def canProcessSafeRows: Boolean = true
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
index 167aea87de077..bb82b562aaaa2 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala
@@ -1997,4 +1997,35 @@ class SQLQuerySuite extends QueryTest with SharedSQLContext {
     sqlContext.setConf("spark.sql.subexpressionElimination.enabled", "true")
     verifyCallCount(df.selectExpr("testUdf(a)", "testUdf(a)"), Row(1, 1), 1)
   }
+
+  test("SPARK-10707: nullability should be correctly propagated through set operations (1)") {
+    // This test produced an incorrect result of 1 before the SPARK-10707 fix because of the
+    // NullPropagation rule: COUNT(v) got replaced with COUNT(1) because the output column of
+    // UNION was incorrectly considered non-nullable:
+    checkAnswer(
+      sql("""SELECT count(v) FROM (
+            |  SELECT v FROM (
+            |    SELECT 'foo' AS v UNION ALL
+            |    SELECT NULL AS v
+            |  ) my_union WHERE isnull(v)
+            |) my_subview""".stripMargin),
+      Seq(Row(0)))
+  }
+
+  test("SPARK-10707: nullability should be correctly propagated through set operations (2)") {
+    // This test uses RAND() to stop column pruning for Union and checks the resulting isnull
+    // value. This would produce an incorrect result before the fix in SPARK-10707 because the "v"
+    // column of the union was considered non-nullable.
+    checkAnswer(
+      sql(
+        """
+          |SELECT a FROM (
+          |  SELECT ISNULL(v) AS a, RAND() FROM (
+          |    SELECT 'foo' AS v UNION ALL SELECT null AS v
+          |  ) my_union
+          |) my_view
+        """.stripMargin),
+      Row(false) :: Row(true) :: Nil)
+  }
+
 }

From 12eea834d7382fbaa9c92182b682b8724049d7c1 Mon Sep 17 00:00:00 2001
From: Xiu Guo <xguo27@gmail.com>
Date: Tue, 24 Nov 2015 00:07:40 -0800
Subject: [PATCH 428/518] [SPARK-11897][SQL] Add @scala.annotations.varargs to
 sql functions

Author: Xiu Guo <xguo27@gmail.com>

Closes #9918 from xguo27/SPARK-11897.
---
 sql/core/src/main/scala/org/apache/spark/sql/functions.scala | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index b27b1340cce46..6137ce3a70fdb 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -689,6 +689,7 @@ object functions extends LegacyFunctions {
    * @group normal_funcs
    * @since 1.4.0
    */
+  @scala.annotation.varargs
   def array(colName: String, colNames: String*): Column = {
     array((colName +: colNames).map(col) : _*)
   }
@@ -871,6 +872,7 @@ object functions extends LegacyFunctions {
    * @group normal_funcs
    * @since 1.4.0
    */
+  @scala.annotation.varargs
   def struct(colName: String, colNames: String*): Column = {
     struct((colName +: colNames).map(col) : _*)
   }

From 800bd799acf7f10a469d8d6537279953129eb2c6 Mon Sep 17 00:00:00 2001
From: Forest Fang <forest.fang@outlook.com>
Date: Tue, 24 Nov 2015 09:03:32 +0000
Subject: [PATCH 429/518] [SPARK-11906][WEB UI] Speculation Tasks Cause
 ProgressBar UI Overflow

When there are speculative tasks in the stage, running progress bar could overflow and goes hidden on a new line:
![image](https://cloud.githubusercontent.com/assets/4317392/11326841/5fd3482e-9142-11e5-8ca5-cb2f0c0c8964.png)
3 completed / 2 running (including 1 speculative) out of 4 total tasks

This is a simple fix by capping the started tasks at `total - completed` tasks
![image](https://cloud.githubusercontent.com/assets/4317392/11326842/6bb67260-9142-11e5-90f0-37f9174878ec.png)

I should note my preferred way to fix it is via css style
```css
.progress { display: flex; }
```
which shifts the correction burden from driver to web browser. However I couldn't get selenium test to measure the position/dimension of the progress bar correctly to get this unit tested.

It also has the side effect that the width will be calibrated so the running occupies 2 / 5 instead of 1 / 4.
![image](https://cloud.githubusercontent.com/assets/4317392/11326848/7b03e9f0-9142-11e5-89ad-bd99cb0647cf.png)

All in all, since this cosmetic bug is minor enough, I suppose the original simple fix should be good enough.

Author: Forest Fang <forest.fang@outlook.com>

Closes #9896 from saurfang/progressbar.
---
 core/src/main/scala/org/apache/spark/ui/UIUtils.scala  |  4 +++-
 .../test/scala/org/apache/spark/ui/UIUtilsSuite.scala  | 10 ++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/ui/UIUtils.scala b/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
index 25dcb604d9e5f..84a1116a5c498 100644
--- a/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
+++ b/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
@@ -319,7 +319,9 @@ private[spark] object UIUtils extends Logging {
       skipped: Int,
       total: Int): Seq[Node] = {
     val completeWidth = "width: %s%%".format((completed.toDouble/total)*100)
-    val startWidth = "width: %s%%".format((started.toDouble/total)*100)
+    // started + completed can be > total when there are speculative tasks
+    val boundedStarted = math.min(started, total - completed)
+    val startWidth = "width: %s%%".format((boundedStarted.toDouble/total)*100)
 
     <div class="progress">
       <span style="text-align:center; position:absolute; width:100%; left:0;">
diff --git a/core/src/test/scala/org/apache/spark/ui/UIUtilsSuite.scala b/core/src/test/scala/org/apache/spark/ui/UIUtilsSuite.scala
index 2b693c165180f..dd8d5ec27f87e 100644
--- a/core/src/test/scala/org/apache/spark/ui/UIUtilsSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ui/UIUtilsSuite.scala
@@ -57,6 +57,16 @@ class UIUtilsSuite extends SparkFunSuite {
     )
   }
 
+  test("SPARK-11906: Progress bar should not overflow because of speculative tasks") {
+    val generated = makeProgressBar(2, 3, 0, 0, 4).head.child.filter(_.label == "div")
+    val expected = Seq(
+      <div class="bar bar-completed" style="width: 75.0%"></div>,
+      <div class="bar bar-running" style="width: 25.0%"></div>
+    )
+    assert(generated.sameElements(expected),
+      s"\nRunning progress bar should round down\n\nExpected:\n$expected\nGenerated:\n$generated")
+  }
+
   private def verify(
       desc: String, expected: Elem, errorMsg: String = "", baseUrl: String = ""): Unit = {
     val generated = makeDescription(desc, baseUrl)

From d4a5e6f719079639ffd38470f4d8d1f6fde3228d Mon Sep 17 00:00:00 2001
From: huangzhaowei <carlmartinmax@gmail.com>
Date: Tue, 24 Nov 2015 23:24:49 +0800
Subject: [PATCH 430/518] [SPARK-11043][SQL] BugFix:Set the operator log in the
 thrift server.

`SessionManager` will set the `operationLog` if the configuration `hive.server2.logging.operation.enabled` is true in version of hive 1.2.1.
But the spark did not adapt to this change, so no matter enabled the configuration or not, spark thrift server will always log the warn message.
PS: if `hive.server2.logging.operation.enabled` is false, it should log the warn message (the same as hive thrift server).

Author: huangzhaowei <carlmartinmax@gmail.com>

Closes #9056 from SaintBacchus/SPARK-11043.
---
 .../SparkExecuteStatementOperation.scala         |  8 ++++----
 .../thriftserver/SparkSQLSessionManager.scala    |  5 +++++
 .../thriftserver/HiveThriftServer2Suites.scala   | 16 +++++++++++++++-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala
index 82fef92dcb73b..e022ee86a763a 100644
--- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala
+++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkExecuteStatementOperation.scala
@@ -134,12 +134,12 @@ private[hive] class SparkExecuteStatementOperation(
 
   def getResultSetSchema: TableSchema = resultSchema
 
-  override def run(): Unit = {
+  override def runInternal(): Unit = {
     setState(OperationState.PENDING)
     setHasResultSet(true) // avoid no resultset for async run
 
     if (!runInBackground) {
-      runInternal()
+      execute()
     } else {
       val sparkServiceUGI = Utils.getUGI()
 
@@ -151,7 +151,7 @@ private[hive] class SparkExecuteStatementOperation(
           val doAsAction = new PrivilegedExceptionAction[Unit]() {
             override def run(): Unit = {
               try {
-                runInternal()
+                execute()
               } catch {
                 case e: HiveSQLException =>
                   setOperationException(e)
@@ -188,7 +188,7 @@ private[hive] class SparkExecuteStatementOperation(
     }
   }
 
-  override def runInternal(): Unit = {
+  private def execute(): Unit = {
     statementId = UUID.randomUUID().toString
     logInfo(s"Running query '$statement' with $statementId")
     setState(OperationState.RUNNING)
diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLSessionManager.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLSessionManager.scala
index af4fcdf021bd4..de4e9c62b57a4 100644
--- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLSessionManager.scala
+++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLSessionManager.scala
@@ -41,6 +41,11 @@ private[hive] class SparkSQLSessionManager(hiveServer: HiveServer2, hiveContext:
   override def init(hiveConf: HiveConf) {
     setSuperField(this, "hiveConf", hiveConf)
 
+    // Create operation log root directory, if operation logging is enabled
+    if (hiveConf.getBoolVar(ConfVars.HIVE_SERVER2_LOGGING_OPERATION_ENABLED)) {
+      invoke(classOf[SessionManager], this, "initOperationLogRootDir")
+    }
+
     val backgroundPoolSize = hiveConf.getIntVar(ConfVars.HIVE_SERVER2_ASYNC_EXEC_THREADS)
     setSuperField(this, "backgroundOperationPool", Executors.newFixedThreadPool(backgroundPoolSize))
     getAncestorField[Log](this, 3, "LOG").info(
diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
index 1dd898aa38350..139d8e897ba1d 100644
--- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
+++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveThriftServer2Suites.scala
@@ -26,6 +26,7 @@ import scala.collection.mutable.ArrayBuffer
 import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.duration._
 import scala.concurrent.{Await, Promise, future}
+import scala.io.Source
 import scala.util.{Random, Try}
 
 import com.google.common.base.Charsets.UTF_8
@@ -507,6 +508,12 @@ class HiveThriftBinaryServerSuite extends HiveThriftJdbcTest {
       assert(rs2.getInt(2) === 500)
     }
   }
+
+  test("SPARK-11043 check operation log root directory") {
+    val expectedLine =
+      "Operation log root directory is created: " + operationLogPath.getAbsoluteFile
+    assert(Source.fromFile(logPath).getLines().exists(_.contains(expectedLine)))
+  }
 }
 
 class SingleSessionSuite extends HiveThriftJdbcTest {
@@ -642,7 +649,8 @@ abstract class HiveThriftServer2Test extends SparkFunSuite with BeforeAndAfterAl
   protected def metastoreJdbcUri = s"jdbc:derby:;databaseName=$metastorePath;create=true"
 
   private val pidDir: File = Utils.createTempDir("thriftserver-pid")
-  private var logPath: File = _
+  protected var logPath: File = _
+  protected var operationLogPath: File = _
   private var logTailingProcess: Process = _
   private var diagnosisBuffer: ArrayBuffer[String] = ArrayBuffer.empty[String]
 
@@ -679,6 +687,7 @@ abstract class HiveThriftServer2Test extends SparkFunSuite with BeforeAndAfterAl
        |  --hiveconf ${ConfVars.METASTOREWAREHOUSE}=$warehousePath
        |  --hiveconf ${ConfVars.HIVE_SERVER2_THRIFT_BIND_HOST}=localhost
        |  --hiveconf ${ConfVars.HIVE_SERVER2_TRANSPORT_MODE}=$mode
+       |  --hiveconf ${ConfVars.HIVE_SERVER2_LOGGING_OPERATION_LOG_LOCATION}=$operationLogPath
        |  --hiveconf $portConf=$port
        |  --driver-class-path $driverClassPath
        |  --driver-java-options -Dlog4j.debug
@@ -706,6 +715,8 @@ abstract class HiveThriftServer2Test extends SparkFunSuite with BeforeAndAfterAl
     warehousePath.delete()
     metastorePath = Utils.createTempDir()
     metastorePath.delete()
+    operationLogPath = Utils.createTempDir()
+    operationLogPath.delete()
     logPath = null
     logTailingProcess = null
 
@@ -782,6 +793,9 @@ abstract class HiveThriftServer2Test extends SparkFunSuite with BeforeAndAfterAl
     metastorePath.delete()
     metastorePath = null
 
+    operationLogPath.delete()
+    operationLogPath = null
+
     Option(logPath).foreach(_.delete())
     logPath = null
 

From 5889880fbe9628681042036892ef7ebd4f0857b4 Mon Sep 17 00:00:00 2001
From: Daoyuan Wang <daoyuan.wang@intel.com>
Date: Tue, 24 Nov 2015 23:32:05 +0800
Subject: [PATCH 431/518] [SPARK-11592][SQL] flush spark-sql command line
 history to history file

Currently, `spark-sql` would not flush command history when exiting.

Author: Daoyuan Wang <daoyuan.wang@intel.com>

Closes #9563 from adrian-wang/jline.
---
 .../hive/thriftserver/SparkSQLCLIDriver.scala    | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala
index 6419002a2aa89..4b928e600b355 100644
--- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala
+++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala
@@ -194,6 +194,22 @@ private[hive] object SparkSQLCLIDriver extends Logging {
         logWarning(e.getMessage)
     }
 
+    // add shutdown hook to flush the history to history file
+    Runtime.getRuntime.addShutdownHook(new Thread(new Runnable() {
+      override def run() = {
+        reader.getHistory match {
+          case h: FileHistory =>
+            try {
+              h.flush()
+            } catch {
+              case e: IOException =>
+                logWarning("WARNING: Failed to write command history file: " + e.getMessage)
+            }
+          case _ =>
+        }
+      }
+    }))
+
     // TODO: missing
 /*
     val clientTransportTSocketField = classOf[CliSessionState].getDeclaredField("transport")

From be9dd1550c1816559d3d418a19c692e715f1c94e Mon Sep 17 00:00:00 2001
From: Jungtaek Lim <kabhwan@gmail.com>
Date: Tue, 24 Nov 2015 09:20:09 -0800
Subject: [PATCH 432/518] =?UTF-8?q?[SPARK-11818][REPL]=20Fix=20ExecutorCla?=
 =?UTF-8?q?ssLoader=20to=20lookup=20resources=20from=20=E2=80=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…parent class loader

Without patch, two additional tests of ExecutorClassLoaderSuite fails.

- "resource from parent"
- "resources from parent"

Detailed explanation is here, https://issues.apache.org/jira/browse/SPARK-11818?focusedCommentId=15011202&page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel#comment-15011202

Author: Jungtaek Lim <kabhwan@gmail.com>

Closes #9812 from HeartSaVioR/SPARK-11818.
---
 .../spark/repl/ExecutorClassLoader.scala      | 12 +++++++-
 .../spark/repl/ExecutorClassLoaderSuite.scala | 29 +++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala b/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
index a976e96809cb8..a8859fcd4584b 100644
--- a/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
+++ b/repl/src/main/scala/org/apache/spark/repl/ExecutorClassLoader.scala
@@ -34,7 +34,9 @@ import org.apache.spark.util.ParentClassLoader
 /**
  * A ClassLoader that reads classes from a Hadoop FileSystem or HTTP URI,
  * used to load classes defined by the interpreter when the REPL is used.
- * Allows the user to specify if user class path should be first
+ * Allows the user to specify if user class path should be first.
+ * This class loader delegates getting/finding resources to parent loader,
+ * which makes sense until REPL never provide resource dynamically.
  */
 class ExecutorClassLoader(conf: SparkConf, classUri: String, parent: ClassLoader,
     userClassPathFirst: Boolean) extends ClassLoader with Logging {
@@ -55,6 +57,14 @@ class ExecutorClassLoader(conf: SparkConf, classUri: String, parent: ClassLoader
     }
   }
 
+  override def getResource(name: String): URL = {
+    parentLoader.getResource(name)
+  }
+
+  override def getResources(name: String): java.util.Enumeration[URL] = {
+    parentLoader.getResources(name)
+  }
+
   override def findClass(name: String): Class[_] = {
     userClassPathFirst match {
       case true => findClassLocally(name).getOrElse(parentLoader.loadClass(name))
diff --git a/repl/src/test/scala/org/apache/spark/repl/ExecutorClassLoaderSuite.scala b/repl/src/test/scala/org/apache/spark/repl/ExecutorClassLoaderSuite.scala
index a58eda12b1120..c1211f7596b9c 100644
--- a/repl/src/test/scala/org/apache/spark/repl/ExecutorClassLoaderSuite.scala
+++ b/repl/src/test/scala/org/apache/spark/repl/ExecutorClassLoaderSuite.scala
@@ -19,8 +19,13 @@ package org.apache.spark.repl
 
 import java.io.File
 import java.net.{URL, URLClassLoader}
+import java.nio.charset.StandardCharsets
+import java.util
+
+import com.google.common.io.Files
 
 import scala.concurrent.duration._
+import scala.io.Source
 import scala.language.implicitConversions
 import scala.language.postfixOps
 
@@ -41,6 +46,7 @@ class ExecutorClassLoaderSuite
 
   val childClassNames = List("ReplFakeClass1", "ReplFakeClass2")
   val parentClassNames = List("ReplFakeClass1", "ReplFakeClass2", "ReplFakeClass3")
+  val parentResourceNames = List("fake-resource.txt")
   var tempDir1: File = _
   var tempDir2: File = _
   var url1: String = _
@@ -54,6 +60,9 @@ class ExecutorClassLoaderSuite
     url1 = "file://" + tempDir1
     urls2 = List(tempDir2.toURI.toURL).toArray
     childClassNames.foreach(TestUtils.createCompiledClass(_, tempDir1, "1"))
+    parentResourceNames.foreach { x =>
+      Files.write("resource".getBytes(StandardCharsets.UTF_8), new File(tempDir2, x))
+    }
     parentClassNames.foreach(TestUtils.createCompiledClass(_, tempDir2, "2"))
   }
 
@@ -99,6 +108,26 @@ class ExecutorClassLoaderSuite
     }
   }
 
+  test("resource from parent") {
+    val parentLoader = new URLClassLoader(urls2, null)
+    val classLoader = new ExecutorClassLoader(new SparkConf(), url1, parentLoader, true)
+    val resourceName: String = parentResourceNames.head
+    val is = classLoader.getResourceAsStream(resourceName)
+    assert(is != null, s"Resource $resourceName not found")
+    val content = Source.fromInputStream(is, "UTF-8").getLines().next()
+    assert(content.contains("resource"), "File doesn't contain 'resource'")
+  }
+
+  test("resources from parent") {
+    val parentLoader = new URLClassLoader(urls2, null)
+    val classLoader = new ExecutorClassLoader(new SparkConf(), url1, parentLoader, true)
+    val resourceName: String = parentResourceNames.head
+    val resources: util.Enumeration[URL] = classLoader.getResources(resourceName)
+    assert(resources.hasMoreElements, s"Resource $resourceName not found")
+    val fileReader = Source.fromInputStream(resources.nextElement().openStream()).bufferedReader()
+    assert(fileReader.readLine().contains("resource"), "File doesn't contain 'resource'")
+  }
+
   test("failing to fetch classes from HTTP server should not leak resources (SPARK-6209)") {
     // This is a regression test for SPARK-6209, a bug where each failed attempt to load a class
     // from the driver's class server would leak a HTTP connection, causing the class server's

From e5aaae6e1145b8c25c4872b2992ab425da9c6f9b Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Tue, 24 Nov 2015 09:28:39 -0800
Subject: [PATCH 433/518] [SPARK-11942][SQL] fix encoder life cycle for CoGroup

we should pass in resolved encodera to logical `CoGroup` and bind them in physical `CoGroup`

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9928 from cloud-fan/cogroup.
---
 .../plans/logical/basicOperators.scala        | 27 ++++++++++---------
 .../org/apache/spark/sql/GroupedDataset.scala |  4 ++-
 .../spark/sql/execution/basicOperators.scala  | 20 +++++++-------
 .../org/apache/spark/sql/DatasetSuite.scala   | 12 +++++++++
 4 files changed, 41 insertions(+), 22 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
index 737e62fd59214..5665fd7e5f419 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala
@@ -553,19 +553,22 @@ case class MapGroups[K, T, U](
 
 /** Factory for constructing new `CoGroup` nodes. */
 object CoGroup {
-  def apply[K : Encoder, Left : Encoder, Right : Encoder, R : Encoder](
-      func: (K, Iterator[Left], Iterator[Right]) => TraversableOnce[R],
+  def apply[Key, Left, Right, Result : Encoder](
+      func: (Key, Iterator[Left], Iterator[Right]) => TraversableOnce[Result],
+      keyEnc: ExpressionEncoder[Key],
+      leftEnc: ExpressionEncoder[Left],
+      rightEnc: ExpressionEncoder[Right],
       leftGroup: Seq[Attribute],
       rightGroup: Seq[Attribute],
       left: LogicalPlan,
-      right: LogicalPlan): CoGroup[K, Left, Right, R] = {
+      right: LogicalPlan): CoGroup[Key, Left, Right, Result] = {
     CoGroup(
       func,
-      encoderFor[K],
-      encoderFor[Left],
-      encoderFor[Right],
-      encoderFor[R],
-      encoderFor[R].schema.toAttributes,
+      keyEnc,
+      leftEnc,
+      rightEnc,
+      encoderFor[Result],
+      encoderFor[Result].schema.toAttributes,
       leftGroup,
       rightGroup,
       left,
@@ -577,12 +580,12 @@ object CoGroup {
  * A relation produced by applying `func` to each grouping key and associated values from left and
  * right children.
  */
-case class CoGroup[K, Left, Right, R](
-    func: (K, Iterator[Left], Iterator[Right]) => TraversableOnce[R],
-    kEncoder: ExpressionEncoder[K],
+case class CoGroup[Key, Left, Right, Result](
+    func: (Key, Iterator[Left], Iterator[Right]) => TraversableOnce[Result],
+    keyEnc: ExpressionEncoder[Key],
     leftEnc: ExpressionEncoder[Left],
     rightEnc: ExpressionEncoder[Right],
-    rEncoder: ExpressionEncoder[R],
+    resultEnc: ExpressionEncoder[Result],
     output: Seq[Attribute],
     leftGroup: Seq[Attribute],
     rightGroup: Seq[Attribute],
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
index 793a86b132907..a10a89342fb5c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedDataset.scala
@@ -304,11 +304,13 @@ class GroupedDataset[K, V] private[sql](
   def cogroup[U, R : Encoder](
       other: GroupedDataset[K, U])(
       f: (K, Iterator[V], Iterator[U]) => TraversableOnce[R]): Dataset[R] = {
-    implicit def uEnc: Encoder[U] = other.unresolvedVEncoder
     new Dataset[R](
       sqlContext,
       CoGroup(
         f,
+        this.resolvedKEncoder,
+        this.resolvedVEncoder,
+        other.resolvedVEncoder,
         this.groupingAttributes,
         other.groupingAttributes,
         this.logicalPlan,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
index d57b8e7a9ed61..a42aea0b96d43 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -375,12 +375,12 @@ case class MapGroups[K, T, U](
  * iterators containing all elements in the group from left and right side.
  * The result of this function is encoded and flattened before being output.
  */
-case class CoGroup[K, Left, Right, R](
-    func: (K, Iterator[Left], Iterator[Right]) => TraversableOnce[R],
-    kEncoder: ExpressionEncoder[K],
+case class CoGroup[Key, Left, Right, Result](
+    func: (Key, Iterator[Left], Iterator[Right]) => TraversableOnce[Result],
+    keyEnc: ExpressionEncoder[Key],
     leftEnc: ExpressionEncoder[Left],
     rightEnc: ExpressionEncoder[Right],
-    rEncoder: ExpressionEncoder[R],
+    resultEnc: ExpressionEncoder[Result],
     output: Seq[Attribute],
     leftGroup: Seq[Attribute],
     rightGroup: Seq[Attribute],
@@ -397,15 +397,17 @@ case class CoGroup[K, Left, Right, R](
     left.execute().zipPartitions(right.execute()) { (leftData, rightData) =>
       val leftGrouped = GroupedIterator(leftData, leftGroup, left.output)
       val rightGrouped = GroupedIterator(rightData, rightGroup, right.output)
-      val groupKeyEncoder = kEncoder.bind(leftGroup)
+      val boundKeyEnc = keyEnc.bind(leftGroup)
+      val boundLeftEnc = leftEnc.bind(left.output)
+      val boundRightEnc = rightEnc.bind(right.output)
 
       new CoGroupedIterator(leftGrouped, rightGrouped, leftGroup).flatMap {
         case (key, leftResult, rightResult) =>
           val result = func(
-            groupKeyEncoder.fromRow(key),
-            leftResult.map(leftEnc.fromRow),
-            rightResult.map(rightEnc.fromRow))
-          result.map(rEncoder.toRow)
+            boundKeyEnc.fromRow(key),
+            leftResult.map(boundLeftEnc.fromRow),
+            rightResult.map(boundRightEnc.fromRow))
+          result.map(resultEnc.toRow)
       }
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index dbdd7ba14a5b7..13eede1b17d8b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -340,6 +340,18 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       1 -> "a#", 2 -> "#q", 3 -> "abcfoo#w", 5 -> "hello#er")
   }
 
+  test("cogroup with complex data") {
+    val ds1 = Seq(1 -> ClassData("a", 1), 2 -> ClassData("b", 2)).toDS()
+    val ds2 = Seq(2 -> ClassData("c", 3), 3 -> ClassData("d", 4)).toDS()
+    val cogrouped = ds1.groupBy(_._1).cogroup(ds2.groupBy(_._1)) { case (key, data1, data2) =>
+      Iterator(key -> (data1.map(_._2.a).mkString + data2.map(_._2.a).mkString))
+    }
+
+    checkAnswer(
+      cogrouped,
+      1 -> "a", 2 -> "bc", 3 -> "d")
+  }
+
   test("SPARK-11436: we should rebind right encoder when join 2 datasets") {
     val ds1 = Seq("1", "2").toDS().as("a")
     val ds2 = Seq(2, 3).toDS().as("b")

From 56a0aba0a60326ba026056c9a23f3f6ec7258c19 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Tue, 24 Nov 2015 09:52:53 -0800
Subject: [PATCH 434/518] [SPARK-11952][ML] Remove duplicate ml examples

Remove duplicate ml examples (only for ml).  mengxr

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9933 from yanboliang/SPARK-11685.
---
 .../main/python/ml/gradient_boosted_trees.py  | 82 -----------------
 .../src/main/python/ml/logistic_regression.py | 66 --------------
 .../main/python/ml/random_forest_example.py   | 87 -------------------
 3 files changed, 235 deletions(-)
 delete mode 100644 examples/src/main/python/ml/gradient_boosted_trees.py
 delete mode 100644 examples/src/main/python/ml/logistic_regression.py
 delete mode 100644 examples/src/main/python/ml/random_forest_example.py

diff --git a/examples/src/main/python/ml/gradient_boosted_trees.py b/examples/src/main/python/ml/gradient_boosted_trees.py
deleted file mode 100644
index c3bf8aa2eb1e6..0000000000000
--- a/examples/src/main/python/ml/gradient_boosted_trees.py
+++ /dev/null
@@ -1,82 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-from __future__ import print_function
-
-import sys
-
-from pyspark import SparkContext
-from pyspark.ml.classification import GBTClassifier
-from pyspark.ml.feature import StringIndexer
-from pyspark.ml.regression import GBTRegressor
-from pyspark.mllib.evaluation import BinaryClassificationMetrics, RegressionMetrics
-from pyspark.sql import Row, SQLContext
-
-"""
-A simple example demonstrating a Gradient Boosted Trees Classification/Regression Pipeline.
-Note: GBTClassifier only supports binary classification currently
-Run with:
-  bin/spark-submit examples/src/main/python/ml/gradient_boosted_trees.py
-"""
-
-
-def testClassification(train, test):
-    # Train a GradientBoostedTrees model.
-
-    rf = GBTClassifier(maxIter=30, maxDepth=4, labelCol="indexedLabel")
-
-    model = rf.fit(train)
-    predictionAndLabels = model.transform(test).select("prediction", "indexedLabel") \
-        .map(lambda x: (x.prediction, x.indexedLabel))
-
-    metrics = BinaryClassificationMetrics(predictionAndLabels)
-    print("AUC %.3f" % metrics.areaUnderROC)
-
-
-def testRegression(train, test):
-    # Train a GradientBoostedTrees model.
-
-    rf = GBTRegressor(maxIter=30, maxDepth=4, labelCol="indexedLabel")
-
-    model = rf.fit(train)
-    predictionAndLabels = model.transform(test).select("prediction", "indexedLabel") \
-        .map(lambda x: (x.prediction, x.indexedLabel))
-
-    metrics = RegressionMetrics(predictionAndLabels)
-    print("rmse %.3f" % metrics.rootMeanSquaredError)
-    print("r2 %.3f" % metrics.r2)
-    print("mae %.3f" % metrics.meanAbsoluteError)
-
-
-if __name__ == "__main__":
-    if len(sys.argv) > 1:
-        print("Usage: gradient_boosted_trees", file=sys.stderr)
-        exit(1)
-    sc = SparkContext(appName="PythonGBTExample")
-    sqlContext = SQLContext(sc)
-
-    # Load the data stored in LIBSVM format as a DataFrame.
-    df = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-    # Map labels into an indexed column of labels in [0, numLabels)
-    stringIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel")
-    si_model = stringIndexer.fit(df)
-    td = si_model.transform(df)
-    [train, test] = td.randomSplit([0.7, 0.3])
-    testClassification(train, test)
-    testRegression(train, test)
-    sc.stop()
diff --git a/examples/src/main/python/ml/logistic_regression.py b/examples/src/main/python/ml/logistic_regression.py
deleted file mode 100644
index 4cd027fdfbe8a..0000000000000
--- a/examples/src/main/python/ml/logistic_regression.py
+++ /dev/null
@@ -1,66 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-from __future__ import print_function
-
-import sys
-
-from pyspark import SparkContext
-from pyspark.ml.classification import LogisticRegression
-from pyspark.mllib.evaluation import MulticlassMetrics
-from pyspark.ml.feature import StringIndexer
-from pyspark.sql import SQLContext
-
-"""
-A simple example demonstrating a logistic regression with elastic net regularization Pipeline.
-Run with:
-  bin/spark-submit examples/src/main/python/ml/logistic_regression.py
-"""
-
-if __name__ == "__main__":
-
-    if len(sys.argv) > 1:
-        print("Usage: logistic_regression", file=sys.stderr)
-        exit(-1)
-
-    sc = SparkContext(appName="PythonLogisticRegressionExample")
-    sqlContext = SQLContext(sc)
-
-    # Load the data stored in LIBSVM format as a DataFrame.
-    df = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-    # Map labels into an indexed column of labels in [0, numLabels)
-    stringIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel")
-    si_model = stringIndexer.fit(df)
-    td = si_model.transform(df)
-    [training, test] = td.randomSplit([0.7, 0.3])
-
-    lr = LogisticRegression(maxIter=100, regParam=0.3).setLabelCol("indexedLabel")
-    lr.setElasticNetParam(0.8)
-
-    # Fit the model
-    lrModel = lr.fit(training)
-
-    predictionAndLabels = lrModel.transform(test).select("prediction", "indexedLabel") \
-        .map(lambda x: (x.prediction, x.indexedLabel))
-
-    metrics = MulticlassMetrics(predictionAndLabels)
-    print("weighted f-measure %.3f" % metrics.weightedFMeasure())
-    print("precision %s" % metrics.precision())
-    print("recall %s" % metrics.recall())
-
-    sc.stop()
diff --git a/examples/src/main/python/ml/random_forest_example.py b/examples/src/main/python/ml/random_forest_example.py
deleted file mode 100644
index dc6a778670193..0000000000000
--- a/examples/src/main/python/ml/random_forest_example.py
+++ /dev/null
@@ -1,87 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-from __future__ import print_function
-
-import sys
-
-from pyspark import SparkContext
-from pyspark.ml.classification import RandomForestClassifier
-from pyspark.ml.feature import StringIndexer
-from pyspark.ml.regression import RandomForestRegressor
-from pyspark.mllib.evaluation import MulticlassMetrics, RegressionMetrics
-from pyspark.mllib.util import MLUtils
-from pyspark.sql import Row, SQLContext
-
-"""
-A simple example demonstrating a RandomForest Classification/Regression Pipeline.
-Run with:
-  bin/spark-submit examples/src/main/python/ml/random_forest_example.py
-"""
-
-
-def testClassification(train, test):
-    # Train a RandomForest model.
-    # Setting featureSubsetStrategy="auto" lets the algorithm choose.
-    # Note: Use larger numTrees in practice.
-
-    rf = RandomForestClassifier(labelCol="indexedLabel", numTrees=3, maxDepth=4)
-
-    model = rf.fit(train)
-    predictionAndLabels = model.transform(test).select("prediction", "indexedLabel") \
-        .map(lambda x: (x.prediction, x.indexedLabel))
-
-    metrics = MulticlassMetrics(predictionAndLabels)
-    print("weighted f-measure %.3f" % metrics.weightedFMeasure())
-    print("precision %s" % metrics.precision())
-    print("recall %s" % metrics.recall())
-
-
-def testRegression(train, test):
-    # Train a RandomForest model.
-    # Note: Use larger numTrees in practice.
-
-    rf = RandomForestRegressor(labelCol="indexedLabel", numTrees=3, maxDepth=4)
-
-    model = rf.fit(train)
-    predictionAndLabels = model.transform(test).select("prediction", "indexedLabel") \
-        .map(lambda x: (x.prediction, x.indexedLabel))
-
-    metrics = RegressionMetrics(predictionAndLabels)
-    print("rmse %.3f" % metrics.rootMeanSquaredError)
-    print("r2 %.3f" % metrics.r2)
-    print("mae %.3f" % metrics.meanAbsoluteError)
-
-
-if __name__ == "__main__":
-    if len(sys.argv) > 1:
-        print("Usage: random_forest_example", file=sys.stderr)
-        exit(1)
-    sc = SparkContext(appName="PythonRandomForestExample")
-    sqlContext = SQLContext(sc)
-
-    # Load the data stored in LIBSVM format as a DataFrame.
-    df = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
-
-    # Map labels into an indexed column of labels in [0, numLabels)
-    stringIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel")
-    si_model = stringIndexer.fit(df)
-    td = si_model.transform(df)
-    [train, test] = td.randomSplit([0.7, 0.3])
-    testClassification(train, test)
-    testRegression(train, test)
-    sc.stop()

From 9e24ba667e43290fbaa3cacb93cf5d9be790f1fd Mon Sep 17 00:00:00 2001
From: "Joseph K. Bradley" <joseph@databricks.com>
Date: Tue, 24 Nov 2015 09:54:55 -0800
Subject: [PATCH 435/518] [SPARK-11521][ML][DOC] Document that Logistic, Linear
 Regression summaries ignore weight col

Doc for 1.6 that the summaries mostly ignore the weight column.
To be corrected for 1.7

CC: mengxr thunterdb

Author: Joseph K. Bradley <joseph@databricks.com>

Closes #9927 from jkbradley/linregsummary-doc.
---
 .../ml/classification/LogisticRegression.scala | 18 ++++++++++++++++++
 .../spark/ml/regression/LinearRegression.scala | 15 +++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
index 418bbdc9a058f..d320d64dd90d0 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
@@ -755,23 +755,35 @@ class BinaryLogisticRegressionSummary private[classification] (
    * Returns the receiver operating characteristic (ROC) curve,
    * which is an Dataframe having two fields (FPR, TPR)
    * with (0.0, 0.0) prepended and (1.0, 1.0) appended to it.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LogisticRegression.weightCol]].
+   *       This will change in later Spark versions.
    * @see http://en.wikipedia.org/wiki/Receiver_operating_characteristic
    */
   @transient lazy val roc: DataFrame = binaryMetrics.roc().toDF("FPR", "TPR")
 
   /**
    * Computes the area under the receiver operating characteristic (ROC) curve.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LogisticRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   lazy val areaUnderROC: Double = binaryMetrics.areaUnderROC()
 
   /**
    * Returns the precision-recall curve, which is an Dataframe containing
    * two fields recall, precision with (0.0, 1.0) prepended to it.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LogisticRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @transient lazy val pr: DataFrame = binaryMetrics.pr().toDF("recall", "precision")
 
   /**
    * Returns a dataframe with two fields (threshold, F-Measure) curve with beta = 1.0.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LogisticRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @transient lazy val fMeasureByThreshold: DataFrame = {
     binaryMetrics.fMeasureByThreshold().toDF("threshold", "F-Measure")
@@ -781,6 +793,9 @@ class BinaryLogisticRegressionSummary private[classification] (
    * Returns a dataframe with two fields (threshold, precision) curve.
    * Every possible probability obtained in transforming the dataset are used
    * as thresholds used in calculating the precision.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LogisticRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @transient lazy val precisionByThreshold: DataFrame = {
     binaryMetrics.precisionByThreshold().toDF("threshold", "precision")
@@ -790,6 +805,9 @@ class BinaryLogisticRegressionSummary private[classification] (
    * Returns a dataframe with two fields (threshold, recall) curve.
    * Every possible probability obtained in transforming the dataset are used
    * as thresholds used in calculating the recall.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LogisticRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @transient lazy val recallByThreshold: DataFrame = {
     binaryMetrics.recallByThreshold().toDF("threshold", "recall")
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
index 70ccec766c471..1db91666f21ab 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala
@@ -540,6 +540,9 @@ class LinearRegressionSummary private[regression] (
    * Returns the explained variance regression score.
    * explainedVariance = 1 - variance(y - \hat{y}) / variance(y)
    * Reference: [[http://en.wikipedia.org/wiki/Explained_variation]]
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LinearRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @Since("1.5.0")
   val explainedVariance: Double = metrics.explainedVariance
@@ -547,6 +550,9 @@ class LinearRegressionSummary private[regression] (
   /**
    * Returns the mean absolute error, which is a risk function corresponding to the
    * expected value of the absolute error loss or l1-norm loss.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LinearRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @Since("1.5.0")
   val meanAbsoluteError: Double = metrics.meanAbsoluteError
@@ -554,6 +560,9 @@ class LinearRegressionSummary private[regression] (
   /**
    * Returns the mean squared error, which is a risk function corresponding to the
    * expected value of the squared error loss or quadratic loss.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LinearRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @Since("1.5.0")
   val meanSquaredError: Double = metrics.meanSquaredError
@@ -561,6 +570,9 @@ class LinearRegressionSummary private[regression] (
   /**
    * Returns the root mean squared error, which is defined as the square root of
    * the mean squared error.
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LinearRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @Since("1.5.0")
   val rootMeanSquaredError: Double = metrics.rootMeanSquaredError
@@ -568,6 +580,9 @@ class LinearRegressionSummary private[regression] (
   /**
    * Returns R^2^, the coefficient of determination.
    * Reference: [[http://en.wikipedia.org/wiki/Coefficient_of_determination]]
+   *
+   * Note: This ignores instance weights (setting all to 1.0) from [[LinearRegression.weightCol]].
+   *       This will change in later Spark versions.
    */
   @Since("1.5.0")
   val r2: Double = metrics.r2

From 52bc25c8e26d4be250d8ff7864067528f4f98592 Mon Sep 17 00:00:00 2001
From: Yuhao Yang <hhbyyh@gmail.com>
Date: Tue, 24 Nov 2015 09:56:17 -0800
Subject: [PATCH 436/518] [SPARK-11847][ML] Model export/import for spark.ml:
 LDA

Add read/write support to LDA, similar to ALS.

save/load for ml.LocalLDAModel is done.
For DistributedLDAModel, I'm not sure if we can invoke save on the mllib.DistributedLDAModel directly. I'll send update after some test.

Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #9894 from hhbyyh/ldaMLsave.
---
 .../org/apache/spark/ml/clustering/LDA.scala  | 110 +++++++++++++++++-
 .../spark/mllib/clustering/LDAModel.scala     |   4 +-
 .../apache/spark/ml/clustering/LDASuite.scala |  44 ++++++-
 3 files changed, 150 insertions(+), 8 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
index 92e05815d6a3d..830510b1698d4 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala
@@ -17,12 +17,13 @@
 
 package org.apache.spark.ml.clustering
 
+import org.apache.hadoop.fs.Path
 import org.apache.spark.Logging
 import org.apache.spark.annotation.{Experimental, Since}
-import org.apache.spark.ml.util.{SchemaUtils, Identifiable}
 import org.apache.spark.ml.{Estimator, Model}
 import org.apache.spark.ml.param.shared.{HasCheckpointInterval, HasFeaturesCol, HasSeed, HasMaxIter}
 import org.apache.spark.ml.param._
+import org.apache.spark.ml.util._
 import org.apache.spark.mllib.clustering.{DistributedLDAModel => OldDistributedLDAModel,
     EMLDAOptimizer => OldEMLDAOptimizer, LDA => OldLDA, LDAModel => OldLDAModel,
     LDAOptimizer => OldLDAOptimizer, LocalLDAModel => OldLocalLDAModel,
@@ -322,7 +323,7 @@ sealed abstract class LDAModel private[ml] (
     @Since("1.6.0") override val uid: String,
     @Since("1.6.0") val vocabSize: Int,
     @Since("1.6.0") @transient protected val sqlContext: SQLContext)
-  extends Model[LDAModel] with LDAParams with Logging {
+  extends Model[LDAModel] with LDAParams with Logging with MLWritable {
 
   // NOTE to developers:
   //  This abstraction should contain all important functionality for basic LDA usage.
@@ -486,6 +487,64 @@ class LocalLDAModel private[ml] (
 
   @Since("1.6.0")
   override def isDistributed: Boolean = false
+
+  @Since("1.6.0")
+  override def write: MLWriter = new LocalLDAModel.LocalLDAModelWriter(this)
+}
+
+
+@Since("1.6.0")
+object LocalLDAModel extends MLReadable[LocalLDAModel] {
+
+  private[LocalLDAModel]
+  class LocalLDAModelWriter(instance: LocalLDAModel) extends MLWriter {
+
+    private case class Data(
+        vocabSize: Int,
+        topicsMatrix: Matrix,
+        docConcentration: Vector,
+        topicConcentration: Double,
+        gammaShape: Double)
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val oldModel = instance.oldLocalModel
+      val data = Data(instance.vocabSize, oldModel.topicsMatrix, oldModel.docConcentration,
+        oldModel.topicConcentration, oldModel.gammaShape)
+      val dataPath = new Path(path, "data").toString
+      sqlContext.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class LocalLDAModelReader extends MLReader[LocalLDAModel] {
+
+    private val className = classOf[LocalLDAModel].getName
+
+    override def load(path: String): LocalLDAModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val data = sqlContext.read.parquet(dataPath)
+        .select("vocabSize", "topicsMatrix", "docConcentration", "topicConcentration",
+          "gammaShape")
+        .head()
+      val vocabSize = data.getAs[Int](0)
+      val topicsMatrix = data.getAs[Matrix](1)
+      val docConcentration = data.getAs[Vector](2)
+      val topicConcentration = data.getAs[Double](3)
+      val gammaShape = data.getAs[Double](4)
+      val oldModel = new OldLocalLDAModel(topicsMatrix, docConcentration, topicConcentration,
+        gammaShape)
+      val model = new LocalLDAModel(metadata.uid, vocabSize, oldModel, sqlContext)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: MLReader[LocalLDAModel] = new LocalLDAModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): LocalLDAModel = super.load(path)
 }
 
 
@@ -562,6 +621,45 @@ class DistributedLDAModel private[ml] (
    */
   @Since("1.6.0")
   lazy val logPrior: Double = oldDistributedModel.logPrior
+
+  @Since("1.6.0")
+  override def write: MLWriter = new DistributedLDAModel.DistributedWriter(this)
+}
+
+
+@Since("1.6.0")
+object DistributedLDAModel extends MLReadable[DistributedLDAModel] {
+
+  private[DistributedLDAModel]
+  class DistributedWriter(instance: DistributedLDAModel) extends MLWriter {
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val modelPath = new Path(path, "oldModel").toString
+      instance.oldDistributedModel.save(sc, modelPath)
+    }
+  }
+
+  private class DistributedLDAModelReader extends MLReader[DistributedLDAModel] {
+
+    private val className = classOf[DistributedLDAModel].getName
+
+    override def load(path: String): DistributedLDAModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val modelPath = new Path(path, "oldModel").toString
+      val oldModel = OldDistributedLDAModel.load(sc, modelPath)
+      val model = new DistributedLDAModel(
+        metadata.uid, oldModel.vocabSize, oldModel, sqlContext, None)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("1.6.0")
+  override def read: MLReader[DistributedLDAModel] = new DistributedLDAModelReader
+
+  @Since("1.6.0")
+  override def load(path: String): DistributedLDAModel = super.load(path)
 }
 
 
@@ -593,7 +691,8 @@ class DistributedLDAModel private[ml] (
 @Since("1.6.0")
 @Experimental
 class LDA @Since("1.6.0") (
-    @Since("1.6.0") override val uid: String) extends Estimator[LDAModel] with LDAParams {
+    @Since("1.6.0") override val uid: String)
+  extends Estimator[LDAModel] with LDAParams with DefaultParamsWritable {
 
   @Since("1.6.0")
   def this() = this(Identifiable.randomUID("lda"))
@@ -695,7 +794,7 @@ class LDA @Since("1.6.0") (
 }
 
 
-private[clustering] object LDA {
+private[clustering] object LDA extends DefaultParamsReadable[LDA] {
 
   /** Get dataset for spark.mllib LDA */
   def getOldDataset(dataset: DataFrame, featuresCol: String): RDD[(Long, Vector)] = {
@@ -706,4 +805,7 @@ private[clustering] object LDA {
         (docId, features)
       }
   }
+
+  @Since("1.6.0")
+  override def load(path: String): LDA = super.load(path)
 }
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala
index cd520f09bd466..7384d065a2ea8 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LDAModel.scala
@@ -187,11 +187,11 @@ abstract class LDAModel private[clustering] extends Saveable {
  * @param topics Inferred topics (vocabSize x k matrix).
  */
 @Since("1.3.0")
-class LocalLDAModel private[clustering] (
+class LocalLDAModel private[spark] (
     @Since("1.3.0") val topics: Matrix,
     @Since("1.5.0") override val docConcentration: Vector,
     @Since("1.5.0") override val topicConcentration: Double,
-    override protected[clustering] val gammaShape: Double = 100)
+    override protected[spark] val gammaShape: Double = 100)
   extends LDAModel with Serializable {
 
   @Since("1.3.0")
diff --git a/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala b/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala
index b634d31cc34f0..97dbfd9a4314a 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala
@@ -18,9 +18,10 @@
 package org.apache.spark.ml.clustering
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.ml.util.MLTestingUtils
+import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
 import org.apache.spark.mllib.linalg.{Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.mllib.util.TestingUtils._
 import org.apache.spark.sql.{DataFrame, Row, SQLContext}
 
 
@@ -39,10 +40,24 @@ object LDASuite {
     }.map(v => new TestRow(v))
     sql.createDataFrame(rdd)
   }
+
+  /**
+   * Mapping from all Params to valid settings which differ from the defaults.
+   * This is useful for tests which need to exercise all Params, such as save/load.
+   * This excludes input columns to simplify some tests.
+   */
+  val allParamSettings: Map[String, Any] = Map(
+    "k" -> 3,
+    "maxIter" -> 2,
+    "checkpointInterval" -> 30,
+    "learningOffset" -> 1023.0,
+    "learningDecay" -> 0.52,
+    "subsamplingRate" -> 0.051
+  )
 }
 
 
-class LDASuite extends SparkFunSuite with MLlibTestSparkContext {
+class LDASuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
 
   val k: Int = 5
   val vocabSize: Int = 30
@@ -218,4 +233,29 @@ class LDASuite extends SparkFunSuite with MLlibTestSparkContext {
     val lp = model.logPrior
     assert(lp <= 0.0 && lp != Double.NegativeInfinity)
   }
+
+  test("read/write LocalLDAModel") {
+    def checkModelData(model: LDAModel, model2: LDAModel): Unit = {
+      assert(model.vocabSize === model2.vocabSize)
+      assert(Vectors.dense(model.topicsMatrix.toArray) ~==
+        Vectors.dense(model2.topicsMatrix.toArray) absTol 1e-6)
+      assert(Vectors.dense(model.getDocConcentration) ~==
+        Vectors.dense(model2.getDocConcentration) absTol 1e-6)
+    }
+    val lda = new LDA()
+    testEstimatorAndModelReadWrite(lda, dataset, LDASuite.allParamSettings, checkModelData)
+  }
+
+  test("read/write DistributedLDAModel") {
+    def checkModelData(model: LDAModel, model2: LDAModel): Unit = {
+      assert(model.vocabSize === model2.vocabSize)
+      assert(Vectors.dense(model.topicsMatrix.toArray) ~==
+        Vectors.dense(model2.topicsMatrix.toArray) absTol 1e-6)
+      assert(Vectors.dense(model.getDocConcentration) ~==
+        Vectors.dense(model2.getDocConcentration) absTol 1e-6)
+    }
+    val lda = new LDA()
+    testEstimatorAndModelReadWrite(lda, dataset,
+      LDASuite.allParamSettings ++ Map("optimizer" -> "em"), checkModelData)
+  }
 }

From 19530da6903fa59b051eec69b9c17e231c68454b Mon Sep 17 00:00:00 2001
From: Wenchen Fan <wenchen@databricks.com>
Date: Tue, 24 Nov 2015 11:09:01 -0800
Subject: [PATCH 437/518] [SPARK-11926][SQL] unify GetStructField and
 GetInternalRowField

Author: Wenchen Fan <wenchen@databricks.com>

Closes #9909 from cloud-fan/get-struct.
---
 .../spark/sql/catalyst/ScalaReflection.scala  |  2 +-
 .../sql/catalyst/analysis/unresolved.scala    |  8 +++----
 .../catalyst/encoders/ExpressionEncoder.scala |  2 +-
 .../sql/catalyst/encoders/RowEncoder.scala    |  2 +-
 .../sql/catalyst/expressions/Expression.scala |  2 +-
 .../expressions/complexTypeExtractors.scala   | 18 ++++++++--------
 .../expressions/namedExpressions.scala        |  4 ++--
 .../sql/catalyst/expressions/objects.scala    | 21 -------------------
 .../expressions/ComplexTypeSuite.scala        |  4 ++--
 9 files changed, 21 insertions(+), 42 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
index 476becec4dd52..d133ad3f0d89d 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala
@@ -130,7 +130,7 @@ object ScalaReflection extends ScalaReflection {
 
     /** Returns the current path with a field at ordinal extracted. */
     def addToPathOrdinal(ordinal: Int, dataType: DataType): Expression = path
-      .map(p => GetInternalRowField(p, ordinal, dataType))
+      .map(p => GetStructField(p, ordinal))
       .getOrElse(BoundReference(ordinal, dataType, false))
 
     /** Returns the current path or `BoundReference`. */
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
index 6485bdfb30234..1b2a8dc4c7f14 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
@@ -201,12 +201,12 @@ case class UnresolvedStar(target: Option[Seq[String]]) extends Star with Unevalu
     if (attribute.isDefined) {
       // This target resolved to an attribute in child. It must be a struct. Expand it.
       attribute.get.dataType match {
-        case s: StructType => {
-          s.fields.map( f => {
-            val extract = GetStructField(attribute.get, f, s.getFieldIndex(f.name).get)
+        case s: StructType => s.zipWithIndex.map {
+          case (f, i) =>
+            val extract = GetStructField(attribute.get, i)
             Alias(extract, target.get + "." + f.name)()
-          })
         }
+
         case _ => {
           throw new AnalysisException("Can only star expand struct data types. Attribute: `" +
             target.get + "`")
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
index 7bc9aed0b204e..0c10a56c555f4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/ExpressionEncoder.scala
@@ -111,7 +111,7 @@ object ExpressionEncoder {
           case UnresolvedAttribute(nameParts) =>
             assert(nameParts.length == 1)
             UnresolvedExtractValue(input, Literal(nameParts.head))
-          case BoundReference(ordinal, dt, _) => GetInternalRowField(input, ordinal, dt)
+          case BoundReference(ordinal, dt, _) => GetStructField(input, ordinal)
         }
       }
     }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
index fa553e7c5324c..67518f52d4a58 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala
@@ -220,7 +220,7 @@ object RowEncoder {
         If(
           Invoke(input, "isNullAt", BooleanType, Literal(i) :: Nil),
           Literal.create(null, externalDataTypeFor(f.dataType)),
-          constructorFor(GetInternalRowField(input, i, f.dataType)))
+          constructorFor(GetStructField(input, i)))
       }
       CreateExternalRow(convertedFields)
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
index 540ed3500616a..169435a10ea2c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
@@ -206,7 +206,7 @@ abstract class Expression extends TreeNode[Expression] {
    */
   def prettyString: String = {
     transform {
-      case a: AttributeReference => PrettyAttribute(a.name)
+      case a: AttributeReference => PrettyAttribute(a.name, a.dataType)
       case u: UnresolvedAttribute => PrettyAttribute(u.name)
     }.toString
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
index f871b737fff3a..10ce10aaf6da2 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeExtractors.scala
@@ -51,7 +51,7 @@ object ExtractValue {
       case (StructType(fields), NonNullLiteral(v, StringType)) =>
         val fieldName = v.toString
         val ordinal = findField(fields, fieldName, resolver)
-        GetStructField(child, fields(ordinal).copy(name = fieldName), ordinal)
+        GetStructField(child, ordinal, Some(fieldName))
 
       case (ArrayType(StructType(fields), containsNull), NonNullLiteral(v, StringType)) =>
         val fieldName = v.toString
@@ -97,18 +97,18 @@ object ExtractValue {
  * Returns the value of fields in the Struct `child`.
  *
  * No need to do type checking since it is handled by [[ExtractValue]].
- * TODO: Unify with [[GetInternalRowField]], remove the need to specify a [[StructField]].
+ *
+ * Note that we can pass in the field name directly to keep case preserving in `toString`.
+ * For example, when get field `yEAr` from `<year: int, month: int>`, we should pass in `yEAr`.
  */
-case class GetStructField(child: Expression, field: StructField, ordinal: Int)
+case class GetStructField(child: Expression, ordinal: Int, name: Option[String] = None)
   extends UnaryExpression {
 
-  override def dataType: DataType = child.dataType match {
-    case s: StructType => s(ordinal).dataType
-    // This is a hack to avoid breaking existing code until we remove the need for the struct field
-    case _ => field.dataType
-  }
+  private lazy val field = child.dataType.asInstanceOf[StructType](ordinal)
+
+  override def dataType: DataType = field.dataType
   override def nullable: Boolean = child.nullable || field.nullable
-  override def toString: String = s"$child.${field.name}"
+  override def toString: String = s"$child.${name.getOrElse(field.name)}"
 
   protected override def nullSafeEval(input: Any): Any =
     input.asInstanceOf[InternalRow].get(ordinal, field.dataType)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
index 00b7970bd16c6..26b6aca79971e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
@@ -273,7 +273,8 @@ case class AttributeReference(
  * A place holder used when printing expressions without debugging information such as the
  * expression id or the unresolved indicator.
  */
-case class PrettyAttribute(name: String) extends Attribute with Unevaluable {
+case class PrettyAttribute(name: String, dataType: DataType = NullType)
+  extends Attribute with Unevaluable {
 
   override def toString: String = name
 
@@ -286,7 +287,6 @@ case class PrettyAttribute(name: String) extends Attribute with Unevaluable {
   override def qualifiers: Seq[String] = throw new UnsupportedOperationException
   override def exprId: ExprId = throw new UnsupportedOperationException
   override def nullable: Boolean = throw new UnsupportedOperationException
-  override def dataType: DataType = NullType
 }
 
 object VirtualColumn {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
index 4a1f419f0ad8d..62d09f0f55105 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects.scala
@@ -517,27 +517,6 @@ case class CreateExternalRow(children: Seq[Expression]) extends Expression {
   }
 }
 
-case class GetInternalRowField(child: Expression, ordinal: Int, dataType: DataType)
-  extends UnaryExpression {
-
-  override def nullable: Boolean = true
-
-  override def eval(input: InternalRow): Any =
-    throw new UnsupportedOperationException("Only code-generated evaluation is supported")
-
-  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    nullSafeCodeGen(ctx, ev, eval => {
-      s"""
-        if ($eval.isNullAt($ordinal)) {
-          ${ev.isNull} = true;
-        } else {
-          ${ev.value} = ${ctx.getValue(eval, dataType, ordinal.toString)};
-        }
-      """
-    })
-  }
-}
-
 /**
  * Serializes an input object using a generic serializer (Kryo or Java).
  * @param kryo if true, use Kryo. Otherwise, use Java.
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala
index e60990aeb423f..62fd47234b33b 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala
@@ -79,8 +79,8 @@ class ComplexTypeSuite extends SparkFunSuite with ExpressionEvalHelper {
     def getStructField(expr: Expression, fieldName: String): GetStructField = {
       expr.dataType match {
         case StructType(fields) =>
-          val field = fields.find(_.name == fieldName).get
-          GetStructField(expr, field, fields.indexOf(field))
+          val index = fields.indexWhere(_.name == fieldName)
+          GetStructField(expr, index)
       }
     }
 

From 81012546ee5a80d2576740af0dad067b0f5962c5 Mon Sep 17 00:00:00 2001
From: tedyu <yuzhihong@gmail.com>
Date: Tue, 24 Nov 2015 12:22:33 -0800
Subject: [PATCH 438/518] [SPARK-11872] Prevent the call to SparkContext#stop()
 in the listener bus's thread

This is continuation of SPARK-11761

Andrew suggested adding this protection. See tail of https://github.com/apache/spark/pull/9741

Author: tedyu <yuzhihong@gmail.com>

Closes #9852 from tedyu/master.
---
 .../scala/org/apache/spark/SparkContext.scala |  4 +++
 .../spark/scheduler/SparkListenerSuite.scala  | 31 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index b153a7b08e590..e19ba113702c6 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -1694,6 +1694,10 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
 
   // Shut down the SparkContext.
   def stop() {
+    if (AsynchronousListenerBus.withinListenerThread.value) {
+      throw new SparkException("Cannot stop SparkContext within listener thread of" +
+        " AsynchronousListenerBus")
+    }
     // Use the stopping variable to ensure no contention for the stop scenario.
     // Still track the stopped variable for use elsewhere in the code.
     if (!stopped.compareAndSet(false, true)) {
diff --git a/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
index 84e545851f49e..f20d5be7c0ee0 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/SparkListenerSuite.scala
@@ -24,6 +24,7 @@ import scala.collection.JavaConverters._
 
 import org.scalatest.Matchers
 
+import org.apache.spark.SparkException
 import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.util.ResetSystemProperties
 import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkFunSuite}
@@ -36,6 +37,21 @@ class SparkListenerSuite extends SparkFunSuite with LocalSparkContext with Match
 
   val jobCompletionTime = 1421191296660L
 
+  test("don't call sc.stop in listener") {
+    sc = new SparkContext("local", "SparkListenerSuite")
+    val listener = new SparkContextStoppingListener(sc)
+    val bus = new LiveListenerBus
+    bus.addListener(listener)
+
+    // Starting listener bus should flush all buffered events
+    bus.start(sc)
+    bus.post(SparkListenerJobEnd(0, jobCompletionTime, JobSucceeded))
+    bus.waitUntilEmpty(WAIT_TIMEOUT_MILLIS)
+
+    bus.stop()
+    assert(listener.sparkExSeen)
+  }
+
   test("basic creation and shutdown of LiveListenerBus") {
     val counter = new BasicJobCounter
     val bus = new LiveListenerBus
@@ -443,6 +459,21 @@ private class BasicJobCounter extends SparkListener {
   override def onJobEnd(job: SparkListenerJobEnd): Unit = count += 1
 }
 
+/**
+ * A simple listener that tries to stop SparkContext.
+ */
+private class SparkContextStoppingListener(val sc: SparkContext) extends SparkListener {
+  @volatile var sparkExSeen = false
+  override def onJobEnd(job: SparkListenerJobEnd): Unit = {
+    try {
+      sc.stop()
+    } catch {
+      case se: SparkException =>
+        sparkExSeen = true
+    }
+  }
+}
+
 private class ListenerThatAcceptsSparkConf(conf: SparkConf) extends SparkListener {
   var count = 0
   override def onJobEnd(job: SparkListenerJobEnd): Unit = count += 1

From f3152722791b163fa66597b3684009058195ba33 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 24 Nov 2015 12:54:37 -0800
Subject: [PATCH 439/518] [SPARK-11946][SQL] Audit pivot API for 1.6.

Currently pivot's signature looks like

```scala
scala.annotation.varargs
def pivot(pivotColumn: Column, values: Column*): GroupedData

scala.annotation.varargs
def pivot(pivotColumn: String, values: Any*): GroupedData
```

I think we can remove the one that takes "Column" types, since callers should always be passing in literals. It'd also be more clear if the values are not varargs, but rather Seq or java.util.List.

I also made similar changes for Python.

Author: Reynold Xin <rxin@databricks.com>

Closes #9929 from rxin/SPARK-11946.
---
 .../apache/spark/scheduler/DAGScheduler.scala |   1 -
 python/pyspark/sql/group.py                   |  12 +-
 .../sql/catalyst/expressions/literals.scala   |   1 +
 .../org/apache/spark/sql/GroupedData.scala    | 154 ++++++++++--------
 .../apache/spark/sql/JavaDataFrameSuite.java  |  16 ++
 .../spark/sql/DataFramePivotSuite.scala       |  21 +--
 .../apache/spark/sql/test/SQLTestData.scala   |   1 +
 7 files changed, 125 insertions(+), 81 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index ae725b467d8c4..77a184dfe4bee 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -1574,7 +1574,6 @@ class DAGScheduler(
   }
 
   def stop() {
-    logInfo("Stopping DAGScheduler")
     messageScheduler.shutdownNow()
     eventProcessLoop.stop()
     taskScheduler.stop()
diff --git a/python/pyspark/sql/group.py b/python/pyspark/sql/group.py
index 227f40bc3cf53..d8ed7eb2dda64 100644
--- a/python/pyspark/sql/group.py
+++ b/python/pyspark/sql/group.py
@@ -168,20 +168,24 @@ def sum(self, *cols):
         """
 
     @since(1.6)
-    def pivot(self, pivot_col, *values):
+    def pivot(self, pivot_col, values=None):
         """Pivots a column of the current DataFrame and preform the specified aggregation.
 
         :param pivot_col: Column to pivot
         :param values: Optional list of values of pivotColumn that will be translated to columns in
             the output data frame. If values are not provided the method with do an immediate call
             to .distinct() on the pivot column.
-        >>> df4.groupBy("year").pivot("course", "dotNET", "Java").sum("earnings").collect()
+
+        >>> df4.groupBy("year").pivot("course", ["dotNET", "Java"]).sum("earnings").collect()
         [Row(year=2012, dotNET=15000, Java=20000), Row(year=2013, dotNET=48000, Java=30000)]
+
         >>> df4.groupBy("year").pivot("course").sum("earnings").collect()
         [Row(year=2012, Java=20000, dotNET=15000), Row(year=2013, Java=30000, dotNET=48000)]
         """
-        jgd = self._jdf.pivot(_to_java_column(pivot_col),
-                              _to_seq(self.sql_ctx._sc, values, _create_column_from_literal))
+        if values is None:
+            jgd = self._jdf.pivot(pivot_col)
+        else:
+            jgd = self._jdf.pivot(pivot_col, values)
         return GroupedData(jgd, self.sql_ctx)
 
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala
index e34fd49be8389..68ec688c99f93 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala
@@ -44,6 +44,7 @@ object Literal {
     case a: Array[Byte] => Literal(a, BinaryType)
     case i: CalendarInterval => Literal(i, CalendarIntervalType)
     case null => Literal(null, NullType)
+    case v: Literal => v
     case _ =>
       throw new RuntimeException("Unsupported literal type " + v.getClass + " " + v)
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
index 63dd7fbcbe9e4..ee7150cbbfbca 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.analysis.{UnresolvedFunction, UnresolvedAli
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical.{Pivot, Rollup, Cube, Aggregate}
-import org.apache.spark.sql.types.{StringType, NumericType}
+import org.apache.spark.sql.types.NumericType
 
 
 /**
@@ -282,74 +282,96 @@ class GroupedData protected[sql](
   }
 
   /**
-    * (Scala-specific) Pivots a column of the current [[DataFrame]] and preform the specified
-    * aggregation.
-    * {{{
-    *   // Compute the sum of earnings for each year by course with each course as a separate column
-    *   df.groupBy($"year").pivot($"course", "dotNET", "Java").agg(sum($"earnings"))
-    *   // Or without specifying column values
-    *   df.groupBy($"year").pivot($"course").agg(sum($"earnings"))
-    * }}}
-    * @param pivotColumn Column to pivot
-    * @param values Optional list of values of pivotColumn that will be translated to columns in the
-    *               output data frame. If values are not provided the method with do an immediate
-    *               call to .distinct() on the pivot column.
-    * @since 1.6.0
-    */
-  @scala.annotation.varargs
-  def pivot(pivotColumn: Column, values: Column*): GroupedData = groupType match {
-    case _: GroupedData.PivotType =>
-      throw new UnsupportedOperationException("repeated pivots are not supported")
-    case GroupedData.GroupByType =>
-      val pivotValues = if (values.nonEmpty) {
-        values.map {
-          case Column(literal: Literal) => literal
-          case other =>
-            throw new UnsupportedOperationException(
-              s"The values of a pivot must be literals, found $other")
-        }
-      } else {
-        // This is to prevent unintended OOM errors when the number of distinct values is large
-        val maxValues = df.sqlContext.conf.getConf(SQLConf.DATAFRAME_PIVOT_MAX_VALUES)
-        // Get the distinct values of the column and sort them so its consistent
-        val values = df.select(pivotColumn)
-          .distinct()
-          .sort(pivotColumn)
-          .map(_.get(0))
-          .take(maxValues + 1)
-          .map(Literal(_)).toSeq
-        if (values.length > maxValues) {
-          throw new RuntimeException(
-            s"The pivot column $pivotColumn has more than $maxValues distinct values, " +
-              "this could indicate an error. " +
-              "If this was intended, set \"" + SQLConf.DATAFRAME_PIVOT_MAX_VALUES.key + "\" " +
-              s"to at least the number of distinct values of the pivot column.")
-        }
-        values
-      }
-      new GroupedData(df, groupingExprs, GroupedData.PivotType(pivotColumn.expr, pivotValues))
-    case _ =>
-      throw new UnsupportedOperationException("pivot is only supported after a groupBy")
+   * Pivots a column of the current [[DataFrame]] and preform the specified aggregation.
+   * There are two versions of pivot function: one that requires the caller to specify the list
+   * of distinct values to pivot on, and one that does not. The latter is more concise but less
+   * efficient, because Spark needs to first compute the list of distinct values internally.
+   *
+   * {{{
+   *   // Compute the sum of earnings for each year by course with each course as a separate column
+   *   df.groupBy("year").pivot("course", Seq("dotNET", "Java")).sum("earnings")
+   *
+   *   // Or without specifying column values (less efficient)
+   *   df.groupBy("year").pivot("course").sum("earnings")
+   * }}}
+   *
+   * @param pivotColumn Name of the column to pivot.
+   * @since 1.6.0
+   */
+  def pivot(pivotColumn: String): GroupedData = {
+    // This is to prevent unintended OOM errors when the number of distinct values is large
+    val maxValues = df.sqlContext.conf.getConf(SQLConf.DATAFRAME_PIVOT_MAX_VALUES)
+    // Get the distinct values of the column and sort them so its consistent
+    val values = df.select(pivotColumn)
+      .distinct()
+      .sort(pivotColumn)
+      .map(_.get(0))
+      .take(maxValues + 1)
+      .toSeq
+
+    if (values.length > maxValues) {
+      throw new AnalysisException(
+        s"The pivot column $pivotColumn has more than $maxValues distinct values, " +
+          "this could indicate an error. " +
+          s"If this was intended, set ${SQLConf.DATAFRAME_PIVOT_MAX_VALUES.key} " +
+          "to at least the number of distinct values of the pivot column.")
+    }
+
+    pivot(pivotColumn, values)
   }
 
   /**
-    * Pivots a column of the current [[DataFrame]] and preform the specified aggregation.
-    * {{{
-    *   // Compute the sum of earnings for each year by course with each course as a separate column
-    *   df.groupBy("year").pivot("course", "dotNET", "Java").sum("earnings")
-    *   // Or without specifying column values
-    *   df.groupBy("year").pivot("course").sum("earnings")
-    * }}}
-    * @param pivotColumn Column to pivot
-    * @param values Optional list of values of pivotColumn that will be translated to columns in the
-    *               output data frame. If values are not provided the method with do an immediate
-    *               call to .distinct() on the pivot column.
-    * @since 1.6.0
-    */
-  @scala.annotation.varargs
-  def pivot(pivotColumn: String, values: Any*): GroupedData = {
-    val resolvedPivotColumn = Column(df.resolve(pivotColumn))
-    pivot(resolvedPivotColumn, values.map(functions.lit): _*)
+   * Pivots a column of the current [[DataFrame]] and preform the specified aggregation.
+   * There are two versions of pivot function: one that requires the caller to specify the list
+   * of distinct values to pivot on, and one that does not. The latter is more concise but less
+   * efficient, because Spark needs to first compute the list of distinct values internally.
+   *
+   * {{{
+   *   // Compute the sum of earnings for each year by course with each course as a separate column
+   *   df.groupBy("year").pivot("course", Seq("dotNET", "Java")).sum("earnings")
+   *
+   *   // Or without specifying column values (less efficient)
+   *   df.groupBy("year").pivot("course").sum("earnings")
+   * }}}
+   *
+   * @param pivotColumn Name of the column to pivot.
+   * @param values List of values that will be translated to columns in the output DataFrame.
+   * @since 1.6.0
+   */
+  def pivot(pivotColumn: String, values: Seq[Any]): GroupedData = {
+    groupType match {
+      case GroupedData.GroupByType =>
+        new GroupedData(
+          df,
+          groupingExprs,
+          GroupedData.PivotType(df.resolve(pivotColumn), values.map(Literal.apply)))
+      case _: GroupedData.PivotType =>
+        throw new UnsupportedOperationException("repeated pivots are not supported")
+      case _ =>
+        throw new UnsupportedOperationException("pivot is only supported after a groupBy")
+    }
+  }
+
+  /**
+   * Pivots a column of the current [[DataFrame]] and preform the specified aggregation.
+   * There are two versions of pivot function: one that requires the caller to specify the list
+   * of distinct values to pivot on, and one that does not. The latter is more concise but less
+   * efficient, because Spark needs to first compute the list of distinct values internally.
+   *
+   * {{{
+   *   // Compute the sum of earnings for each year by course with each course as a separate column
+   *   df.groupBy("year").pivot("course", Arrays.<Object>asList("dotNET", "Java")).sum("earnings");
+   *
+   *   // Or without specifying column values (less efficient)
+   *   df.groupBy("year").pivot("course").sum("earnings");
+   * }}}
+   *
+   * @param pivotColumn Name of the column to pivot.
+   * @param values List of values that will be translated to columns in the output DataFrame.
+   * @since 1.6.0
+   */
+  def pivot(pivotColumn: String, values: java.util.List[Any]): GroupedData = {
+    pivot(pivotColumn, values.asScala)
   }
 }
 
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index 567bdddece80e..a12fed3c0c6af 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -282,4 +282,20 @@ public void testSampleBy() {
     Assert.assertEquals(1, actual[1].getLong(0));
     Assert.assertTrue(2 <= actual[1].getLong(1) && actual[1].getLong(1) <= 13);
   }
+
+  @Test
+  public void pivot() {
+    DataFrame df = context.table("courseSales");
+    Row[] actual = df.groupBy("year")
+      .pivot("course", Arrays.<Object>asList("dotNET", "Java"))
+      .agg(sum("earnings")).orderBy("year").collect();
+
+    Assert.assertEquals(2012, actual[0].getInt(0));
+    Assert.assertEquals(15000.0, actual[0].getDouble(1), 0.01);
+    Assert.assertEquals(20000.0, actual[0].getDouble(2), 0.01);
+
+    Assert.assertEquals(2013, actual[1].getInt(0));
+    Assert.assertEquals(48000.0, actual[1].getDouble(1), 0.01);
+    Assert.assertEquals(30000.0, actual[1].getDouble(2), 0.01);
+  }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala
index 0c23d142670c1..fc53aba68ebb7 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala
@@ -25,7 +25,7 @@ class DataFramePivotSuite extends QueryTest with SharedSQLContext{
 
   test("pivot courses with literals") {
     checkAnswer(
-      courseSales.groupBy($"year").pivot($"course", lit("dotNET"), lit("Java"))
+      courseSales.groupBy("year").pivot("course", Seq("dotNET", "Java"))
         .agg(sum($"earnings")),
       Row(2012, 15000.0, 20000.0) :: Row(2013, 48000.0, 30000.0) :: Nil
     )
@@ -33,14 +33,15 @@ class DataFramePivotSuite extends QueryTest with SharedSQLContext{
 
   test("pivot year with literals") {
     checkAnswer(
-      courseSales.groupBy($"course").pivot($"year", lit(2012), lit(2013)).agg(sum($"earnings")),
+      courseSales.groupBy("course").pivot("year", Seq(2012, 2013)).agg(sum($"earnings")),
       Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
     )
   }
 
   test("pivot courses with literals and multiple aggregations") {
     checkAnswer(
-      courseSales.groupBy($"year").pivot($"course", lit("dotNET"), lit("Java"))
+      courseSales.groupBy($"year")
+        .pivot("course", Seq("dotNET", "Java"))
         .agg(sum($"earnings"), avg($"earnings")),
       Row(2012, 15000.0, 7500.0, 20000.0, 20000.0) ::
         Row(2013, 48000.0, 48000.0, 30000.0, 30000.0) :: Nil
@@ -49,14 +50,14 @@ class DataFramePivotSuite extends QueryTest with SharedSQLContext{
 
   test("pivot year with string values (cast)") {
     checkAnswer(
-      courseSales.groupBy("course").pivot("year", "2012", "2013").sum("earnings"),
+      courseSales.groupBy("course").pivot("year", Seq("2012", "2013")).sum("earnings"),
       Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
     )
   }
 
   test("pivot year with int values") {
     checkAnswer(
-      courseSales.groupBy("course").pivot("year", 2012, 2013).sum("earnings"),
+      courseSales.groupBy("course").pivot("year", Seq(2012, 2013)).sum("earnings"),
       Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
     )
   }
@@ -64,22 +65,22 @@ class DataFramePivotSuite extends QueryTest with SharedSQLContext{
   test("pivot courses with no values") {
     // Note Java comes before dotNet in sorted order
     checkAnswer(
-      courseSales.groupBy($"year").pivot($"course").agg(sum($"earnings")),
+      courseSales.groupBy("year").pivot("course").agg(sum($"earnings")),
       Row(2012, 20000.0, 15000.0) :: Row(2013, 30000.0, 48000.0) :: Nil
     )
   }
 
   test("pivot year with no values") {
     checkAnswer(
-      courseSales.groupBy($"course").pivot($"year").agg(sum($"earnings")),
+      courseSales.groupBy("course").pivot("year").agg(sum($"earnings")),
       Row("dotNET", 15000.0, 48000.0) :: Row("Java", 20000.0, 30000.0) :: Nil
     )
   }
 
-  test("pivot max values inforced") {
+  test("pivot max values enforced") {
     sqlContext.conf.setConf(SQLConf.DATAFRAME_PIVOT_MAX_VALUES, 1)
-    intercept[RuntimeException](
-      courseSales.groupBy($"year").pivot($"course")
+    intercept[AnalysisException](
+      courseSales.groupBy("year").pivot("course")
     )
     sqlContext.conf.setConf(SQLConf.DATAFRAME_PIVOT_MAX_VALUES,
       SQLConf.DATAFRAME_PIVOT_MAX_VALUES.defaultValue.get)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala
index abad0d7eaaedf..83c63e04f344a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SQLTestData.scala
@@ -281,6 +281,7 @@ private[sql] trait SQLTestData { self =>
     person
     salary
     complexData
+    courseSales
   }
 }
 

From e6dd237463d2de8c506f0735dfdb3f43e8122513 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Tue, 24 Nov 2015 15:08:02 -0600
Subject: [PATCH 440/518] [SPARK-11929][CORE] Make the repl log4j configuration
 override the root logger.

In the default Spark distribution, there are currently two separate
log4j config files, with different default values for the root logger,
so that when running the shell you have a different default log level.
This makes the shell more usable, since the logs don't overwhelm the
output.

But if you install a custom log4j.properties, you lose that, because
then it's going to be used no matter whether you're running a regular
app or the shell.

With this change, the overriding of the log level is done differently;
the log level repl's main class (org.apache.spark.repl.Main) is used
to define the root logger's level when running the shell, defaulting
to WARN if it's not set explicitly.

On a somewhat related change, the shell output about the "sc" variable
was changed a bit to contain a little more useful information about
the application, since when the root logger's log level is WARN, that
information is never shown to the user.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9816 from vanzin/shell-logging.
---
 conf/log4j.properties.template                |  5 +++
 .../spark/log4j-defaults-repl.properties      | 33 --------------
 .../apache/spark/log4j-defaults.properties    |  5 +++
 .../main/scala/org/apache/spark/Logging.scala | 45 ++++++++++---------
 .../apache/spark/repl/SparkILoopInit.scala    | 21 ++++-----
 .../org/apache/spark/repl/SparkILoop.scala    | 25 ++++++-----
 6 files changed, 57 insertions(+), 77 deletions(-)
 delete mode 100644 core/src/main/resources/org/apache/spark/log4j-defaults-repl.properties

diff --git a/conf/log4j.properties.template b/conf/log4j.properties.template
index f3046be54d7c6..9809b0c828487 100644
--- a/conf/log4j.properties.template
+++ b/conf/log4j.properties.template
@@ -22,6 +22,11 @@ log4j.appender.console.target=System.err
 log4j.appender.console.layout=org.apache.log4j.PatternLayout
 log4j.appender.console.layout.ConversionPattern=%d{yy/MM/dd HH:mm:ss} %p %c{1}: %m%n
 
+# Set the default spark-shell log level to WARN. When running the spark-shell, the
+# log level for this class is used to overwrite the root logger's log level, so that
+# the user can have different defaults for the shell and regular Spark apps.
+log4j.logger.org.apache.spark.repl.Main=WARN
+
 # Settings to quiet third party logs that are too verbose
 log4j.logger.org.spark-project.jetty=WARN
 log4j.logger.org.spark-project.jetty.util.component.AbstractLifeCycle=ERROR
diff --git a/core/src/main/resources/org/apache/spark/log4j-defaults-repl.properties b/core/src/main/resources/org/apache/spark/log4j-defaults-repl.properties
deleted file mode 100644
index c85abc35b93bf..0000000000000
--- a/core/src/main/resources/org/apache/spark/log4j-defaults-repl.properties
+++ /dev/null
@@ -1,33 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-# Set everything to be logged to the console
-log4j.rootCategory=WARN, console
-log4j.appender.console=org.apache.log4j.ConsoleAppender
-log4j.appender.console.target=System.err
-log4j.appender.console.layout=org.apache.log4j.PatternLayout
-log4j.appender.console.layout.ConversionPattern=%d{yy/MM/dd HH:mm:ss} %p %c{1}: %m%n
-
-# Settings to quiet third party logs that are too verbose
-log4j.logger.org.spark-project.jetty=WARN
-log4j.logger.org.spark-project.jetty.util.component.AbstractLifeCycle=ERROR
-log4j.logger.org.apache.spark.repl.SparkIMain$exprTyper=INFO
-log4j.logger.org.apache.spark.repl.SparkILoop$SparkILoopInterpreter=INFO
-
-# SPARK-9183: Settings to avoid annoying messages when looking up nonexistent UDFs in SparkSQL with Hive support
-log4j.logger.org.apache.hadoop.hive.metastore.RetryingHMSHandler=FATAL
-log4j.logger.org.apache.hadoop.hive.ql.exec.FunctionRegistry=ERROR
diff --git a/core/src/main/resources/org/apache/spark/log4j-defaults.properties b/core/src/main/resources/org/apache/spark/log4j-defaults.properties
index d44cc85dcbd82..0750488e4adf9 100644
--- a/core/src/main/resources/org/apache/spark/log4j-defaults.properties
+++ b/core/src/main/resources/org/apache/spark/log4j-defaults.properties
@@ -22,6 +22,11 @@ log4j.appender.console.target=System.err
 log4j.appender.console.layout=org.apache.log4j.PatternLayout
 log4j.appender.console.layout.ConversionPattern=%d{yy/MM/dd HH:mm:ss} %p %c{1}: %m%n
 
+# Set the default spark-shell log level to WARN. When running the spark-shell, the
+# log level for this class is used to overwrite the root logger's log level, so that
+# the user can have different defaults for the shell and regular Spark apps.
+log4j.logger.org.apache.spark.repl.Main=WARN
+
 # Settings to quiet third party logs that are too verbose
 log4j.logger.org.spark-project.jetty=WARN
 log4j.logger.org.spark-project.jetty.util.component.AbstractLifeCycle=ERROR
diff --git a/core/src/main/scala/org/apache/spark/Logging.scala b/core/src/main/scala/org/apache/spark/Logging.scala
index 69f6e06ee0057..e35e158c7e8a6 100644
--- a/core/src/main/scala/org/apache/spark/Logging.scala
+++ b/core/src/main/scala/org/apache/spark/Logging.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark
 
-import org.apache.log4j.{LogManager, PropertyConfigurator}
+import org.apache.log4j.{Level, LogManager, PropertyConfigurator}
 import org.slf4j.{Logger, LoggerFactory}
 import org.slf4j.impl.StaticLoggerBinder
 
@@ -119,30 +119,31 @@ trait Logging {
     val usingLog4j12 = "org.slf4j.impl.Log4jLoggerFactory".equals(binderClass)
     if (usingLog4j12) {
       val log4j12Initialized = LogManager.getRootLogger.getAllAppenders.hasMoreElements
+      // scalastyle:off println
       if (!log4j12Initialized) {
-        // scalastyle:off println
-        if (Utils.isInInterpreter) {
-          val replDefaultLogProps = "org/apache/spark/log4j-defaults-repl.properties"
-          Option(Utils.getSparkClassLoader.getResource(replDefaultLogProps)) match {
-            case Some(url) =>
-              PropertyConfigurator.configure(url)
-              System.err.println(s"Using Spark's repl log4j profile: $replDefaultLogProps")
-              System.err.println("To adjust logging level use sc.setLogLevel(\"INFO\")")
-            case None =>
-              System.err.println(s"Spark was unable to load $replDefaultLogProps")
-          }
-        } else {
-          val defaultLogProps = "org/apache/spark/log4j-defaults.properties"
-          Option(Utils.getSparkClassLoader.getResource(defaultLogProps)) match {
-            case Some(url) =>
-              PropertyConfigurator.configure(url)
-              System.err.println(s"Using Spark's default log4j profile: $defaultLogProps")
-            case None =>
-              System.err.println(s"Spark was unable to load $defaultLogProps")
-          }
+        val defaultLogProps = "org/apache/spark/log4j-defaults.properties"
+        Option(Utils.getSparkClassLoader.getResource(defaultLogProps)) match {
+          case Some(url) =>
+            PropertyConfigurator.configure(url)
+            System.err.println(s"Using Spark's default log4j profile: $defaultLogProps")
+          case None =>
+            System.err.println(s"Spark was unable to load $defaultLogProps")
         }
-        // scalastyle:on println
       }
+
+      if (Utils.isInInterpreter) {
+        // Use the repl's main class to define the default log level when running the shell,
+        // overriding the root logger's config if they're different.
+        val rootLogger = LogManager.getRootLogger()
+        val replLogger = LogManager.getLogger("org.apache.spark.repl.Main")
+        val replLevel = Option(replLogger.getLevel()).getOrElse(Level.WARN)
+        if (replLevel != rootLogger.getEffectiveLevel()) {
+          System.err.printf("Setting default log level to \"%s\".\n", replLevel)
+          System.err.println("To adjust logging level use sc.setLogLevel(newLevel).")
+          rootLogger.setLevel(replLevel)
+        }
+      }
+      // scalastyle:on println
     }
     Logging.initialized = true
 
diff --git a/repl/scala-2.10/src/main/scala/org/apache/spark/repl/SparkILoopInit.scala b/repl/scala-2.10/src/main/scala/org/apache/spark/repl/SparkILoopInit.scala
index bd3314d94eed6..99e1e1df33fd8 100644
--- a/repl/scala-2.10/src/main/scala/org/apache/spark/repl/SparkILoopInit.scala
+++ b/repl/scala-2.10/src/main/scala/org/apache/spark/repl/SparkILoopInit.scala
@@ -123,18 +123,19 @@ private[repl] trait SparkILoopInit {
   def initializeSpark() {
     intp.beQuietDuring {
       command("""
-         @transient val sc = {
-           val _sc = org.apache.spark.repl.Main.interp.createSparkContext()
-           println("Spark context available as sc.")
-           _sc
-         }
+        @transient val sc = {
+          val _sc = org.apache.spark.repl.Main.interp.createSparkContext()
+          println("Spark context available as sc " +
+            s"(master = ${_sc.master}, app id = ${_sc.applicationId}).")
+          _sc
+        }
         """)
       command("""
-         @transient val sqlContext = {
-           val _sqlContext = org.apache.spark.repl.Main.interp.createSQLContext()
-           println("SQL context available as sqlContext.")
-           _sqlContext
-         }
+        @transient val sqlContext = {
+          val _sqlContext = org.apache.spark.repl.Main.interp.createSQLContext()
+          println("SQL context available as sqlContext.")
+          _sqlContext
+        }
         """)
       command("import org.apache.spark.SparkContext._")
       command("import sqlContext.implicits._")
diff --git a/repl/scala-2.11/src/main/scala/org/apache/spark/repl/SparkILoop.scala b/repl/scala-2.11/src/main/scala/org/apache/spark/repl/SparkILoop.scala
index 33d262558b1fc..e91139fb29f69 100644
--- a/repl/scala-2.11/src/main/scala/org/apache/spark/repl/SparkILoop.scala
+++ b/repl/scala-2.11/src/main/scala/org/apache/spark/repl/SparkILoop.scala
@@ -37,18 +37,19 @@ class SparkILoop(in0: Option[BufferedReader], out: JPrintWriter)
   def initializeSpark() {
     intp.beQuietDuring {
       processLine("""
-         @transient val sc = {
-           val _sc = org.apache.spark.repl.Main.createSparkContext()
-           println("Spark context available as sc.")
-           _sc
-         }
+        @transient val sc = {
+          val _sc = org.apache.spark.repl.Main.createSparkContext()
+          println("Spark context available as sc " +
+            s"(master = ${_sc.master}, app id = ${_sc.applicationId}).")
+          _sc
+        }
         """)
       processLine("""
-         @transient val sqlContext = {
-           val _sqlContext = org.apache.spark.repl.Main.createSQLContext()
-           println("SQL context available as sqlContext.")
-           _sqlContext
-         }
+        @transient val sqlContext = {
+          val _sqlContext = org.apache.spark.repl.Main.createSQLContext()
+          println("SQL context available as sqlContext.")
+          _sqlContext
+        }
         """)
       processLine("import org.apache.spark.SparkContext._")
       processLine("import sqlContext.implicits._")
@@ -85,7 +86,7 @@ class SparkILoop(in0: Option[BufferedReader], out: JPrintWriter)
   /** Available commands */
   override def commands: List[LoopCommand] = sparkStandardCommands
 
-  /** 
+  /**
    * We override `loadFiles` because we need to initialize Spark *before* the REPL
    * sees any files, so that the Spark context is visible in those files. This is a bit of a
    * hack, but there isn't another hook available to us at this point.
@@ -98,7 +99,7 @@ class SparkILoop(in0: Option[BufferedReader], out: JPrintWriter)
 
 object SparkILoop {
 
-  /** 
+  /**
    * Creates an interpreter loop with default settings and feeds
    * the given code to it as input.
    */

From 58d9b260556a89a3d0832d583acafba1df7c6751 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Tue, 24 Nov 2015 14:33:28 -0800
Subject: [PATCH 441/518] [SPARK-11805] free the array in UnsafeExternalSorter
 during spilling

After calling spill() on SortedIterator, the array inside InMemorySorter is not needed, it should be freed during spilling, this could help to join multiple tables with limited memory.

Author: Davies Liu <davies@databricks.com>

Closes #9793 from davies/free_array.
---
 .../unsafe/sort/UnsafeExternalSorter.java     | 10 +++---
 .../unsafe/sort/UnsafeInMemorySorter.java     | 31 ++++++++-----------
 2 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 9a7b2ad06cab6..2e40312674737 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -468,6 +468,12 @@ public long spill() throws IOException {
           }
           allocatedPages.clear();
         }
+
+        // in-memory sorter will not be used after spilling
+        assert(inMemSorter != null);
+        released += inMemSorter.getMemoryUsage();
+        inMemSorter.free();
+        inMemSorter = null;
         return released;
       }
     }
@@ -489,10 +495,6 @@ public void loadNext() throws IOException {
           }
           upstream = nextUpstream;
           nextUpstream = null;
-
-          assert(inMemSorter != null);
-          inMemSorter.free();
-          inMemSorter = null;
         }
         numRecords--;
         upstream.loadNext();
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
index a218ad4623f46..dce1f15a2963c 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -108,6 +108,7 @@ public UnsafeInMemorySorter(
    */
   public void free() {
     consumer.freeArray(array);
+    array = null;
   }
 
   public void reset() {
@@ -160,28 +161,22 @@ public void insertRecord(long recordPointer, long keyPrefix) {
     pos++;
   }
 
-  public static final class SortedIterator extends UnsafeSorterIterator {
+  public final class SortedIterator extends UnsafeSorterIterator {
 
-    private final TaskMemoryManager memoryManager;
-    private final int sortBufferInsertPosition;
-    private final LongArray sortBuffer;
-    private int position = 0;
+    private final int numRecords;
+    private int position;
     private Object baseObject;
     private long baseOffset;
     private long keyPrefix;
     private int recordLength;
 
-    private SortedIterator(
-        TaskMemoryManager memoryManager,
-        int sortBufferInsertPosition,
-        LongArray sortBuffer) {
-      this.memoryManager = memoryManager;
-      this.sortBufferInsertPosition = sortBufferInsertPosition;
-      this.sortBuffer = sortBuffer;
+    private SortedIterator(int numRecords) {
+      this.numRecords = numRecords;
+      this.position = 0;
     }
 
     public SortedIterator clone () {
-      SortedIterator iter = new SortedIterator(memoryManager, sortBufferInsertPosition, sortBuffer);
+      SortedIterator iter = new SortedIterator(numRecords);
       iter.position = position;
       iter.baseObject = baseObject;
       iter.baseOffset = baseOffset;
@@ -192,21 +187,21 @@ public SortedIterator clone () {
 
     @Override
     public boolean hasNext() {
-      return position < sortBufferInsertPosition;
+      return position / 2 < numRecords;
     }
 
     public int numRecordsLeft() {
-      return (sortBufferInsertPosition - position) / 2;
+      return numRecords - position / 2;
     }
 
     @Override
     public void loadNext() {
       // This pointer points to a 4-byte record length, followed by the record's bytes
-      final long recordPointer = sortBuffer.get(position);
+      final long recordPointer = array.get(position);
       baseObject = memoryManager.getPage(recordPointer);
       baseOffset = memoryManager.getOffsetInPage(recordPointer) + 4;  // Skip over record length
       recordLength = Platform.getInt(baseObject, baseOffset - 4);
-      keyPrefix = sortBuffer.get(position + 1);
+      keyPrefix = array.get(position + 1);
       position += 2;
     }
 
@@ -229,6 +224,6 @@ public void loadNext() {
    */
   public SortedIterator getSortedIterator() {
     sorter.sort(array, 0, pos / 2, sortComparator);
-    return new SortedIterator(memoryManager, pos, array);
+    return new SortedIterator(pos / 2);
   }
 }

From 34ca392da7097a1fbe48cd6c3ebff51453ca26ca Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 24 Nov 2015 14:51:01 -0800
Subject: [PATCH 442/518] Added a line of comment to explain why the extra sort
 exists in pivot.

---
 sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
index ee7150cbbfbca..abd531c4ba541 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
@@ -304,7 +304,7 @@ class GroupedData protected[sql](
     // Get the distinct values of the column and sort them so its consistent
     val values = df.select(pivotColumn)
       .distinct()
-      .sort(pivotColumn)
+      .sort(pivotColumn)  // ensure that the output columns are in a consistent logical order
       .map(_.get(0))
       .take(maxValues + 1)
       .toSeq

From c7f95df5c6d8eb2e6f11cf58b704fea34326a5f2 Mon Sep 17 00:00:00 2001
From: Cheng Lian <lian@databricks.com>
Date: Tue, 24 Nov 2015 14:59:14 -0800
Subject: [PATCH 443/518] [SPARK-11783][SQL] Fixes execution Hive client when
 using remote Hive metastore

When using remote Hive metastore, `hive.metastore.uris` is set to the metastore URI.  However, it overrides `javax.jdo.option.ConnectionURL` unexpectedly, thus the execution Hive client connects to the actual remote Hive metastore instead of the Derby metastore created in the temporary directory.  Cleaning this configuration for the execution Hive client fixes this issue.

Author: Cheng Lian <lian@databricks.com>

Closes #9895 from liancheng/spark-11783.clean-remote-metastore-config.
---
 .../org/apache/spark/sql/hive/HiveContext.scala   | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index c0bb5af7d5c85..8a4264194ae8d 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -736,6 +736,21 @@ private[hive] object HiveContext {
       s"jdbc:derby:;databaseName=${localMetastore.getAbsolutePath};create=true")
     propMap.put("datanucleus.rdbms.datastoreAdapterClassName",
       "org.datanucleus.store.rdbms.adapter.DerbyAdapter")
+
+    // SPARK-11783: When "hive.metastore.uris" is set, the metastore connection mode will be
+    // remote (https://cwiki.apache.org/confluence/display/Hive/AdminManual+MetastoreAdmin
+    // mentions that "If hive.metastore.uris is empty local mode is assumed, remote otherwise").
+    // Remote means that the metastore server is running in its own process.
+    // When the mode is remote, configurations like "javax.jdo.option.ConnectionURL" will not be
+    // used (because they are used by remote metastore server that talks to the database).
+    // Because execution Hive should always connects to a embedded derby metastore.
+    // We have to remove the value of hive.metastore.uris. So, the execution Hive client connects
+    // to the actual embedded derby metastore instead of the remote metastore.
+    // You can search HiveConf.ConfVars.METASTOREURIS in the code of HiveConf (in Hive's repo).
+    // Then, you will find that the local metastore mode is only set to true when
+    // hive.metastore.uris is not set.
+    propMap.put(ConfVars.METASTOREURIS.varname, "")
+
     propMap.toMap
   }
 

From 238ae51b66ac12d15fba6aff061804004c5ca6cb Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Tue, 24 Nov 2015 15:54:10 -0800
Subject: [PATCH 444/518] [SPARK-11914][SQL] Support coalesce and repartition
 in Dataset APIs

This PR is to provide two common `coalesce` and `repartition` in Dataset APIs.

After reading the comments of SPARK-9999, I am unclear about the plan for supporting re-partitioning in Dataset APIs. Currently, both RDD APIs and Dataframe APIs provide users such a flexibility to control the number of partitions.

In most traditional RDBMS, they expose the number of partitions, the partitioning columns, the table partitioning methods to DBAs for performance tuning and storage planning. Normally, these parameters could largely affect the query performance. Since the actual performance depends on the workload types, I think it is almost impossible to automate the discovery of the best partitioning strategy for all the scenarios.

I am wondering if Dataset APIs are planning to hide these APIs from users? Feel free to reject my PR if it does not match the plan.

Thank you for your answers. marmbrus rxin cloud-fan

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9899 from gatorsmile/coalesce.
---
 .../scala/org/apache/spark/sql/Dataset.scala  | 19 +++++++++++++++++++
 .../org/apache/spark/sql/DatasetSuite.scala   | 15 +++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 07647508421a4..17e2611790d5a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -152,6 +152,25 @@ class Dataset[T] private[sql](
    */
   def count(): Long = toDF().count()
 
+  /**
+    * Returns a new [[Dataset]] that has exactly `numPartitions` partitions.
+    * @since 1.6.0
+    */
+  def repartition(numPartitions: Int): Dataset[T] = withPlan {
+    Repartition(numPartitions, shuffle = true, _)
+  }
+
+  /**
+    * Returns a new [[Dataset]] that has exactly `numPartitions` partitions.
+    * Similar to coalesce defined on an [[RDD]], this operation results in a narrow dependency, e.g.
+    * if you go from 1000 partitions to 100 partitions, there will not be a shuffle, instead each of
+    * the 100 new partitions will claim 10 of the current partitions.
+    * @since 1.6.0
+    */
+  def coalesce(numPartitions: Int): Dataset[T] = withPlan {
+    Repartition(numPartitions, shuffle = false, _)
+  }
+
   /* *********************** *
    *  Functional Operations  *
    * *********************** */
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 13eede1b17d8b..c253fdbb8c99e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -52,6 +52,21 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     assert(ds.takeAsList(1).get(0) == item)
   }
 
+  test("coalesce, repartition") {
+    val data = (1 to 100).map(i => ClassData(i.toString, i))
+    val ds = data.toDS()
+
+    assert(ds.repartition(10).rdd.partitions.length == 10)
+    checkAnswer(
+      ds.repartition(10),
+      data: _*)
+
+    assert(ds.coalesce(1).rdd.partitions.length == 1)
+    checkAnswer(
+      ds.coalesce(1),
+      data: _*)
+  }
+
   test("as tuple") {
     val data = Seq(("a", 1), ("b", 2)).toDF("a", "b")
     checkAnswer(

From 25bbd3c16e8e8be4d2c43000223d54650e9a3696 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 24 Nov 2015 18:16:07 -0800
Subject: [PATCH 445/518] [SPARK-11967][SQL] Consistent use of varargs for
 multiple paths in DataFrameReader

This patch makes it consistent to use varargs in all DataFrameReader methods, including Parquet, JSON, text, and the generic load function.

Also added a few more API tests for the Java API.

Author: Reynold Xin <rxin@databricks.com>

Closes #9945 from rxin/SPARK-11967.
---
 python/pyspark/sql/readwriter.py              | 19 ++++++----
 .../apache/spark/sql/DataFrameReader.scala    | 36 +++++++++++++++----
 .../apache/spark/sql/JavaDataFrameSuite.java  | 23 ++++++++++++
 sql/core/src/test/resources/text-suite2.txt   |  1 +
 .../org/apache/spark/sql/DataFrameSuite.scala |  2 +-
 5 files changed, 66 insertions(+), 15 deletions(-)
 create mode 100644 sql/core/src/test/resources/text-suite2.txt

diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py
index e8f0d7ec77035..2e75f0c8a1827 100644
--- a/python/pyspark/sql/readwriter.py
+++ b/python/pyspark/sql/readwriter.py
@@ -109,7 +109,7 @@ def options(self, **options):
     def load(self, path=None, format=None, schema=None, **options):
         """Loads data from a data source and returns it as a :class`DataFrame`.
 
-        :param path: optional string for file-system backed data sources.
+        :param path: optional string or a list of string for file-system backed data sources.
         :param format: optional string for format of the data source. Default to 'parquet'.
         :param schema: optional :class:`StructType` for the input schema.
         :param options: all other string options
@@ -118,6 +118,7 @@ def load(self, path=None, format=None, schema=None, **options):
         ...     opt2=1, opt3='str')
         >>> df.dtypes
         [('name', 'string'), ('year', 'int'), ('month', 'int'), ('day', 'int')]
+
         >>> df = sqlContext.read.format('json').load(['python/test_support/sql/people.json',
         ...     'python/test_support/sql/people1.json'])
         >>> df.dtypes
@@ -130,10 +131,8 @@ def load(self, path=None, format=None, schema=None, **options):
         self.options(**options)
         if path is not None:
             if type(path) == list:
-                paths = path
-                gateway = self._sqlContext._sc._gateway
-                jpaths = utils.toJArray(gateway, gateway.jvm.java.lang.String, paths)
-                return self._df(self._jreader.load(jpaths))
+                return self._df(
+                    self._jreader.load(self._sqlContext._sc._jvm.PythonUtils.toSeq(path)))
             else:
                 return self._df(self._jreader.load(path))
         else:
@@ -175,6 +174,8 @@ def json(self, path, schema=None):
             self.schema(schema)
         if isinstance(path, basestring):
             return self._df(self._jreader.json(path))
+        elif type(path) == list:
+            return self._df(self._jreader.json(self._sqlContext._sc._jvm.PythonUtils.toSeq(path)))
         elif isinstance(path, RDD):
             return self._df(self._jreader.json(path._jrdd))
         else:
@@ -205,16 +206,20 @@ def parquet(self, *paths):
 
     @ignore_unicode_prefix
     @since(1.6)
-    def text(self, path):
+    def text(self, paths):
         """Loads a text file and returns a [[DataFrame]] with a single string column named "text".
 
         Each line in the text file is a new row in the resulting DataFrame.
 
+        :param paths: string, or list of strings, for input path(s).
+
         >>> df = sqlContext.read.text('python/test_support/sql/text-test.txt')
         >>> df.collect()
         [Row(value=u'hello'), Row(value=u'this')]
         """
-        return self._df(self._jreader.text(path))
+        if isinstance(paths, basestring):
+            paths = [paths]
+        return self._df(self._jreader.text(self._sqlContext._sc._jvm.PythonUtils.toSeq(paths)))
 
     @since(1.5)
     def orc(self, path):
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
index dcb3737b70fbf..3ed1e55adec6d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
@@ -24,17 +24,17 @@ import scala.collection.JavaConverters._
 import org.apache.hadoop.fs.Path
 import org.apache.hadoop.util.StringUtils
 
+import org.apache.spark.{Logging, Partition}
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.JavaRDD
 import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.SqlParser
 import org.apache.spark.sql.execution.datasources.jdbc.{JDBCPartition, JDBCPartitioningInfo, JDBCRelation}
-import org.apache.spark.sql.execution.datasources.json.{JSONOptions, JSONRelation}
+import org.apache.spark.sql.execution.datasources.json.JSONRelation
 import org.apache.spark.sql.execution.datasources.parquet.ParquetRelation
 import org.apache.spark.sql.execution.datasources.{LogicalRelation, ResolvedDataSource}
 import org.apache.spark.sql.types.StructType
-import org.apache.spark.{Logging, Partition}
-import org.apache.spark.sql.catalyst.{SqlParser, TableIdentifier}
 
 /**
  * :: Experimental ::
@@ -104,6 +104,7 @@ class DataFrameReader private[sql](sqlContext: SQLContext) extends Logging {
    *
    * @since 1.4.0
    */
+  // TODO: Remove this one in Spark 2.0.
   def load(path: String): DataFrame = {
     option("path", path).load()
   }
@@ -130,7 +131,8 @@ class DataFrameReader private[sql](sqlContext: SQLContext) extends Logging {
    *
    * @since 1.6.0
    */
-  def load(paths: Array[String]): DataFrame = {
+  @scala.annotation.varargs
+  def load(paths: String*): DataFrame = {
     option("paths", paths.map(StringUtils.escapeString(_, '\\', ',')).mkString(",")).load()
   }
 
@@ -236,11 +238,30 @@ class DataFrameReader private[sql](sqlContext: SQLContext) extends Logging {
    * <li>`allowNumericLeadingZeros` (default `false`): allows leading zeros in numbers
    * (e.g. 00012)</li>
    *
-   * @param path input path
    * @since 1.4.0
    */
+  // TODO: Remove this one in Spark 2.0.
   def json(path: String): DataFrame = format("json").load(path)
 
+  /**
+   * Loads a JSON file (one object per line) and returns the result as a [[DataFrame]].
+   *
+   * This function goes through the input once to determine the input schema. If you know the
+   * schema in advance, use the version that specifies the schema to avoid the extra scan.
+   *
+   * You can set the following JSON-specific options to deal with non-standard JSON files:
+   * <li>`primitivesAsString` (default `false`): infers all primitive values as a string type</li>
+   * <li>`allowComments` (default `false`): ignores Java/C++ style comment in JSON records</li>
+   * <li>`allowUnquotedFieldNames` (default `false`): allows unquoted JSON field names</li>
+   * <li>`allowSingleQuotes` (default `true`): allows single quotes in addition to double quotes
+   * </li>
+   * <li>`allowNumericLeadingZeros` (default `false`): allows leading zeros in numbers
+   * (e.g. 00012)</li>
+   *
+   * @since 1.6.0
+   */
+  def json(paths: String*): DataFrame = format("json").load(paths : _*)
+
   /**
    * Loads an `JavaRDD[String]` storing JSON objects (one object per record) and
    * returns the result as a [[DataFrame]].
@@ -328,10 +349,11 @@ class DataFrameReader private[sql](sqlContext: SQLContext) extends Logging {
    *   sqlContext.read().text("/path/to/spark/README.md")
    * }}}
    *
-   * @param path input path
+   * @param paths input path
    * @since 1.6.0
    */
-  def text(path: String): DataFrame = format("text").load(path)
+  @scala.annotation.varargs
+  def text(paths: String*): DataFrame = format("text").load(paths : _*)
 
   ///////////////////////////////////////////////////////////////////////////////////////
   // Builder pattern config options
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index a12fed3c0c6af..8e0b2dbca4a98 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -298,4 +298,27 @@ public void pivot() {
     Assert.assertEquals(48000.0, actual[1].getDouble(1), 0.01);
     Assert.assertEquals(30000.0, actual[1].getDouble(2), 0.01);
   }
+
+  public void testGenericLoad() {
+    DataFrame df1 = context.read().format("text").load(
+      Thread.currentThread().getContextClassLoader().getResource("text-suite.txt").toString());
+    Assert.assertEquals(4L, df1.count());
+
+    DataFrame df2 = context.read().format("text").load(
+      Thread.currentThread().getContextClassLoader().getResource("text-suite.txt").toString(),
+      Thread.currentThread().getContextClassLoader().getResource("text-suite2.txt").toString());
+    Assert.assertEquals(5L, df2.count());
+  }
+
+  @Test
+  public void testTextLoad() {
+    DataFrame df1 = context.read().text(
+      Thread.currentThread().getContextClassLoader().getResource("text-suite.txt").toString());
+    Assert.assertEquals(4L, df1.count());
+
+    DataFrame df2 = context.read().text(
+      Thread.currentThread().getContextClassLoader().getResource("text-suite.txt").toString(),
+      Thread.currentThread().getContextClassLoader().getResource("text-suite2.txt").toString());
+    Assert.assertEquals(5L, df2.count());
+  }
 }
diff --git a/sql/core/src/test/resources/text-suite2.txt b/sql/core/src/test/resources/text-suite2.txt
new file mode 100644
index 0000000000000..f9d498c80493c
--- /dev/null
+++ b/sql/core/src/test/resources/text-suite2.txt
@@ -0,0 +1 @@
+This is another file for testing multi path loading.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
index dd6d06512ff60..76e9648aa7533 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala
@@ -897,7 +897,7 @@ class DataFrameSuite extends QueryTest with SharedSQLContext {
       val dir2 = new File(dir, "dir2").getCanonicalPath
       df2.write.format("json").save(dir2)
 
-      checkAnswer(sqlContext.read.format("json").load(Array(dir1, dir2)),
+      checkAnswer(sqlContext.read.format("json").load(dir1, dir2),
         Row(1, 22) :: Row(2, 23) :: Nil)
 
       checkAnswer(sqlContext.read.format("json").load(dir1),

From 4d6bbbc03ddb6650b00eb638e4876a196014c19c Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 24 Nov 2015 18:58:55 -0800
Subject: [PATCH 446/518] [SPARK-11947][SQL] Mark deprecated methods with "This
 will be removed in Spark 2.0."

Also fixed some documentation as I saw them.

Author: Reynold Xin <rxin@databricks.com>

Closes #9930 from rxin/SPARK-11947.
---
 project/MimaExcludes.scala                    |  3 +-
 .../scala/org/apache/spark/sql/Column.scala   | 20 +++--
 .../org/apache/spark/sql/DataFrame.scala      | 72 +++++++++------
 .../scala/org/apache/spark/sql/Dataset.scala  |  1 +
 .../org/apache/spark/sql/SQLContext.scala     | 88 ++++++++++---------
 .../org/apache/spark/sql/SQLImplicits.scala   | 25 +++++-
 .../sql/{ => execution}/SparkSQLParser.scala  | 15 ++--
 .../org/apache/spark/sql/functions.scala      | 52 ++++++-----
 .../SimpleTextHadoopFsRelationSuite.scala     |  6 +-
 9 files changed, 172 insertions(+), 110 deletions(-)
 rename sql/core/src/main/scala/org/apache/spark/sql/{ => execution}/SparkSQLParser.scala (89%)

diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index bb45d1bb12146..54a9ad956d119 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -108,7 +108,8 @@ object MimaExcludes {
         ProblemFilters.exclude[MissingClassProblem](
           "org.apache.spark.rdd.MapPartitionsWithPreparationRDD"),
         ProblemFilters.exclude[MissingClassProblem](
-          "org.apache.spark.rdd.MapPartitionsWithPreparationRDD$")
+          "org.apache.spark.rdd.MapPartitionsWithPreparationRDD$"),
+        ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.SparkSQLParser")
       ) ++ Seq(
         // SPARK-11485
         ProblemFilters.exclude[MissingMethodProblem]("org.apache.spark.sql.DataFrameHolder.df"),
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
index 30c554a85e693..b3cd9e1eff142 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Column.scala
@@ -42,7 +42,8 @@ private[sql] object Column {
 
 /**
  * A [[Column]] where an [[Encoder]] has been given for the expected input and return type.
- * @since 1.6.0
+ * To create a [[TypedColumn]], use the `as` function on a [[Column]].
+ *
  * @tparam T The input type expected for this expression.  Can be `Any` if the expression is type
  *           checked by the analyzer instead of the compiler (i.e. `expr("sum(...)")`).
  * @tparam U The output type of this column.
@@ -51,7 +52,8 @@ private[sql] object Column {
  */
 class TypedColumn[-T, U](
     expr: Expression,
-    private[sql] val encoder: ExpressionEncoder[U]) extends Column(expr) {
+    private[sql] val encoder: ExpressionEncoder[U])
+  extends Column(expr) {
 
   /**
    * Inserts the specific input type and schema into any expressions that are expected to operate
@@ -61,12 +63,11 @@ class TypedColumn[-T, U](
       inputEncoder: ExpressionEncoder[_],
       schema: Seq[Attribute]): TypedColumn[T, U] = {
     val boundEncoder = inputEncoder.bind(schema).asInstanceOf[ExpressionEncoder[Any]]
-    new TypedColumn[T, U] (expr transform {
-      case ta: TypedAggregateExpression if ta.aEncoder.isEmpty =>
-        ta.copy(
-          aEncoder = Some(boundEncoder),
-          children = schema)
-    }, encoder)
+    new TypedColumn[T, U](
+      expr transform { case ta: TypedAggregateExpression if ta.aEncoder.isEmpty =>
+        ta.copy(aEncoder = Some(boundEncoder), children = schema)
+      },
+      encoder)
   }
 }
 
@@ -691,8 +692,9 @@ class Column(protected[sql] val expr: Expression) extends Logging {
    *
    * @group expr_ops
    * @since 1.3.0
+   * @deprecated As of 1.5.0. Use isin. This will be removed in Spark 2.0.
    */
-  @deprecated("use isin", "1.5.0")
+  @deprecated("use isin. This will be removed in Spark 2.0.", "1.5.0")
   @scala.annotation.varargs
   def in(list: Any*): Column = isin(list : _*)
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 5586fc994b98a..5eca1db9525ec 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -1713,9 +1713,9 @@ class DataFrame private[sql](
   ////////////////////////////////////////////////////////////////////////////
 
   /**
-   * @deprecated As of 1.3.0, replaced by `toDF()`.
+   * @deprecated As of 1.3.0, replaced by `toDF()`. This will be removed in Spark 2.0.
    */
-  @deprecated("use toDF", "1.3.0")
+  @deprecated("Use toDF. This will be removed in Spark 2.0.", "1.3.0")
   def toSchemaRDD: DataFrame = this
 
   /**
@@ -1725,9 +1725,9 @@ class DataFrame private[sql](
    * given name; if you pass `false`, it will throw if the table already
    * exists.
    * @group output
-   * @deprecated As of 1.340, replaced by `write().jdbc()`.
+   * @deprecated As of 1.340, replaced by `write().jdbc()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.jdbc()", "1.4.0")
+  @deprecated("Use write.jdbc(). This will be removed in Spark 2.0.", "1.4.0")
   def createJDBCTable(url: String, table: String, allowExisting: Boolean): Unit = {
     val w = if (allowExisting) write.mode(SaveMode.Overwrite) else write
     w.jdbc(url, table, new Properties)
@@ -1744,9 +1744,9 @@ class DataFrame private[sql](
    * the RDD in order via the simple statement
    * `INSERT INTO table VALUES (?, ?, ..., ?)` should not fail.
    * @group output
-   * @deprecated As of 1.4.0, replaced by `write().jdbc()`.
+   * @deprecated As of 1.4.0, replaced by `write().jdbc()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.jdbc()", "1.4.0")
+  @deprecated("Use write.jdbc(). This will be removed in Spark 2.0.", "1.4.0")
   def insertIntoJDBC(url: String, table: String, overwrite: Boolean): Unit = {
     val w = if (overwrite) write.mode(SaveMode.Overwrite) else write.mode(SaveMode.Append)
     w.jdbc(url, table, new Properties)
@@ -1757,9 +1757,9 @@ class DataFrame private[sql](
    * Files that are written out using this method can be read back in as a [[DataFrame]]
    * using the `parquetFile` function in [[SQLContext]].
    * @group output
-   * @deprecated As of 1.4.0, replaced by `write().parquet()`.
+   * @deprecated As of 1.4.0, replaced by `write().parquet()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.parquet(path)", "1.4.0")
+  @deprecated("Use write.parquet(path). This will be removed in Spark 2.0.", "1.4.0")
   def saveAsParquetFile(path: String): Unit = {
     write.format("parquet").mode(SaveMode.ErrorIfExists).save(path)
   }
@@ -1782,8 +1782,9 @@ class DataFrame private[sql](
    *
    * @group output
    * @deprecated As of 1.4.0, replaced by `write().saveAsTable(tableName)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.saveAsTable(tableName)", "1.4.0")
+  @deprecated("Use write.saveAsTable(tableName). This will be removed in Spark 2.0.", "1.4.0")
   def saveAsTable(tableName: String): Unit = {
     write.mode(SaveMode.ErrorIfExists).saveAsTable(tableName)
   }
@@ -1805,8 +1806,10 @@ class DataFrame private[sql](
    *
    * @group output
    * @deprecated As of 1.4.0, replaced by `write().mode(mode).saveAsTable(tableName)`.
+   *              This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.mode(mode).saveAsTable(tableName)", "1.4.0")
+  @deprecated("Use write.mode(mode).saveAsTable(tableName). This will be removed in Spark 2.0.",
+    "1.4.0")
   def saveAsTable(tableName: String, mode: SaveMode): Unit = {
     write.mode(mode).saveAsTable(tableName)
   }
@@ -1829,8 +1832,10 @@ class DataFrame private[sql](
    *
    * @group output
    * @deprecated As of 1.4.0, replaced by `write().format(source).saveAsTable(tableName)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.format(source).saveAsTable(tableName)", "1.4.0")
+  @deprecated("Use write.format(source).saveAsTable(tableName). This will be removed in Spark 2.0.",
+    "1.4.0")
   def saveAsTable(tableName: String, source: String): Unit = {
     write.format(source).saveAsTable(tableName)
   }
@@ -1853,8 +1858,10 @@ class DataFrame private[sql](
    *
    * @group output
    * @deprecated As of 1.4.0, replaced by `write().mode(mode).saveAsTable(tableName)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.format(source).mode(mode).saveAsTable(tableName)", "1.4.0")
+  @deprecated("Use write.format(source).mode(mode).saveAsTable(tableName). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def saveAsTable(tableName: String, source: String, mode: SaveMode): Unit = {
     write.format(source).mode(mode).saveAsTable(tableName)
   }
@@ -1877,9 +1884,10 @@ class DataFrame private[sql](
    * @group output
    * @deprecated As of 1.4.0, replaced by
    *            `write().format(source).mode(mode).options(options).saveAsTable(tableName)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.format(source).mode(mode).options(options).saveAsTable(tableName)",
-    "1.4.0")
+  @deprecated("Use write.format(source).mode(mode).options(options).saveAsTable(tableName). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def saveAsTable(
       tableName: String,
       source: String,
@@ -1907,9 +1915,10 @@ class DataFrame private[sql](
    * @group output
    * @deprecated As of 1.4.0, replaced by
    *            `write().format(source).mode(mode).options(options).saveAsTable(tableName)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.format(source).mode(mode).options(options).saveAsTable(tableName)",
-    "1.4.0")
+  @deprecated("Use write.format(source).mode(mode).options(options).saveAsTable(tableName). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def saveAsTable(
       tableName: String,
       source: String,
@@ -1923,9 +1932,9 @@ class DataFrame private[sql](
    * using the default data source configured by spark.sql.sources.default and
    * [[SaveMode.ErrorIfExists]] as the save mode.
    * @group output
-   * @deprecated As of 1.4.0, replaced by `write().save(path)`.
+   * @deprecated As of 1.4.0, replaced by `write().save(path)`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.save(path)", "1.4.0")
+  @deprecated("Use write.save(path). This will be removed in Spark 2.0.", "1.4.0")
   def save(path: String): Unit = {
     write.save(path)
   }
@@ -1935,8 +1944,9 @@ class DataFrame private[sql](
    * using the default data source configured by spark.sql.sources.default.
    * @group output
    * @deprecated As of 1.4.0, replaced by `write().mode(mode).save(path)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.mode(mode).save(path)", "1.4.0")
+  @deprecated("Use write.mode(mode).save(path). This will be removed in Spark 2.0.", "1.4.0")
   def save(path: String, mode: SaveMode): Unit = {
     write.mode(mode).save(path)
   }
@@ -1946,8 +1956,9 @@ class DataFrame private[sql](
    * using [[SaveMode.ErrorIfExists]] as the save mode.
    * @group output
    * @deprecated As of 1.4.0, replaced by `write().format(source).save(path)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.format(source).save(path)", "1.4.0")
+  @deprecated("Use write.format(source).save(path). This will be removed in Spark 2.0.", "1.4.0")
   def save(path: String, source: String): Unit = {
     write.format(source).save(path)
   }
@@ -1957,8 +1968,10 @@ class DataFrame private[sql](
    * [[SaveMode]] specified by mode.
    * @group output
    * @deprecated As of 1.4.0, replaced by `write().format(source).mode(mode).save(path)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.format(source).mode(mode).save(path)", "1.4.0")
+  @deprecated("Use write.format(source).mode(mode).save(path). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def save(path: String, source: String, mode: SaveMode): Unit = {
     write.format(source).mode(mode).save(path)
   }
@@ -1969,8 +1982,10 @@ class DataFrame private[sql](
    * @group output
    * @deprecated As of 1.4.0, replaced by
    *            `write().format(source).mode(mode).options(options).save(path)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.format(source).mode(mode).options(options).save()", "1.4.0")
+  @deprecated("Use write.format(source).mode(mode).options(options).save(). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def save(
       source: String,
       mode: SaveMode,
@@ -1985,8 +2000,10 @@ class DataFrame private[sql](
    * @group output
    * @deprecated As of 1.4.0, replaced by
    *            `write().format(source).mode(mode).options(options).save(path)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.format(source).mode(mode).options(options).save()", "1.4.0")
+  @deprecated("Use write.format(source).mode(mode).options(options).save(). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def save(
       source: String,
       mode: SaveMode,
@@ -1994,14 +2011,15 @@ class DataFrame private[sql](
     write.format(source).mode(mode).options(options).save()
   }
 
-
   /**
    * Adds the rows from this RDD to the specified table, optionally overwriting the existing data.
    * @group output
    * @deprecated As of 1.4.0, replaced by
    *            `write().mode(SaveMode.Append|SaveMode.Overwrite).saveAsTable(tableName)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.mode(SaveMode.Append|SaveMode.Overwrite).saveAsTable(tableName)", "1.4.0")
+  @deprecated("Use write.mode(SaveMode.Append|SaveMode.Overwrite).saveAsTable(tableName). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def insertInto(tableName: String, overwrite: Boolean): Unit = {
     write.mode(if (overwrite) SaveMode.Overwrite else SaveMode.Append).insertInto(tableName)
   }
@@ -2012,8 +2030,10 @@ class DataFrame private[sql](
    * @group output
    * @deprecated As of 1.4.0, replaced by
    *            `write().mode(SaveMode.Append).saveAsTable(tableName)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use write.mode(SaveMode.Append).saveAsTable(tableName)", "1.4.0")
+  @deprecated("Use write.mode(SaveMode.Append).saveAsTable(tableName). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def insertInto(tableName: String): Unit = {
     write.mode(SaveMode.Append).insertInto(tableName)
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 17e2611790d5a..dd84b8bc11e2b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql.execution.{Queryable, QueryExecution}
 import org.apache.spark.sql.types.StructType
 
 /**
+ * :: Experimental ::
  * A [[Dataset]] is a strongly typed collection of objects that can be transformed in parallel
  * using functional or relational operations.
  *
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index 39471d2fb79a7..46bf544fd885f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -942,33 +942,33 @@ class SQLContext private[sql](
   ////////////////////////////////////////////////////////////////////////////
 
   /**
-   * @deprecated As of 1.3.0, replaced by `createDataFrame()`.
+   * @deprecated As of 1.3.0, replaced by `createDataFrame()`. This will be removed in Spark 2.0.
    */
-  @deprecated("use createDataFrame", "1.3.0")
+  @deprecated("Use createDataFrame. This will be removed in Spark 2.0.", "1.3.0")
   def applySchema(rowRDD: RDD[Row], schema: StructType): DataFrame = {
     createDataFrame(rowRDD, schema)
   }
 
   /**
-   * @deprecated As of 1.3.0, replaced by `createDataFrame()`.
+   * @deprecated As of 1.3.0, replaced by `createDataFrame()`. This will be removed in Spark 2.0.
    */
-  @deprecated("use createDataFrame", "1.3.0")
+  @deprecated("Use createDataFrame. This will be removed in Spark 2.0.", "1.3.0")
   def applySchema(rowRDD: JavaRDD[Row], schema: StructType): DataFrame = {
     createDataFrame(rowRDD, schema)
   }
 
   /**
-   * @deprecated As of 1.3.0, replaced by `createDataFrame()`.
+   * @deprecated As of 1.3.0, replaced by `createDataFrame()`. This will be removed in Spark 2.0.
    */
-  @deprecated("use createDataFrame", "1.3.0")
+  @deprecated("Use createDataFrame. This will be removed in Spark 2.0.", "1.3.0")
   def applySchema(rdd: RDD[_], beanClass: Class[_]): DataFrame = {
     createDataFrame(rdd, beanClass)
   }
 
   /**
-   * @deprecated As of 1.3.0, replaced by `createDataFrame()`.
+   * @deprecated As of 1.3.0, replaced by `createDataFrame()`. This will be removed in Spark 2.0.
    */
-  @deprecated("use createDataFrame", "1.3.0")
+  @deprecated("Use createDataFrame. This will be removed in Spark 2.0.", "1.3.0")
   def applySchema(rdd: JavaRDD[_], beanClass: Class[_]): DataFrame = {
     createDataFrame(rdd, beanClass)
   }
@@ -978,9 +978,9 @@ class SQLContext private[sql](
    * [[DataFrame]] if no paths are passed in.
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().parquet()`.
+   * @deprecated As of 1.4.0, replaced by `read().parquet()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.parquet()", "1.4.0")
+  @deprecated("Use read.parquet(). This will be removed in Spark 2.0.", "1.4.0")
   @scala.annotation.varargs
   def parquetFile(paths: String*): DataFrame = {
     if (paths.isEmpty) {
@@ -995,9 +995,9 @@ class SQLContext private[sql](
    * It goes through the entire dataset once to determine the schema.
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonFile(path: String): DataFrame = {
     read.json(path)
   }
@@ -1007,18 +1007,18 @@ class SQLContext private[sql](
    * returning the result as a [[DataFrame]].
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonFile(path: String, schema: StructType): DataFrame = {
     read.schema(schema).json(path)
   }
 
   /**
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonFile(path: String, samplingRatio: Double): DataFrame = {
     read.option("samplingRatio", samplingRatio.toString).json(path)
   }
@@ -1029,9 +1029,9 @@ class SQLContext private[sql](
    * It goes through the entire dataset once to determine the schema.
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonRDD(json: RDD[String]): DataFrame = read.json(json)
 
   /**
@@ -1040,9 +1040,9 @@ class SQLContext private[sql](
    * It goes through the entire dataset once to determine the schema.
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonRDD(json: JavaRDD[String]): DataFrame = read.json(json)
 
   /**
@@ -1050,9 +1050,9 @@ class SQLContext private[sql](
    * returning the result as a [[DataFrame]].
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonRDD(json: RDD[String], schema: StructType): DataFrame = {
     read.schema(schema).json(json)
   }
@@ -1062,9 +1062,9 @@ class SQLContext private[sql](
    * schema, returning the result as a [[DataFrame]].
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonRDD(json: JavaRDD[String], schema: StructType): DataFrame = {
     read.schema(schema).json(json)
   }
@@ -1074,9 +1074,9 @@ class SQLContext private[sql](
    * schema, returning the result as a [[DataFrame]].
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonRDD(json: RDD[String], samplingRatio: Double): DataFrame = {
     read.option("samplingRatio", samplingRatio.toString).json(json)
   }
@@ -1086,9 +1086,9 @@ class SQLContext private[sql](
    * schema, returning the result as a [[DataFrame]].
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().json()`.
+   * @deprecated As of 1.4.0, replaced by `read().json()`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.json()", "1.4.0")
+  @deprecated("Use read.json(). This will be removed in Spark 2.0.", "1.4.0")
   def jsonRDD(json: JavaRDD[String], samplingRatio: Double): DataFrame = {
     read.option("samplingRatio", samplingRatio.toString).json(json)
   }
@@ -1098,9 +1098,9 @@ class SQLContext private[sql](
    * using the default data source configured by spark.sql.sources.default.
    *
    * @group genericdata
-   * @deprecated As of 1.4.0, replaced by `read().load(path)`.
+   * @deprecated As of 1.4.0, replaced by `read().load(path)`. This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.load(path)", "1.4.0")
+  @deprecated("Use read.load(path). This will be removed in Spark 2.0.", "1.4.0")
   def load(path: String): DataFrame = {
     read.load(path)
   }
@@ -1110,8 +1110,9 @@ class SQLContext private[sql](
    *
    * @group genericdata
    * @deprecated As of 1.4.0, replaced by `read().format(source).load(path)`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.format(source).load(path)", "1.4.0")
+  @deprecated("Use read.format(source).load(path). This will be removed in Spark 2.0.", "1.4.0")
   def load(path: String, source: String): DataFrame = {
     read.format(source).load(path)
   }
@@ -1122,8 +1123,10 @@ class SQLContext private[sql](
    *
    * @group genericdata
    * @deprecated As of 1.4.0, replaced by `read().format(source).options(options).load()`.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use read.format(source).options(options).load()", "1.4.0")
+  @deprecated("Use read.format(source).options(options).load(). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def load(source: String, options: java.util.Map[String, String]): DataFrame = {
     read.options(options).format(source).load()
   }
@@ -1135,7 +1138,8 @@ class SQLContext private[sql](
    * @group genericdata
    * @deprecated As of 1.4.0, replaced by `read().format(source).options(options).load()`.
    */
-  @deprecated("Use read.format(source).options(options).load()", "1.4.0")
+  @deprecated("Use read.format(source).options(options).load(). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def load(source: String, options: Map[String, String]): DataFrame = {
     read.options(options).format(source).load()
   }
@@ -1148,7 +1152,8 @@ class SQLContext private[sql](
    * @deprecated As of 1.4.0, replaced by
    *            `read().format(source).schema(schema).options(options).load()`.
    */
-  @deprecated("Use read.format(source).schema(schema).options(options).load()", "1.4.0")
+  @deprecated("Use read.format(source).schema(schema).options(options).load(). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def load(source: String, schema: StructType, options: java.util.Map[String, String]): DataFrame =
   {
     read.format(source).schema(schema).options(options).load()
@@ -1162,7 +1167,8 @@ class SQLContext private[sql](
    * @deprecated As of 1.4.0, replaced by
    *            `read().format(source).schema(schema).options(options).load()`.
    */
-  @deprecated("Use read.format(source).schema(schema).options(options).load()", "1.4.0")
+  @deprecated("Use read.format(source).schema(schema).options(options).load(). " +
+    "This will be removed in Spark 2.0.", "1.4.0")
   def load(source: String, schema: StructType, options: Map[String, String]): DataFrame = {
     read.format(source).schema(schema).options(options).load()
   }
@@ -1172,9 +1178,9 @@ class SQLContext private[sql](
    * url named table.
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().jdbc()`.
+   * @deprecated As of 1.4.0, replaced by `read().jdbc()`. This will be removed in Spark 2.0.
    */
-  @deprecated("use read.jdbc()", "1.4.0")
+  @deprecated("Use read.jdbc(). This will be removed in Spark 2.0.", "1.4.0")
   def jdbc(url: String, table: String): DataFrame = {
     read.jdbc(url, table, new Properties)
   }
@@ -1190,9 +1196,9 @@ class SQLContext private[sql](
    * @param numPartitions the number of partitions.  the range `minValue`-`maxValue` will be split
    *                      evenly into this many partitions
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().jdbc()`.
+   * @deprecated As of 1.4.0, replaced by `read().jdbc()`. This will be removed in Spark 2.0.
    */
-  @deprecated("use read.jdbc()", "1.4.0")
+  @deprecated("Use read.jdbc(). This will be removed in Spark 2.0.", "1.4.0")
   def jdbc(
       url: String,
       table: String,
@@ -1210,9 +1216,9 @@ class SQLContext private[sql](
    * of the [[DataFrame]].
    *
    * @group specificdata
-   * @deprecated As of 1.4.0, replaced by `read().jdbc()`.
+   * @deprecated As of 1.4.0, replaced by `read().jdbc()`. This will be removed in Spark 2.0.
    */
-  @deprecated("use read.jdbc()", "1.4.0")
+  @deprecated("Use read.jdbc(). This will be removed in Spark 2.0.", "1.4.0")
   def jdbc(url: String, table: String, theParts: Array[String]): DataFrame = {
     read.jdbc(url, table, theParts, new Properties)
   }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
index 25ffdcde17717..6735d02954b8d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLImplicits.scala
@@ -30,19 +30,38 @@ import org.apache.spark.unsafe.types.UTF8String
 
 /**
  * A collection of implicit methods for converting common Scala objects into [[DataFrame]]s.
+ *
+ * @since 1.6.0
  */
 abstract class SQLImplicits {
+
   protected def _sqlContext: SQLContext
 
+  /** @since 1.6.0 */
   implicit def newProductEncoder[T <: Product : TypeTag]: Encoder[T] = ExpressionEncoder()
 
+  /** @since 1.6.0 */
   implicit def newIntEncoder: Encoder[Int] = ExpressionEncoder()
+
+  /** @since 1.6.0 */
   implicit def newLongEncoder: Encoder[Long] = ExpressionEncoder()
+
+  /** @since 1.6.0 */
   implicit def newDoubleEncoder: Encoder[Double] = ExpressionEncoder()
+
+  /** @since 1.6.0 */
   implicit def newFloatEncoder: Encoder[Float] = ExpressionEncoder()
+
+  /** @since 1.6.0 */
   implicit def newByteEncoder: Encoder[Byte] = ExpressionEncoder()
+
+  /** @since 1.6.0 */
   implicit def newShortEncoder: Encoder[Short] = ExpressionEncoder()
+  /** @since 1.6.0 */
+
   implicit def newBooleanEncoder: Encoder[Boolean] = ExpressionEncoder()
+
+  /** @since 1.6.0 */
   implicit def newStringEncoder: Encoder[String] = ExpressionEncoder()
 
   /**
@@ -84,9 +103,9 @@ abstract class SQLImplicits {
     DataFrameHolder(_sqlContext.createDataFrame(data))
   }
 
-  // Do NOT add more implicit conversions. They are likely to break source compatibility by
-  // making existing implicit conversions ambiguous. In particular, RDD[Double] is dangerous
-  // because of [[DoubleRDDFunctions]].
+  // Do NOT add more implicit conversions for primitive types.
+  // They are likely to break source compatibility by making existing implicit conversions
+  // ambiguous. In particular, RDD[Double] is dangerous because of [[DoubleRDDFunctions]].
 
   /**
    * Creates a single column DataFrame from an RDD[Int].
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SparkSQLParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSQLParser.scala
similarity index 89%
rename from sql/core/src/main/scala/org/apache/spark/sql/SparkSQLParser.scala
rename to sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSQLParser.scala
index ea8fce6ca9cf2..b3e8d0d84937e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSQLParser.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSQLParser.scala
@@ -15,24 +15,23 @@
  * limitations under the License.
  */
 
-package org.apache.spark.sql
+package org.apache.spark.sql.execution
 
 import scala.util.parsing.combinator.RegexParsers
 
 import org.apache.spark.sql.catalyst.AbstractSparkSQLParser
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference}
-import org.apache.spark.sql.catalyst.plans.logical.{DescribeFunction, LogicalPlan, ShowFunctions}
-import org.apache.spark.sql.execution._
+import org.apache.spark.sql.catalyst.plans.logical
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.types.StringType
 
-
 /**
  * The top level Spark SQL parser. This parser recognizes syntaxes that are available for all SQL
  * dialects supported by Spark SQL, and delegates all the other syntaxes to the `fallback` parser.
  *
  * @param fallback A function that parses an input string to a logical plan
  */
-private[sql] class SparkSQLParser(fallback: String => LogicalPlan) extends AbstractSparkSQLParser {
+class SparkSQLParser(fallback: String => LogicalPlan) extends AbstractSparkSQLParser {
 
   // A parser for the key-value part of the "SET [key = [value ]]" syntax
   private object SetCommandParser extends RegexParsers {
@@ -100,14 +99,14 @@ private[sql] class SparkSQLParser(fallback: String => LogicalPlan) extends Abstr
         case _ ~ dbName => ShowTablesCommand(dbName)
       }
     | SHOW ~ FUNCTIONS ~> ((ident <~ ".").? ~ (ident | stringLit)).? ^^ {
-        case Some(f) => ShowFunctions(f._1, Some(f._2))
-        case None => ShowFunctions(None, None)
+        case Some(f) => logical.ShowFunctions(f._1, Some(f._2))
+        case None => logical.ShowFunctions(None, None)
       }
     )
 
   private lazy val desc: Parser[LogicalPlan] =
     DESCRIBE ~ FUNCTION ~> EXTENDED.? ~ (ident | stringLit) ^^ {
-      case isExtended ~ functionName => DescribeFunction(functionName, isExtended.isDefined)
+      case isExtended ~ functionName => logical.DescribeFunction(functionName, isExtended.isDefined)
     }
 
   private lazy val others: Parser[LogicalPlan] =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 6137ce3a70fdb..77dd5bc72508b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -17,8 +17,6 @@
 
 package org.apache.spark.sql
 
-
-
 import scala.language.implicitConversions
 import scala.reflect.runtime.universe.{TypeTag, typeTag}
 import scala.util.Try
@@ -39,11 +37,11 @@ import org.apache.spark.util.Utils
  * "bridge" methods due to the use of covariant return types.
  *
  * {{{
- * In LegacyFunctions:
- * public abstract org.apache.spark.sql.Column avg(java.lang.String);
+ *   // In LegacyFunctions:
+ *   public abstract org.apache.spark.sql.Column avg(java.lang.String);
  *
- * In functions:
- * public static org.apache.spark.sql.TypedColumn<java.lang.Object, java.lang.Object> avg(...);
+ *   // In functions:
+ *   public static org.apache.spark.sql.TypedColumn<java.lang.Object, java.lang.Object> avg(...);
  * }}}
  *
  * This allows us to use the same functions both in typed [[Dataset]] operations and untyped
@@ -2528,8 +2526,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function0[_], returnType: DataType): Column = withExpr {
     ScalaUDF(f, returnType, Seq())
   }
@@ -2541,8 +2540,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function1[_, _], returnType: DataType, arg1: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr))
   }
@@ -2554,8 +2554,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function2[_, _, _], returnType: DataType, arg1: Column, arg2: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr))
   }
@@ -2567,8 +2568,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function3[_, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr))
   }
@@ -2580,8 +2582,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function4[_, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr))
   }
@@ -2593,8 +2596,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function5[_, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr))
   }
@@ -2606,8 +2610,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function6[_, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr))
   }
@@ -2619,8 +2624,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function7[_, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr, arg7.expr))
   }
@@ -2632,8 +2638,9 @@ object functions extends LegacyFunctions {
     * @group udf_funcs
     * @since 1.3.0
     * @deprecated As of 1.5.0, since it's redundant with udf()
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function8[_, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr, arg7.expr, arg8.expr))
   }
@@ -2644,9 +2651,10 @@ object functions extends LegacyFunctions {
     *
     * @group udf_funcs
     * @since 1.3.0
-    * @deprecated As of 1.5.0, since it's redundant with udf()
+    * @deprecated As of 1.5.0, since it's redundant with udf().
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function9[_, _, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column, arg9: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr, arg7.expr, arg8.expr, arg9.expr))
   }
@@ -2657,9 +2665,10 @@ object functions extends LegacyFunctions {
     *
     * @group udf_funcs
     * @since 1.3.0
-    * @deprecated As of 1.5.0, since it's redundant with udf()
+    * @deprecated As of 1.5.0, since it's redundant with udf().
+   *              This will be removed in Spark 2.0.
     */
-  @deprecated("Use udf", "1.5.0")
+  @deprecated("Use udf. This will be removed in Spark 2.0.", "1.5.0")
   def callUDF(f: Function10[_, _, _, _, _, _, _, _, _, _, _], returnType: DataType, arg1: Column, arg2: Column, arg3: Column, arg4: Column, arg5: Column, arg6: Column, arg7: Column, arg8: Column, arg9: Column, arg10: Column): Column = withExpr {
     ScalaUDF(f, returnType, Seq(arg1.expr, arg2.expr, arg3.expr, arg4.expr, arg5.expr, arg6.expr, arg7.expr, arg8.expr, arg9.expr, arg10.expr))
   }
@@ -2700,9 +2709,10 @@ object functions extends LegacyFunctions {
    *
    * @group udf_funcs
    * @since 1.4.0
-   * @deprecated As of 1.5.0, since it was not coherent to have two functions callUdf and callUDF
+   * @deprecated As of 1.5.0, since it was not coherent to have two functions callUdf and callUDF.
+   *             This will be removed in Spark 2.0.
    */
-  @deprecated("Use callUDF", "1.5.0")
+  @deprecated("Use callUDF. This will be removed in Spark 2.0.", "1.5.0")
   def callUdf(udfName: String, cols: Column*): Column = withExpr {
     // Note: we avoid using closures here because on file systems that are case-insensitive, the
     // compiled class file for the closure here will conflict with the one in callUDF (upper case).
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
index 81af684ba0bf1..b554d135e4b5c 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/SimpleTextHadoopFsRelationSuite.scala
@@ -80,7 +80,11 @@ class SimpleTextHadoopFsRelationSuite extends HadoopFsRelationTest with Predicat
 
   private var partitionedDF: DataFrame = _
 
-  private val partitionedDataSchema: StructType = StructType('a.int :: 'b.int :: 'c.string :: Nil)
+  private val partitionedDataSchema: StructType =
+    new StructType()
+      .add("a", IntegerType)
+      .add("b", IntegerType)
+      .add("c", StringType)
 
   protected override def beforeAll(): Unit = {
     this.tempPath = Utils.createTempDir()

From a5d988763319f63a8e2b58673dd4f9098f17c835 Mon Sep 17 00:00:00 2001
From: Burak Yavuz <brkyvz@gmail.com>
Date: Tue, 24 Nov 2015 20:58:47 -0800
Subject: [PATCH 447/518] [STREAMING][FLAKY-TEST] Catch execution context race
 condition in `FileBasedWriteAheadLog.close()`

There is a race condition in `FileBasedWriteAheadLog.close()`, where if delete's of old log files are in progress, the write ahead log may close, and result in a `RejectedExecutionException`. This is okay, and should be handled gracefully.

Example test failures:
https://amplab.cs.berkeley.edu/jenkins/job/Spark-1.6-SBT/AMPLAB_JENKINS_BUILD_PROFILE=hadoop1.0,label=spark-test/95/testReport/junit/org.apache.spark.streaming.util/BatchedWriteAheadLogWithCloseFileAfterWriteSuite/BatchedWriteAheadLog___clean_old_logs/

The reason the test fails is in `afterEach`, `writeAheadLog.close` is called, and there may still be async deletes in flight.

tdas zsxwing

Author: Burak Yavuz <brkyvz@gmail.com>

Closes #9953 from brkyvz/flaky-ss.
---
 .../streaming/util/FileBasedWriteAheadLog.scala  | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala
index 72705f1a9c010..f5165f7c39122 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/FileBasedWriteAheadLog.scala
@@ -17,7 +17,7 @@
 package org.apache.spark.streaming.util
 
 import java.nio.ByteBuffer
-import java.util.concurrent.ThreadPoolExecutor
+import java.util.concurrent.{RejectedExecutionException, ThreadPoolExecutor}
 import java.util.{Iterator => JIterator}
 
 import scala.collection.JavaConverters._
@@ -176,10 +176,16 @@ private[streaming] class FileBasedWriteAheadLog(
     }
     oldLogFiles.foreach { logInfo =>
       if (!executionContext.isShutdown) {
-        val f = Future { deleteFile(logInfo) }(executionContext)
-        if (waitForCompletion) {
-          import scala.concurrent.duration._
-          Await.ready(f, 1 second)
+        try {
+          val f = Future { deleteFile(logInfo) }(executionContext)
+          if (waitForCompletion) {
+            import scala.concurrent.duration._
+            Await.ready(f, 1 second)
+          }
+        } catch {
+          case e: RejectedExecutionException =>
+            logWarning("Execution context shutdown before deleting old WriteAheadLogs. " +
+              "This would not affect recovery correctness.", e)
         }
       }
     }

From 151d7c2baf18403e6e59e97c80c8bcded6148038 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Tue, 24 Nov 2015 21:30:53 -0800
Subject: [PATCH 448/518] [SPARK-10621][SQL] Consistent naming for functions in
 SQL, Python, Scala

Author: Reynold Xin <rxin@databricks.com>

Closes #9948 from rxin/SPARK-10621.
---
 python/pyspark/sql/functions.py               | 111 +++++++++++++---
 .../org/apache/spark/sql/functions.scala      | 124 ++++++++++++++----
 2 files changed, 196 insertions(+), 39 deletions(-)

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index a1ca723bbd7ab..e3786e0fa5fb2 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -150,18 +150,18 @@ def _():
 
 _window_functions = {
     'rowNumber':
-        """returns a sequential number starting at 1 within a window partition.
-
-        This is equivalent to the ROW_NUMBER function in SQL.""",
+        """.. note:: Deprecated in 1.6, use row_number instead.""",
+    'row_number':
+        """returns a sequential number starting at 1 within a window partition.""",
     'denseRank':
+        """.. note:: Deprecated in 1.6, use dense_rank instead.""",
+    'dense_rank':
         """returns the rank of rows within a window partition, without any gaps.
 
         The difference between rank and denseRank is that denseRank leaves no gaps in ranking
         sequence when there are ties. That is, if you were ranking a competition using denseRank
         and had three people tie for second place, you would say that all three were in second
-        place and that the next person came in third.
-
-        This is equivalent to the DENSE_RANK function in SQL.""",
+        place and that the next person came in third.""",
     'rank':
         """returns the rank of rows within a window partition.
 
@@ -172,14 +172,14 @@ def _():
 
         This is equivalent to the RANK function in SQL.""",
     'cumeDist':
+        """.. note:: Deprecated in 1.6, use cume_dist instead.""",
+    'cume_dist':
         """returns the cumulative distribution of values within a window partition,
-        i.e. the fraction of rows that are below the current row.
-
-        This is equivalent to the CUME_DIST function in SQL.""",
+        i.e. the fraction of rows that are below the current row.""",
     'percentRank':
-        """returns the relative rank (i.e. percentile) of rows within a window partition.
-
-        This is equivalent to the PERCENT_RANK function in SQL.""",
+        """.. note:: Deprecated in 1.6, use percent_rank instead.""",
+    'percent_rank':
+        """returns the relative rank (i.e. percentile) of rows within a window partition.""",
 }
 
 for _name, _doc in _functions.items():
@@ -189,7 +189,7 @@ def _():
 for _name, _doc in _binary_mathfunctions.items():
     globals()[_name] = since(1.4)(_create_binary_mathfunction(_name, _doc))
 for _name, _doc in _window_functions.items():
-    globals()[_name] = since(1.4)(_create_window_function(_name, _doc))
+    globals()[_name] = since(1.6)(_create_window_function(_name, _doc))
 for _name, _doc in _functions_1_6.items():
     globals()[_name] = since(1.6)(_create_function(_name, _doc))
 del _name, _doc
@@ -288,6 +288,38 @@ def countDistinct(col, *cols):
 
 @since(1.4)
 def monotonicallyIncreasingId():
+    """
+    .. note:: Deprecated in 1.6, use monotonically_increasing_id instead.
+    """
+    return monotonically_increasing_id()
+
+
+@since(1.6)
+def input_file_name():
+    """Creates a string column for the file name of the current Spark task.
+    """
+    sc = SparkContext._active_spark_context
+    return Column(sc._jvm.functions.input_file_name())
+
+
+@since(1.6)
+def isnan(col):
+    """An expression that returns true iff the column is NaN.
+    """
+    sc = SparkContext._active_spark_context
+    return Column(sc._jvm.functions.isnan(_to_java_column(col)))
+
+
+@since(1.6)
+def isnull(col):
+    """An expression that returns true iff the column is null.
+    """
+    sc = SparkContext._active_spark_context
+    return Column(sc._jvm.functions.isnull(_to_java_column(col)))
+
+
+@since(1.6)
+def monotonically_increasing_id():
     """A column that generates monotonically increasing 64-bit integers.
 
     The generated ID is guaranteed to be monotonically increasing and unique, but not consecutive.
@@ -300,11 +332,21 @@ def monotonicallyIncreasingId():
     0, 1, 2, 8589934592 (1L << 33), 8589934593, 8589934594.
 
     >>> df0 = sc.parallelize(range(2), 2).mapPartitions(lambda x: [(1,), (2,), (3,)]).toDF(['col1'])
-    >>> df0.select(monotonicallyIncreasingId().alias('id')).collect()
+    >>> df0.select(monotonically_increasing_id().alias('id')).collect()
     [Row(id=0), Row(id=1), Row(id=2), Row(id=8589934592), Row(id=8589934593), Row(id=8589934594)]
     """
     sc = SparkContext._active_spark_context
-    return Column(sc._jvm.functions.monotonicallyIncreasingId())
+    return Column(sc._jvm.functions.monotonically_increasing_id())
+
+
+@since(1.6)
+def nanvl(col1, col2):
+    """Returns col1 if it is not NaN, or col2 if col1 is NaN.
+
+    Both inputs should be floating point columns (DoubleType or FloatType).
+    """
+    sc = SparkContext._active_spark_context
+    return Column(sc._jvm.functions.nanvl(_to_java_column(col1), _to_java_column(col2)))
 
 
 @since(1.4)
@@ -382,15 +424,23 @@ def shiftRightUnsigned(col, numBits):
 
 @since(1.4)
 def sparkPartitionId():
+    """
+    .. note:: Deprecated in 1.6, use spark_partition_id instead.
+    """
+    return spark_partition_id()
+
+
+@since(1.6)
+def spark_partition_id():
     """A column for partition ID of the Spark task.
 
     Note that this is indeterministic because it depends on data partitioning and task scheduling.
 
-    >>> df.repartition(1).select(sparkPartitionId().alias("pid")).collect()
+    >>> df.repartition(1).select(spark_partition_id().alias("pid")).collect()
     [Row(pid=0), Row(pid=0)]
     """
     sc = SparkContext._active_spark_context
-    return Column(sc._jvm.functions.sparkPartitionId())
+    return Column(sc._jvm.functions.spark_partition_id())
 
 
 @since(1.5)
@@ -1410,6 +1460,33 @@ def explode(col):
     return Column(jc)
 
 
+@since(1.6)
+def get_json_object(col, path):
+    """
+    Extracts json object from a json string based on json path specified, and returns json string
+    of the extracted json object. It will return null if the input json string is invalid.
+
+    :param col: string column in json format
+    :param path: path to the json object to extract
+    """
+    sc = SparkContext._active_spark_context
+    jc = sc._jvm.functions.get_json_object(_to_java_column(col), path)
+    return Column(jc)
+
+
+@since(1.6)
+def json_tuple(col, fields):
+    """Creates a new row for a json column according to the given field names.
+
+    :param col: string column in json format
+    :param fields: list of fields to extract
+
+    """
+    sc = SparkContext._active_spark_context
+    jc = sc._jvm.functions.json_tuple(_to_java_column(col), fields)
+    return Column(jc)
+
+
 @since(1.5)
 def size(col):
     """
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 77dd5bc72508b..276c5dfc8b062 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -472,6 +472,13 @@ object functions extends LegacyFunctions {
   // Window functions
   //////////////////////////////////////////////////////////////////////////////////////////////
 
+  /**
+   * @group window_funcs
+   * @deprecated As of 1.6.0, replaced by `cume_dist`. This will be removed in Spark 2.0.
+   */
+  @deprecated("Use cume_dist. This will be removed in Spark 2.0.", "1.6.0")
+  def cumeDist(): Column = cume_dist()
+
   /**
    * Window function: returns the cumulative distribution of values within a window partition,
    * i.e. the fraction of rows that are below the current row.
@@ -481,13 +488,17 @@ object functions extends LegacyFunctions {
    *   cumeDist(x) = number of values before (and including) x / N
    * }}}
    *
-   *
-   * This is equivalent to the CUME_DIST function in SQL.
-   *
    * @group window_funcs
-   * @since 1.4.0
+   * @since 1.6.0
    */
-  def cumeDist(): Column = withExpr { UnresolvedWindowFunction("cume_dist", Nil) }
+  def cume_dist(): Column = withExpr { UnresolvedWindowFunction("cume_dist", Nil) }
+
+  /**
+   * @group window_funcs
+   * @deprecated As of 1.6.0, replaced by `dense_rank`. This will be removed in Spark 2.0.
+   */
+  @deprecated("Use dense_rank. This will be removed in Spark 2.0.", "1.6.0")
+  def denseRank(): Column = dense_rank()
 
   /**
    * Window function: returns the rank of rows within a window partition, without any gaps.
@@ -497,12 +508,10 @@ object functions extends LegacyFunctions {
    * and had three people tie for second place, you would say that all three were in second
    * place and that the next person came in third.
    *
-   * This is equivalent to the DENSE_RANK function in SQL.
-   *
    * @group window_funcs
-   * @since 1.4.0
+   * @since 1.6.0
    */
-  def denseRank(): Column = withExpr { UnresolvedWindowFunction("dense_rank", Nil) }
+  def dense_rank(): Column = withExpr { UnresolvedWindowFunction("dense_rank", Nil) }
 
   /**
    * Window function: returns the value that is `offset` rows before the current row, and
@@ -620,6 +629,13 @@ object functions extends LegacyFunctions {
    */
   def ntile(n: Int): Column = withExpr { UnresolvedWindowFunction("ntile", lit(n).expr :: Nil) }
 
+  /**
+   * @group window_funcs
+   * @deprecated As of 1.6.0, replaced by `percent_rank`. This will be removed in Spark 2.0.
+   */
+  @deprecated("Use percent_rank. This will be removed in Spark 2.0.", "1.6.0")
+  def percentRank(): Column = percent_rank()
+
   /**
    * Window function: returns the relative rank (i.e. percentile) of rows within a window partition.
    *
@@ -631,9 +647,9 @@ object functions extends LegacyFunctions {
    * This is equivalent to the PERCENT_RANK function in SQL.
    *
    * @group window_funcs
-   * @since 1.4.0
+   * @since 1.6.0
    */
-  def percentRank(): Column = withExpr { UnresolvedWindowFunction("percent_rank", Nil) }
+  def percent_rank(): Column = withExpr { UnresolvedWindowFunction("percent_rank", Nil) }
 
   /**
    * Window function: returns the rank of rows within a window partition.
@@ -650,15 +666,20 @@ object functions extends LegacyFunctions {
    */
   def rank(): Column = withExpr { UnresolvedWindowFunction("rank", Nil) }
 
+  /**
+   * @group window_funcs
+   * @deprecated As of 1.6.0, replaced by `row_number`. This will be removed in Spark 2.0.
+   */
+  @deprecated("Use row_number. This will be removed in Spark 2.0.", "1.6.0")
+  def rowNumber(): Column = row_number()
+
   /**
    * Window function: returns a sequential number starting at 1 within a window partition.
    *
-   * This is equivalent to the ROW_NUMBER function in SQL.
-   *
    * @group window_funcs
-   * @since 1.4.0
+   * @since 1.6.0
    */
-  def rowNumber(): Column = withExpr { UnresolvedWindowFunction("row_number", Nil) }
+  def row_number(): Column = withExpr { UnresolvedWindowFunction("row_number", Nil) }
 
   //////////////////////////////////////////////////////////////////////////////////////////////
   // Non-aggregate functions
@@ -720,20 +741,43 @@ object functions extends LegacyFunctions {
   @scala.annotation.varargs
   def coalesce(e: Column*): Column = withExpr { Coalesce(e.map(_.expr)) }
 
+  /**
+   * @group normal_funcs
+   * @deprecated As of 1.6.0, replaced by `input_file_name`. This will be removed in Spark 2.0.
+   */
+  @deprecated("Use input_file_name. This will be removed in Spark 2.0.", "1.6.0")
+  def inputFileName(): Column = input_file_name()
+
   /**
    * Creates a string column for the file name of the current Spark task.
    *
    * @group normal_funcs
+   * @since 1.6.0
    */
-  def inputFileName(): Column = withExpr { InputFileName() }
+  def input_file_name(): Column = withExpr { InputFileName() }
+
+  /**
+   * @group normal_funcs
+   * @deprecated As of 1.6.0, replaced by `isnan`. This will be removed in Spark 2.0.
+   */
+  @deprecated("Use isnan. This will be removed in Spark 2.0.", "1.6.0")
+  def isNaN(e: Column): Column = isnan(e)
 
   /**
    * Return true iff the column is NaN.
    *
    * @group normal_funcs
-   * @since 1.5.0
+   * @since 1.6.0
+   */
+  def isnan(e: Column): Column = withExpr { IsNaN(e.expr) }
+
+  /**
+   * Return true iff the column is null.
+   *
+   * @group normal_funcs
+   * @since 1.6.0
    */
-  def isNaN(e: Column): Column = withExpr { IsNaN(e.expr) }
+  def isnull(e: Column): Column = withExpr { IsNull(e.expr) }
 
   /**
    * A column expression that generates monotonically increasing 64-bit integers.
@@ -750,7 +794,24 @@ object functions extends LegacyFunctions {
    * @group normal_funcs
    * @since 1.4.0
    */
-  def monotonicallyIncreasingId(): Column = withExpr { MonotonicallyIncreasingID() }
+  def monotonicallyIncreasingId(): Column = monotonically_increasing_id()
+
+  /**
+   * A column expression that generates monotonically increasing 64-bit integers.
+   *
+   * The generated ID is guaranteed to be monotonically increasing and unique, but not consecutive.
+   * The current implementation puts the partition ID in the upper 31 bits, and the record number
+   * within each partition in the lower 33 bits. The assumption is that the data frame has
+   * less than 1 billion partitions, and each partition has less than 8 billion records.
+   *
+   * As an example, consider a [[DataFrame]] with two partitions, each with 3 records.
+   * This expression would return the following IDs:
+   * 0, 1, 2, 8589934592 (1L << 33), 8589934593, 8589934594.
+   *
+   * @group normal_funcs
+   * @since 1.6.0
+   */
+  def monotonically_increasing_id(): Column = withExpr { MonotonicallyIncreasingID() }
 
   /**
    * Returns col1 if it is not NaN, or col2 if col1 is NaN.
@@ -825,15 +886,23 @@ object functions extends LegacyFunctions {
    */
   def randn(): Column = randn(Utils.random.nextLong)
 
+  /**
+   * @group normal_funcs
+   * @since 1.4.0
+   * @deprecated As of 1.6.0, replaced by `spark_partition_id`. This will be removed in Spark 2.0.
+   */
+  @deprecated("Use cume_dist. This will be removed in Spark 2.0.", "1.6.0")
+  def sparkPartitionId(): Column = spark_partition_id()
+
   /**
    * Partition ID of the Spark task.
    *
    * Note that this is indeterministic because it depends on data partitioning and task scheduling.
    *
    * @group normal_funcs
-   * @since 1.4.0
+   * @since 1.6.0
    */
-  def sparkPartitionId(): Column = withExpr { SparkPartitionID() }
+  def spark_partition_id(): Column = withExpr { SparkPartitionID() }
 
   /**
    * Computes the square root of the specified float value.
@@ -2305,6 +2374,17 @@ object functions extends LegacyFunctions {
    */
   def explode(e: Column): Column = withExpr { Explode(e.expr) }
 
+  /**
+   * Extracts json object from a json string based on json path specified, and returns json string
+   * of the extracted json object. It will return null if the input json string is invalid.
+   *
+   * @group collection_funcs
+   * @since 1.6.0
+   */
+  def get_json_object(e: Column, path: String): Column = withExpr {
+    GetJsonObject(e.expr, lit(path).expr)
+  }
+
   /**
    * Creates a new row for a json column according to the given field names.
    *
@@ -2313,7 +2393,7 @@ object functions extends LegacyFunctions {
    */
   @scala.annotation.varargs
   def json_tuple(json: Column, fields: String*): Column = withExpr {
-    require(fields.length > 0, "at least 1 field name should be given.")
+    require(fields.nonEmpty, "at least 1 field name should be given.")
     JsonTuple(json.expr +: fields.map(Literal.apply))
   }
 

From 2169886883d33b33acf378ac42a626576b342df1 Mon Sep 17 00:00:00 2001
From: Tathagata Das <tathagata.das1565@gmail.com>
Date: Tue, 24 Nov 2015 23:13:01 -0800
Subject: [PATCH 449/518] [SPARK-11979][STREAMING] Empty TrackStateRDD cannot
 be checkpointed and recovered from checkpoint file

This solves the following exception caused when empty state RDD is checkpointed and recovered. The root cause is that an empty OpenHashMapBasedStateMap cannot be deserialized as the initialCapacity is set to zero.
```
Job aborted due to stage failure: Task 0 in stage 6.0 failed 1 times, most recent failure: Lost task 0.0 in stage 6.0 (TID 20, localhost): java.lang.IllegalArgumentException: requirement failed: Invalid initial capacity
	at scala.Predef$.require(Predef.scala:233)
	at org.apache.spark.streaming.util.OpenHashMapBasedStateMap.<init>(StateMap.scala:96)
	at org.apache.spark.streaming.util.OpenHashMapBasedStateMap.<init>(StateMap.scala:86)
	at org.apache.spark.streaming.util.OpenHashMapBasedStateMap.readObject(StateMap.scala:291)
	at sun.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
	at sun.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:57)
	at sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
	at java.lang.reflect.Method.invoke(Method.java:606)
	at java.io.ObjectStreamClass.invokeReadObject(ObjectStreamClass.java:1017)
	at java.io.ObjectInputStream.readSerialData(ObjectInputStream.java:1893)
	at java.io.ObjectInputStream.readOrdinaryObject(ObjectInputStream.java:1798)
	at java.io.ObjectInputStream.readObject0(ObjectInputStream.java:1350)
	at java.io.ObjectInputStream.defaultReadFields(ObjectInputStream.java:1990)
	at java.io.ObjectInputStream.readSerialData(ObjectInputStream.java:1915)
	at java.io.ObjectInputStream.readOrdinaryObject(ObjectInputStream.java:1798)
	at java.io.ObjectInputStream.readObject0(ObjectInputStream.java:1350)
	at java.io.ObjectInputStream.readObject(ObjectInputStream.java:370)
	at org.apache.spark.serializer.JavaDeserializationStream.readObject(JavaSerializer.scala:76)
	at org.apache.spark.serializer.DeserializationStream$$anon$1.getNext(Serializer.scala:181)
	at org.apache.spark.util.NextIterator.hasNext(NextIterator.scala:73)
	at scala.collection.Iterator$$anon$13.hasNext(Iterator.scala:371)
	at scala.collection.Iterator$class.foreach(Iterator.scala:727)
	at scala.collection.AbstractIterator.foreach(Iterator.scala:1157)
	at scala.collection.generic.Growable$class.$plus$plus$eq(Growable.scala:48)
	at scala.collection.mutable.ArrayBuffer.$plus$plus$eq(ArrayBuffer.scala:103)
	at scala.collection.mutable.ArrayBuffer.$plus$plus$eq(ArrayBuffer.scala:47)
	at scala.collection.TraversableOnce$class.to(TraversableOnce.scala:273)
	at scala.collection.AbstractIterator.to(Iterator.scala:1157)
	at scala.collection.TraversableOnce$class.toBuffer(TraversableOnce.scala:265)
	at scala.collection.AbstractIterator.toBuffer(Iterator.scala:1157)
	at scala.collection.TraversableOnce$class.toArray(TraversableOnce.scala:252)
	at scala.collection.AbstractIterator.toArray(Iterator.scala:1157)
	at org.apache.spark.rdd.RDD$$anonfun$collect$1$$anonfun$12.apply(RDD.scala:921)
	at org.apache.spark.rdd.RDD$$anonfun$collect$1$$anonfun$12.apply(RDD.scala:921)
	at org.apache.spark.SparkContext$$anonfun$runJob$5.apply(SparkContext.scala:1858)
	at org.apache.spark.SparkContext$$anonfun$runJob$5.apply(SparkContext.scala:1858)
	at org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:66)
	at org.apache.spark.scheduler.Task.run(Task.scala:88)
	at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:214)
	at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1145)
	at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:615)
	at java.lang.Thread.run(Thread.java:744)
```

Author: Tathagata Das <tathagata.das1565@gmail.com>

Closes #9958 from tdas/SPARK-11979.
---
 .../spark/streaming/util/StateMap.scala       | 19 +++++++-----
 .../spark/streaming/StateMapSuite.scala       | 30 ++++++++++++-------
 .../streaming/rdd/TrackStateRDDSuite.scala    | 10 +++++++
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala b/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
index 34287c3e00908..3f139ad138c88 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/util/StateMap.scala
@@ -59,7 +59,7 @@ private[streaming] object StateMap {
   def create[K: ClassTag, S: ClassTag](conf: SparkConf): StateMap[K, S] = {
     val deltaChainThreshold = conf.getInt("spark.streaming.sessionByKey.deltaChainThreshold",
       DELTA_CHAIN_LENGTH_THRESHOLD)
-    new OpenHashMapBasedStateMap[K, S](64, deltaChainThreshold)
+    new OpenHashMapBasedStateMap[K, S](deltaChainThreshold)
   }
 }
 
@@ -79,7 +79,7 @@ private[streaming] class EmptyStateMap[K: ClassTag, S: ClassTag] extends StateMa
 /** Implementation of StateMap based on Spark's [[org.apache.spark.util.collection.OpenHashMap]] */
 private[streaming] class OpenHashMapBasedStateMap[K: ClassTag, S: ClassTag](
     @transient @volatile var parentStateMap: StateMap[K, S],
-    initialCapacity: Int = 64,
+    initialCapacity: Int = DEFAULT_INITIAL_CAPACITY,
     deltaChainThreshold: Int = DELTA_CHAIN_LENGTH_THRESHOLD
   ) extends StateMap[K, S] { self =>
 
@@ -89,12 +89,14 @@ private[streaming] class OpenHashMapBasedStateMap[K: ClassTag, S: ClassTag](
     deltaChainThreshold = deltaChainThreshold)
 
   def this(deltaChainThreshold: Int) = this(
-    initialCapacity = 64, deltaChainThreshold = deltaChainThreshold)
+    initialCapacity = DEFAULT_INITIAL_CAPACITY, deltaChainThreshold = deltaChainThreshold)
 
   def this() = this(DELTA_CHAIN_LENGTH_THRESHOLD)
 
-  @transient @volatile private var deltaMap =
-    new OpenHashMap[K, StateInfo[S]](initialCapacity)
+  require(initialCapacity >= 1, "Invalid initial capacity")
+  require(deltaChainThreshold >= 1, "Invalid delta chain threshold")
+
+  @transient @volatile private var deltaMap = new OpenHashMap[K, StateInfo[S]](initialCapacity)
 
   /** Get the session data if it exists */
   override def get(key: K): Option[S] = {
@@ -284,9 +286,10 @@ private[streaming] class OpenHashMapBasedStateMap[K: ClassTag, S: ClassTag](
     // Read the data of the parent map. Keep reading records, until the limiter is reached
     // First read the approximate number of records to expect and allocate properly size
     // OpenHashMap
-    val parentSessionStoreSizeHint = inputStream.readInt()
+    val parentStateMapSizeHint = inputStream.readInt()
+    val newStateMapInitialCapacity = math.max(parentStateMapSizeHint, DEFAULT_INITIAL_CAPACITY)
     val newParentSessionStore = new OpenHashMapBasedStateMap[K, S](
-      initialCapacity = parentSessionStoreSizeHint, deltaChainThreshold)
+      initialCapacity = newStateMapInitialCapacity, deltaChainThreshold)
 
     // Read the records until the limit marking object has been reached
     var parentSessionLoopDone = false
@@ -338,4 +341,6 @@ private[streaming] object OpenHashMapBasedStateMap {
   class LimitMarker(val num: Int) extends Serializable
 
   val DELTA_CHAIN_LENGTH_THRESHOLD = 20
+
+  val DEFAULT_INITIAL_CAPACITY = 64
 }
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala
index 48d3b41b66cbf..c4a01eaea739e 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/StateMapSuite.scala
@@ -122,23 +122,27 @@ class StateMapSuite extends SparkFunSuite {
 
   test("OpenHashMapBasedStateMap - serializing and deserializing") {
     val map1 = new OpenHashMapBasedStateMap[Int, Int]()
+    testSerialization(map1, "error deserializing and serialized empty map")
+
     map1.put(1, 100, 1)
     map1.put(2, 200, 2)
+    testSerialization(map1, "error deserializing and serialized map with data + no delta")
 
     val map2 = map1.copy()
+    // Do not test compaction
+    assert(map2.asInstanceOf[OpenHashMapBasedStateMap[_, _]].shouldCompact === false)
+    testSerialization(map2, "error deserializing and serialized map with 1 delta + no new data")
+
     map2.put(3, 300, 3)
     map2.put(4, 400, 4)
+    testSerialization(map2, "error deserializing and serialized map with 1 delta + new data")
 
     val map3 = map2.copy()
+    assert(map3.asInstanceOf[OpenHashMapBasedStateMap[_, _]].shouldCompact === false)
+    testSerialization(map3, "error deserializing and serialized map with 2 delta + no new data")
     map3.put(3, 600, 3)
     map3.remove(2)
-
-    // Do not test compaction
-    assert(map3.asInstanceOf[OpenHashMapBasedStateMap[_, _]].shouldCompact === false)
-
-    val deser_map3 = Utils.deserialize[StateMap[Int, Int]](
-      Utils.serialize(map3), Thread.currentThread().getContextClassLoader)
-    assertMap(deser_map3, map3, 1, "Deserialized map not same as original map")
+    testSerialization(map3, "error deserializing and serialized map with 2 delta + new data")
   }
 
   test("OpenHashMapBasedStateMap - serializing and deserializing with compaction") {
@@ -156,11 +160,9 @@ class StateMapSuite extends SparkFunSuite {
     assert(map.deltaChainLength > deltaChainThreshold)
     assert(map.shouldCompact === true)
 
-    val deser_map = Utils.deserialize[OpenHashMapBasedStateMap[Int, Int]](
-      Utils.serialize(map), Thread.currentThread().getContextClassLoader)
+    val deser_map = testSerialization(map, "Deserialized + compacted map not same as original map")
     assert(deser_map.deltaChainLength < deltaChainThreshold)
     assert(deser_map.shouldCompact === false)
-    assertMap(deser_map, map, 1, "Deserialized + compacted map not same as original map")
   }
 
   test("OpenHashMapBasedStateMap - all possible sequences of operations with copies ") {
@@ -265,6 +267,14 @@ class StateMapSuite extends SparkFunSuite {
     assertMap(stateMap, refMap.toMap, time, "Final state map does not match reference map")
   }
 
+  private def testSerialization[MapType <: StateMap[Int, Int]](
+    map: MapType, msg: String): MapType = {
+    val deserMap = Utils.deserialize[MapType](
+      Utils.serialize(map), Thread.currentThread().getContextClassLoader)
+    assertMap(deserMap, map, 1, msg)
+    deserMap
+  }
+
   // Assert whether all the data and operations on a state map matches that of a reference state map
   private def assertMap(
       mapToTest: StateMap[Int, Int],
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
index 0feb3af1abb0f..3b2d43f2ce581 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/rdd/TrackStateRDDSuite.scala
@@ -332,6 +332,16 @@ class TrackStateRDDSuite extends SparkFunSuite with RDDCheckpointTester with Bef
       makeStateRDDWithLongLineageParenttateRDD, reliableCheckpoint = true, rddCollectFunc _)
   }
 
+  test("checkpointing empty state RDD") {
+    val emptyStateRDD = TrackStateRDD.createFromPairRDD[Int, Int, Int, Int](
+      sc.emptyRDD[(Int, Int)], new HashPartitioner(10), Time(0))
+    emptyStateRDD.checkpoint()
+    assert(emptyStateRDD.flatMap { _.stateMap.getAll() }.collect().isEmpty)
+    val cpRDD = sc.checkpointFile[TrackStateRDDRecord[Int, Int, Int]](
+      emptyStateRDD.getCheckpointFile.get)
+    assert(cpRDD.flatMap { _.stateMap.getAll() }.collect().isEmpty)
+  }
+
   /** Assert whether the `trackStateByKey` operation generates expected results */
   private def assertOperation[K: ClassTag, V: ClassTag, S: ClassTag, T: ClassTag](
       testStateRDD: TrackStateRDD[K, V, S, T],

From 2610e06124c7fc0b2b1cfb2e3050a35ab492fb71 Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Wed, 25 Nov 2015 01:02:36 -0800
Subject: [PATCH 450/518] [SPARK-11970][SQL] Adding JoinType into JoinWith and
 support Sample in Dataset API

Except inner join, maybe the other join types are also useful when users are using the joinWith function. Thus, added the joinType into the existing joinWith call in Dataset APIs.

Also providing another joinWith interface for the cartesian-join-like functionality.

Please provide your opinions. marmbrus rxin cloud-fan Thank you!

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9921 from gatorsmile/joinWith.
---
 .../scala/org/apache/spark/sql/Dataset.scala  | 45 +++++++++++++++----
 .../org/apache/spark/sql/DatasetSuite.scala   | 36 ++++++++++++---
 2 files changed, 65 insertions(+), 16 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index dd84b8bc11e2b..97eb5b969280d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -20,16 +20,16 @@ package org.apache.spark.sql
 import scala.collection.JavaConverters._
 
 import org.apache.spark.annotation.Experimental
-import org.apache.spark.rdd.RDD
 import org.apache.spark.api.java.function._
-
+import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.encoders._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAlias
-import org.apache.spark.sql.catalyst.plans.Inner
+import org.apache.spark.sql.catalyst.plans.JoinType
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.execution.{Queryable, QueryExecution}
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.Utils
 
 /**
  * :: Experimental ::
@@ -83,7 +83,6 @@ class Dataset[T] private[sql](
 
   /**
    * Returns the schema of the encoded form of the objects in this [[Dataset]].
-   *
    * @since 1.6.0
    */
   def schema: StructType = resolvedTEncoder.schema
@@ -185,7 +184,6 @@ class Dataset[T] private[sql](
    *     .transform(featurize)
    *     .transform(...)
    * }}}
-   *
    * @since 1.6.0
    */
   def transform[U](t: Dataset[T] => Dataset[U]): Dataset[U] = t(this)
@@ -453,6 +451,21 @@ class Dataset[T] private[sql](
       c5: TypedColumn[T, U5]): Dataset[(U1, U2, U3, U4, U5)] =
     selectUntyped(c1, c2, c3, c4, c5).asInstanceOf[Dataset[(U1, U2, U3, U4, U5)]]
 
+  /**
+   * Returns a new [[Dataset]] by sampling a fraction of records.
+   * @since 1.6.0
+   */
+  def sample(withReplacement: Boolean, fraction: Double, seed: Long) : Dataset[T] =
+    withPlan(Sample(0.0, fraction, withReplacement, seed, _))
+
+  /**
+   * Returns a new [[Dataset]] by sampling a fraction of records, using a random seed.
+   * @since 1.6.0
+   */
+  def sample(withReplacement: Boolean, fraction: Double) : Dataset[T] = {
+    sample(withReplacement, fraction, Utils.random.nextLong)
+  }
+
   /* **************** *
    *  Set operations  *
    * **************** */
@@ -511,13 +524,17 @@ class Dataset[T] private[sql](
    * types as well as working with relational data where either side of the join has column
    * names in common.
    *
+   * @param other Right side of the join.
+   * @param condition Join expression.
+   * @param joinType One of: `inner`, `outer`, `left_outer`, `right_outer`, `leftsemi`.
    * @since 1.6.0
    */
-  def joinWith[U](other: Dataset[U], condition: Column): Dataset[(T, U)] = {
+  def joinWith[U](other: Dataset[U], condition: Column, joinType: String): Dataset[(T, U)] = {
     val left = this.logicalPlan
     val right = other.logicalPlan
 
-    val joined = sqlContext.executePlan(Join(left, right, Inner, Some(condition.expr)))
+    val joined = sqlContext.executePlan(Join(left, right, joinType =
+      JoinType(joinType), Some(condition.expr)))
     val leftOutput = joined.analyzed.output.take(left.output.length)
     val rightOutput = joined.analyzed.output.takeRight(right.output.length)
 
@@ -540,6 +557,18 @@ class Dataset[T] private[sql](
     }
   }
 
+  /**
+   * Using inner equi-join to join this [[Dataset]] returning a [[Tuple2]] for each pair
+   * where `condition` evaluates to true.
+   *
+   * @param other Right side of the join.
+   * @param condition Join expression.
+   * @since 1.6.0
+   */
+  def joinWith[U](other: Dataset[U], condition: Column): Dataset[(T, U)] = {
+    joinWith(other, condition, "inner")
+  }
+
   /* ************************** *
    *  Gather to Driver Actions  *
    * ************************** */
@@ -584,7 +613,6 @@ class Dataset[T] private[sql](
    *
    * Running take requires moving data into the application's driver process, and doing so with
    * a very large `n` can crash the driver process with OutOfMemoryError.
-   *
    * @since 1.6.0
    */
   def take(num: Int): Array[T] = withPlan(Limit(Literal(num), _)).collect()
@@ -594,7 +622,6 @@ class Dataset[T] private[sql](
    *
    * Running take requires moving data into the application's driver process, and doing so with
    * a very large `n` can crash the driver process with OutOfMemoryError.
-   *
    * @since 1.6.0
    */
   def takeAsList(num: Int): java.util.List[T] = java.util.Arrays.asList(take(num) : _*)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index c253fdbb8c99e..7d539180ded9e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -185,17 +185,23 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     val ds2 = Seq(1, 2).toDS().as("b")
 
     checkAnswer(
-      ds1.joinWith(ds2, $"a.value" === $"b.value"),
+      ds1.joinWith(ds2, $"a.value" === $"b.value", "inner"),
       (1, 1), (2, 2))
   }
 
-  test("joinWith, expression condition") {
-    val ds1 = Seq(ClassData("a", 1), ClassData("b", 2)).toDS()
-    val ds2 = Seq(("a", 1), ("b", 2)).toDS()
+  test("joinWith, expression condition, outer join") {
+    val nullInteger = null.asInstanceOf[Integer]
+    val nullString = null.asInstanceOf[String]
+    val ds1 = Seq(ClassNullableData("a", 1),
+      ClassNullableData("c", 3)).toDS()
+    val ds2 = Seq(("a", new Integer(1)),
+      ("b", new Integer(2))).toDS()
 
     checkAnswer(
-      ds1.joinWith(ds2, $"_1" === $"a"),
-      (ClassData("a", 1), ("a", 1)), (ClassData("b", 2), ("b", 2)))
+      ds1.joinWith(ds2, $"_1" === $"a", "outer"),
+      (ClassNullableData("a", 1), ("a", new Integer(1))),
+      (ClassNullableData("c", 3), (nullString, nullInteger)),
+      (ClassNullableData(nullString, nullInteger), ("b", new Integer(2))))
   }
 
   test("joinWith tuple with primitive, expression") {
@@ -225,7 +231,6 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       ds1.joinWith(ds2, $"a._2" === $"b._2").as("ab").joinWith(ds3, $"ab._1._2" === $"c._2"),
       ((("a", 1), ("a", 1)), ("a", 1)),
       ((("b", 2), ("b", 2)), ("b", 2)))
-
   }
 
   test("groupBy function, keys") {
@@ -367,6 +372,22 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
       1 -> "a", 2 -> "bc", 3 -> "d")
   }
 
+  test("sample with replacement") {
+    val n = 100
+    val data = sparkContext.parallelize(1 to n, 2).toDS()
+    checkAnswer(
+      data.sample(withReplacement = true, 0.05, seed = 13),
+      5, 10, 52, 73)
+  }
+
+  test("sample without replacement") {
+    val n = 100
+    val data = sparkContext.parallelize(1 to n, 2).toDS()
+    checkAnswer(
+      data.sample(withReplacement = false, 0.05, seed = 13),
+      3, 17, 27, 58, 62)
+  }
+
   test("SPARK-11436: we should rebind right encoder when join 2 datasets") {
     val ds1 = Seq("1", "2").toDS().as("a")
     val ds2 = Seq(2, 3).toDS().as("b")
@@ -440,6 +461,7 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
 
 
 case class ClassData(a: String, b: Int)
+case class ClassNullableData(a: String, b: Integer)
 
 /**
  * A class used to test serialization using encoders. This class throws exceptions when using

From a0f1a11837bfffb76582499d36fbaf21a1d628cb Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 25 Nov 2015 01:03:18 -0800
Subject: [PATCH 451/518] [SPARK-11981][SQL] Move implementations of methods
 back to DataFrame from Queryable

Also added show methods to Dataset.

Author: Reynold Xin <rxin@databricks.com>

Closes #9964 from rxin/SPARK-11981.
---
 .../org/apache/spark/sql/DataFrame.scala      | 35 ++++++++-
 .../scala/org/apache/spark/sql/Dataset.scala  | 77 ++++++++++++++++++-
 .../spark/sql/execution/Queryable.scala       | 32 ++------
 3 files changed, 111 insertions(+), 33 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index 5eca1db9525ec..d8319b9a97fcf 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.plans.{Inner, JoinType}
 import org.apache.spark.sql.catalyst.{CatalystTypeConverters, ScalaReflection, SqlParser}
-import org.apache.spark.sql.execution.{EvaluatePython, FileRelation, LogicalRDD, QueryExecution, Queryable, SQLExecution}
+import org.apache.spark.sql.execution.{EvaluatePython, ExplainCommand, FileRelation, LogicalRDD, QueryExecution, Queryable, SQLExecution}
 import org.apache.spark.sql.execution.datasources.{CreateTableUsingAsSelect, LogicalRelation}
 import org.apache.spark.sql.execution.datasources.json.JacksonGenerator
 import org.apache.spark.sql.sources.HadoopFsRelation
@@ -112,8 +112,8 @@ private[sql] object DataFrame {
  */
 @Experimental
 class DataFrame private[sql](
-    @transient val sqlContext: SQLContext,
-    @DeveloperApi @transient val queryExecution: QueryExecution)
+    @transient override val sqlContext: SQLContext,
+    @DeveloperApi @transient override val queryExecution: QueryExecution)
   extends Queryable with Serializable {
 
   // Note for Spark contributors: if adding or updating any action in `DataFrame`, please make sure
@@ -282,6 +282,35 @@ class DataFrame private[sql](
    */
   def schema: StructType = queryExecution.analyzed.schema
 
+  /**
+   * Prints the schema to the console in a nice tree format.
+   * @group basic
+   * @since 1.3.0
+   */
+  // scalastyle:off println
+  override def printSchema(): Unit = println(schema.treeString)
+  // scalastyle:on println
+
+  /**
+   * Prints the plans (logical and physical) to the console for debugging purposes.
+   * @group basic
+   * @since 1.3.0
+   */
+  override def explain(extended: Boolean): Unit = {
+    val explain = ExplainCommand(queryExecution.logical, extended = extended)
+    sqlContext.executePlan(explain).executedPlan.executeCollect().foreach {
+      // scalastyle:off println
+      r => println(r.getString(0))
+      // scalastyle:on println
+    }
+  }
+
+  /**
+   * Prints the physical plan to the console for debugging purposes.
+   * @since 1.3.0
+   */
+  override def explain(): Unit = explain(extended = false)
+
   /**
    * Returns all column names and their data types as an array.
    * @group basic
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
index 97eb5b969280d..da4600133290f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -61,8 +61,8 @@ import org.apache.spark.util.Utils
  */
 @Experimental
 class Dataset[T] private[sql](
-    @transient val sqlContext: SQLContext,
-    @transient val queryExecution: QueryExecution,
+    @transient override val sqlContext: SQLContext,
+    @transient override val queryExecution: QueryExecution,
     tEncoder: Encoder[T]) extends Queryable with Serializable {
 
   /**
@@ -85,7 +85,25 @@ class Dataset[T] private[sql](
    * Returns the schema of the encoded form of the objects in this [[Dataset]].
    * @since 1.6.0
    */
-  def schema: StructType = resolvedTEncoder.schema
+  override def schema: StructType = resolvedTEncoder.schema
+
+  /**
+   * Prints the schema of the underlying [[DataFrame]] to the console in a nice tree format.
+   * @since 1.6.0
+   */
+  override def printSchema(): Unit = toDF().printSchema()
+
+  /**
+   * Prints the plans (logical and physical) to the console for debugging purposes.
+   * @since 1.6.0
+   */
+  override def explain(extended: Boolean): Unit = toDF().explain(extended)
+
+  /**
+   * Prints the physical plan to the console for debugging purposes.
+   * @since 1.6.0
+   */
+  override def explain(): Unit = toDF().explain()
 
   /* ************* *
    *  Conversions  *
@@ -152,6 +170,59 @@ class Dataset[T] private[sql](
    */
   def count(): Long = toDF().count()
 
+  /**
+   * Displays the content of this [[Dataset]] in a tabular form. Strings more than 20 characters
+   * will be truncated, and all cells will be aligned right. For example:
+   * {{{
+   *   year  month AVG('Adj Close) MAX('Adj Close)
+   *   1980  12    0.503218        0.595103
+   *   1981  01    0.523289        0.570307
+   *   1982  02    0.436504        0.475256
+   *   1983  03    0.410516        0.442194
+   *   1984  04    0.450090        0.483521
+   * }}}
+   * @param numRows Number of rows to show
+   *
+   * @since 1.6.0
+   */
+  def show(numRows: Int): Unit = show(numRows, truncate = true)
+
+  /**
+   * Displays the top 20 rows of [[DataFrame]] in a tabular form. Strings more than 20 characters
+   * will be truncated, and all cells will be aligned right.
+   *
+   * @since 1.6.0
+   */
+  def show(): Unit = show(20)
+
+  /**
+   * Displays the top 20 rows of [[DataFrame]] in a tabular form.
+   *
+   * @param truncate Whether truncate long strings. If true, strings more than 20 characters will
+   *              be truncated and all cells will be aligned right
+   *
+   * @since 1.6.0
+   */
+  def show(truncate: Boolean): Unit = show(20, truncate)
+
+  /**
+   * Displays the [[DataFrame]] in a tabular form. For example:
+   * {{{
+   *   year  month AVG('Adj Close) MAX('Adj Close)
+   *   1980  12    0.503218        0.595103
+   *   1981  01    0.523289        0.570307
+   *   1982  02    0.436504        0.475256
+   *   1983  03    0.410516        0.442194
+   *   1984  04    0.450090        0.483521
+   * }}}
+   * @param numRows Number of rows to show
+   * @param truncate Whether truncate long strings. If true, strings more than 20 characters will
+   *              be truncated and all cells will be aligned right
+   *
+   * @since 1.6.0
+   */
+  def show(numRows: Int, truncate: Boolean): Unit = toDF().show(numRows, truncate)
+
   /**
     * Returns a new [[Dataset]] that has exactly `numPartitions` partitions.
     * @since 1.6.0
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
index 321e2c783537f..f2f5997d1b7c6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Queryable.scala
@@ -17,11 +17,11 @@
 
 package org.apache.spark.sql.execution
 
+import scala.util.control.NonFatal
+
 import org.apache.spark.sql.SQLContext
 import org.apache.spark.sql.types.StructType
 
-import scala.util.control.NonFatal
-
 /** A trait that holds shared code between DataFrames and Datasets. */
 private[sql] trait Queryable {
   def schema: StructType
@@ -37,31 +37,9 @@ private[sql] trait Queryable {
     }
   }
 
-  /**
-   * Prints the schema to the console in a nice tree format.
-   * @group basic
-   * @since 1.3.0
-   */
-  // scalastyle:off println
-  def printSchema(): Unit = println(schema.treeString)
-  // scalastyle:on println
+  def printSchema(): Unit
 
-  /**
-   * Prints the plans (logical and physical) to the console for debugging purposes.
-   * @since 1.3.0
-   */
-  def explain(extended: Boolean): Unit = {
-    val explain = ExplainCommand(queryExecution.logical, extended = extended)
-    sqlContext.executePlan(explain).executedPlan.executeCollect().foreach {
-      // scalastyle:off println
-      r => println(r.getString(0))
-      // scalastyle:on println
-    }
-  }
+  def explain(extended: Boolean): Unit
 
-  /**
-   * Only prints the physical plan to the console for debugging purposes.
-   * @since 1.3.0
-   */
-  def explain(): Unit = explain(extended = false)
+  def explain(): Unit
 }

From 63850026576b3ea7783f9d4b975171dc3cff6e4c Mon Sep 17 00:00:00 2001
From: Ashwin Swaroop <Ashwin Swaroop>
Date: Wed, 25 Nov 2015 13:41:14 +0000
Subject: [PATCH 452/518] [SPARK-11686][CORE] Issue WARN when dynamic
 allocation is disabled due to spark.dynamicAllocation.enabled and
 spark.executor.instances both set

Changed the log type to a 'warning' instead of 'info' as required.

Author: Ashwin Swaroop <Ashwin Swaroop>

Closes #9926 from ashwinswaroop/master.
---
 core/src/main/scala/org/apache/spark/SparkContext.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index e19ba113702c6..2c10779f2b893 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -556,7 +556,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
     // Optionally scale number of executors dynamically based on workload. Exposed for testing.
     val dynamicAllocationEnabled = Utils.isDynamicAllocationEnabled(_conf)
     if (!dynamicAllocationEnabled && _conf.getBoolean("spark.dynamicAllocation.enabled", false)) {
-      logInfo("Dynamic Allocation and num executors both set, thus dynamic allocation disabled.")
+      logWarning("Dynamic Allocation and num executors both set, thus dynamic allocation disabled.")
     }
 
     _executorAllocationManager =

From b9b6fbe89b6d1a890faa02c1a53bb670a6255362 Mon Sep 17 00:00:00 2001
From: Jeff Zhang <zjffdu@apache.org>
Date: Wed, 25 Nov 2015 13:49:58 +0000
Subject: [PATCH 453/518] =?UTF-8?q?[SPARK-11860][PYSAPRK][DOCUMENTATION]?=
 =?UTF-8?q?=20Invalid=20argument=20specification=20=E2=80=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

…for registerFunction [Python]

Straightforward change on the python doc

Author: Jeff Zhang <zjffdu@apache.org>

Closes #9901 from zjffdu/SPARK-11860.
---
 python/pyspark/sql/context.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/python/pyspark/sql/context.py b/python/pyspark/sql/context.py
index 5a85ac31025e8..a49c1b58d0180 100644
--- a/python/pyspark/sql/context.py
+++ b/python/pyspark/sql/context.py
@@ -195,14 +195,15 @@ def range(self, start, end=None, step=1, numPartitions=None):
     @ignore_unicode_prefix
     @since(1.2)
     def registerFunction(self, name, f, returnType=StringType()):
-        """Registers a lambda function as a UDF so it can be used in SQL statements.
+        """Registers a python function (including lambda function) as a UDF
+        so it can be used in SQL statements.
 
         In addition to a name and the function itself, the return type can be optionally specified.
         When the return type is not given it default to a string and conversion will automatically
         be done.  For any other return type, the produced object must match the specified type.
 
         :param name: name of the UDF
-        :param samplingRatio: lambda function
+        :param f: python function
         :param returnType: a :class:`DataType` object
 
         >>> sqlContext.registerFunction("stringLengthString", lambda x: len(x))

From 0a5aef753e70e93d7e56054f354a52e4d4e18932 Mon Sep 17 00:00:00 2001
From: Mark Hamstra <markhamstra@gmail.com>
Date: Wed, 25 Nov 2015 09:34:34 -0600
Subject: [PATCH 454/518] [SPARK-10666][SPARK-6880][CORE] Use properties from
 ActiveJob associated with a Stage

This issue was addressed in https://github.com/apache/spark/pull/5494, but the fix in that PR, while safe in the sense that it will prevent the SparkContext from shutting down, misses the actual bug.  The intent of `submitMissingTasks` should be understood as "submit the Tasks that are missing for the Stage, and run them as part of the ActiveJob identified by jobId".  Because of a long-standing bug, the `jobId` parameter was never being used.  Instead, we were trying to use the jobId with which the Stage was created -- which may no longer exist as an ActiveJob, hence the crash reported in SPARK-6880.

The correct fix is to use the ActiveJob specified by the supplied jobId parameter, which is guaranteed to exist at the call sites of submitMissingTasks.

This fix should be applied to all maintenance branches, since it has existed since 1.0.

kayousterhout pankajarora12

Author: Mark Hamstra <markhamstra@gmail.com>
Author: Imran Rashid <irashid@cloudera.com>

Closes #6291 from markhamstra/SPARK-6880.
---
 .../apache/spark/scheduler/DAGScheduler.scala |   6 +-
 .../spark/scheduler/DAGSchedulerSuite.scala   | 107 +++++++++++++++++-
 2 files changed, 109 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
index 77a184dfe4bee..e01a9609b9a0d 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala
@@ -946,7 +946,9 @@ class DAGScheduler(
       stage.resetInternalAccumulators()
     }
 
-    val properties = jobIdToActiveJob.get(stage.firstJobId).map(_.properties).orNull
+    // Use the scheduling pool, job group, description, etc. from an ActiveJob associated
+    // with this Stage
+    val properties = jobIdToActiveJob(jobId).properties
 
     runningStages += stage
     // SparkListenerStageSubmitted should be posted before testing whether tasks are
@@ -1047,7 +1049,7 @@ class DAGScheduler(
       stage.pendingPartitions ++= tasks.map(_.partitionId)
       logDebug("New pending partitions: " + stage.pendingPartitions)
       taskScheduler.submitTasks(new TaskSet(
-        tasks.toArray, stage.id, stage.latestInfo.attemptId, stage.firstJobId, properties))
+        tasks.toArray, stage.id, stage.latestInfo.attemptId, jobId, properties))
       stage.latestInfo.submissionTime = Some(clock.getTimeMillis())
     } else {
       // Because we posted SparkListenerStageSubmitted earlier, we should mark
diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
index 4d6b25455226f..653d41fc053c9 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.scheduler
 
+import java.util.Properties
+
 import scala.collection.mutable.{ArrayBuffer, HashSet, HashMap, Map}
 import scala.language.reflectiveCalls
 import scala.util.control.NonFatal
@@ -262,9 +264,10 @@ class DAGSchedulerSuite
       rdd: RDD[_],
       partitions: Array[Int],
       func: (TaskContext, Iterator[_]) => _ = jobComputeFunc,
-      listener: JobListener = jobListener): Int = {
+      listener: JobListener = jobListener,
+      properties: Properties = null): Int = {
     val jobId = scheduler.nextJobId.getAndIncrement()
-    runEvent(JobSubmitted(jobId, rdd, func, partitions, CallSite("", ""), listener))
+    runEvent(JobSubmitted(jobId, rdd, func, partitions, CallSite("", ""), listener, properties))
     jobId
   }
 
@@ -1322,6 +1325,106 @@ class DAGSchedulerSuite
     assertDataStructuresEmpty()
   }
 
+  def checkJobPropertiesAndPriority(taskSet: TaskSet, expected: String, priority: Int): Unit = {
+    assert(taskSet.properties != null)
+    assert(taskSet.properties.getProperty("testProperty") === expected)
+    assert(taskSet.priority === priority)
+  }
+
+  def launchJobsThatShareStageAndCancelFirst(): ShuffleDependency[Int, Int, Nothing] = {
+    val baseRdd = new MyRDD(sc, 1, Nil)
+    val shuffleDep1 = new ShuffleDependency(baseRdd, new HashPartitioner(1))
+    val intermediateRdd = new MyRDD(sc, 1, List(shuffleDep1))
+    val shuffleDep2 = new ShuffleDependency(intermediateRdd, new HashPartitioner(1))
+    val finalRdd1 = new MyRDD(sc, 1, List(shuffleDep2))
+    val finalRdd2 = new MyRDD(sc, 1, List(shuffleDep2))
+    val job1Properties = new Properties()
+    val job2Properties = new Properties()
+    job1Properties.setProperty("testProperty", "job1")
+    job2Properties.setProperty("testProperty", "job2")
+
+    // Run jobs 1 & 2, both referencing the same stage, then cancel job1.
+    // Note that we have to submit job2 before we cancel job1 to have them actually share
+    // *Stages*, and not just shuffle dependencies, due to skipped stages (at least until
+    // we address SPARK-10193.)
+    val jobId1 = submit(finalRdd1, Array(0), properties = job1Properties)
+    val jobId2 = submit(finalRdd2, Array(0), properties = job2Properties)
+    assert(scheduler.activeJobs.nonEmpty)
+    val testProperty1 = scheduler.jobIdToActiveJob(jobId1).properties.getProperty("testProperty")
+
+    // remove job1 as an ActiveJob
+    cancel(jobId1)
+
+    // job2 should still be running
+    assert(scheduler.activeJobs.nonEmpty)
+    val testProperty2 = scheduler.jobIdToActiveJob(jobId2).properties.getProperty("testProperty")
+    assert(testProperty1 != testProperty2)
+    // NB: This next assert isn't necessarily the "desired" behavior; it's just to document
+    // the current behavior.  We've already submitted the TaskSet for stage 0 based on job1, but
+    // even though we have cancelled that job and are now running it because of job2, we haven't
+    // updated the TaskSet's properties.  Changing the properties to "job2" is likely the more
+    // correct behavior.
+    val job1Id = 0  // TaskSet priority for Stages run with "job1" as the ActiveJob
+    checkJobPropertiesAndPriority(taskSets(0), "job1", job1Id)
+    complete(taskSets(0), Seq((Success, makeMapStatus("hostA", 1))))
+
+    shuffleDep1
+  }
+
+  /**
+   * Makes sure that tasks for a stage used by multiple jobs are submitted with the properties of a
+   * later, active job if they were previously run under a job that is no longer active
+   */
+  test("stage used by two jobs, the first no longer active (SPARK-6880)") {
+    launchJobsThatShareStageAndCancelFirst()
+
+    // The next check is the key for SPARK-6880.  For the stage which was shared by both job1 and
+    // job2 but never had any tasks submitted for job1, the properties of job2 are now used to run
+    // the stage.
+    checkJobPropertiesAndPriority(taskSets(1), "job2", 1)
+
+    complete(taskSets(1), Seq((Success, makeMapStatus("hostA", 1))))
+    assert(taskSets(2).properties != null)
+    complete(taskSets(2), Seq((Success, 42)))
+    assert(results === Map(0 -> 42))
+    assert(scheduler.activeJobs.isEmpty)
+
+    assertDataStructuresEmpty()
+  }
+
+  /**
+   * Makes sure that tasks for a stage used by multiple jobs are submitted with the properties of a
+   * later, active job if they were previously run under a job that is no longer active, even when
+   * there are fetch failures
+   */
+  test("stage used by two jobs, some fetch failures, and the first job no longer active " +
+    "(SPARK-6880)") {
+    val shuffleDep1 = launchJobsThatShareStageAndCancelFirst()
+    val job2Id = 1  // TaskSet priority for Stages run with "job2" as the ActiveJob
+
+    // lets say there is a fetch failure in this task set, which makes us go back and
+    // run stage 0, attempt 1
+    complete(taskSets(1), Seq(
+      (FetchFailed(makeBlockManagerId("hostA"), shuffleDep1.shuffleId, 0, 0, "ignored"), null)))
+    scheduler.resubmitFailedStages()
+
+    // stage 0, attempt 1 should have the properties of job2
+    assert(taskSets(2).stageId === 0)
+    assert(taskSets(2).stageAttemptId === 1)
+    checkJobPropertiesAndPriority(taskSets(2), "job2", job2Id)
+
+    // run the rest of the stages normally, checking that they have the correct properties
+    complete(taskSets(2), Seq((Success, makeMapStatus("hostA", 1))))
+    checkJobPropertiesAndPriority(taskSets(3), "job2", job2Id)
+    complete(taskSets(3), Seq((Success, makeMapStatus("hostA", 1))))
+    checkJobPropertiesAndPriority(taskSets(4), "job2", job2Id)
+    complete(taskSets(4), Seq((Success, 42)))
+    assert(results === Map(0 -> 42))
+    assert(scheduler.activeJobs.isEmpty)
+
+    assertDataStructuresEmpty()
+  }
+
   test("run trivial shuffle with out-of-band failure and retry") {
     val shuffleMapRdd = new MyRDD(sc, 2, Nil)
     val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(2))

From c1f85fc71e71e07534b89c84677d977bb20994f8 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Wed, 25 Nov 2015 09:47:20 -0800
Subject: [PATCH 455/518] [SPARK-11956][CORE] Fix a few bugs in network
 lib-based file transfer.

- NettyRpcEnv::openStream() now correctly propagates errors to
  the read side of the pipe.
- NettyStreamManager now throws if the file being transferred does
  not exist.
- The network library now correctly handles zero-sized streams.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9941 from vanzin/SPARK-11956.
---
 .../apache/spark/rpc/netty/NettyRpcEnv.scala  | 19 +++++++++----
 .../spark/rpc/netty/NettyStreamManager.scala  |  2 +-
 .../org/apache/spark/rpc/RpcEnvSuite.scala    | 27 +++++++++++++-----
 .../client/TransportResponseHandler.java      | 28 ++++++++++++-------
 .../org/apache/spark/network/StreamSuite.java | 23 ++++++++++++++-
 5 files changed, 75 insertions(+), 24 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
index 68701f609f77a..c8fa870f50e68 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
@@ -27,7 +27,7 @@ import javax.annotation.Nullable
 
 import scala.concurrent.{Future, Promise}
 import scala.reflect.ClassTag
-import scala.util.{DynamicVariable, Failure, Success}
+import scala.util.{DynamicVariable, Failure, Success, Try}
 import scala.util.control.NonFatal
 
 import org.apache.spark.{Logging, SecurityManager, SparkConf}
@@ -368,13 +368,22 @@ private[netty] class NettyRpcEnv(
 
     @volatile private var error: Throwable = _
 
-    def setError(e: Throwable): Unit = error = e
+    def setError(e: Throwable): Unit = {
+      error = e
+      source.close()
+    }
 
     override def read(dst: ByteBuffer): Int = {
-      if (error != null) {
-        throw error
+      val result = if (error == null) {
+        Try(source.read(dst))
+      } else {
+        Failure(error)
+      }
+
+      result match {
+        case Success(bytesRead) => bytesRead
+        case Failure(error) => throw error
       }
-      source.read(dst)
     }
 
     override def close(): Unit = source.close()
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala
index eb1d2604fb235..a2768b4252dcb 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala
@@ -44,7 +44,7 @@ private[netty] class NettyStreamManager(rpcEnv: NettyRpcEnv)
       case _ => throw new IllegalArgumentException(s"Invalid file type: $ftype")
     }
 
-    require(file != null, s"File not found: $streamId")
+    require(file != null && file.isFile(), s"File not found: $streamId")
     new FileSegmentManagedBuffer(rpcEnv.transportConf, file, 0, file.length())
   }
 
diff --git a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
index 2b664c6313efa..6cc958a5f6bc8 100644
--- a/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/RpcEnvSuite.scala
@@ -729,23 +729,36 @@ abstract class RpcEnvSuite extends SparkFunSuite with BeforeAndAfterAll {
     val tempDir = Utils.createTempDir()
     val file = new File(tempDir, "file")
     Files.write(UUID.randomUUID().toString(), file, UTF_8)
+    val empty = new File(tempDir, "empty")
+    Files.write("", empty, UTF_8);
     val jar = new File(tempDir, "jar")
     Files.write(UUID.randomUUID().toString(), jar, UTF_8)
 
     val fileUri = env.fileServer.addFile(file)
+    val emptyUri = env.fileServer.addFile(empty)
     val jarUri = env.fileServer.addJar(jar)
 
     val destDir = Utils.createTempDir()
-    val destFile = new File(destDir, file.getName())
-    val destJar = new File(destDir, jar.getName())
-
     val sm = new SecurityManager(conf)
     val hc = SparkHadoopUtil.get.conf
-    Utils.fetchFile(fileUri, destDir, conf, sm, hc, 0L, false)
-    Utils.fetchFile(jarUri, destDir, conf, sm, hc, 0L, false)
 
-    assert(Files.equal(file, destFile))
-    assert(Files.equal(jar, destJar))
+    val files = Seq(
+      (file, fileUri),
+      (empty, emptyUri),
+      (jar, jarUri))
+    files.foreach { case (f, uri) =>
+      val destFile = new File(destDir, f.getName())
+      Utils.fetchFile(uri, destDir, conf, sm, hc, 0L, false)
+      assert(Files.equal(f, destFile))
+    }
+
+    // Try to download files that do not exist.
+    Seq("files", "jars").foreach { root =>
+      intercept[Exception] {
+        val uri = env.address.toSparkURL + s"/$root/doesNotExist"
+        Utils.fetchFile(uri, destDir, conf, sm, hc, 0L, false)
+      }
+    }
   }
 
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
index be181e0660826..4c15045363b84 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
@@ -185,16 +185,24 @@ public void handle(ResponseMessage message) {
       StreamResponse resp = (StreamResponse) message;
       StreamCallback callback = streamCallbacks.poll();
       if (callback != null) {
-        StreamInterceptor interceptor = new StreamInterceptor(this, resp.streamId, resp.byteCount,
-          callback);
-        try {
-          TransportFrameDecoder frameDecoder = (TransportFrameDecoder)
-            channel.pipeline().get(TransportFrameDecoder.HANDLER_NAME);
-          frameDecoder.setInterceptor(interceptor);
-          streamActive = true;
-        } catch (Exception e) {
-          logger.error("Error installing stream handler.", e);
-          deactivateStream();
+        if (resp.byteCount > 0) {
+          StreamInterceptor interceptor = new StreamInterceptor(this, resp.streamId, resp.byteCount,
+            callback);
+          try {
+            TransportFrameDecoder frameDecoder = (TransportFrameDecoder)
+              channel.pipeline().get(TransportFrameDecoder.HANDLER_NAME);
+            frameDecoder.setInterceptor(interceptor);
+            streamActive = true;
+          } catch (Exception e) {
+            logger.error("Error installing stream handler.", e);
+            deactivateStream();
+          }
+        } else {
+          try {
+            callback.onComplete(resp.streamId);
+          } catch (Exception e) {
+            logger.warn("Error in stream handler onComplete().", e);
+          }
         }
       } else {
         logger.error("Could not find callback for StreamResponse.");
diff --git a/network/common/src/test/java/org/apache/spark/network/StreamSuite.java b/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
index 00158fd081626..538f3efe8d6f2 100644
--- a/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
@@ -51,13 +51,14 @@
 import org.apache.spark.network.util.TransportConf;
 
 public class StreamSuite {
-  private static final String[] STREAMS = { "largeBuffer", "smallBuffer", "file" };
+  private static final String[] STREAMS = { "largeBuffer", "smallBuffer", "emptyBuffer", "file" };
 
   private static TransportServer server;
   private static TransportClientFactory clientFactory;
   private static File testFile;
   private static File tempDir;
 
+  private static ByteBuffer emptyBuffer;
   private static ByteBuffer smallBuffer;
   private static ByteBuffer largeBuffer;
 
@@ -73,6 +74,7 @@ private static ByteBuffer createBuffer(int bufSize) {
   @BeforeClass
   public static void setUp() throws Exception {
     tempDir = Files.createTempDir();
+    emptyBuffer = createBuffer(0);
     smallBuffer = createBuffer(100);
     largeBuffer = createBuffer(100000);
 
@@ -103,6 +105,8 @@ public ManagedBuffer openStream(String streamId) {
             return new NioManagedBuffer(largeBuffer);
           case "smallBuffer":
             return new NioManagedBuffer(smallBuffer);
+          case "emptyBuffer":
+            return new NioManagedBuffer(emptyBuffer);
           case "file":
             return new FileSegmentManagedBuffer(conf, testFile, 0, testFile.length());
           default:
@@ -138,6 +142,18 @@ public static void tearDown() {
     }
   }
 
+  @Test
+  public void testZeroLengthStream() throws Throwable {
+    TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
+    try {
+      StreamTask task = new StreamTask(client, "emptyBuffer", TimeUnit.SECONDS.toMillis(5));
+      task.run();
+      task.check();
+    } finally {
+      client.close();
+    }
+  }
+
   @Test
   public void testSingleStream() throws Throwable {
     TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
@@ -226,6 +242,11 @@ public void run() {
             outFile = File.createTempFile("data", ".tmp", tempDir);
             out = new FileOutputStream(outFile);
             break;
+          case "emptyBuffer":
+            baos = new ByteArrayOutputStream();
+            out = baos;
+            srcBuffer = emptyBuffer;
+            break;
           default:
             throw new IllegalArgumentException(streamId);
         }

From faabdfa2bd416ae514961535f1953e8e9e8b1f3f Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Wed, 25 Nov 2015 10:36:35 -0800
Subject: [PATCH 456/518] [SPARK-11984][SQL][PYTHON] Fix typos in doc for pivot
 for scala and python

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9967 from felixcheung/pypivotdoc.
---
 python/pyspark/sql/group.py                                 | 6 +++---
 .../src/main/scala/org/apache/spark/sql/GroupedData.scala   | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/python/pyspark/sql/group.py b/python/pyspark/sql/group.py
index d8ed7eb2dda64..1911588309aff 100644
--- a/python/pyspark/sql/group.py
+++ b/python/pyspark/sql/group.py
@@ -169,11 +169,11 @@ def sum(self, *cols):
 
     @since(1.6)
     def pivot(self, pivot_col, values=None):
-        """Pivots a column of the current DataFrame and preform the specified aggregation.
+        """Pivots a column of the current DataFrame and perform the specified aggregation.
 
         :param pivot_col: Column to pivot
-        :param values: Optional list of values of pivotColumn that will be translated to columns in
-            the output data frame. If values are not provided the method with do an immediate call
+        :param values: Optional list of values of pivot column that will be translated to columns in
+            the output DataFrame. If values are not provided the method will do an immediate call
             to .distinct() on the pivot column.
 
         >>> df4.groupBy("year").pivot("course", ["dotNET", "Java"]).sum("earnings").collect()
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
index abd531c4ba541..13341a88a6b74 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/GroupedData.scala
@@ -282,7 +282,7 @@ class GroupedData protected[sql](
   }
 
   /**
-   * Pivots a column of the current [[DataFrame]] and preform the specified aggregation.
+   * Pivots a column of the current [[DataFrame]] and perform the specified aggregation.
    * There are two versions of pivot function: one that requires the caller to specify the list
    * of distinct values to pivot on, and one that does not. The latter is more concise but less
    * efficient, because Spark needs to first compute the list of distinct values internally.
@@ -321,7 +321,7 @@ class GroupedData protected[sql](
   }
 
   /**
-   * Pivots a column of the current [[DataFrame]] and preform the specified aggregation.
+   * Pivots a column of the current [[DataFrame]] and perform the specified aggregation.
    * There are two versions of pivot function: one that requires the caller to specify the list
    * of distinct values to pivot on, and one that does not. The latter is more concise but less
    * efficient, because Spark needs to first compute the list of distinct values internally.
@@ -353,7 +353,7 @@ class GroupedData protected[sql](
   }
 
   /**
-   * Pivots a column of the current [[DataFrame]] and preform the specified aggregation.
+   * Pivots a column of the current [[DataFrame]] and perform the specified aggregation.
    * There are two versions of pivot function: one that requires the caller to specify the list
    * of distinct values to pivot on, and one that does not. The latter is more concise but less
    * efficient, because Spark needs to first compute the list of distinct values internally.

From 6b781576a15d8d5c5fbed8bef1c5bda95b3d44ac Mon Sep 17 00:00:00 2001
From: Zhongshuai Pei <peizhongshuai@huawei.com>
Date: Wed, 25 Nov 2015 10:37:34 -0800
Subject: [PATCH 457/518] [SPARK-11974][CORE] Not all the temp dirs had been
 deleted when the JVM exits

deleting the temp dir like that

```

scala> import scala.collection.mutable
import scala.collection.mutable

scala> val a = mutable.Set(1,2,3,4,7,0,8,98,9)
a: scala.collection.mutable.Set[Int] = Set(0, 9, 1, 2, 3, 7, 4, 8, 98)

scala> a.foreach(x => {a.remove(x) })

scala> a.foreach(println(_))
98
```

You may not modify a collection while traversing or iterating over it.This can not delete all element of the collection

Author: Zhongshuai Pei <peizhongshuai@huawei.com>

Closes #9951 from DoingDone9/Bug_RemainDir.
---
 .../scala/org/apache/spark/util/ShutdownHookManager.scala     | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/util/ShutdownHookManager.scala b/core/src/main/scala/org/apache/spark/util/ShutdownHookManager.scala
index db4a8b304ec3e..4012dca3ecdf8 100644
--- a/core/src/main/scala/org/apache/spark/util/ShutdownHookManager.scala
+++ b/core/src/main/scala/org/apache/spark/util/ShutdownHookManager.scala
@@ -57,7 +57,9 @@ private[spark] object ShutdownHookManager extends Logging {
   // Add a shutdown hook to delete the temp dirs when the JVM exits
   addShutdownHook(TEMP_DIR_SHUTDOWN_PRIORITY) { () =>
     logInfo("Shutdown hook called")
-    shutdownDeletePaths.foreach { dirPath =>
+    // we need to materialize the paths to delete because deleteRecursively removes items from
+    // shutdownDeletePaths as we are traversing through it.
+    shutdownDeletePaths.toArray.foreach { dirPath =>
       try {
         logInfo("Deleting directory " + dirPath)
         Utils.deleteRecursively(new File(dirPath))

From dc1d324fdf83e9f4b1debfb277533b002691d71f Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Wed, 25 Nov 2015 11:11:39 -0800
Subject: [PATCH 458/518] [SPARK-11969] [SQL] [PYSPARK] visualization of SQL
 query for pyspark

Currently, we does not have visualization for SQL query from Python, this PR fix that.

cc zsxwing

Author: Davies Liu <davies@databricks.com>

Closes #9949 from davies/pyspark_sql_ui.
---
 python/pyspark/sql/dataframe.py                      |  2 +-
 .../main/scala/org/apache/spark/sql/DataFrame.scala  |  7 +++++++
 .../org/apache/spark/sql/execution/python.scala      | 12 +++++++-----
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index 0dd75ba7ca820..746bb55e14f22 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -277,7 +277,7 @@ def collect(self):
         [Row(age=2, name=u'Alice'), Row(age=5, name=u'Bob')]
         """
         with SCCallSiteSync(self._sc) as css:
-            port = self._sc._jvm.PythonRDD.collectAndServe(self._jdf.javaToPython().rdd())
+            port = self._jdf.collectToPython()
         return list(_load_from_socket(port, BatchedSerializer(PickleSerializer())))
 
     @ignore_unicode_prefix
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
index d8319b9a97fcf..6197f10813a3b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala
@@ -29,6 +29,7 @@ import org.apache.commons.lang3.StringUtils
 
 import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.api.java.JavaRDD
+import org.apache.spark.api.python.PythonRDD
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis._
@@ -1735,6 +1736,12 @@ class DataFrame private[sql](
     EvaluatePython.javaToPython(rdd)
   }
 
+  protected[sql] def collectToPython(): Int = {
+    withNewExecutionId {
+      PythonRDD.collectAndServe(javaToPython.rdd)
+    }
+  }
+
   ////////////////////////////////////////////////////////////////////////////
   ////////////////////////////////////////////////////////////////////////////
   // Deprecated methods
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python.scala
index d611b0011da16..defcec95fb555 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python.scala
@@ -121,11 +121,13 @@ object EvaluatePython {
 
   def takeAndServe(df: DataFrame, n: Int): Int = {
     registerPicklers()
-    val iter = new SerDeUtil.AutoBatchedPickler(
-      df.queryExecution.executedPlan.executeTake(n).iterator.map { row =>
-        EvaluatePython.toJava(row, df.schema)
-      })
-    PythonRDD.serveIterator(iter, s"serve-DataFrame")
+    df.withNewExecutionId {
+      val iter = new SerDeUtil.AutoBatchedPickler(
+        df.queryExecution.executedPlan.executeTake(n).iterator.map { row =>
+          EvaluatePython.toJava(row, df.schema)
+        })
+      PythonRDD.serveIterator(iter, s"serve-DataFrame")
+    }
   }
 
   /**

From 0dee44a6646daae0cc03dbc32125e080dff0f4ae Mon Sep 17 00:00:00 2001
From: Yu ISHIKAWA <yuu.ishikawa@gmail.com>
Date: Wed, 25 Nov 2015 11:35:52 -0800
Subject: [PATCH 459/518] [MINOR] Remove unnecessary spaces in
 `include_example.rb`

Author: Yu ISHIKAWA <yuu.ishikawa@gmail.com>

Closes #9960 from yu-iskw/minor-remove-spaces.
---
 docs/_plugins/include_example.rb | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/_plugins/include_example.rb b/docs/_plugins/include_example.rb
index 549f81fe1b1bc..564c86680f68e 100644
--- a/docs/_plugins/include_example.rb
+++ b/docs/_plugins/include_example.rb
@@ -20,12 +20,12 @@
 
 module Jekyll
   class IncludeExampleTag < Liquid::Tag
-    
+
     def initialize(tag_name, markup, tokens)
       @markup = markup
       super
     end
- 
+
     def render(context)
       site = context.registers[:site]
       config_dir = '../examples/src/main'
@@ -37,7 +37,7 @@ def render(context)
 
       code = File.open(@file).read.encode("UTF-8")
       code = select_lines(code)
- 
+
       rendered_code = Pygments.highlight(code, :lexer => @lang)
 
       hint = "<div><small>Find full example code at " \
@@ -45,7 +45,7 @@ def render(context)
 
       rendered_code + hint
     end
- 
+
     # Trim the code block so as to have the same indention, regardless of their positions in the
     # code file.
     def trim_codeblock(lines)

From 67b67320884282ccf3102e2af96f877e9b186517 Mon Sep 17 00:00:00 2001
From: Jeff Zhang <zjffdu@apache.org>
Date: Wed, 25 Nov 2015 11:37:42 -0800
Subject: [PATCH 460/518] [DOCUMENTATION] Fix minor doc error

Author: Jeff Zhang <zjffdu@apache.org>

Closes #9956 from zjffdu/dev_typo.
---
 docs/configuration.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/configuration.md b/docs/configuration.md
index 4de202d7f7631..741d6b2b37a87 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -35,7 +35,7 @@ val sc = new SparkContext(conf)
 {% endhighlight %}
 
 Note that we can have more than 1 thread in local mode, and in cases like Spark Streaming, we may
-actually require one to prevent any sort of starvation issues.
+actually require more than 1 thread to prevent any sort of starvation issues.
 
 Properties that specify some time duration should be configured with a unit of time.
 The following format is accepted:

From 83653ac5e71996c5a366a42170bed316b208f1b5 Mon Sep 17 00:00:00 2001
From: Alex Bozarth <ajbozart@us.ibm.com>
Date: Wed, 25 Nov 2015 11:39:00 -0800
Subject: [PATCH 461/518] [SPARK-10864][WEB UI] app name is hidden if window is
 resized

Currently the Web UI navbar has a minimum width of 1200px; so if a window is resized smaller than that the app name goes off screen. The 1200px width seems to have been chosen since it fits the longest example app name without wrapping.

To work with smaller window widths I made the tabs wrap since it looked better than wrapping the app name. This is a distinct change in how the navbar looks and I'm not sure if it's what we actually want to do.

Other notes:
- min-width set to 600px to keep the tabs from wrapping individually (will need to be adjusted if tabs are added)
- app name will also wrap (making three levels) if a really really long app name is used

Author: Alex Bozarth <ajbozart@us.ibm.com>

Closes #9874 from ajbozarth/spark10864.
---
 .../main/resources/org/apache/spark/ui/static/webui.css   | 8 ++------
 core/src/main/scala/org/apache/spark/ui/UIUtils.scala     | 2 +-
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/core/src/main/resources/org/apache/spark/ui/static/webui.css b/core/src/main/resources/org/apache/spark/ui/static/webui.css
index 04f3070d25b4a..c628a0c706553 100644
--- a/core/src/main/resources/org/apache/spark/ui/static/webui.css
+++ b/core/src/main/resources/org/apache/spark/ui/static/webui.css
@@ -16,14 +16,9 @@
  */
 
 .navbar {
-  height: 50px;
   font-size: 15px;
   margin-bottom: 15px;
-  min-width: 1200px
-}
-
-.navbar .navbar-inner {
-  height: 50px;
+  min-width: 600px;
 }
 
 .navbar .brand {
@@ -46,6 +41,7 @@
 .navbar-text {
   height: 50px;
   line-height: 3.3;
+  white-space: nowrap;
 }
 
 table.sortable thead {
diff --git a/core/src/main/scala/org/apache/spark/ui/UIUtils.scala b/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
index 84a1116a5c498..1e8194f57888e 100644
--- a/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
+++ b/core/src/main/scala/org/apache/spark/ui/UIUtils.scala
@@ -210,10 +210,10 @@ private[spark] object UIUtils extends Logging {
                 <span class="version">{org.apache.spark.SPARK_VERSION}</span>
               </a>
             </div>
-            <ul class="nav">{header}</ul>
             <p class="navbar-text pull-right">
               <strong title={appName}>{shortAppName}</strong> application UI
             </p>
+            <ul class="nav">{header}</ul>
           </div>
         </div>
         <div class="container-fluid">

From 9f3e59a16822fb61d60cf103bd4f7823552939c6 Mon Sep 17 00:00:00 2001
From: wangt <wangtao.upc@gmail.com>
Date: Wed, 25 Nov 2015 11:41:05 -0800
Subject: [PATCH 462/518] [SPARK-11880][WINDOWS][SPARK SUBMIT]
 bin/load-spark-env.cmd loads spark-env.cmd from wrong directory

* On windows the `bin/load-spark-env.cmd` tries to load `spark-env.cmd` from `%~dp0..\..\conf`, where `~dp0` points to `bin` and `conf` is only one level up.
* Updated `bin/load-spark-env.cmd` to load `spark-env.cmd` from `%~dp0..\conf`, instead of `%~dp0..\..\conf`

Author: wangt <wangtao.upc@gmail.com>

Closes #9863 from toddwan/master.
---
 bin/load-spark-env.cmd | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bin/load-spark-env.cmd b/bin/load-spark-env.cmd
index 36d932c453b6f..59080edd294f2 100644
--- a/bin/load-spark-env.cmd
+++ b/bin/load-spark-env.cmd
@@ -27,7 +27,7 @@ if [%SPARK_ENV_LOADED%] == [] (
   if not [%SPARK_CONF_DIR%] == [] (
     set user_conf_dir=%SPARK_CONF_DIR%
   ) else (
-    set user_conf_dir=%~dp0..\..\conf
+    set user_conf_dir=%~dp0..\conf
   )
 
   call :LoadSparkEnv

From 88875d9413ec7d64a88d40857ffcf97b5853a7f2 Mon Sep 17 00:00:00 2001
From: jerryshao <sshao@hortonworks.com>
Date: Wed, 25 Nov 2015 11:42:53 -0800
Subject: [PATCH 463/518] [SPARK-10558][CORE] Fix wrong executor state in
 Master

`ExecutorAdded` can only be sent to `AppClient` when worker report back the executor state as `LOADING`, otherwise because of concurrency issue, `AppClient` will possibly receive `ExectuorAdded` at first, then `ExecutorStateUpdated` with `LOADING` state.

Also Master will change the executor state from `LAUNCHING` to `RUNNING` (`AppClient` report back the state as `RUNNING`), then to `LOADING` (worker report back to state as `LOADING`), it should be `LAUNCHING` -> `LOADING` -> `RUNNING`.

Also it is wrongly shown in master UI, the state of executor should be `RUNNING` rather than `LOADING`:

![screen shot 2015-09-11 at 2 30 28 pm](https://cloud.githubusercontent.com/assets/850797/9809254/3155d840-5899-11e5-8cdf-ad06fef75762.png)

Author: jerryshao <sshao@hortonworks.com>

Closes #8714 from jerryshao/SPARK-10558.
---
 .../org/apache/spark/deploy/ExecutorState.scala    |  2 +-
 .../org/apache/spark/deploy/client/AppClient.scala |  3 ---
 .../org/apache/spark/deploy/master/Master.scala    | 14 +++++++++++---
 .../org/apache/spark/deploy/worker/Worker.scala    |  2 +-
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/ExecutorState.scala b/core/src/main/scala/org/apache/spark/deploy/ExecutorState.scala
index efa88c62e1f5d..69c98e28931d7 100644
--- a/core/src/main/scala/org/apache/spark/deploy/ExecutorState.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/ExecutorState.scala
@@ -19,7 +19,7 @@ package org.apache.spark.deploy
 
 private[deploy] object ExecutorState extends Enumeration {
 
-  val LAUNCHING, LOADING, RUNNING, KILLED, FAILED, LOST, EXITED = Value
+  val LAUNCHING, RUNNING, KILLED, FAILED, LOST, EXITED = Value
 
   type ExecutorState = Value
 
diff --git a/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala b/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
index afab362e213b5..df6ba7d669ce9 100644
--- a/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/client/AppClient.scala
@@ -178,9 +178,6 @@ private[spark] class AppClient(
         val fullId = appId + "/" + id
         logInfo("Executor added: %s on %s (%s) with %d cores".format(fullId, workerId, hostPort,
           cores))
-        // FIXME if changing master and `ExecutorAdded` happen at the same time (the order is not
-        // guaranteed), `ExecutorStateChanged` may be sent to a dead master.
-        sendToMaster(ExecutorStateChanged(appId.get, id, ExecutorState.RUNNING, None, None))
         listener.executorAdded(fullId, workerId, hostPort, cores, memory)
 
       case ExecutorUpdated(id, state, message, exitStatus) =>
diff --git a/core/src/main/scala/org/apache/spark/deploy/master/Master.scala b/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
index b25a487806c7f..9952c97dbdffc 100644
--- a/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
@@ -253,9 +253,17 @@ private[deploy] class Master(
       execOption match {
         case Some(exec) => {
           val appInfo = idToApp(appId)
+          val oldState = exec.state
           exec.state = state
-          if (state == ExecutorState.RUNNING) { appInfo.resetRetryCount() }
+
+          if (state == ExecutorState.RUNNING) {
+            assert(oldState == ExecutorState.LAUNCHING,
+              s"executor $execId state transfer from $oldState to RUNNING is illegal")
+            appInfo.resetRetryCount()
+          }
+
           exec.application.driver.send(ExecutorUpdated(execId, state, message, exitStatus))
+
           if (ExecutorState.isFinished(state)) {
             // Remove this executor from the worker and app
             logInfo(s"Removing executor ${exec.fullId} because it is $state")
@@ -702,8 +710,8 @@ private[deploy] class Master(
     worker.addExecutor(exec)
     worker.endpoint.send(LaunchExecutor(masterUrl,
       exec.application.id, exec.id, exec.application.desc, exec.cores, exec.memory))
-    exec.application.driver.send(ExecutorAdded(
-      exec.id, worker.id, worker.hostPort, exec.cores, exec.memory))
+    exec.application.driver.send(
+      ExecutorAdded(exec.id, worker.id, worker.hostPort, exec.cores, exec.memory))
   }
 
   private def registerWorker(worker: WorkerInfo): Boolean = {
diff --git a/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala b/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala
index a45867e7680ec..418faf8fc967f 100755
--- a/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala
@@ -469,7 +469,7 @@ private[deploy] class Worker(
             executorDir,
             workerUri,
             conf,
-            appLocalDirs, ExecutorState.LOADING)
+            appLocalDirs, ExecutorState.RUNNING)
           executors(appId + "/" + execId) = manager
           manager.start()
           coresUsed += cores_

From d29e2ef4cf43c7f7c5aa40d305cf02be44ce19e0 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Wed, 25 Nov 2015 11:47:21 -0800
Subject: [PATCH 464/518] [SPARK-11935][PYSPARK] Send the Python exceptions in
 TransformFunction and TransformFunctionSerializer to Java

The Python exception track in TransformFunction and TransformFunctionSerializer is not sent back to Java. Py4j just throws a very general exception, which is hard to debug.

This PRs adds `getFailure` method to get the failure message in Java side.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9922 from zsxwing/SPARK-11935.
---
 python/pyspark/streaming/tests.py             | 82 ++++++++++++++++++-
 python/pyspark/streaming/util.py              | 29 +++++--
 .../streaming/api/python/PythonDStream.scala  | 52 ++++++++++--
 3 files changed, 144 insertions(+), 19 deletions(-)

diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index a0e0267cafa58..d380d697bc51c 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -404,17 +404,69 @@ def func(dstream):
         self._test_func(input, func, expected)
 
     def test_failed_func(self):
+        # Test failure in
+        # TransformFunction.apply(rdd: Option[RDD[_]], time: Time)
         input = [self.sc.parallelize([d], 1) for d in range(4)]
         input_stream = self.ssc.queueStream(input)
 
         def failed_func(i):
-            raise ValueError("failed")
+            raise ValueError("This is a special error")
 
         input_stream.map(failed_func).pprint()
         self.ssc.start()
         try:
             self.ssc.awaitTerminationOrTimeout(10)
         except:
+            import traceback
+            failure = traceback.format_exc()
+            self.assertTrue("This is a special error" in failure)
+            return
+
+        self.fail("a failed func should throw an error")
+
+    def test_failed_func2(self):
+        # Test failure in
+        # TransformFunction.apply(rdd: Option[RDD[_]], rdd2: Option[RDD[_]], time: Time)
+        input = [self.sc.parallelize([d], 1) for d in range(4)]
+        input_stream1 = self.ssc.queueStream(input)
+        input_stream2 = self.ssc.queueStream(input)
+
+        def failed_func(rdd1, rdd2):
+            raise ValueError("This is a special error")
+
+        input_stream1.transformWith(failed_func, input_stream2, True).pprint()
+        self.ssc.start()
+        try:
+            self.ssc.awaitTerminationOrTimeout(10)
+        except:
+            import traceback
+            failure = traceback.format_exc()
+            self.assertTrue("This is a special error" in failure)
+            return
+
+        self.fail("a failed func should throw an error")
+
+    def test_failed_func_with_reseting_failure(self):
+        input = [self.sc.parallelize([d], 1) for d in range(4)]
+        input_stream = self.ssc.queueStream(input)
+
+        def failed_func(i):
+            if i == 1:
+                # Make it fail in the second batch
+                raise ValueError("This is a special error")
+            else:
+                return i
+
+        # We should be able to see the results of the 3rd and 4th batches even if the second batch
+        # fails
+        expected = [[0], [2], [3]]
+        self.assertEqual(expected, self._collect(input_stream.map(failed_func), 3))
+        try:
+            self.ssc.awaitTerminationOrTimeout(10)
+        except:
+            import traceback
+            failure = traceback.format_exc()
+            self.assertTrue("This is a special error" in failure)
             return
 
         self.fail("a failed func should throw an error")
@@ -780,6 +832,34 @@ def tearDown(self):
         if self.cpd is not None:
             shutil.rmtree(self.cpd)
 
+    def test_transform_function_serializer_failure(self):
+        inputd = tempfile.mkdtemp()
+        self.cpd = tempfile.mkdtemp("test_transform_function_serializer_failure")
+
+        def setup():
+            conf = SparkConf().set("spark.default.parallelism", 1)
+            sc = SparkContext(conf=conf)
+            ssc = StreamingContext(sc, 0.5)
+
+            # A function that cannot be serialized
+            def process(time, rdd):
+                sc.parallelize(range(1, 10))
+
+            ssc.textFileStream(inputd).foreachRDD(process)
+            return ssc
+
+        self.ssc = StreamingContext.getOrCreate(self.cpd, setup)
+        try:
+            self.ssc.start()
+        except:
+            import traceback
+            failure = traceback.format_exc()
+            self.assertTrue(
+                "It appears that you are attempting to reference SparkContext" in failure)
+            return
+
+        self.fail("using SparkContext in process should fail because it's not Serializable")
+
     def test_get_or_create_and_get_active_or_create(self):
         inputd = tempfile.mkdtemp()
         outputd = tempfile.mkdtemp() + "/"
diff --git a/python/pyspark/streaming/util.py b/python/pyspark/streaming/util.py
index 767c732eb90b4..c7f02bca2ae38 100644
--- a/python/pyspark/streaming/util.py
+++ b/python/pyspark/streaming/util.py
@@ -38,12 +38,15 @@ def __init__(self, ctx, func, *deserializers):
         self.func = func
         self.deserializers = deserializers
         self._rdd_wrapper = lambda jrdd, ctx, ser: RDD(jrdd, ctx, ser)
+        self.failure = None
 
     def rdd_wrapper(self, func):
         self._rdd_wrapper = func
         return self
 
     def call(self, milliseconds, jrdds):
+        # Clear the failure
+        self.failure = None
         try:
             if self.ctx is None:
                 self.ctx = SparkContext._active_spark_context
@@ -62,9 +65,11 @@ def call(self, milliseconds, jrdds):
             r = self.func(t, *rdds)
             if r:
                 return r._jrdd
-        except Exception:
-            traceback.print_exc()
-            raise
+        except:
+            self.failure = traceback.format_exc()
+
+    def getLastFailure(self):
+        return self.failure
 
     def __repr__(self):
         return "TransformFunction(%s)" % self.func
@@ -89,22 +94,28 @@ def __init__(self, ctx, serializer, gateway=None):
         self.serializer = serializer
         self.gateway = gateway or self.ctx._gateway
         self.gateway.jvm.PythonDStream.registerSerializer(self)
+        self.failure = None
 
     def dumps(self, id):
+        # Clear the failure
+        self.failure = None
         try:
             func = self.gateway.gateway_property.pool[id]
             return bytearray(self.serializer.dumps((func.func, func.deserializers)))
-        except Exception:
-            traceback.print_exc()
-            raise
+        except:
+            self.failure = traceback.format_exc()
 
     def loads(self, data):
+        # Clear the failure
+        self.failure = None
         try:
             f, deserializers = self.serializer.loads(bytes(data))
             return TransformFunction(self.ctx, f, *deserializers)
-        except Exception:
-            traceback.print_exc()
-            raise
+        except:
+            self.failure = traceback.format_exc()
+
+    def getLastFailure(self):
+        return self.failure
 
     def __repr__(self):
         return "TransformFunctionSerializer(%s)" % self.serializer
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/api/python/PythonDStream.scala b/streaming/src/main/scala/org/apache/spark/streaming/api/python/PythonDStream.scala
index dfc569451df86..994309ddd0a3e 100644
--- a/streaming/src/main/scala/org/apache/spark/streaming/api/python/PythonDStream.scala
+++ b/streaming/src/main/scala/org/apache/spark/streaming/api/python/PythonDStream.scala
@@ -26,6 +26,7 @@ import scala.language.existentials
 
 import py4j.GatewayServer
 
+import org.apache.spark.SparkException
 import org.apache.spark.api.java._
 import org.apache.spark.rdd.RDD
 import org.apache.spark.storage.StorageLevel
@@ -40,6 +41,13 @@ import org.apache.spark.util.Utils
  */
 private[python] trait PythonTransformFunction {
   def call(time: Long, rdds: JList[_]): JavaRDD[Array[Byte]]
+
+  /**
+   * Get the failure, if any, in the last call to `call`.
+   *
+   * @return the failure message if there was a failure, or `null` if there was no failure.
+   */
+  def getLastFailure: String
 }
 
 /**
@@ -48,6 +56,13 @@ private[python] trait PythonTransformFunction {
 private[python] trait PythonTransformFunctionSerializer {
   def dumps(id: String): Array[Byte]
   def loads(bytes: Array[Byte]): PythonTransformFunction
+
+  /**
+   * Get the failure, if any, in the last call to `dumps` or `loads`.
+   *
+   * @return the failure message if there was a failure, or `null` if there was no failure.
+   */
+  def getLastFailure: String
 }
 
 /**
@@ -59,18 +74,27 @@ private[python] class TransformFunction(@transient var pfunc: PythonTransformFun
   extends function.Function2[JList[JavaRDD[_]], Time, JavaRDD[Array[Byte]]] {
 
   def apply(rdd: Option[RDD[_]], time: Time): Option[RDD[Array[Byte]]] = {
-    Option(pfunc.call(time.milliseconds, List(rdd.map(JavaRDD.fromRDD(_)).orNull).asJava))
-      .map(_.rdd)
+    val rdds = List(rdd.map(JavaRDD.fromRDD(_)).orNull).asJava
+    Option(callPythonTransformFunction(time.milliseconds, rdds)).map(_.rdd)
   }
 
   def apply(rdd: Option[RDD[_]], rdd2: Option[RDD[_]], time: Time): Option[RDD[Array[Byte]]] = {
     val rdds = List(rdd.map(JavaRDD.fromRDD(_)).orNull, rdd2.map(JavaRDD.fromRDD(_)).orNull).asJava
-    Option(pfunc.call(time.milliseconds, rdds)).map(_.rdd)
+    Option(callPythonTransformFunction(time.milliseconds, rdds)).map(_.rdd)
   }
 
   // for function.Function2
   def call(rdds: JList[JavaRDD[_]], time: Time): JavaRDD[Array[Byte]] = {
-    pfunc.call(time.milliseconds, rdds)
+    callPythonTransformFunction(time.milliseconds, rdds)
+  }
+
+  private def callPythonTransformFunction(time: Long, rdds: JList[_]): JavaRDD[Array[Byte]] = {
+    val resultRDD = pfunc.call(time, rdds)
+    val failure = pfunc.getLastFailure
+    if (failure != null) {
+      throw new SparkException("An exception was raised by Python:\n" + failure)
+    }
+    resultRDD
   }
 
   private def writeObject(out: ObjectOutputStream): Unit = Utils.tryOrIOException {
@@ -103,23 +127,33 @@ private[python] object PythonTransformFunctionSerializer {
   /*
    * Register a serializer from Python, should be called during initialization
    */
-  def register(ser: PythonTransformFunctionSerializer): Unit = {
+  def register(ser: PythonTransformFunctionSerializer): Unit = synchronized {
     serializer = ser
   }
 
-  def serialize(func: PythonTransformFunction): Array[Byte] = {
+  def serialize(func: PythonTransformFunction): Array[Byte] = synchronized {
     require(serializer != null, "Serializer has not been registered!")
     // get the id of PythonTransformFunction in py4j
     val h = Proxy.getInvocationHandler(func.asInstanceOf[Proxy])
     val f = h.getClass().getDeclaredField("id")
     f.setAccessible(true)
     val id = f.get(h).asInstanceOf[String]
-    serializer.dumps(id)
+    val results = serializer.dumps(id)
+    val failure = serializer.getLastFailure
+    if (failure != null) {
+      throw new SparkException("An exception was raised by Python:\n" + failure)
+    }
+    results
   }
 
-  def deserialize(bytes: Array[Byte]): PythonTransformFunction = {
+  def deserialize(bytes: Array[Byte]): PythonTransformFunction = synchronized {
     require(serializer != null, "Serializer has not been registered!")
-    serializer.loads(bytes)
+    val pfunc = serializer.loads(bytes)
+    val failure = serializer.getLastFailure
+    if (failure != null) {
+      throw new SparkException("An exception was raised by Python:\n" + failure)
+    }
+    pfunc
   }
 }
 

From 4e81783e92f464d479baaf93eccc3adb1496989a Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Wed, 25 Nov 2015 12:58:18 -0800
Subject: [PATCH 465/518] [SPARK-11866][NETWORK][CORE] Make sure timed out RPCs
 are cleaned up.

This change does a couple of different things to make sure that the RpcEnv-level
code and the network library agree about the status of outstanding RPCs.

For RPCs that do not expect a reply ("RpcEnv.send"), support for one way
messages (hello CORBA!) was added to the network layer. This is a
"fire and forget" message that does not require any state to be kept
by the TransportClient; as a result, the RpcEnv 'Ack' message is not needed
anymore.

For RPCs that do expect a reply ("RpcEnv.ask"), the network library now
returns the internal RPC id; if the RpcEnv layer decides to time out the
RPC before the network layer does, it now asks the TransportClient to
forget about the RPC, so that if the network-level timeout occurs, the
client is not killed.

As part of implementing the above, I cleaned up some of the code in the
netty rpc backend, removing types that were not necessary and factoring
out some common code. Of interest is a slight change in the exceptions
when posting messages to a stopped RpcEnv; that's mostly to avoid nasty
error messages from the local-cluster backend when shutting down, which
pollutes the terminal output.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9917 from vanzin/SPARK-11866.
---
 .../spark/deploy/worker/ExecutorRunner.scala  |   6 +-
 .../apache/spark/rpc/netty/Dispatcher.scala   |  55 +++----
 .../org/apache/spark/rpc/netty/Inbox.scala    |  28 ++--
 .../spark/rpc/netty/NettyRpcCallContext.scala |  35 +---
 .../apache/spark/rpc/netty/NettyRpcEnv.scala  | 153 +++++++-----------
 .../org/apache/spark/rpc/netty/Outbox.scala   |  64 ++++++--
 .../apache/spark/rpc/netty/InboxSuite.scala   |   6 +-
 .../rpc/netty/NettyRpcHandlerSuite.scala      |   2 +-
 .../spark/network/client/TransportClient.java |  34 +++-
 .../spark/network/protocol/Message.java       |   4 +-
 .../network/protocol/MessageDecoder.java      |   3 +
 .../spark/network/protocol/OneWayMessage.java |  75 +++++++++
 .../spark/network/sasl/SaslRpcHandler.java    |   5 +
 .../spark/network/server/RpcHandler.java      |  36 +++++
 .../server/TransportRequestHandler.java       |  18 ++-
 .../apache/spark/network/ProtocolSuite.java   |   2 +
 .../spark/network/RpcIntegrationSuite.java    |  31 ++++
 .../spark/network/sasl/SparkSaslSuite.java    |   9 ++
 18 files changed, 374 insertions(+), 192 deletions(-)
 create mode 100644 network/common/src/main/java/org/apache/spark/network/protocol/OneWayMessage.java

diff --git a/core/src/main/scala/org/apache/spark/deploy/worker/ExecutorRunner.scala b/core/src/main/scala/org/apache/spark/deploy/worker/ExecutorRunner.scala
index 3aef0515cbf6e..25a17473e4b53 100644
--- a/core/src/main/scala/org/apache/spark/deploy/worker/ExecutorRunner.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/worker/ExecutorRunner.scala
@@ -92,7 +92,11 @@ private[deploy] class ExecutorRunner(
       process.destroy()
       exitCode = Some(process.waitFor())
     }
-    worker.send(ExecutorStateChanged(appId, execId, state, message, exitCode))
+    try {
+      worker.send(ExecutorStateChanged(appId, execId, state, message, exitCode))
+    } catch {
+      case e: IllegalStateException => logWarning(e.getMessage(), e)
+    }
   }
 
   /** Stop this executor runner, including killing the process it launched */
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala b/core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala
index eb25d6c7b721b..533c9847661b6 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/Dispatcher.scala
@@ -106,44 +106,30 @@ private[netty] class Dispatcher(nettyEnv: NettyRpcEnv) extends Logging {
     val iter = endpoints.keySet().iterator()
     while (iter.hasNext) {
       val name = iter.next
-      postMessage(
-        name,
-        _ => message,
-        () => { logWarning(s"Drop $message because $name has been stopped") })
+      postMessage(name, message, (e) => logWarning(s"Message $message dropped.", e))
     }
   }
 
   /** Posts a message sent by a remote endpoint. */
   def postRemoteMessage(message: RequestMessage, callback: RpcResponseCallback): Unit = {
-    def createMessage(sender: NettyRpcEndpointRef): InboxMessage = {
-      val rpcCallContext =
-        new RemoteNettyRpcCallContext(
-          nettyEnv, sender, callback, message.senderAddress, message.needReply)
-      ContentMessage(message.senderAddress, message.content, message.needReply, rpcCallContext)
-    }
-
-    def onEndpointStopped(): Unit = {
-      callback.onFailure(
-        new SparkException(s"Could not find ${message.receiver.name} or it has been stopped"))
-    }
-
-    postMessage(message.receiver.name, createMessage, onEndpointStopped)
+    val rpcCallContext =
+      new RemoteNettyRpcCallContext(nettyEnv, callback, message.senderAddress)
+    val rpcMessage = RpcMessage(message.senderAddress, message.content, rpcCallContext)
+    postMessage(message.receiver.name, rpcMessage, (e) => callback.onFailure(e))
   }
 
   /** Posts a message sent by a local endpoint. */
   def postLocalMessage(message: RequestMessage, p: Promise[Any]): Unit = {
-    def createMessage(sender: NettyRpcEndpointRef): InboxMessage = {
-      val rpcCallContext =
-        new LocalNettyRpcCallContext(sender, message.senderAddress, message.needReply, p)
-      ContentMessage(message.senderAddress, message.content, message.needReply, rpcCallContext)
-    }
-
-    def onEndpointStopped(): Unit = {
-      p.tryFailure(
-        new SparkException(s"Could not find ${message.receiver.name} or it has been stopped"))
-    }
+    val rpcCallContext =
+      new LocalNettyRpcCallContext(message.senderAddress, p)
+    val rpcMessage = RpcMessage(message.senderAddress, message.content, rpcCallContext)
+    postMessage(message.receiver.name, rpcMessage, (e) => p.tryFailure(e))
+  }
 
-    postMessage(message.receiver.name, createMessage, onEndpointStopped)
+  /** Posts a one-way message. */
+  def postOneWayMessage(message: RequestMessage): Unit = {
+    postMessage(message.receiver.name, OneWayMessage(message.senderAddress, message.content),
+      (e) => throw e)
   }
 
   /**
@@ -155,21 +141,26 @@ private[netty] class Dispatcher(nettyEnv: NettyRpcEnv) extends Logging {
    */
   private def postMessage(
       endpointName: String,
-      createMessageFn: NettyRpcEndpointRef => InboxMessage,
-      callbackIfStopped: () => Unit): Unit = {
+      message: InboxMessage,
+      callbackIfStopped: (Exception) => Unit): Unit = {
     val shouldCallOnStop = synchronized {
       val data = endpoints.get(endpointName)
       if (stopped || data == null) {
         true
       } else {
-        data.inbox.post(createMessageFn(data.ref))
+        data.inbox.post(message)
         receivers.offer(data)
         false
       }
     }
     if (shouldCallOnStop) {
       // We don't need to call `onStop` in the `synchronized` block
-      callbackIfStopped()
+      val error = if (stopped) {
+          new IllegalStateException("RpcEnv already stopped.")
+        } else {
+          new SparkException(s"Could not find $endpointName or it has been stopped.")
+        }
+      callbackIfStopped(error)
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala b/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala
index 464027f07cc88..175463cc10319 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/Inbox.scala
@@ -27,10 +27,13 @@ import org.apache.spark.rpc.{RpcAddress, RpcEndpoint, ThreadSafeRpcEndpoint}
 
 private[netty] sealed trait InboxMessage
 
-private[netty] case class ContentMessage(
+private[netty] case class OneWayMessage(
+    senderAddress: RpcAddress,
+    content: Any) extends InboxMessage
+
+private[netty] case class RpcMessage(
     senderAddress: RpcAddress,
     content: Any,
-    needReply: Boolean,
     context: NettyRpcCallContext) extends InboxMessage
 
 private[netty] case object OnStart extends InboxMessage
@@ -96,29 +99,24 @@ private[netty] class Inbox(
     while (true) {
       safelyCall(endpoint) {
         message match {
-          case ContentMessage(_sender, content, needReply, context) =>
-            // The partial function to call
-            val pf = if (needReply) endpoint.receiveAndReply(context) else endpoint.receive
+          case RpcMessage(_sender, content, context) =>
             try {
-              pf.applyOrElse[Any, Unit](content, { msg =>
+              endpoint.receiveAndReply(context).applyOrElse[Any, Unit](content, { msg =>
                 throw new SparkException(s"Unsupported message $message from ${_sender}")
               })
-              if (!needReply) {
-                context.finish()
-              }
             } catch {
               case NonFatal(e) =>
-                if (needReply) {
-                  // If the sender asks a reply, we should send the error back to the sender
-                  context.sendFailure(e)
-                } else {
-                  context.finish()
-                }
+                context.sendFailure(e)
                 // Throw the exception -- this exception will be caught by the safelyCall function.
                 // The endpoint's onError function will be called.
                 throw e
             }
 
+          case OneWayMessage(_sender, content) =>
+            endpoint.receive.applyOrElse[Any, Unit](content, { msg =>
+              throw new SparkException(s"Unsupported message $message from ${_sender}")
+            })
+
           case OnStart =>
             endpoint.onStart()
             if (!endpoint.isInstanceOf[ThreadSafeRpcEndpoint]) {
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcCallContext.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcCallContext.scala
index 21d5bb4923d1b..6637e2321f673 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcCallContext.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcCallContext.scala
@@ -23,49 +23,28 @@ import org.apache.spark.Logging
 import org.apache.spark.network.client.RpcResponseCallback
 import org.apache.spark.rpc.{RpcAddress, RpcCallContext}
 
-private[netty] abstract class NettyRpcCallContext(
-    endpointRef: NettyRpcEndpointRef,
-    override val senderAddress: RpcAddress,
-    needReply: Boolean)
+private[netty] abstract class NettyRpcCallContext(override val senderAddress: RpcAddress)
   extends RpcCallContext with Logging {
 
   protected def send(message: Any): Unit
 
   override def reply(response: Any): Unit = {
-    if (needReply) {
-      send(AskResponse(endpointRef, response))
-    } else {
-      throw new IllegalStateException(
-        s"Cannot send $response to the sender because the sender does not expect a reply")
-    }
+    send(response)
   }
 
   override def sendFailure(e: Throwable): Unit = {
-    if (needReply) {
-      send(AskResponse(endpointRef, RpcFailure(e)))
-    } else {
-      logError(e.getMessage, e)
-      throw new IllegalStateException(
-        "Cannot send reply to the sender because the sender won't handle it")
-    }
+    send(RpcFailure(e))
   }
 
-  def finish(): Unit = {
-    if (!needReply) {
-      send(Ack(endpointRef))
-    }
-  }
 }
 
 /**
  * If the sender and the receiver are in the same process, the reply can be sent back via `Promise`.
  */
 private[netty] class LocalNettyRpcCallContext(
-    endpointRef: NettyRpcEndpointRef,
     senderAddress: RpcAddress,
-    needReply: Boolean,
     p: Promise[Any])
-  extends NettyRpcCallContext(endpointRef, senderAddress, needReply) {
+  extends NettyRpcCallContext(senderAddress) {
 
   override protected def send(message: Any): Unit = {
     p.success(message)
@@ -77,11 +56,9 @@ private[netty] class LocalNettyRpcCallContext(
  */
 private[netty] class RemoteNettyRpcCallContext(
     nettyEnv: NettyRpcEnv,
-    endpointRef: NettyRpcEndpointRef,
     callback: RpcResponseCallback,
-    senderAddress: RpcAddress,
-    needReply: Boolean)
-  extends NettyRpcCallContext(endpointRef, senderAddress, needReply) {
+    senderAddress: RpcAddress)
+  extends NettyRpcCallContext(senderAddress) {
 
   override protected def send(message: Any): Unit = {
     val reply = nettyEnv.serialize(message)
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
index c8fa870f50e68..c7d74fa1d9195 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
@@ -150,7 +150,7 @@ private[netty] class NettyRpcEnv(
 
   private def postToOutbox(receiver: NettyRpcEndpointRef, message: OutboxMessage): Unit = {
     if (receiver.client != null) {
-      receiver.client.sendRpc(message.content, message.createCallback(receiver.client));
+      message.sendWith(receiver.client)
     } else {
       require(receiver.address != null,
         "Cannot send message to client endpoint with no listen address.")
@@ -182,25 +182,10 @@ private[netty] class NettyRpcEnv(
     val remoteAddr = message.receiver.address
     if (remoteAddr == address) {
       // Message to a local RPC endpoint.
-      val promise = Promise[Any]()
-      dispatcher.postLocalMessage(message, promise)
-      promise.future.onComplete {
-        case Success(response) =>
-          val ack = response.asInstanceOf[Ack]
-          logTrace(s"Received ack from ${ack.sender}")
-        case Failure(e) =>
-          logWarning(s"Exception when sending $message", e)
-      }(ThreadUtils.sameThread)
+      dispatcher.postOneWayMessage(message)
     } else {
       // Message to a remote RPC endpoint.
-      postToOutbox(message.receiver, OutboxMessage(serialize(message),
-        (e) => {
-          logWarning(s"Exception when sending $message", e)
-        },
-        (client, response) => {
-          val ack = deserialize[Ack](client, response)
-          logDebug(s"Receive ack from ${ack.sender}")
-        }))
+      postToOutbox(message.receiver, OneWayOutboxMessage(serialize(message)))
     }
   }
 
@@ -208,46 +193,52 @@ private[netty] class NettyRpcEnv(
     clientFactory.createClient(address.host, address.port)
   }
 
-  private[netty] def ask(message: RequestMessage): Future[Any] = {
+  private[netty] def ask[T: ClassTag](message: RequestMessage, timeout: RpcTimeout): Future[T] = {
     val promise = Promise[Any]()
     val remoteAddr = message.receiver.address
+
+    def onFailure(e: Throwable): Unit = {
+      if (!promise.tryFailure(e)) {
+        logWarning(s"Ignored failure: $e")
+      }
+    }
+
+    def onSuccess(reply: Any): Unit = reply match {
+      case RpcFailure(e) => onFailure(e)
+      case rpcReply =>
+        if (!promise.trySuccess(rpcReply)) {
+          logWarning(s"Ignored message: $reply")
+        }
+    }
+
     if (remoteAddr == address) {
       val p = Promise[Any]()
-      dispatcher.postLocalMessage(message, p)
       p.future.onComplete {
-        case Success(response) =>
-          val reply = response.asInstanceOf[AskResponse]
-          if (reply.reply.isInstanceOf[RpcFailure]) {
-            if (!promise.tryFailure(reply.reply.asInstanceOf[RpcFailure].e)) {
-              logWarning(s"Ignore failure: ${reply.reply}")
-            }
-          } else if (!promise.trySuccess(reply.reply)) {
-            logWarning(s"Ignore message: ${reply}")
-          }
-        case Failure(e) =>
-          if (!promise.tryFailure(e)) {
-            logWarning("Ignore Exception", e)
-          }
+        case Success(response) => onSuccess(response)
+        case Failure(e) => onFailure(e)
       }(ThreadUtils.sameThread)
+      dispatcher.postLocalMessage(message, p)
     } else {
-      postToOutbox(message.receiver, OutboxMessage(serialize(message),
-        (e) => {
-          if (!promise.tryFailure(e)) {
-            logWarning("Ignore Exception", e)
-          }
-        },
-        (client, response) => {
-          val reply = deserialize[AskResponse](client, response)
-          if (reply.reply.isInstanceOf[RpcFailure]) {
-            if (!promise.tryFailure(reply.reply.asInstanceOf[RpcFailure].e)) {
-              logWarning(s"Ignore failure: ${reply.reply}")
-            }
-          } else if (!promise.trySuccess(reply.reply)) {
-            logWarning(s"Ignore message: ${reply}")
-          }
-        }))
+      val rpcMessage = RpcOutboxMessage(serialize(message),
+        onFailure,
+        (client, response) => onSuccess(deserialize[Any](client, response)))
+      postToOutbox(message.receiver, rpcMessage)
+      promise.future.onFailure {
+        case _: TimeoutException => rpcMessage.onTimeout()
+        case _ =>
+      }(ThreadUtils.sameThread)
     }
-    promise.future
+
+    val timeoutCancelable = timeoutScheduler.schedule(new Runnable {
+      override def run(): Unit = {
+        promise.tryFailure(
+          new TimeoutException("Cannot receive any reply in ${timeout.duration}"))
+      }
+    }, timeout.duration.toNanos, TimeUnit.NANOSECONDS)
+    promise.future.onComplete { v =>
+      timeoutCancelable.cancel(true)
+    }(ThreadUtils.sameThread)
+    promise.future.mapTo[T].recover(timeout.addMessageIfTimeout)(ThreadUtils.sameThread)
   }
 
   private[netty] def serialize(content: Any): Array[Byte] = {
@@ -512,25 +503,12 @@ private[netty] class NettyRpcEndpointRef(
   override def name: String = _name
 
   override def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T] = {
-    val promise = Promise[Any]()
-    val timeoutCancelable = nettyEnv.timeoutScheduler.schedule(new Runnable {
-      override def run(): Unit = {
-        promise.tryFailure(new TimeoutException("Cannot receive any reply in " + timeout.duration))
-      }
-    }, timeout.duration.toNanos, TimeUnit.NANOSECONDS)
-    val f = nettyEnv.ask(RequestMessage(nettyEnv.address, this, message, true))
-    f.onComplete { v =>
-      timeoutCancelable.cancel(true)
-      if (!promise.tryComplete(v)) {
-        logWarning(s"Ignore message $v")
-      }
-    }(ThreadUtils.sameThread)
-    promise.future.mapTo[T].recover(timeout.addMessageIfTimeout)(ThreadUtils.sameThread)
+    nettyEnv.ask(RequestMessage(nettyEnv.address, this, message), timeout)
   }
 
   override def send(message: Any): Unit = {
     require(message != null, "Message is null")
-    nettyEnv.send(RequestMessage(nettyEnv.address, this, message, false))
+    nettyEnv.send(RequestMessage(nettyEnv.address, this, message))
   }
 
   override def toString: String = s"NettyRpcEndpointRef(${_address})"
@@ -549,24 +527,7 @@ private[netty] class NettyRpcEndpointRef(
  * The message that is sent from the sender to the receiver.
  */
 private[netty] case class RequestMessage(
-    senderAddress: RpcAddress, receiver: NettyRpcEndpointRef, content: Any, needReply: Boolean)
-
-/**
- * The base trait for all messages that are sent back from the receiver to the sender.
- */
-private[netty] trait ResponseMessage
-
-/**
- * The reply for `ask` from the receiver side.
- */
-private[netty] case class AskResponse(sender: NettyRpcEndpointRef, reply: Any)
-  extends ResponseMessage
-
-/**
- * A message to send back to the receiver side. It's necessary because [[TransportClient]] only
- * clean the resources when it receives a reply.
- */
-private[netty] case class Ack(sender: NettyRpcEndpointRef) extends ResponseMessage
+    senderAddress: RpcAddress, receiver: NettyRpcEndpointRef, content: Any)
 
 /**
  * A response that indicates some failure happens in the receiver side.
@@ -598,6 +559,18 @@ private[netty] class NettyRpcHandler(
       client: TransportClient,
       message: Array[Byte],
       callback: RpcResponseCallback): Unit = {
+    val messageToDispatch = internalReceive(client, message)
+    dispatcher.postRemoteMessage(messageToDispatch, callback)
+  }
+
+  override def receive(
+      client: TransportClient,
+      message: Array[Byte]): Unit = {
+    val messageToDispatch = internalReceive(client, message)
+    dispatcher.postOneWayMessage(messageToDispatch)
+  }
+
+  private def internalReceive(client: TransportClient, message: Array[Byte]): RequestMessage = {
     val addr = client.getChannel().remoteAddress().asInstanceOf[InetSocketAddress]
     assert(addr != null)
     val clientAddr = RpcAddress(addr.getHostName, addr.getPort)
@@ -605,14 +578,12 @@ private[netty] class NettyRpcHandler(
       dispatcher.postToAll(RemoteProcessConnected(clientAddr))
     }
     val requestMessage = nettyEnv.deserialize[RequestMessage](client, message)
-    val messageToDispatch = if (requestMessage.senderAddress == null) {
-        // Create a new message with the socket address of the client as the sender.
-        RequestMessage(clientAddr, requestMessage.receiver, requestMessage.content,
-          requestMessage.needReply)
-      } else {
-        requestMessage
-      }
-    dispatcher.postRemoteMessage(messageToDispatch, callback)
+    if (requestMessage.senderAddress == null) {
+      // Create a new message with the socket address of the client as the sender.
+      RequestMessage(clientAddr, requestMessage.receiver, requestMessage.content)
+    } else {
+      requestMessage
+    }
   }
 
   override def getStreamManager: StreamManager = streamManager
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala b/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
index 2f6817f2eb935..36fdd00bbc4c2 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
@@ -22,22 +22,56 @@ import javax.annotation.concurrent.GuardedBy
 
 import scala.util.control.NonFatal
 
-import org.apache.spark.SparkException
+import org.apache.spark.{Logging, SparkException}
 import org.apache.spark.network.client.{RpcResponseCallback, TransportClient}
 import org.apache.spark.rpc.RpcAddress
 
-private[netty] case class OutboxMessage(content: Array[Byte],
-  _onFailure: (Throwable) => Unit,
-  _onSuccess: (TransportClient, Array[Byte]) => Unit) {
+private[netty] sealed trait OutboxMessage {
 
-  def createCallback(client: TransportClient): RpcResponseCallback = new RpcResponseCallback() {
-    override def onFailure(e: Throwable): Unit = {
-      _onFailure(e)
-    }
+  def sendWith(client: TransportClient): Unit
 
-    override def onSuccess(response: Array[Byte]): Unit = {
-      _onSuccess(client, response)
-    }
+  def onFailure(e: Throwable): Unit
+
+}
+
+private[netty] case class OneWayOutboxMessage(content: Array[Byte]) extends OutboxMessage
+  with Logging {
+
+  override def sendWith(client: TransportClient): Unit = {
+    client.send(content)
+  }
+
+  override def onFailure(e: Throwable): Unit = {
+    logWarning(s"Failed to send one-way RPC.", e)
+  }
+
+}
+
+private[netty] case class RpcOutboxMessage(
+    content: Array[Byte],
+    _onFailure: (Throwable) => Unit,
+    _onSuccess: (TransportClient, Array[Byte]) => Unit)
+  extends OutboxMessage with RpcResponseCallback {
+
+  private var client: TransportClient = _
+  private var requestId: Long = _
+
+  override def sendWith(client: TransportClient): Unit = {
+    this.client = client
+    this.requestId = client.sendRpc(content, this)
+  }
+
+  def onTimeout(): Unit = {
+    require(client != null, "TransportClient has not yet been set.")
+    client.removeRpcRequest(requestId)
+  }
+
+  override def onFailure(e: Throwable): Unit = {
+    _onFailure(e)
+  }
+
+  override def onSuccess(response: Array[Byte]): Unit = {
+    _onSuccess(client, response)
   }
 
 }
@@ -82,7 +116,7 @@ private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
       }
     }
     if (dropped) {
-      message._onFailure(new SparkException("Message is dropped because Outbox is stopped"))
+      message.onFailure(new SparkException("Message is dropped because Outbox is stopped"))
     } else {
       drainOutbox()
     }
@@ -122,7 +156,7 @@ private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
       try {
         val _client = synchronized { client }
         if (_client != null) {
-          _client.sendRpc(message.content, message.createCallback(_client))
+          message.sendWith(_client)
         } else {
           assert(stopped == true)
         }
@@ -195,7 +229,7 @@ private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
     // update messages and it's safe to just drain the queue.
     var message = messages.poll()
     while (message != null) {
-      message._onFailure(e)
+      message.onFailure(e)
       message = messages.poll()
     }
     assert(messages.isEmpty)
@@ -229,7 +263,7 @@ private[netty] class Outbox(nettyEnv: NettyRpcEnv, val address: RpcAddress) {
     // update messages and it's safe to just drain the queue.
     var message = messages.poll()
     while (message != null) {
-      message._onFailure(new SparkException("Message is dropped because Outbox is stopped"))
+      message.onFailure(new SparkException("Message is dropped because Outbox is stopped"))
       message = messages.poll()
     }
   }
diff --git a/core/src/test/scala/org/apache/spark/rpc/netty/InboxSuite.scala b/core/src/test/scala/org/apache/spark/rpc/netty/InboxSuite.scala
index 276c077b3d13e..2136795b18813 100644
--- a/core/src/test/scala/org/apache/spark/rpc/netty/InboxSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/netty/InboxSuite.scala
@@ -35,7 +35,7 @@ class InboxSuite extends SparkFunSuite {
     val dispatcher = mock(classOf[Dispatcher])
 
     val inbox = new Inbox(endpointRef, endpoint)
-    val message = ContentMessage(null, "hi", false, null)
+    val message = OneWayMessage(null, "hi")
     inbox.post(message)
     inbox.process(dispatcher)
     assert(inbox.isEmpty)
@@ -55,7 +55,7 @@ class InboxSuite extends SparkFunSuite {
     val dispatcher = mock(classOf[Dispatcher])
 
     val inbox = new Inbox(endpointRef, endpoint)
-    val message = ContentMessage(null, "hi", true, null)
+    val message = RpcMessage(null, "hi", null)
     inbox.post(message)
     inbox.process(dispatcher)
     assert(inbox.isEmpty)
@@ -83,7 +83,7 @@ class InboxSuite extends SparkFunSuite {
       new Thread {
         override def run(): Unit = {
           for (_ <- 0 until 100) {
-            val message = ContentMessage(null, "hi", false, null)
+            val message = OneWayMessage(null, "hi")
             inbox.post(message)
           }
           exitLatch.countDown()
diff --git a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
index ccca795683da3..323184cdd9b6e 100644
--- a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
@@ -33,7 +33,7 @@ class NettyRpcHandlerSuite extends SparkFunSuite {
   val env = mock(classOf[NettyRpcEnv])
   val sm = mock(classOf[StreamManager])
   when(env.deserialize(any(classOf[TransportClient]), any(classOf[Array[Byte]]))(any()))
-    .thenReturn(RequestMessage(RpcAddress("localhost", 12345), null, null, false))
+    .thenReturn(RequestMessage(RpcAddress("localhost", 12345), null, null))
 
   test("receive") {
     val dispatcher = mock(classOf[Dispatcher])
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java b/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
index 876fcd846791c..8a58e7b24585b 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
@@ -25,6 +25,7 @@
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
@@ -36,6 +37,7 @@
 import org.slf4j.LoggerFactory;
 
 import org.apache.spark.network.protocol.ChunkFetchRequest;
+import org.apache.spark.network.protocol.OneWayMessage;
 import org.apache.spark.network.protocol.RpcRequest;
 import org.apache.spark.network.protocol.StreamChunkId;
 import org.apache.spark.network.protocol.StreamRequest;
@@ -205,8 +207,12 @@ public void operationComplete(ChannelFuture future) throws Exception {
   /**
    * Sends an opaque message to the RpcHandler on the server-side. The callback will be invoked
    * with the server's response or upon any failure.
+   *
+   * @param message The message to send.
+   * @param callback Callback to handle the RPC's reply.
+   * @return The RPC's id.
    */
-  public void sendRpc(byte[] message, final RpcResponseCallback callback) {
+  public long sendRpc(byte[] message, final RpcResponseCallback callback) {
     final String serverAddr = NettyUtils.getRemoteAddress(channel);
     final long startTime = System.currentTimeMillis();
     logger.trace("Sending RPC to {}", serverAddr);
@@ -235,6 +241,8 @@ public void operationComplete(ChannelFuture future) throws Exception {
           }
         }
       });
+
+    return requestId;
   }
 
   /**
@@ -265,11 +273,35 @@ public void onFailure(Throwable e) {
     }
   }
 
+  /**
+   * Sends an opaque message to the RpcHandler on the server-side. No reply is expected for the
+   * message, and no delivery guarantees are made.
+   *
+   * @param message The message to send.
+   */
+  public void send(byte[] message) {
+    channel.writeAndFlush(new OneWayMessage(message));
+  }
+
+  /**
+   * Removes any state associated with the given RPC.
+   *
+   * @param requestId The RPC id returned by {@link #sendRpc(byte[], RpcResponseCallback)}.
+   */
+  public void removeRpcRequest(long requestId) {
+    handler.removeRpcRequest(requestId);
+  }
+
   /** Mark this channel as having timed out. */
   public void timeOut() {
     this.timedOut = true;
   }
 
+  @VisibleForTesting
+  public TransportResponseHandler getHandler() {
+    return handler;
+  }
+
   @Override
   public void close() {
     // close is a local operation and should finish with milliseconds; timeout just to be safe
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/Message.java b/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
index d01598c20f16f..39afd03db60ee 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
@@ -28,7 +28,8 @@ public interface Message extends Encodable {
   public static enum Type implements Encodable {
     ChunkFetchRequest(0), ChunkFetchSuccess(1), ChunkFetchFailure(2),
     RpcRequest(3), RpcResponse(4), RpcFailure(5),
-    StreamRequest(6), StreamResponse(7), StreamFailure(8);
+    StreamRequest(6), StreamResponse(7), StreamFailure(8),
+    OneWayMessage(9);
 
     private final byte id;
 
@@ -55,6 +56,7 @@ public static Type decode(ByteBuf buf) {
         case 6: return StreamRequest;
         case 7: return StreamResponse;
         case 8: return StreamFailure;
+        case 9: return OneWayMessage;
         default: throw new IllegalArgumentException("Unknown message type: " + id);
       }
     }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/MessageDecoder.java b/network/common/src/main/java/org/apache/spark/network/protocol/MessageDecoder.java
index 3c04048f3821a..074780f2b95ce 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/MessageDecoder.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/MessageDecoder.java
@@ -63,6 +63,9 @@ private Message decode(Message.Type msgType, ByteBuf in) {
       case RpcFailure:
         return RpcFailure.decode(in);
 
+      case OneWayMessage:
+        return OneWayMessage.decode(in);
+
       case StreamRequest:
         return StreamRequest.decode(in);
 
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/OneWayMessage.java b/network/common/src/main/java/org/apache/spark/network/protocol/OneWayMessage.java
new file mode 100644
index 0000000000000..95a0270be3da9
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/OneWayMessage.java
@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.protocol;
+
+import java.util.Arrays;
+
+import com.google.common.base.Objects;
+import io.netty.buffer.ByteBuf;
+
+/**
+ * A RPC that does not expect a reply, which is handled by a remote
+ * {@link org.apache.spark.network.server.RpcHandler}.
+ */
+public final class OneWayMessage implements RequestMessage {
+  /** Serialized message to send to remote RpcHandler. */
+  public final byte[] message;
+
+  public OneWayMessage(byte[] message) {
+    this.message = message;
+  }
+
+  @Override
+  public Type type() { return Type.OneWayMessage; }
+
+  @Override
+  public int encodedLength() {
+    return Encoders.ByteArrays.encodedLength(message);
+  }
+
+  @Override
+  public void encode(ByteBuf buf) {
+    Encoders.ByteArrays.encode(buf, message);
+  }
+
+  public static OneWayMessage decode(ByteBuf buf) {
+    byte[] message = Encoders.ByteArrays.decode(buf);
+    return new OneWayMessage(message);
+  }
+
+  @Override
+  public int hashCode() {
+    return Arrays.hashCode(message);
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (other instanceof OneWayMessage) {
+      OneWayMessage o = (OneWayMessage) other;
+      return Arrays.equals(message, o.message);
+    }
+    return false;
+  }
+
+  @Override
+  public String toString() {
+    return Objects.toStringHelper(this)
+      .add("message", message)
+      .toString();
+  }
+}
diff --git a/network/common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java b/network/common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
index 7033adb9cae6f..830db94b890c5 100644
--- a/network/common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
@@ -108,6 +108,11 @@ public void receive(TransportClient client, byte[] message, RpcResponseCallback
     }
   }
 
+  @Override
+  public void receive(TransportClient client, byte[] message) {
+    delegate.receive(client, message);
+  }
+
   @Override
   public StreamManager getStreamManager() {
     return delegate.getStreamManager();
diff --git a/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java b/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
index dbb7f95f55bc0..65109ddfe13b9 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
@@ -17,6 +17,9 @@
 
 package org.apache.spark.network.server;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
 
@@ -24,6 +27,9 @@
  * Handler for sendRPC() messages sent by {@link org.apache.spark.network.client.TransportClient}s.
  */
 public abstract class RpcHandler {
+
+  private static final RpcResponseCallback ONE_WAY_CALLBACK = new OneWayRpcCallback();
+
   /**
    * Receive a single RPC message. Any exception thrown while in this method will be sent back to
    * the client in string form as a standard RPC failure.
@@ -47,6 +53,19 @@ public abstract void receive(
    */
   public abstract StreamManager getStreamManager();
 
+  /**
+   * Receives an RPC message that does not expect a reply. The default implementation will
+   * call "{@link receive(TransportClient, byte[], RpcResponseCallback}" and log a warning if
+   * any of the callback methods are called.
+   *
+   * @param client A channel client which enables the handler to make requests back to the sender
+   *               of this RPC. This will always be the exact same object for a particular channel.
+   * @param message The serialized bytes of the RPC.
+   */
+  public void receive(TransportClient client, byte[] message) {
+    receive(client, message, ONE_WAY_CALLBACK);
+  }
+
   /**
    * Invoked when the connection associated with the given client has been invalidated.
    * No further requests will come from this client.
@@ -54,4 +73,21 @@ public abstract void receive(
   public void connectionTerminated(TransportClient client) { }
 
   public void exceptionCaught(Throwable cause, TransportClient client) { }
+
+  private static class OneWayRpcCallback implements RpcResponseCallback {
+
+    private final Logger logger = LoggerFactory.getLogger(OneWayRpcCallback.class);
+
+    @Override
+    public void onSuccess(byte[] response) {
+      logger.warn("Response provided for one-way RPC.");
+    }
+
+    @Override
+    public void onFailure(Throwable e) {
+      logger.error("Error response provided for one-way RPC.", e);
+    }
+
+  }
+
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java b/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
index 4f67bd573be21..db18ea77d1073 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
@@ -17,6 +17,7 @@
 
 package org.apache.spark.network.server;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelFuture;
@@ -27,13 +28,14 @@
 import org.apache.spark.network.buffer.ManagedBuffer;
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
-import org.apache.spark.network.protocol.Encodable;
-import org.apache.spark.network.protocol.RequestMessage;
 import org.apache.spark.network.protocol.ChunkFetchRequest;
-import org.apache.spark.network.protocol.RpcRequest;
 import org.apache.spark.network.protocol.ChunkFetchFailure;
 import org.apache.spark.network.protocol.ChunkFetchSuccess;
+import org.apache.spark.network.protocol.Encodable;
+import org.apache.spark.network.protocol.OneWayMessage;
+import org.apache.spark.network.protocol.RequestMessage;
 import org.apache.spark.network.protocol.RpcFailure;
+import org.apache.spark.network.protocol.RpcRequest;
 import org.apache.spark.network.protocol.RpcResponse;
 import org.apache.spark.network.protocol.StreamFailure;
 import org.apache.spark.network.protocol.StreamRequest;
@@ -95,6 +97,8 @@ public void handle(RequestMessage request) {
       processFetchRequest((ChunkFetchRequest) request);
     } else if (request instanceof RpcRequest) {
       processRpcRequest((RpcRequest) request);
+    } else if (request instanceof OneWayMessage) {
+      processOneWayMessage((OneWayMessage) request);
     } else if (request instanceof StreamRequest) {
       processStreamRequest((StreamRequest) request);
     } else {
@@ -156,6 +160,14 @@ public void onFailure(Throwable e) {
     }
   }
 
+  private void processOneWayMessage(OneWayMessage req) {
+    try {
+      rpcHandler.receive(reverseClient, req.message);
+    } catch (Exception e) {
+      logger.error("Error while invoking RpcHandler#receive() for one-way message.", e);
+    }
+  }
+
   /**
    * Responds to a single message with some Encodable object. If a failure occurs while sending,
    * it will be logged and the channel closed.
diff --git a/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java b/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
index 22b451fc0e60e..1aa20900ffe74 100644
--- a/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
@@ -35,6 +35,7 @@
 import org.apache.spark.network.protocol.Message;
 import org.apache.spark.network.protocol.MessageDecoder;
 import org.apache.spark.network.protocol.MessageEncoder;
+import org.apache.spark.network.protocol.OneWayMessage;
 import org.apache.spark.network.protocol.RpcFailure;
 import org.apache.spark.network.protocol.RpcRequest;
 import org.apache.spark.network.protocol.RpcResponse;
@@ -84,6 +85,7 @@ public void requests() {
     testClientToServer(new RpcRequest(12345, new byte[0]));
     testClientToServer(new RpcRequest(12345, new byte[100]));
     testClientToServer(new StreamRequest("abcde"));
+    testClientToServer(new OneWayMessage(new byte[100]));
   }
 
   @Test
diff --git a/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java b/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
index 8eb56bdd9846f..88fa2258bb794 100644
--- a/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
@@ -17,9 +17,11 @@
 
 package org.apache.spark.network;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Set;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
@@ -46,6 +48,7 @@ public class RpcIntegrationSuite {
   static TransportServer server;
   static TransportClientFactory clientFactory;
   static RpcHandler rpcHandler;
+  static List<String> oneWayMsgs;
 
   @BeforeClass
   public static void setUp() throws Exception {
@@ -64,12 +67,19 @@ public void receive(TransportClient client, byte[] message, RpcResponseCallback
         }
       }
 
+      @Override
+      public void receive(TransportClient client, byte[] message) {
+        String msg = new String(message, Charsets.UTF_8);
+        oneWayMsgs.add(msg);
+      }
+
       @Override
       public StreamManager getStreamManager() { return new OneForOneStreamManager(); }
     };
     TransportContext context = new TransportContext(conf, rpcHandler);
     server = context.createServer();
     clientFactory = context.createClientFactory();
+    oneWayMsgs = new ArrayList<>();
   }
 
   @AfterClass
@@ -158,6 +168,27 @@ public void sendSuccessAndFailure() throws Exception {
     assertErrorsContain(res.errorMessages, Sets.newHashSet("Thrown: the", "Returned: !"));
   }
 
+  @Test
+  public void sendOneWayMessage() throws Exception {
+    final String message = "no reply";
+    TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
+    try {
+      client.send(message.getBytes(Charsets.UTF_8));
+      assertEquals(0, client.getHandler().numOutstandingRequests());
+
+      // Make sure the message arrives.
+      long deadline = System.nanoTime() + TimeUnit.NANOSECONDS.convert(10, TimeUnit.SECONDS);
+      while (System.nanoTime() < deadline && oneWayMsgs.size() == 0) {
+        TimeUnit.MILLISECONDS.sleep(10);
+      }
+
+      assertEquals(1, oneWayMsgs.size());
+      assertEquals(message, oneWayMsgs.get(0));
+    } finally {
+      client.close();
+    }
+  }
+
   private void assertErrorsContain(Set<String> errors, Set<String> contains) {
     assertEquals(contains.size(), errors.size());
 
diff --git a/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java b/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
index b146899670180..a6f180bc40c9a 100644
--- a/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
@@ -21,6 +21,7 @@
 import static org.mockito.Mockito.*;
 
 import java.io.File;
+import java.lang.reflect.Method;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.List;
@@ -353,6 +354,14 @@ public void testRpcHandlerDelegate() throws Exception {
     verify(handler).exceptionCaught(any(Throwable.class), any(TransportClient.class));
   }
 
+  @Test
+  public void testDelegates() throws Exception {
+    Method[] rpcHandlerMethods = RpcHandler.class.getDeclaredMethods();
+    for (Method m : rpcHandlerMethods) {
+      SaslRpcHandler.class.getDeclaredMethod(m.getName(), m.getParameterTypes());
+    }
+  }
+
   private static class SaslTestCtx {
 
     final TransportClient client;

From ecac2835458bbf73fe59413d5bf921500c5b987d Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Wed, 25 Nov 2015 13:45:41 -0800
Subject: [PATCH 466/518] Fix Aggregator documentation (rename present to
 finish).

---
 .../scala/org/apache/spark/sql/expressions/Aggregator.scala     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
index b0cd32b5f73e6..65117d5824755 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/Aggregator.scala
@@ -34,7 +34,7 @@ import org.apache.spark.sql.{DataFrame, Dataset, Encoder, TypedColumn}
  *     def zero: Int = 0
  *     def reduce(b: Int, a: Data): Int = b + a.i
  *     def merge(b1: Int, b2: Int): Int = b1 + b2
- *     def present(r: Int): Int = r
+ *     def finish(r: Int): Int = r
  *   }.toColumn()
  *
  *   val ds: Dataset[Data] = ...

From 21e5606419c4b7462d30580c549e9bfa0123ae23 Mon Sep 17 00:00:00 2001
From: Daoyuan Wang <daoyuan.wang@intel.com>
Date: Wed, 25 Nov 2015 13:51:30 -0800
Subject: [PATCH 467/518] [SPARK-11983][SQL] remove all unused codegen fallback
 trait

Author: Daoyuan Wang <daoyuan.wang@intel.com>

Closes #9966 from adrian-wang/removeFallback.
---
 .../org/apache/spark/sql/catalyst/expressions/Cast.scala      | 3 +--
 .../spark/sql/catalyst/expressions/regexpExpressions.scala    | 4 ++--
 .../spark/sql/catalyst/expressions/NonFoldableLiteral.scala   | 3 +--
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
index 533d17ea5c172..a2c6c39fd8ce2 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala
@@ -104,8 +104,7 @@ object Cast {
 }
 
 /** Cast the child expression to the target data type. */
-case class Cast(child: Expression, dataType: DataType)
-  extends UnaryExpression with CodegenFallback {
+case class Cast(child: Expression, dataType: DataType) extends UnaryExpression {
 
   override def toString: String = s"cast($child as ${dataType.simpleString})"
 
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala
index 9e484c5ed83bf..adef6050c3565 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala
@@ -66,7 +66,7 @@ trait StringRegexExpression extends ImplicitCastInputTypes {
  * Simple RegEx pattern matching function
  */
 case class Like(left: Expression, right: Expression)
-  extends BinaryExpression with StringRegexExpression with CodegenFallback {
+  extends BinaryExpression with StringRegexExpression {
 
   override def escape(v: String): String = StringUtils.escapeLikeRegex(v)
 
@@ -117,7 +117,7 @@ case class Like(left: Expression, right: Expression)
 
 
 case class RLike(left: Expression, right: Expression)
-  extends BinaryExpression with StringRegexExpression with CodegenFallback {
+  extends BinaryExpression with StringRegexExpression {
 
   override def escape(v: String): String = v
   override def matches(regex: Pattern, str: String): Boolean = regex.matcher(str).find(0)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/NonFoldableLiteral.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/NonFoldableLiteral.scala
index 31ecf4a9e810a..118fd695fe2f5 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/NonFoldableLiteral.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/NonFoldableLiteral.scala
@@ -26,8 +26,7 @@ import org.apache.spark.sql.types._
  * A literal value that is not foldable. Used in expression codegen testing to test code path
  * that behave differently based on foldable values.
  */
-case class NonFoldableLiteral(value: Any, dataType: DataType)
-  extends LeafExpression with CodegenFallback {
+case class NonFoldableLiteral(value: Any, dataType: DataType) extends LeafExpression {
 
   override def foldable: Boolean = false
   override def nullable: Boolean = true

From cc243a079b1c039d6e7f0b410d1654d94a090e14 Mon Sep 17 00:00:00 2001
From: Carson Wang <carson.wang@intel.com>
Date: Wed, 25 Nov 2015 15:13:13 -0800
Subject: [PATCH 468/518] [SPARK-11206] Support SQL UI on the history server

On the live web UI, there is a SQL tab which provides valuable information for the SQL query. But once the workload is finished, we won't see the SQL tab on the history server. It will be helpful if we support SQL UI on the history server so we can analyze it even after its execution.

To support SQL UI on the history server:
1. I added an `onOtherEvent` method to the `SparkListener` trait and post all SQL related events to the same event bus.
2. Two SQL events `SparkListenerSQLExecutionStart` and `SparkListenerSQLExecutionEnd` are defined in the sql module.
3. The new SQL events are written to event log using Jackson.
4.  A new trait `SparkHistoryListenerFactory` is added to allow the history server to feed events to the SQL history listener. The SQL implementation is loaded at runtime using `java.util.ServiceLoader`.

Author: Carson Wang <carson.wang@intel.com>

Closes #9297 from carsonwang/SqlHistoryUI.
---
 .rat-excludes                                 |   1 +
 .../org/apache/spark/JavaSparkListener.java   |   3 +
 .../apache/spark/SparkFirehoseListener.java   |   4 +
 .../scheduler/EventLoggingListener.scala      |   4 +
 .../spark/scheduler/SparkListener.scala       |  24 ++-
 .../spark/scheduler/SparkListenerBus.scala    |   1 +
 .../scala/org/apache/spark/ui/SparkUI.scala   |  16 +-
 .../org/apache/spark/util/JsonProtocol.scala  |  11 +-
 ...park.scheduler.SparkHistoryListenerFactory |   1 +
 .../org/apache/spark/sql/SQLContext.scala     |  18 ++-
 .../spark/sql/execution/SQLExecution.scala    |  24 +--
 .../spark/sql/execution/SparkPlanInfo.scala   |  46 ++++++
 .../sql/execution/metric/SQLMetricInfo.scala  |  30 ++++
 .../sql/execution/metric/SQLMetrics.scala     |  56 ++++---
 .../sql/execution/ui/ExecutionPage.scala      |   4 +-
 .../spark/sql/execution/ui/SQLListener.scala  | 139 ++++++++++++------
 .../spark/sql/execution/ui/SQLTab.scala       |  12 +-
 .../sql/execution/ui/SparkPlanGraph.scala     |  20 +--
 .../execution/metric/SQLMetricsSuite.scala    |   4 +-
 .../sql/execution/ui/SQLListenerSuite.scala   |  43 +++---
 .../spark/sql/test/SharedSQLContext.scala     |   1 +
 21 files changed, 327 insertions(+), 135 deletions(-)
 create mode 100644 sql/core/src/main/resources/META-INF/services/org.apache.spark.scheduler.SparkHistoryListenerFactory
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
 create mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetricInfo.scala

diff --git a/.rat-excludes b/.rat-excludes
index 08fba6d351d6a..7262c960ed6bb 100644
--- a/.rat-excludes
+++ b/.rat-excludes
@@ -82,4 +82,5 @@ INDEX
 gen-java.*
 .*avpr
 org.apache.spark.sql.sources.DataSourceRegister
+org.apache.spark.scheduler.SparkHistoryListenerFactory
 .*parquet
diff --git a/core/src/main/java/org/apache/spark/JavaSparkListener.java b/core/src/main/java/org/apache/spark/JavaSparkListener.java
index fa9acf0a15b88..23bc9a2e81727 100644
--- a/core/src/main/java/org/apache/spark/JavaSparkListener.java
+++ b/core/src/main/java/org/apache/spark/JavaSparkListener.java
@@ -82,4 +82,7 @@ public void onExecutorRemoved(SparkListenerExecutorRemoved executorRemoved) { }
   @Override
   public void onBlockUpdated(SparkListenerBlockUpdated blockUpdated) { }
 
+  @Override
+  public void onOtherEvent(SparkListenerEvent event) { }
+
 }
diff --git a/core/src/main/java/org/apache/spark/SparkFirehoseListener.java b/core/src/main/java/org/apache/spark/SparkFirehoseListener.java
index 1214d05ba6063..e6b24afd88ad4 100644
--- a/core/src/main/java/org/apache/spark/SparkFirehoseListener.java
+++ b/core/src/main/java/org/apache/spark/SparkFirehoseListener.java
@@ -118,4 +118,8 @@ public void onBlockUpdated(SparkListenerBlockUpdated blockUpdated) {
         onEvent(blockUpdated);
     }
 
+    @Override
+    public void onOtherEvent(SparkListenerEvent event) {
+        onEvent(event);
+    }
 }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala b/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala
index 000a021a528cf..eaa07acc5132e 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala
@@ -207,6 +207,10 @@ private[spark] class EventLoggingListener(
   // No-op because logging every update would be overkill
   override def onExecutorMetricsUpdate(event: SparkListenerExecutorMetricsUpdate): Unit = { }
 
+  override def onOtherEvent(event: SparkListenerEvent): Unit = {
+    logEvent(event, flushLogger = true)
+  }
+
   /**
    * Stop logging events. The event log file will be renamed so that it loses the
    * ".inprogress" suffix.
diff --git a/core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala b/core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala
index 896f1743332f1..075a7f13172de 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala
@@ -22,15 +22,19 @@ import java.util.Properties
 import scala.collection.Map
 import scala.collection.mutable
 
-import org.apache.spark.{Logging, TaskEndReason}
+import com.fasterxml.jackson.annotation.JsonTypeInfo
+
+import org.apache.spark.{Logging, SparkConf, TaskEndReason}
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.scheduler.cluster.ExecutorInfo
 import org.apache.spark.storage.{BlockManagerId, BlockUpdatedInfo}
 import org.apache.spark.util.{Distribution, Utils}
+import org.apache.spark.ui.SparkUI
 
 @DeveloperApi
-sealed trait SparkListenerEvent
+@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "Event")
+trait SparkListenerEvent
 
 @DeveloperApi
 case class SparkListenerStageSubmitted(stageInfo: StageInfo, properties: Properties = null)
@@ -130,6 +134,17 @@ case class SparkListenerApplicationEnd(time: Long) extends SparkListenerEvent
  */
 private[spark] case class SparkListenerLogStart(sparkVersion: String) extends SparkListenerEvent
 
+/**
+ * Interface for creating history listeners defined in other modules like SQL, which are used to
+ * rebuild the history UI.
+ */
+private[spark] trait SparkHistoryListenerFactory {
+  /**
+   * Create listeners used to rebuild the history UI.
+   */
+  def createListeners(conf: SparkConf, sparkUI: SparkUI): Seq[SparkListener]
+}
+
 /**
  * :: DeveloperApi ::
  * Interface for listening to events from the Spark scheduler. Note that this is an internal
@@ -223,6 +238,11 @@ trait SparkListener {
    * Called when the driver receives a block update info.
    */
   def onBlockUpdated(blockUpdated: SparkListenerBlockUpdated) { }
+
+  /**
+   * Called when other events like SQL-specific events are posted.
+   */
+  def onOtherEvent(event: SparkListenerEvent) { }
 }
 
 /**
diff --git a/core/src/main/scala/org/apache/spark/scheduler/SparkListenerBus.scala b/core/src/main/scala/org/apache/spark/scheduler/SparkListenerBus.scala
index 04afde33f5aad..95722a07144ec 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/SparkListenerBus.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/SparkListenerBus.scala
@@ -61,6 +61,7 @@ private[spark] trait SparkListenerBus extends ListenerBus[SparkListener, SparkLi
       case blockUpdated: SparkListenerBlockUpdated =>
         listener.onBlockUpdated(blockUpdated)
       case logStart: SparkListenerLogStart => // ignore event log metadata
+      case _ => listener.onOtherEvent(event)
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/ui/SparkUI.scala b/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
index 4608bce202ec8..8da6884a38535 100644
--- a/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
+++ b/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
@@ -17,10 +17,13 @@
 
 package org.apache.spark.ui
 
-import java.util.Date
+import java.util.{Date, ServiceLoader}
+
+import scala.collection.JavaConverters._
 
 import org.apache.spark.status.api.v1.{ApiRootResource, ApplicationAttemptInfo, ApplicationInfo,
   UIRoot}
+import org.apache.spark.util.Utils
 import org.apache.spark.{Logging, SecurityManager, SparkConf, SparkContext}
 import org.apache.spark.scheduler._
 import org.apache.spark.storage.StorageStatusListener
@@ -154,7 +157,16 @@ private[spark] object SparkUI {
       appName: String,
       basePath: String,
       startTime: Long): SparkUI = {
-    create(None, conf, listenerBus, securityManager, appName, basePath, startTime = startTime)
+    val sparkUI = create(
+      None, conf, listenerBus, securityManager, appName, basePath, startTime = startTime)
+
+    val listenerFactories = ServiceLoader.load(classOf[SparkHistoryListenerFactory],
+      Utils.getContextOrSparkClassLoader).asScala
+    listenerFactories.foreach { listenerFactory =>
+      val listeners = listenerFactory.createListeners(conf, sparkUI)
+      listeners.foreach(listenerBus.addListener)
+    }
+    sparkUI
   }
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
index c9beeb25e05af..7f5d713ec6505 100644
--- a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
+++ b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
@@ -19,19 +19,21 @@ package org.apache.spark.util
 
 import java.util.{Properties, UUID}
 
-import org.apache.spark.scheduler.cluster.ExecutorInfo
-
 import scala.collection.JavaConverters._
 import scala.collection.Map
 
+import com.fasterxml.jackson.databind.ObjectMapper
+import com.fasterxml.jackson.module.scala.DefaultScalaModule
 import org.json4s.DefaultFormats
 import org.json4s.JsonDSL._
 import org.json4s.JsonAST._
+import org.json4s.jackson.JsonMethods._
 
 import org.apache.spark._
 import org.apache.spark.executor._
 import org.apache.spark.rdd.RDDOperationScope
 import org.apache.spark.scheduler._
+import org.apache.spark.scheduler.cluster.ExecutorInfo
 import org.apache.spark.storage._
 
 /**
@@ -54,6 +56,8 @@ private[spark] object JsonProtocol {
 
   private implicit val format = DefaultFormats
 
+  private val mapper = new ObjectMapper().registerModule(DefaultScalaModule)
+
   /** ------------------------------------------------- *
    * JSON serialization methods for SparkListenerEvents |
    * -------------------------------------------------- */
@@ -96,6 +100,7 @@ private[spark] object JsonProtocol {
         executorMetricsUpdateToJson(metricsUpdate)
       case blockUpdated: SparkListenerBlockUpdated =>
         throw new MatchError(blockUpdated)  // TODO(ekl) implement this
+      case _ => parse(mapper.writeValueAsString(event))
     }
   }
 
@@ -511,6 +516,8 @@ private[spark] object JsonProtocol {
       case `executorRemoved` => executorRemovedFromJson(json)
       case `logStart` => logStartFromJson(json)
       case `metricsUpdate` => executorMetricsUpdateFromJson(json)
+      case other => mapper.readValue(compact(render(json)), Utils.classForName(other))
+        .asInstanceOf[SparkListenerEvent]
     }
   }
 
diff --git a/sql/core/src/main/resources/META-INF/services/org.apache.spark.scheduler.SparkHistoryListenerFactory b/sql/core/src/main/resources/META-INF/services/org.apache.spark.scheduler.SparkHistoryListenerFactory
new file mode 100644
index 0000000000000..507100be90967
--- /dev/null
+++ b/sql/core/src/main/resources/META-INF/services/org.apache.spark.scheduler.SparkHistoryListenerFactory
@@ -0,0 +1 @@
+org.apache.spark.sql.execution.ui.SQLHistoryListenerFactory
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index 46bf544fd885f..1c2ac5f6f11bf 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -1263,6 +1263,8 @@ object SQLContext {
    */
   @transient private val instantiatedContext = new AtomicReference[SQLContext]()
 
+  @transient private val sqlListener = new AtomicReference[SQLListener]()
+
   /**
    * Get the singleton SQLContext if it exists or create a new one using the given SparkContext.
    *
@@ -1307,6 +1309,10 @@ object SQLContext {
     Option(instantiatedContext.get())
   }
 
+  private[sql] def clearSqlListener(): Unit = {
+    sqlListener.set(null)
+  }
+
   /**
    * Changes the SQLContext that will be returned in this thread and its children when
    * SQLContext.getOrCreate() is called. This can be used to ensure that a given thread receives
@@ -1355,9 +1361,13 @@ object SQLContext {
    * Create a SQLListener then add it into SparkContext, and create an SQLTab if there is SparkUI.
    */
   private[sql] def createListenerAndUI(sc: SparkContext): SQLListener = {
-    val listener = new SQLListener(sc.conf)
-    sc.addSparkListener(listener)
-    sc.ui.foreach(new SQLTab(listener, _))
-    listener
+    if (sqlListener.get() == null) {
+      val listener = new SQLListener(sc.conf)
+      if (sqlListener.compareAndSet(null, listener)) {
+        sc.addSparkListener(listener)
+        sc.ui.foreach(new SQLTab(listener, _))
+      }
+    }
+    sqlListener.get()
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
index 1422e15549c94..34971986261c2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
@@ -21,7 +21,8 @@ import java.util.concurrent.atomic.AtomicLong
 
 import org.apache.spark.SparkContext
 import org.apache.spark.sql.SQLContext
-import org.apache.spark.sql.execution.ui.SparkPlanGraph
+import org.apache.spark.sql.execution.ui.{SparkListenerSQLExecutionStart,
+  SparkListenerSQLExecutionEnd}
 import org.apache.spark.util.Utils
 
 private[sql] object SQLExecution {
@@ -45,25 +46,14 @@ private[sql] object SQLExecution {
       sc.setLocalProperty(EXECUTION_ID_KEY, executionId.toString)
       val r = try {
         val callSite = Utils.getCallSite()
-        sqlContext.listener.onExecutionStart(
-          executionId,
-          callSite.shortForm,
-          callSite.longForm,
-          queryExecution.toString,
-          SparkPlanGraph(queryExecution.executedPlan),
-          System.currentTimeMillis())
+        sqlContext.sparkContext.listenerBus.post(SparkListenerSQLExecutionStart(
+          executionId, callSite.shortForm, callSite.longForm, queryExecution.toString,
+          SparkPlanInfo.fromSparkPlan(queryExecution.executedPlan), System.currentTimeMillis()))
         try {
           body
         } finally {
-          // Ideally, we need to make sure onExecutionEnd happens after onJobStart and onJobEnd.
-          // However, onJobStart and onJobEnd run in the listener thread. Because we cannot add new
-          // SQL event types to SparkListener since it's a public API, we cannot guarantee that.
-          //
-          // SQLListener should handle the case that onExecutionEnd happens before onJobEnd.
-          //
-          // The worst case is onExecutionEnd may happen before onJobStart when the listener thread
-          // is very busy. If so, we cannot track the jobs for the execution. It seems acceptable.
-          sqlContext.listener.onExecutionEnd(executionId, System.currentTimeMillis())
+          sqlContext.sparkContext.listenerBus.post(SparkListenerSQLExecutionEnd(
+            executionId, System.currentTimeMillis()))
         }
       } finally {
         sc.setLocalProperty(EXECUTION_ID_KEY, null)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
new file mode 100644
index 0000000000000..486ce34064e43
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import org.apache.spark.annotation.DeveloperApi
+import org.apache.spark.sql.execution.metric.SQLMetricInfo
+import org.apache.spark.util.Utils
+
+/**
+ * :: DeveloperApi ::
+ * Stores information about a SQL SparkPlan.
+ */
+@DeveloperApi
+class SparkPlanInfo(
+    val nodeName: String,
+    val simpleString: String,
+    val children: Seq[SparkPlanInfo],
+    val metrics: Seq[SQLMetricInfo])
+
+private[sql] object SparkPlanInfo {
+
+  def fromSparkPlan(plan: SparkPlan): SparkPlanInfo = {
+    val metrics = plan.metrics.toSeq.map { case (key, metric) =>
+      new SQLMetricInfo(metric.name.getOrElse(key), metric.id,
+        Utils.getFormattedClassName(metric.param))
+    }
+    val children = plan.children.map(fromSparkPlan)
+
+    new SparkPlanInfo(plan.nodeName, plan.simpleString, children, metrics)
+  }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetricInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetricInfo.scala
new file mode 100644
index 0000000000000..2708219ad3485
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetricInfo.scala
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.metric
+
+import org.apache.spark.annotation.DeveloperApi
+
+/**
+ * :: DeveloperApi ::
+ * Stores information about a SQL Metric.
+ */
+@DeveloperApi
+class SQLMetricInfo(
+    val name: String,
+    val accumulatorId: Long,
+    val metricParam: String)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala
index 1c253e3942e95..6c0f6f8a52dc5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala
@@ -104,21 +104,39 @@ private class LongSQLMetricParam(val stringValue: Seq[Long] => String, initialVa
   override def zero: LongSQLMetricValue = new LongSQLMetricValue(initialValue)
 }
 
+private object LongSQLMetricParam extends LongSQLMetricParam(_.sum.toString, 0L)
+
+private object StaticsLongSQLMetricParam extends LongSQLMetricParam(
+  (values: Seq[Long]) => {
+    // This is a workaround for SPARK-11013.
+    // We use -1 as initial value of the accumulator, if the accumulator is valid, we will update
+    // it at the end of task and the value will be at least 0.
+    val validValues = values.filter(_ >= 0)
+    val Seq(sum, min, med, max) = {
+      val metric = if (validValues.length == 0) {
+        Seq.fill(4)(0L)
+      } else {
+        val sorted = validValues.sorted
+        Seq(sorted.sum, sorted(0), sorted(validValues.length / 2), sorted(validValues.length - 1))
+      }
+      metric.map(Utils.bytesToString)
+    }
+    s"\n$sum ($min, $med, $max)"
+  }, -1L)
+
 private[sql] object SQLMetrics {
 
   private def createLongMetric(
       sc: SparkContext,
       name: String,
-      stringValue: Seq[Long] => String,
-      initialValue: Long): LongSQLMetric = {
-    val param = new LongSQLMetricParam(stringValue, initialValue)
+      param: LongSQLMetricParam): LongSQLMetric = {
     val acc = new LongSQLMetric(name, param)
     sc.cleaner.foreach(_.registerAccumulatorForCleanup(acc))
     acc
   }
 
   def createLongMetric(sc: SparkContext, name: String): LongSQLMetric = {
-    createLongMetric(sc, name, _.sum.toString, 0L)
+    createLongMetric(sc, name, LongSQLMetricParam)
   }
 
   /**
@@ -126,31 +144,25 @@ private[sql] object SQLMetrics {
    * spill size, etc.
    */
   def createSizeMetric(sc: SparkContext, name: String): LongSQLMetric = {
-    val stringValue = (values: Seq[Long]) => {
-      // This is a workaround for SPARK-11013.
-      // We use -1 as initial value of the accumulator, if the accumulator is valid, we will update
-      // it at the end of task and the value will be at least 0.
-      val validValues = values.filter(_ >= 0)
-      val Seq(sum, min, med, max) = {
-        val metric = if (validValues.length == 0) {
-          Seq.fill(4)(0L)
-        } else {
-          val sorted = validValues.sorted
-          Seq(sorted.sum, sorted(0), sorted(validValues.length / 2), sorted(validValues.length - 1))
-        }
-        metric.map(Utils.bytesToString)
-      }
-      s"\n$sum ($min, $med, $max)"
-    }
     // The final result of this metric in physical operator UI may looks like:
     // data size total (min, med, max):
     // 100GB (100MB, 1GB, 10GB)
-    createLongMetric(sc, s"$name total (min, med, max)", stringValue, -1L)
+    createLongMetric(sc, s"$name total (min, med, max)", StaticsLongSQLMetricParam)
+  }
+
+  def getMetricParam(metricParamName: String): SQLMetricParam[SQLMetricValue[Any], Any] = {
+    val longSQLMetricParam = Utils.getFormattedClassName(LongSQLMetricParam)
+    val staticsSQLMetricParam = Utils.getFormattedClassName(StaticsLongSQLMetricParam)
+    val metricParam = metricParamName match {
+      case `longSQLMetricParam` => LongSQLMetricParam
+      case `staticsSQLMetricParam` => StaticsLongSQLMetricParam
+    }
+    metricParam.asInstanceOf[SQLMetricParam[SQLMetricValue[Any], Any]]
   }
 
   /**
    * A metric that its value will be ignored. Use this one when we need a metric parameter but don't
    * care about the value.
    */
-  val nullLongMetric = new LongSQLMetric("null", new LongSQLMetricParam(_.sum.toString, 0L))
+  val nullLongMetric = new LongSQLMetric("null", LongSQLMetricParam)
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala
index e74d6fb396e1c..c74ad40406992 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala
@@ -19,9 +19,7 @@ package org.apache.spark.sql.execution.ui
 
 import javax.servlet.http.HttpServletRequest
 
-import scala.xml.{Node, Unparsed}
-
-import org.apache.commons.lang3.StringEscapeUtils
+import scala.xml.Node
 
 import org.apache.spark.Logging
 import org.apache.spark.ui.{UIUtils, WebUIPage}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala
index 5a072de400b6a..e19a1e3e5851f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala
@@ -19,11 +19,34 @@ package org.apache.spark.sql.execution.ui
 
 import scala.collection.mutable
 
-import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.scheduler._
 import org.apache.spark.sql.execution.SQLExecution
-import org.apache.spark.sql.execution.metric.{SQLMetricParam, SQLMetricValue}
+import org.apache.spark.sql.execution.SparkPlanInfo
+import org.apache.spark.sql.execution.metric.{LongSQLMetricValue, SQLMetricValue, SQLMetricParam}
 import org.apache.spark.{JobExecutionStatus, Logging, SparkConf}
+import org.apache.spark.ui.SparkUI
+
+@DeveloperApi
+case class SparkListenerSQLExecutionStart(
+    executionId: Long,
+    description: String,
+    details: String,
+    physicalPlanDescription: String,
+    sparkPlanInfo: SparkPlanInfo,
+    time: Long)
+  extends SparkListenerEvent
+
+@DeveloperApi
+case class SparkListenerSQLExecutionEnd(executionId: Long, time: Long)
+  extends SparkListenerEvent
+
+private[sql] class SQLHistoryListenerFactory extends SparkHistoryListenerFactory {
+
+  override def createListeners(conf: SparkConf, sparkUI: SparkUI): Seq[SparkListener] = {
+    List(new SQLHistoryListener(conf, sparkUI))
+  }
+}
 
 private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Logging {
 
@@ -118,7 +141,8 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
   override def onExecutorMetricsUpdate(
       executorMetricsUpdate: SparkListenerExecutorMetricsUpdate): Unit = synchronized {
     for ((taskId, stageId, stageAttemptID, metrics) <- executorMetricsUpdate.taskMetrics) {
-      updateTaskAccumulatorValues(taskId, stageId, stageAttemptID, metrics, finishTask = false)
+      updateTaskAccumulatorValues(taskId, stageId, stageAttemptID, metrics.accumulatorUpdates(),
+        finishTask = false)
     }
   }
 
@@ -140,7 +164,7 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
       taskEnd.taskInfo.taskId,
       taskEnd.stageId,
       taskEnd.stageAttemptId,
-      taskEnd.taskMetrics,
+      taskEnd.taskMetrics.accumulatorUpdates(),
       finishTask = true)
   }
 
@@ -148,15 +172,12 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
    * Update the accumulator values of a task with the latest metrics for this task. This is called
    * every time we receive an executor heartbeat or when a task finishes.
    */
-  private def updateTaskAccumulatorValues(
+  protected def updateTaskAccumulatorValues(
       taskId: Long,
       stageId: Int,
       stageAttemptID: Int,
-      metrics: TaskMetrics,
+      accumulatorUpdates: Map[Long, Any],
       finishTask: Boolean): Unit = {
-    if (metrics == null) {
-      return
-    }
 
     _stageIdToStageMetrics.get(stageId) match {
       case Some(stageMetrics) =>
@@ -174,9 +195,9 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
             case Some(taskMetrics) =>
               if (finishTask) {
                 taskMetrics.finished = true
-                taskMetrics.accumulatorUpdates = metrics.accumulatorUpdates()
+                taskMetrics.accumulatorUpdates = accumulatorUpdates
               } else if (!taskMetrics.finished) {
-                taskMetrics.accumulatorUpdates = metrics.accumulatorUpdates()
+                taskMetrics.accumulatorUpdates = accumulatorUpdates
               } else {
                 // If a task is finished, we should not override with accumulator updates from
                 // heartbeat reports
@@ -185,7 +206,7 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
               // TODO Now just set attemptId to 0. Should fix here when we can get the attempt
               // id from SparkListenerExecutorMetricsUpdate
               stageMetrics.taskIdToMetricUpdates(taskId) = new SQLTaskMetrics(
-                  attemptId = 0, finished = finishTask, metrics.accumulatorUpdates())
+                  attemptId = 0, finished = finishTask, accumulatorUpdates)
           }
         }
       case None =>
@@ -193,38 +214,40 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
     }
   }
 
-  def onExecutionStart(
-      executionId: Long,
-      description: String,
-      details: String,
-      physicalPlanDescription: String,
-      physicalPlanGraph: SparkPlanGraph,
-      time: Long): Unit = {
-    val sqlPlanMetrics = physicalPlanGraph.nodes.flatMap { node =>
-      node.metrics.map(metric => metric.accumulatorId -> metric)
-    }
-
-    val executionUIData = new SQLExecutionUIData(executionId, description, details,
-      physicalPlanDescription, physicalPlanGraph, sqlPlanMetrics.toMap, time)
-    synchronized {
-      activeExecutions(executionId) = executionUIData
-      _executionIdToData(executionId) = executionUIData
-    }
-  }
-
-  def onExecutionEnd(executionId: Long, time: Long): Unit = synchronized {
-    _executionIdToData.get(executionId).foreach { executionUIData =>
-      executionUIData.completionTime = Some(time)
-      if (!executionUIData.hasRunningJobs) {
-        // onExecutionEnd happens after all "onJobEnd"s
-        // So we should update the execution lists.
-        markExecutionFinished(executionId)
-      } else {
-        // There are some running jobs, onExecutionEnd happens before some "onJobEnd"s.
-        // Then we don't if the execution is successful, so let the last onJobEnd updates the
-        // execution lists.
+  override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
+    case SparkListenerSQLExecutionStart(executionId, description, details,
+      physicalPlanDescription, sparkPlanInfo, time) =>
+      val physicalPlanGraph = SparkPlanGraph(sparkPlanInfo)
+      val sqlPlanMetrics = physicalPlanGraph.nodes.flatMap { node =>
+        node.metrics.map(metric => metric.accumulatorId -> metric)
+      }
+      val executionUIData = new SQLExecutionUIData(
+        executionId,
+        description,
+        details,
+        physicalPlanDescription,
+        physicalPlanGraph,
+        sqlPlanMetrics.toMap,
+        time)
+      synchronized {
+        activeExecutions(executionId) = executionUIData
+        _executionIdToData(executionId) = executionUIData
+      }
+    case SparkListenerSQLExecutionEnd(executionId, time) => synchronized {
+      _executionIdToData.get(executionId).foreach { executionUIData =>
+        executionUIData.completionTime = Some(time)
+        if (!executionUIData.hasRunningJobs) {
+          // onExecutionEnd happens after all "onJobEnd"s
+          // So we should update the execution lists.
+          markExecutionFinished(executionId)
+        } else {
+          // There are some running jobs, onExecutionEnd happens before some "onJobEnd"s.
+          // Then we don't if the execution is successful, so let the last onJobEnd updates the
+          // execution lists.
+        }
       }
     }
+    case _ => // Ignore
   }
 
   private def markExecutionFinished(executionId: Long): Unit = {
@@ -289,6 +312,38 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
 
 }
 
+private[spark] class SQLHistoryListener(conf: SparkConf, sparkUI: SparkUI)
+  extends SQLListener(conf) {
+
+  private var sqlTabAttached = false
+
+  override def onExecutorMetricsUpdate(
+      executorMetricsUpdate: SparkListenerExecutorMetricsUpdate): Unit = synchronized {
+    // Do nothing
+  }
+
+  override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = synchronized {
+    updateTaskAccumulatorValues(
+      taskEnd.taskInfo.taskId,
+      taskEnd.stageId,
+      taskEnd.stageAttemptId,
+      taskEnd.taskInfo.accumulables.map { acc =>
+        (acc.id, new LongSQLMetricValue(acc.update.getOrElse("0").toLong))
+      }.toMap,
+      finishTask = true)
+  }
+
+  override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
+    case _: SparkListenerSQLExecutionStart =>
+      if (!sqlTabAttached) {
+        new SQLTab(this, sparkUI)
+        sqlTabAttached = true
+      }
+      super.onOtherEvent(event)
+    case _ => super.onOtherEvent(event)
+  }
+}
+
 /**
  * Represent all necessary data for an execution that will be used in Web UI.
  */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLTab.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLTab.scala
index 9c27944d42fc6..4f50b2ecdc8f8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLTab.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLTab.scala
@@ -17,13 +17,11 @@
 
 package org.apache.spark.sql.execution.ui
 
-import java.util.concurrent.atomic.AtomicInteger
-
 import org.apache.spark.Logging
 import org.apache.spark.ui.{SparkUI, SparkUITab}
 
 private[sql] class SQLTab(val listener: SQLListener, sparkUI: SparkUI)
-  extends SparkUITab(sparkUI, SQLTab.nextTabName) with Logging {
+  extends SparkUITab(sparkUI, "SQL") with Logging {
 
   val parent = sparkUI
 
@@ -35,13 +33,5 @@ private[sql] class SQLTab(val listener: SQLListener, sparkUI: SparkUI)
 }
 
 private[sql] object SQLTab {
-
   private val STATIC_RESOURCE_DIR = "org/apache/spark/sql/execution/ui/static"
-
-  private val nextTabId = new AtomicInteger(0)
-
-  private def nextTabName: String = {
-    val nextId = nextTabId.getAndIncrement()
-    if (nextId == 0) "SQL" else s"SQL$nextId"
-  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SparkPlanGraph.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SparkPlanGraph.scala
index f1fce5478a3fe..7af0ff09c5c6d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SparkPlanGraph.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SparkPlanGraph.scala
@@ -21,8 +21,8 @@ import java.util.concurrent.atomic.AtomicLong
 
 import scala.collection.mutable
 
-import org.apache.spark.sql.execution.SparkPlan
-import org.apache.spark.sql.execution.metric.{SQLMetricParam, SQLMetricValue}
+import org.apache.spark.sql.execution.SparkPlanInfo
+import org.apache.spark.sql.execution.metric.SQLMetrics
 
 /**
  * A graph used for storing information of an executionPlan of DataFrame.
@@ -48,27 +48,27 @@ private[sql] object SparkPlanGraph {
   /**
    * Build a SparkPlanGraph from the root of a SparkPlan tree.
    */
-  def apply(plan: SparkPlan): SparkPlanGraph = {
+  def apply(planInfo: SparkPlanInfo): SparkPlanGraph = {
     val nodeIdGenerator = new AtomicLong(0)
     val nodes = mutable.ArrayBuffer[SparkPlanGraphNode]()
     val edges = mutable.ArrayBuffer[SparkPlanGraphEdge]()
-    buildSparkPlanGraphNode(plan, nodeIdGenerator, nodes, edges)
+    buildSparkPlanGraphNode(planInfo, nodeIdGenerator, nodes, edges)
     new SparkPlanGraph(nodes, edges)
   }
 
   private def buildSparkPlanGraphNode(
-      plan: SparkPlan,
+      planInfo: SparkPlanInfo,
       nodeIdGenerator: AtomicLong,
       nodes: mutable.ArrayBuffer[SparkPlanGraphNode],
       edges: mutable.ArrayBuffer[SparkPlanGraphEdge]): SparkPlanGraphNode = {
-    val metrics = plan.metrics.toSeq.map { case (key, metric) =>
-      SQLPlanMetric(metric.name.getOrElse(key), metric.id,
-        metric.param.asInstanceOf[SQLMetricParam[SQLMetricValue[Any], Any]])
+    val metrics = planInfo.metrics.map { metric =>
+      SQLPlanMetric(metric.name, metric.accumulatorId,
+        SQLMetrics.getMetricParam(metric.metricParam))
     }
     val node = SparkPlanGraphNode(
-      nodeIdGenerator.getAndIncrement(), plan.nodeName, plan.simpleString, metrics)
+      nodeIdGenerator.getAndIncrement(), planInfo.nodeName, planInfo.simpleString, metrics)
     nodes += node
-    val childrenNodes = plan.children.map(
+    val childrenNodes = planInfo.children.map(
       child => buildSparkPlanGraphNode(child, nodeIdGenerator, nodes, edges))
     for (child <- childrenNodes) {
       edges += SparkPlanGraphEdge(child.id, node.id)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index 5e2b4154dd7ce..ebfa1eaf3e5bc 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -26,6 +26,7 @@ import org.apache.xbean.asm5.Opcodes._
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql._
+import org.apache.spark.sql.execution.SparkPlanInfo
 import org.apache.spark.sql.execution.ui.SparkPlanGraph
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
@@ -82,7 +83,8 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
     if (jobs.size == expectedNumOfJobs) {
       // If we can track all jobs, check the metric values
       val metricValues = sqlContext.listener.getExecutionMetrics(executionId)
-      val actualMetrics = SparkPlanGraph(df.queryExecution.executedPlan).nodes.filter { node =>
+      val actualMetrics = SparkPlanGraph(SparkPlanInfo.fromSparkPlan(
+        df.queryExecution.executedPlan)).nodes.filter { node =>
         expectedMetrics.contains(node.id)
       }.map { node =>
         val nodeMetrics = node.metrics.map { metric =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
index c15aac775096c..f93d081d0c30e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
@@ -21,10 +21,10 @@ import java.util.Properties
 
 import org.apache.spark.{SparkException, SparkContext, SparkConf, SparkFunSuite}
 import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.sql.execution.metric.LongSQLMetricValue
 import org.apache.spark.scheduler._
 import org.apache.spark.sql.{DataFrame, SQLContext}
-import org.apache.spark.sql.execution.SQLExecution
+import org.apache.spark.sql.execution.{SparkPlanInfo, SQLExecution}
+import org.apache.spark.sql.execution.metric.LongSQLMetricValue
 import org.apache.spark.sql.test.SharedSQLContext
 
 class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
@@ -82,7 +82,8 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
     val executionId = 0
     val df = createTestDataFrame
     val accumulatorIds =
-      SparkPlanGraph(df.queryExecution.executedPlan).nodes.flatMap(_.metrics.map(_.accumulatorId))
+      SparkPlanGraph(SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan))
+        .nodes.flatMap(_.metrics.map(_.accumulatorId))
     // Assume all accumulators are long
     var accumulatorValue = 0L
     val accumulatorUpdates = accumulatorIds.map { id =>
@@ -90,13 +91,13 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
       (id, accumulatorValue)
     }.toMap
 
-    listener.onExecutionStart(
+    listener.onOtherEvent(SparkListenerSQLExecutionStart(
       executionId,
       "test",
       "test",
       df.queryExecution.toString,
-      SparkPlanGraph(df.queryExecution.executedPlan),
-      System.currentTimeMillis())
+      SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
+      System.currentTimeMillis()))
 
     val executionUIData = listener.executionIdToData(0)
 
@@ -206,7 +207,8 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
       time = System.currentTimeMillis(),
       JobSucceeded
     ))
-    listener.onExecutionEnd(executionId, System.currentTimeMillis())
+    listener.onOtherEvent(SparkListenerSQLExecutionEnd(
+      executionId, System.currentTimeMillis()))
 
     assert(executionUIData.runningJobs.isEmpty)
     assert(executionUIData.succeededJobs === Seq(0))
@@ -219,19 +221,20 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
     val listener = new SQLListener(sqlContext.sparkContext.conf)
     val executionId = 0
     val df = createTestDataFrame
-    listener.onExecutionStart(
+    listener.onOtherEvent(SparkListenerSQLExecutionStart(
       executionId,
       "test",
       "test",
       df.queryExecution.toString,
-      SparkPlanGraph(df.queryExecution.executedPlan),
-      System.currentTimeMillis())
+      SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
+      System.currentTimeMillis()))
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
       stageInfos = Nil,
       createProperties(executionId)))
-    listener.onExecutionEnd(executionId, System.currentTimeMillis())
+    listener.onOtherEvent(SparkListenerSQLExecutionEnd(
+      executionId, System.currentTimeMillis()))
     listener.onJobEnd(SparkListenerJobEnd(
       jobId = 0,
       time = System.currentTimeMillis(),
@@ -248,13 +251,13 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
     val listener = new SQLListener(sqlContext.sparkContext.conf)
     val executionId = 0
     val df = createTestDataFrame
-    listener.onExecutionStart(
+    listener.onOtherEvent(SparkListenerSQLExecutionStart(
       executionId,
       "test",
       "test",
       df.queryExecution.toString,
-      SparkPlanGraph(df.queryExecution.executedPlan),
-      System.currentTimeMillis())
+      SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
+      System.currentTimeMillis()))
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
@@ -271,7 +274,8 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
       time = System.currentTimeMillis(),
       stageInfos = Nil,
       createProperties(executionId)))
-    listener.onExecutionEnd(executionId, System.currentTimeMillis())
+    listener.onOtherEvent(SparkListenerSQLExecutionEnd(
+      executionId, System.currentTimeMillis()))
     listener.onJobEnd(SparkListenerJobEnd(
       jobId = 1,
       time = System.currentTimeMillis(),
@@ -288,19 +292,20 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
     val listener = new SQLListener(sqlContext.sparkContext.conf)
     val executionId = 0
     val df = createTestDataFrame
-    listener.onExecutionStart(
+    listener.onOtherEvent(SparkListenerSQLExecutionStart(
       executionId,
       "test",
       "test",
       df.queryExecution.toString,
-      SparkPlanGraph(df.queryExecution.executedPlan),
-      System.currentTimeMillis())
+      SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
+      System.currentTimeMillis()))
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
       stageInfos = Seq.empty,
       createProperties(executionId)))
-    listener.onExecutionEnd(executionId, System.currentTimeMillis())
+    listener.onOtherEvent(SparkListenerSQLExecutionEnd(
+      executionId, System.currentTimeMillis()))
     listener.onJobEnd(SparkListenerJobEnd(
       jobId = 0,
       time = System.currentTimeMillis(),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
index 963d10eed62ed..e7b376548787c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
@@ -42,6 +42,7 @@ trait SharedSQLContext extends SQLTestUtils {
    * Initialize the [[TestSQLContext]].
    */
   protected override def beforeAll(): Unit = {
+    SQLContext.clearSqlListener()
     if (_ctx == null) {
       _ctx = new TestSQLContext
     }

From d1930ec01ab5a9d83f801f8ae8d4f15a38d98b76 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Wed, 25 Nov 2015 21:25:20 -0800
Subject: [PATCH 469/518] [SPARK-12003] [SQL] remove the prefix for name after
 expanded star

Right now, the expended start will include the name of expression as prefix for column, that's not better than without expending, we should not have the prefix.

Author: Davies Liu <davies@databricks.com>

Closes #9984 from davies/expand_star.
---
 .../org/apache/spark/sql/catalyst/analysis/unresolved.scala     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
index 1b2a8dc4c7f14..4f89b462a6ce3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala
@@ -204,7 +204,7 @@ case class UnresolvedStar(target: Option[Seq[String]]) extends Star with Unevalu
         case s: StructType => s.zipWithIndex.map {
           case (f, i) =>
             val extract = GetStructField(attribute.get, i)
-            Alias(extract, target.get + "." + f.name)()
+            Alias(extract, f.name)()
         }
 
         case _ => {

From 068b6438d6886ce5b4aa698383866f466d913d66 Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Wed, 25 Nov 2015 23:24:33 -0800
Subject: [PATCH 470/518] [SPARK-11980][SPARK-10621][SQL] Fix json_tuple and
 add test cases for

Added Python test cases for the function `isnan`, `isnull`, `nanvl` and `json_tuple`.

Fixed a bug in the function `json_tuple`

rxin , could you help me review my changes? Please let me know anything is missing.

Thank you! Have a good Thanksgiving day!

Author: gatorsmile <gatorsmile@gmail.com>

Closes #9977 from gatorsmile/json_tuple.
---
 python/pyspark/sql/functions.py | 44 +++++++++++++++++++++++++--------
 1 file changed, 34 insertions(+), 10 deletions(-)

diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index e3786e0fa5fb2..90625949f747a 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -286,14 +286,6 @@ def countDistinct(col, *cols):
     return Column(jc)
 
 
-@since(1.4)
-def monotonicallyIncreasingId():
-    """
-    .. note:: Deprecated in 1.6, use monotonically_increasing_id instead.
-    """
-    return monotonically_increasing_id()
-
-
 @since(1.6)
 def input_file_name():
     """Creates a string column for the file name of the current Spark task.
@@ -305,6 +297,10 @@ def input_file_name():
 @since(1.6)
 def isnan(col):
     """An expression that returns true iff the column is NaN.
+
+    >>> df = sqlContext.createDataFrame([(1.0, float('nan')), (float('nan'), 2.0)], ("a", "b"))
+    >>> df.select(isnan("a").alias("r1"), isnan(df.a).alias("r2")).collect()
+    [Row(r1=False, r2=False), Row(r1=True, r2=True)]
     """
     sc = SparkContext._active_spark_context
     return Column(sc._jvm.functions.isnan(_to_java_column(col)))
@@ -313,11 +309,23 @@ def isnan(col):
 @since(1.6)
 def isnull(col):
     """An expression that returns true iff the column is null.
+
+    >>> df = sqlContext.createDataFrame([(1, None), (None, 2)], ("a", "b"))
+    >>> df.select(isnull("a").alias("r1"), isnull(df.a).alias("r2")).collect()
+    [Row(r1=False, r2=False), Row(r1=True, r2=True)]
     """
     sc = SparkContext._active_spark_context
     return Column(sc._jvm.functions.isnull(_to_java_column(col)))
 
 
+@since(1.4)
+def monotonicallyIncreasingId():
+    """
+    .. note:: Deprecated in 1.6, use monotonically_increasing_id instead.
+    """
+    return monotonically_increasing_id()
+
+
 @since(1.6)
 def monotonically_increasing_id():
     """A column that generates monotonically increasing 64-bit integers.
@@ -344,6 +352,10 @@ def nanvl(col1, col2):
     """Returns col1 if it is not NaN, or col2 if col1 is NaN.
 
     Both inputs should be floating point columns (DoubleType or FloatType).
+
+    >>> df = sqlContext.createDataFrame([(1.0, float('nan')), (float('nan'), 2.0)], ("a", "b"))
+    >>> df.select(nanvl("a", "b").alias("r1"), nanvl(df.a, df.b).alias("r2")).collect()
+    [Row(r1=1.0, r2=1.0), Row(r1=2.0, r2=2.0)]
     """
     sc = SparkContext._active_spark_context
     return Column(sc._jvm.functions.nanvl(_to_java_column(col1), _to_java_column(col2)))
@@ -1460,6 +1472,7 @@ def explode(col):
     return Column(jc)
 
 
+@ignore_unicode_prefix
 @since(1.6)
 def get_json_object(col, path):
     """
@@ -1468,22 +1481,33 @@ def get_json_object(col, path):
 
     :param col: string column in json format
     :param path: path to the json object to extract
+
+    >>> data = [("1", '''{"f1": "value1", "f2": "value2"}'''), ("2", '''{"f1": "value12"}''')]
+    >>> df = sqlContext.createDataFrame(data, ("key", "jstring"))
+    >>> df.select(df.key, get_json_object(df.jstring, '$.f1').alias("c0"), \
+                          get_json_object(df.jstring, '$.f2').alias("c1") ).collect()
+    [Row(key=u'1', c0=u'value1', c1=u'value2'), Row(key=u'2', c0=u'value12', c1=None)]
     """
     sc = SparkContext._active_spark_context
     jc = sc._jvm.functions.get_json_object(_to_java_column(col), path)
     return Column(jc)
 
 
+@ignore_unicode_prefix
 @since(1.6)
-def json_tuple(col, fields):
+def json_tuple(col, *fields):
     """Creates a new row for a json column according to the given field names.
 
     :param col: string column in json format
     :param fields: list of fields to extract
 
+    >>> data = [("1", '''{"f1": "value1", "f2": "value2"}'''), ("2", '''{"f1": "value12"}''')]
+    >>> df = sqlContext.createDataFrame(data, ("key", "jstring"))
+    >>> df.select(df.key, json_tuple(df.jstring, 'f1', 'f2')).collect()
+    [Row(key=u'1', c0=u'value1', c1=u'value2'), Row(key=u'2', c0=u'value12', c1=None)]
     """
     sc = SparkContext._active_spark_context
-    jc = sc._jvm.functions.json_tuple(_to_java_column(col), fields)
+    jc = sc._jvm.functions.json_tuple(_to_java_column(col), _to_seq(sc, fields))
     return Column(jc)
 
 

From d3ef693325f91a1ed340c9756c81244a80398eb2 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Wed, 25 Nov 2015 23:31:21 -0800
Subject: [PATCH 471/518] [SPARK-11999][CORE] Fix the issue that
 ThreadUtils.newDaemonCachedThreadPool doesn't cache any task

In the previous codes, `newDaemonCachedThreadPool` uses `SynchronousQueue`, which is wrong. `SynchronousQueue` is an empty queue that cannot cache any task. This patch uses `LinkedBlockingQueue` to fix it along with other fixes to make sure `newDaemonCachedThreadPool` can use at most `maxThreadNumber` threads, and after that, cache tasks to `LinkedBlockingQueue`.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9978 from zsxwing/cached-threadpool.
---
 .../org/apache/spark/util/ThreadUtils.scala   | 14 ++++--
 .../apache/spark/util/ThreadUtilsSuite.scala  | 45 +++++++++++++++++++
 2 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/util/ThreadUtils.scala b/core/src/main/scala/org/apache/spark/util/ThreadUtils.scala
index 53283448c87b1..f9fbe2ff858ce 100644
--- a/core/src/main/scala/org/apache/spark/util/ThreadUtils.scala
+++ b/core/src/main/scala/org/apache/spark/util/ThreadUtils.scala
@@ -56,10 +56,18 @@ private[spark] object ThreadUtils {
    * Create a cached thread pool whose max number of threads is `maxThreadNumber`. Thread names
    * are formatted as prefix-ID, where ID is a unique, sequentially assigned integer.
    */
-  def newDaemonCachedThreadPool(prefix: String, maxThreadNumber: Int): ThreadPoolExecutor = {
+  def newDaemonCachedThreadPool(
+      prefix: String, maxThreadNumber: Int, keepAliveSeconds: Int = 60): ThreadPoolExecutor = {
     val threadFactory = namedThreadFactory(prefix)
-    new ThreadPoolExecutor(
-      0, maxThreadNumber, 60L, TimeUnit.SECONDS, new SynchronousQueue[Runnable], threadFactory)
+    val threadPool = new ThreadPoolExecutor(
+      maxThreadNumber, // corePoolSize: the max number of threads to create before queuing the tasks
+      maxThreadNumber, // maximumPoolSize: because we use LinkedBlockingDeque, this one is not used
+      keepAliveSeconds,
+      TimeUnit.SECONDS,
+      new LinkedBlockingQueue[Runnable],
+      threadFactory)
+    threadPool.allowCoreThreadTimeOut(true)
+    threadPool
   }
 
   /**
diff --git a/core/src/test/scala/org/apache/spark/util/ThreadUtilsSuite.scala b/core/src/test/scala/org/apache/spark/util/ThreadUtilsSuite.scala
index 620e4debf4e08..92ae038967528 100644
--- a/core/src/test/scala/org/apache/spark/util/ThreadUtilsSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/ThreadUtilsSuite.scala
@@ -24,6 +24,8 @@ import scala.concurrent.duration._
 import scala.concurrent.{Await, Future}
 import scala.util.Random
 
+import org.scalatest.concurrent.Eventually._
+
 import org.apache.spark.SparkFunSuite
 
 class ThreadUtilsSuite extends SparkFunSuite {
@@ -59,6 +61,49 @@ class ThreadUtilsSuite extends SparkFunSuite {
     }
   }
 
+  test("newDaemonCachedThreadPool") {
+    val maxThreadNumber = 10
+    val startThreadsLatch = new CountDownLatch(maxThreadNumber)
+    val latch = new CountDownLatch(1)
+    val cachedThreadPool = ThreadUtils.newDaemonCachedThreadPool(
+      "ThreadUtilsSuite-newDaemonCachedThreadPool",
+      maxThreadNumber,
+      keepAliveSeconds = 2)
+    try {
+      for (_ <- 1 to maxThreadNumber) {
+        cachedThreadPool.execute(new Runnable {
+          override def run(): Unit = {
+            startThreadsLatch.countDown()
+            latch.await(10, TimeUnit.SECONDS)
+          }
+        })
+      }
+      startThreadsLatch.await(10, TimeUnit.SECONDS)
+      assert(cachedThreadPool.getActiveCount === maxThreadNumber)
+      assert(cachedThreadPool.getQueue.size === 0)
+
+      // Submit a new task and it should be put into the queue since the thread number reaches the
+      // limitation
+      cachedThreadPool.execute(new Runnable {
+        override def run(): Unit = {
+          latch.await(10, TimeUnit.SECONDS)
+        }
+      })
+
+      assert(cachedThreadPool.getActiveCount === maxThreadNumber)
+      assert(cachedThreadPool.getQueue.size === 1)
+
+      latch.countDown()
+      eventually(timeout(10.seconds)) {
+        // All threads should be stopped after keepAliveSeconds
+        assert(cachedThreadPool.getActiveCount === 0)
+        assert(cachedThreadPool.getPoolSize === 0)
+      }
+    } finally {
+      cachedThreadPool.shutdownNow()
+    }
+  }
+
   test("sameThread") {
     val callerThreadName = Thread.currentThread().getName()
     val f = Future {

From 27d69a0573ed55e916a464e268dcfd5ecc6ed849 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Thu, 26 Nov 2015 00:19:42 -0800
Subject: [PATCH 472/518] [SPARK-11973] [SQL] push filter through aggregation
 with alias and literals

Currently, filter can't be pushed through aggregation with alias or literals, this patch fix that.

After this patch, the time of TPC-DS query 4 go down to 13 seconds from 141 seconds (10x improvements).

cc nongli  yhuai

Author: Davies Liu <davies@databricks.com>

Closes #9959 from davies/push_filter2.
---
 .../sql/catalyst/expressions/predicates.scala |  9 ++++
 .../sql/catalyst/optimizer/Optimizer.scala    | 28 ++++++----
 .../optimizer/FilterPushdownSuite.scala       | 53 +++++++++++++++++++
 3 files changed, 79 insertions(+), 11 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/predicates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/predicates.scala
index 68557479a9591..304b438c84ba4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/predicates.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/predicates.scala
@@ -65,6 +65,15 @@ trait PredicateHelper {
     }
   }
 
+  // Substitute any known alias from a map.
+  protected def replaceAlias(
+      condition: Expression,
+      aliases: AttributeMap[Expression]): Expression = {
+    condition.transform {
+      case a: Attribute => aliases.getOrElse(a, a)
+    }
+  }
+
   /**
    * Returns true if `expr` can be evaluated using only the output of `plan`.  This method
    * can be used to determine when it is acceptable to move expression evaluation within a query
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index f4dba67f13b54..52f609bc158ca 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -640,20 +640,14 @@ object PushPredicateThroughProject extends Rule[LogicalPlan] with PredicateHelpe
           filter
         } else {
           // Push down the small conditions without nondeterministic expressions.
-          val pushedCondition = deterministic.map(replaceAlias(_, aliasMap)).reduce(And)
+          val pushedCondition =
+            deterministic.map(replaceAlias(_, aliasMap)).reduce(And)
           Filter(nondeterministic.reduce(And),
             project.copy(child = Filter(pushedCondition, grandChild)))
         }
       }
   }
 
-  // Substitute any attributes that are produced by the child projection, so that we safely
-  // eliminate it.
-  private def replaceAlias(condition: Expression, sourceAliases: AttributeMap[Expression]) = {
-    condition.transform {
-      case a: Attribute => sourceAliases.getOrElse(a, a)
-    }
-  }
 }
 
 /**
@@ -690,12 +684,24 @@ object PushPredicateThroughAggregate extends Rule[LogicalPlan] with PredicateHel
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case filter @ Filter(condition,
         aggregate @ Aggregate(groupingExpressions, aggregateExpressions, grandChild)) =>
-      val (pushDown, stayUp) = splitConjunctivePredicates(condition).partition {
-        conjunct => conjunct.references subsetOf AttributeSet(groupingExpressions)
+
+      def hasAggregate(expression: Expression): Boolean = expression match {
+        case agg: AggregateExpression => true
+        case other => expression.children.exists(hasAggregate)
+      }
+      // Create a map of Alias for expressions that does not have AggregateExpression
+      val aliasMap = AttributeMap(aggregateExpressions.collect {
+        case a: Alias if !hasAggregate(a.child) => (a.toAttribute, a.child)
+      })
+
+      val (pushDown, stayUp) = splitConjunctivePredicates(condition).partition { conjunct =>
+        val replaced = replaceAlias(conjunct, aliasMap)
+        replaced.references.subsetOf(grandChild.outputSet) && replaced.deterministic
       }
       if (pushDown.nonEmpty) {
         val pushDownPredicate = pushDown.reduce(And)
-        val withPushdown = aggregate.copy(child = Filter(pushDownPredicate, grandChild))
+        val replaced = replaceAlias(pushDownPredicate, aliasMap)
+        val withPushdown = aggregate.copy(child = Filter(replaced, grandChild))
         stayUp.reduceOption(And).map(Filter(_, withPushdown)).getOrElse(withPushdown)
       } else {
         filter
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
index 0290fafe879f6..0128c220baaca 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
@@ -697,4 +697,57 @@ class FilterPushdownSuite extends PlanTest {
 
     comparePlans(optimized, correctAnswer)
   }
+
+  test("aggregate: push down filters with alias") {
+    val originalQuery = testRelation
+      .select('a, 'b)
+      .groupBy('a)(('a + 1) as 'aa, count('b) as 'c)
+      .where(('c === 2L || 'aa > 4) && 'aa < 3)
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+
+    val correctAnswer = testRelation
+      .select('a, 'b)
+      .where('a + 1 < 3)
+      .groupBy('a)(('a + 1) as 'aa, count('b) as 'c)
+      .where('c === 2L || 'aa > 4)
+      .analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("aggregate: push down filters with literal") {
+    val originalQuery = testRelation
+      .select('a, 'b)
+      .groupBy('a)('a, count('b) as 'c, "s" as 'd)
+      .where('c === 2L && 'd === "s")
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+
+    val correctAnswer = testRelation
+      .select('a, 'b)
+      .where("s" === "s")
+      .groupBy('a)('a, count('b) as 'c, "s" as 'd)
+      .where('c === 2L)
+      .analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("aggregate: don't push down filters which is nondeterministic") {
+    val originalQuery = testRelation
+      .select('a, 'b)
+      .groupBy('a)('a + Rand(10) as 'aa, count('b) as 'c, Rand(11).as("rnd"))
+      .where('c === 2L && 'aa + Rand(10).as("rnd") === 3 && 'rnd === 5)
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+
+    val correctAnswer = testRelation
+      .select('a, 'b)
+      .groupBy('a)('a + Rand(10) as 'aa, count('b) as 'c, Rand(11).as("rnd"))
+      .where('c === 2L && 'aa + Rand(10).as("rnd") === 3 && 'rnd === 5)
+      .analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
 }

From 001f0528a851ac314b390e65eb0583f89e69a949 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Thu, 26 Nov 2015 01:15:05 -0800
Subject: [PATCH 473/518] [SPARK-12005][SQL] Work around VerifyError in
 HyperLogLogPlusPlus.

Just move the code around a bit; that seems to make the JVM happy.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9985 from vanzin/SPARK-12005.
---
 .../expressions/aggregate/HyperLogLogPlusPlus.scala | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
index 8a95c541f1e86..e1fd22e36764e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala
@@ -63,11 +63,7 @@ case class HyperLogLogPlusPlus(
   def this(child: Expression, relativeSD: Expression) = {
     this(
       child = child,
-      relativeSD = relativeSD match {
-        case Literal(d: Double, DoubleType) => d
-        case _ =>
-          throw new AnalysisException("The second argument should be a double literal.")
-      },
+      relativeSD = HyperLogLogPlusPlus.validateDoubleLiteral(relativeSD),
       mutableAggBufferOffset = 0,
       inputAggBufferOffset = 0)
   }
@@ -448,4 +444,11 @@ object HyperLogLogPlusPlus {
     Array(189083, 185696.913, 182348.774, 179035.946, 175762.762, 172526.444, 169329.754, 166166.099, 163043.269, 159958.91, 156907.912, 153906.845, 150924.199, 147996.568, 145093.457, 142239.233, 139421.475, 136632.27, 133889.588, 131174.2, 128511.619, 125868.621, 123265.385, 120721.061, 118181.769, 115709.456, 113252.446, 110840.198, 108465.099, 106126.164, 103823.469, 101556.618, 99308.004, 97124.508, 94937.803, 92833.731, 90745.061, 88677.627, 86617.47, 84650.442, 82697.833, 80769.132, 78879.629, 77014.432, 75215.626, 73384.587, 71652.482, 69895.93, 68209.301, 66553.669, 64921.981, 63310.323, 61742.115, 60205.018, 58698.658, 57190.657, 55760.865, 54331.169, 52908.167, 51550.273, 50225.254, 48922.421, 47614.533, 46362.049, 45098.569, 43926.083, 42736.03, 41593.473, 40425.26, 39316.237, 38243.651, 37170.617, 36114.609, 35084.19, 34117.233, 33206.509, 32231.505, 31318.728, 30403.404, 29540.0550000001, 28679.236, 27825.862, 26965.216, 26179.148, 25462.08, 24645.952, 23922.523, 23198.144, 22529.128, 21762.4179999999, 21134.779, 20459.117, 19840.818, 19187.04, 18636.3689999999, 17982.831, 17439.7389999999, 16874.547, 16358.2169999999, 15835.684, 15352.914, 14823.681, 14329.313, 13816.897, 13342.874, 12880.882, 12491.648, 12021.254, 11625.392, 11293.7610000001, 10813.697, 10456.209, 10099.074, 9755.39000000001, 9393.18500000006, 9047.57900000003, 8657.98499999999, 8395.85900000005, 8033, 7736.95900000003, 7430.59699999995, 7258.47699999996, 6924.58200000005, 6691.29399999999, 6357.92500000005, 6202.05700000003, 5921.19700000004, 5628.28399999999, 5404.96799999999, 5226.71100000001, 4990.75600000005, 4799.77399999998, 4622.93099999998, 4472.478, 4171.78700000001, 3957.46299999999, 3868.95200000005, 3691.14300000004, 3474.63100000005, 3341.67200000002, 3109.14000000001, 3071.97400000005, 2796.40399999998, 2756.17799999996, 2611.46999999997, 2471.93000000005, 2382.26399999997, 2209.22400000005, 2142.28399999999, 2013.96100000001, 1911.18999999994, 1818.27099999995, 1668.47900000005, 1519.65800000005, 1469.67599999998, 1367.13800000004, 1248.52899999998, 1181.23600000003, 1022.71900000004, 1088.20700000005, 959.03600000008, 876.095999999903, 791.183999999892, 703.337000000058, 731.949999999953, 586.86400000006, 526.024999999907, 323.004999999888, 320.448000000091, 340.672999999952, 309.638999999966, 216.601999999955, 102.922999999952, 19.2399999999907, -0.114000000059605, -32.6240000000689, -89.3179999999702, -153.497999999905, -64.2970000000205, -143.695999999996, -259.497999999905, -253.017999999924, -213.948000000091, -397.590000000084, -434.006000000052, -403.475000000093, -297.958000000101, -404.317000000039, -528.898999999976, -506.621000000043, -513.205000000075, -479.351000000024, -596.139999999898, -527.016999999993, -664.681000000099, -680.306000000099, -704.050000000047, -850.486000000034, -757.43200000003, -713.308999999892)
   )
   // scalastyle:on
+
+  private def validateDoubleLiteral(exp: Expression): Double = exp match {
+    case Literal(d: Double, DoubleType) => d
+    case _ =>
+      throw new AnalysisException("The second argument should be a double literal.")
+  }
+
 }

From bc16a67562560c732833260cbc34825f7e9dcb8f Mon Sep 17 00:00:00 2001
From: Dilip Biswal <dbiswal@us.ibm.com>
Date: Thu, 26 Nov 2015 11:31:28 -0800
Subject: [PATCH 474/518] [SPARK-11863][SQL] Unable to resolve order by if it
 contains mixture of aliases and real columns

this is based on https://github.com/apache/spark/pull/9844, with some bug fix and clean up.

The problems is that, normal operator should be resolved based on its child, but `Sort` operator can also be resolved based on its grandchild. So we have 3 rules that can resolve `Sort`: `ResolveReferences`, `ResolveSortReferences`(if grandchild is `Project`) and `ResolveAggregateFunctions`(if grandchild is `Aggregate`).
For example, `select c1 as a , c2 as b from tab group by c1, c2 order by a, c2`, we need to resolve `a` and `c2` for `Sort`. Firstly `a` will be resolved in `ResolveReferences` based on its child, and when we reach `ResolveAggregateFunctions`, we will try to resolve both `a` and `c2` based on its grandchild, but failed because `a` is not a legal aggregate expression.

whoever merge this PR, please give the credit to dilipbiswal

Author: Dilip Biswal <dbiswal@us.ibm.com>
Author: Wenchen Fan <wenchen@databricks.com>

Closes #9961 from cloud-fan/sort.
---
 .../spark/sql/catalyst/analysis/Analyzer.scala | 13 ++++++++++---
 .../sql/catalyst/analysis/AnalysisSuite.scala  | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index 47962ebe6ef82..94ffbbb2e5c65 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -630,7 +630,8 @@ class Analyzer(
 
         // Try resolving the ordering as though it is in the aggregate clause.
         try {
-          val aliasedOrdering = sortOrder.map(o => Alias(o.child, "aggOrder")())
+          val unresolvedSortOrders = sortOrder.filter(s => !s.resolved || containsAggregate(s))
+          val aliasedOrdering = unresolvedSortOrders.map(o => Alias(o.child, "aggOrder")())
           val aggregatedOrdering = aggregate.copy(aggregateExpressions = aliasedOrdering)
           val resolvedAggregate: Aggregate = execute(aggregatedOrdering).asInstanceOf[Aggregate]
           val resolvedAliasedOrdering: Seq[Alias] =
@@ -663,13 +664,19 @@ class Analyzer(
               }
           }
 
+          val sortOrdersMap = unresolvedSortOrders
+            .map(new TreeNodeRef(_))
+            .zip(evaluatedOrderings)
+            .toMap
+          val finalSortOrders = sortOrder.map(s => sortOrdersMap.getOrElse(new TreeNodeRef(s), s))
+
           // Since we don't rely on sort.resolved as the stop condition for this rule,
           // we need to check this and prevent applying this rule multiple times
-          if (sortOrder == evaluatedOrderings) {
+          if (sortOrder == finalSortOrders) {
             sort
           } else {
             Project(aggregate.output,
-              Sort(evaluatedOrderings, global,
+              Sort(finalSortOrders, global,
                 aggregate.copy(aggregateExpressions = originalAggExprs ++ needsPushDown)))
           }
         } catch {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
index e051069951887..aeeca802d8bb3 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala
@@ -220,6 +220,24 @@ class AnalysisSuite extends AnalysisTest {
     // checkUDF(udf4, expected4)
   }
 
+  test("SPARK-11863 mixture of aliases and real columns in order by clause - tpcds 19,55,71") {
+    val a = testRelation2.output(0)
+    val c = testRelation2.output(2)
+    val alias1 = a.as("a1")
+    val alias2 = c.as("a2")
+    val alias3 = count(a).as("a3")
+
+    val plan = testRelation2
+      .groupBy('a, 'c)('a.as("a1"), 'c.as("a2"), count('a).as("a3"))
+      .orderBy('a1.asc, 'c.asc)
+
+    val expected = testRelation2
+      .groupBy(a, c)(alias1, alias2, alias3)
+      .orderBy(alias1.toAttribute.asc, alias2.toAttribute.asc)
+      .select(alias1.toAttribute, alias2.toAttribute, alias3.toAttribute)
+    checkAnalysis(plan, expected)
+  }
+
   test("analyzer should replace current_timestamp with literals") {
     val in = Project(Seq(Alias(CurrentTimestamp(), "a")(), Alias(CurrentTimestamp(), "b")()),
       LocalRelation())

From ad76562390b81207f8f32491c0bd8ad0e020141f Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Thu, 26 Nov 2015 16:20:08 -0800
Subject: [PATCH 475/518] [SPARK-11998][SQL][TEST-HADOOP2.0] When downloading
 Hadoop artifacts from maven, we need to try to download the version that is
 used by Spark

If we need to download Hive/Hadoop artifacts, try to download a Hadoop that matches the Hadoop used by Spark. If the Hadoop artifact cannot be resolved (e.g. Hadoop version is a vendor specific version like 2.0.0-cdh4.1.1), we will use Hadoop 2.4.0 (we used to hard code this version as the hadoop that we will download from maven) and we will not share Hadoop classes.

I tested this match in my laptop with the following confs (these confs are used by our builds). All tests are good.
```
build/sbt -Phadoop-1 -Dhadoop.version=1.2.1 -Pkinesis-asl -Phive-thriftserver -Phive
build/sbt -Phadoop-1 -Dhadoop.version=2.0.0-mr1-cdh4.1.1 -Pkinesis-asl -Phive-thriftserver -Phive
build/sbt -Pyarn -Phadoop-2.2 -Pkinesis-asl -Phive-thriftserver -Phive
build/sbt -Pyarn -Phadoop-2.3 -Dhadoop.version=2.3.0 -Pkinesis-asl -Phive-thriftserver -Phive
```

Author: Yin Huai <yhuai@databricks.com>

Closes #9979 from yhuai/versionsSuite.
---
 .../apache/spark/sql/hive/HiveContext.scala   |  4 +-
 .../hive/client/IsolatedClientLoader.scala    | 62 +++++++++++++++----
 .../spark/sql/hive/client/VersionsSuite.scala | 23 +++++--
 3 files changed, 72 insertions(+), 17 deletions(-)

diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
index 8a4264194ae8d..e83941c2ecf66 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveContext.scala
@@ -35,6 +35,7 @@ import org.apache.hadoop.hive.conf.HiveConf.ConfVars
 import org.apache.hadoop.hive.ql.metadata.Table
 import org.apache.hadoop.hive.ql.parse.VariableSubstitution
 import org.apache.hadoop.hive.serde2.io.{DateWritable, TimestampWritable}
+import org.apache.hadoop.util.VersionInfo
 
 import org.apache.spark.api.java.JavaSparkContext
 import org.apache.spark.sql.SQLConf.SQLConfEntry
@@ -288,7 +289,8 @@ class HiveContext private[hive](
       logInfo(
         s"Initializing HiveMetastoreConnection version $hiveMetastoreVersion using maven.")
       IsolatedClientLoader.forVersion(
-        version = hiveMetastoreVersion,
+        hiveMetastoreVersion = hiveMetastoreVersion,
+        hadoopVersion = VersionInfo.getVersion,
         config = allConfig,
         barrierPrefixes = hiveMetastoreBarrierPrefixes,
         sharedPrefixes = hiveMetastoreSharedPrefixes)
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/IsolatedClientLoader.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/IsolatedClientLoader.scala
index e041e0d8e5ae8..010051d255fdc 100644
--- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/IsolatedClientLoader.scala
+++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/IsolatedClientLoader.scala
@@ -34,23 +34,51 @@ import org.apache.spark.sql.hive.HiveContext
 import org.apache.spark.util.{MutableURLClassLoader, Utils}
 
 /** Factory for `IsolatedClientLoader` with specific versions of hive. */
-private[hive] object IsolatedClientLoader {
+private[hive] object IsolatedClientLoader extends Logging {
   /**
    * Creates isolated Hive client loaders by downloading the requested version from maven.
    */
   def forVersion(
-      version: String,
+      hiveMetastoreVersion: String,
+      hadoopVersion: String,
       config: Map[String, String] = Map.empty,
       ivyPath: Option[String] = None,
       sharedPrefixes: Seq[String] = Seq.empty,
       barrierPrefixes: Seq[String] = Seq.empty): IsolatedClientLoader = synchronized {
-    val resolvedVersion = hiveVersion(version)
-    val files = resolvedVersions.getOrElseUpdate(resolvedVersion,
-      downloadVersion(resolvedVersion, ivyPath))
+    val resolvedVersion = hiveVersion(hiveMetastoreVersion)
+    // We will first try to share Hadoop classes. If we cannot resolve the Hadoop artifact
+    // with the given version, we will use Hadoop 2.4.0 and then will not share Hadoop classes.
+    var sharesHadoopClasses = true
+    val files = if (resolvedVersions.contains((resolvedVersion, hadoopVersion))) {
+      resolvedVersions((resolvedVersion, hadoopVersion))
+    } else {
+      val (downloadedFiles, actualHadoopVersion) =
+        try {
+          (downloadVersion(resolvedVersion, hadoopVersion, ivyPath), hadoopVersion)
+        } catch {
+          case e: RuntimeException if e.getMessage.contains("hadoop") =>
+            // If the error message contains hadoop, it is probably because the hadoop
+            // version cannot be resolved (e.g. it is a vendor specific version like
+            // 2.0.0-cdh4.1.1). If it is the case, we will try just
+            // "org.apache.hadoop:hadoop-client:2.4.0". "org.apache.hadoop:hadoop-client:2.4.0"
+            // is used just because we used to hard code it as the hadoop artifact to download.
+            logWarning(s"Failed to resolve Hadoop artifacts for the version ${hadoopVersion}. " +
+              s"We will change the hadoop version from ${hadoopVersion} to 2.4.0 and try again. " +
+              "Hadoop classes will not be shared between Spark and Hive metastore client. " +
+              "It is recommended to set jars used by Hive metastore client through " +
+              "spark.sql.hive.metastore.jars in the production environment.")
+            sharesHadoopClasses = false
+            (downloadVersion(resolvedVersion, "2.4.0", ivyPath), "2.4.0")
+        }
+      resolvedVersions.put((resolvedVersion, actualHadoopVersion), downloadedFiles)
+      resolvedVersions((resolvedVersion, actualHadoopVersion))
+    }
+
     new IsolatedClientLoader(
-      version = hiveVersion(version),
+      version = hiveVersion(hiveMetastoreVersion),
       execJars = files,
       config = config,
+      sharesHadoopClasses = sharesHadoopClasses,
       sharedPrefixes = sharedPrefixes,
       barrierPrefixes = barrierPrefixes)
   }
@@ -64,12 +92,15 @@ private[hive] object IsolatedClientLoader {
     case "1.2" | "1.2.0" | "1.2.1" => hive.v1_2
   }
 
-  private def downloadVersion(version: HiveVersion, ivyPath: Option[String]): Seq[URL] = {
+  private def downloadVersion(
+      version: HiveVersion,
+      hadoopVersion: String,
+      ivyPath: Option[String]): Seq[URL] = {
     val hiveArtifacts = version.extraDeps ++
       Seq("hive-metastore", "hive-exec", "hive-common", "hive-serde")
         .map(a => s"org.apache.hive:$a:${version.fullVersion}") ++
       Seq("com.google.guava:guava:14.0.1",
-        "org.apache.hadoop:hadoop-client:2.4.0")
+        s"org.apache.hadoop:hadoop-client:$hadoopVersion")
 
     val classpath = quietly {
       SparkSubmitUtils.resolveMavenCoordinates(
@@ -86,7 +117,10 @@ private[hive] object IsolatedClientLoader {
     tempDir.listFiles().map(_.toURI.toURL)
   }
 
-  private def resolvedVersions = new scala.collection.mutable.HashMap[HiveVersion, Seq[URL]]
+  // A map from a given pair of HiveVersion and Hadoop version to jar files.
+  // It is only used by forVersion.
+  private val resolvedVersions =
+    new scala.collection.mutable.HashMap[(HiveVersion, String), Seq[URL]]
 }
 
 /**
@@ -106,6 +140,7 @@ private[hive] object IsolatedClientLoader {
  * @param config   A set of options that will be added to the HiveConf of the constructed client.
  * @param isolationOn When true, custom versions of barrier classes will be constructed.  Must be
  *                    true unless loading the version of hive that is on Sparks classloader.
+ * @param sharesHadoopClasses When true, we will share Hadoop classes between Spark and
  * @param rootClassLoader The system root classloader. Must not know about Hive classes.
  * @param baseClassLoader The spark classloader that is used to load shared classes.
  */
@@ -114,6 +149,7 @@ private[hive] class IsolatedClientLoader(
     val execJars: Seq[URL] = Seq.empty,
     val config: Map[String, String] = Map.empty,
     val isolationOn: Boolean = true,
+    val sharesHadoopClasses: Boolean = true,
     val rootClassLoader: ClassLoader = ClassLoader.getSystemClassLoader.getParent.getParent,
     val baseClassLoader: ClassLoader = Thread.currentThread().getContextClassLoader,
     val sharedPrefixes: Seq[String] = Seq.empty,
@@ -126,16 +162,20 @@ private[hive] class IsolatedClientLoader(
   /** All jars used by the hive specific classloader. */
   protected def allJars = execJars.toArray
 
-  protected def isSharedClass(name: String): Boolean =
+  protected def isSharedClass(name: String): Boolean = {
+    val isHadoopClass =
+      name.startsWith("org.apache.hadoop.") && !name.startsWith("org.apache.hadoop.hive.")
+
     name.contains("slf4j") ||
     name.contains("log4j") ||
     name.startsWith("org.apache.spark.") ||
-    (name.startsWith("org.apache.hadoop.") && !name.startsWith("org.apache.hadoop.hive.")) ||
+    (sharesHadoopClasses && isHadoopClass) ||
     name.startsWith("scala.") ||
     (name.startsWith("com.google") && !name.startsWith("com.google.cloud")) ||
     name.startsWith("java.lang.") ||
     name.startsWith("java.net") ||
     sharedPrefixes.exists(name.startsWith)
+  }
 
   /** True if `name` refers to a spark class that must see specific version of Hive. */
   protected def isBarrierClass(name: String): Boolean =
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala
index 7bc13bc60d30e..502b240f3650f 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/client/VersionsSuite.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.hive.client
 
 import java.io.File
 
+import org.apache.hadoop.util.VersionInfo
+
 import org.apache.spark.sql.hive.HiveContext
 import org.apache.spark.{Logging, SparkFunSuite}
 import org.apache.spark.sql.catalyst.expressions.{NamedExpression, Literal, AttributeReference, EqualTo}
@@ -53,9 +55,11 @@ class VersionsSuite extends SparkFunSuite with Logging {
   }
 
   test("success sanity check") {
-    val badClient = IsolatedClientLoader.forVersion(HiveContext.hiveExecutionVersion,
-      buildConf(),
-      ivyPath).createClient()
+    val badClient = IsolatedClientLoader.forVersion(
+      hiveMetastoreVersion = HiveContext.hiveExecutionVersion,
+      hadoopVersion = VersionInfo.getVersion,
+      config = buildConf(),
+      ivyPath = ivyPath).createClient()
     val db = new HiveDatabase("default", "")
     badClient.createDatabase(db)
   }
@@ -85,7 +89,11 @@ class VersionsSuite extends SparkFunSuite with Logging {
   ignore("failure sanity check") {
     val e = intercept[Throwable] {
       val badClient = quietly {
-        IsolatedClientLoader.forVersion("13", buildConf(), ivyPath).createClient()
+        IsolatedClientLoader.forVersion(
+          hiveMetastoreVersion = "13",
+          hadoopVersion = VersionInfo.getVersion,
+          config = buildConf(),
+          ivyPath = ivyPath).createClient()
       }
     }
     assert(getNestedMessages(e) contains "Unknown column 'A0.OWNER_NAME' in 'field list'")
@@ -99,7 +107,12 @@ class VersionsSuite extends SparkFunSuite with Logging {
     test(s"$version: create client") {
       client = null
       System.gc() // Hack to avoid SEGV on some JVM versions.
-      client = IsolatedClientLoader.forVersion(version, buildConf(), ivyPath).createClient()
+      client =
+        IsolatedClientLoader.forVersion(
+          hiveMetastoreVersion = version,
+          hadoopVersion = VersionInfo.getVersion,
+          config = buildConf(),
+          ivyPath = ivyPath).createClient()
     }
 
     test(s"$version: createDatabase") {

From de28e4d4deca385b7c40b3a6a1efcd6e2fec2f9b Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 26 Nov 2015 18:47:54 -0800
Subject: [PATCH 476/518] [SPARK-11973][SQL] Improve optimizer code
 readability.

This is a followup for https://github.com/apache/spark/pull/9959.

I added more documentation and rewrote some monadic code into simpler ifs.

Author: Reynold Xin <rxin@databricks.com>

Closes #9995 from rxin/SPARK-11973.
---
 .../sql/catalyst/optimizer/Optimizer.scala    | 50 +++++++++----------
 .../optimizer/FilterPushdownSuite.scala       |  2 +-
 2 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 52f609bc158ca..2901d8f2efddf 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -59,7 +59,7 @@ object DefaultOptimizer extends Optimizer {
       ConstantFolding,
       LikeSimplification,
       BooleanSimplification,
-      RemoveDispensable,
+      RemoveDispensableExpressions,
       SimplifyFilters,
       SimplifyCasts,
       SimplifyCaseConversionExpressions) ::
@@ -660,14 +660,14 @@ object PushPredicateThroughGenerate extends Rule[LogicalPlan] with PredicateHelp
     case filter @ Filter(condition, g: Generate) =>
       // Predicates that reference attributes produced by the `Generate` operator cannot
       // be pushed below the operator.
-      val (pushDown, stayUp) = splitConjunctivePredicates(condition).partition {
-        conjunct => conjunct.references subsetOf g.child.outputSet
+      val (pushDown, stayUp) = splitConjunctivePredicates(condition).partition { cond =>
+        cond.references subsetOf g.child.outputSet
       }
       if (pushDown.nonEmpty) {
         val pushDownPredicate = pushDown.reduce(And)
-        val withPushdown = Generate(g.generator, join = g.join, outer = g.outer,
+        val newGenerate = Generate(g.generator, join = g.join, outer = g.outer,
           g.qualifier, g.generatorOutput, Filter(pushDownPredicate, g.child))
-        stayUp.reduceOption(And).map(Filter(_, withPushdown)).getOrElse(withPushdown)
+        if (stayUp.isEmpty) newGenerate else Filter(stayUp.reduce(And), newGenerate)
       } else {
         filter
       }
@@ -675,34 +675,34 @@ object PushPredicateThroughGenerate extends Rule[LogicalPlan] with PredicateHelp
 }
 
 /**
- * Push [[Filter]] operators through [[Aggregate]] operators. Parts of the predicate that reference
- * attributes which are subset of group by attribute set of [[Aggregate]] will be pushed beneath,
- * and the rest should remain above.
+ * Push [[Filter]] operators through [[Aggregate]] operators, iff the filters reference only
+ * non-aggregate attributes (typically literals or grouping expressions).
  */
 object PushPredicateThroughAggregate extends Rule[LogicalPlan] with PredicateHelper {
 
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
-    case filter @ Filter(condition,
-        aggregate @ Aggregate(groupingExpressions, aggregateExpressions, grandChild)) =>
-
-      def hasAggregate(expression: Expression): Boolean = expression match {
-        case agg: AggregateExpression => true
-        case other => expression.children.exists(hasAggregate)
-      }
-      // Create a map of Alias for expressions that does not have AggregateExpression
-      val aliasMap = AttributeMap(aggregateExpressions.collect {
-        case a: Alias if !hasAggregate(a.child) => (a.toAttribute, a.child)
+    case filter @ Filter(condition, aggregate: Aggregate) =>
+      // Find all the aliased expressions in the aggregate list that don't include any actual
+      // AggregateExpression, and create a map from the alias to the expression
+      val aliasMap = AttributeMap(aggregate.aggregateExpressions.collect {
+        case a: Alias if a.child.find(_.isInstanceOf[AggregateExpression]).isEmpty =>
+          (a.toAttribute, a.child)
       })
 
-      val (pushDown, stayUp) = splitConjunctivePredicates(condition).partition { conjunct =>
-        val replaced = replaceAlias(conjunct, aliasMap)
-        replaced.references.subsetOf(grandChild.outputSet) && replaced.deterministic
+      // For each filter, expand the alias and check if the filter can be evaluated using
+      // attributes produced by the aggregate operator's child operator.
+      val (pushDown, stayUp) = splitConjunctivePredicates(condition).partition { cond =>
+        val replaced = replaceAlias(cond, aliasMap)
+        replaced.references.subsetOf(aggregate.child.outputSet) && replaced.deterministic
       }
+
       if (pushDown.nonEmpty) {
         val pushDownPredicate = pushDown.reduce(And)
         val replaced = replaceAlias(pushDownPredicate, aliasMap)
-        val withPushdown = aggregate.copy(child = Filter(replaced, grandChild))
-        stayUp.reduceOption(And).map(Filter(_, withPushdown)).getOrElse(withPushdown)
+        val newAggregate = aggregate.copy(child = Filter(replaced, aggregate.child))
+        // If there is no more filter to stay up, just eliminate the filter.
+        // Otherwise, create "Filter(stayUp) <- Aggregate <- Filter(pushDownPredicate)".
+        if (stayUp.isEmpty) newAggregate else Filter(stayUp.reduce(And), newAggregate)
       } else {
         filter
       }
@@ -714,7 +714,7 @@ object PushPredicateThroughAggregate extends Rule[LogicalPlan] with PredicateHel
  * evaluated using only the attributes of the left or right side of a join.  Other
  * [[Filter]] conditions are moved into the `condition` of the [[Join]].
  *
- * And also Pushes down the join filter, where the `condition` can be evaluated using only the
+ * And also pushes down the join filter, where the `condition` can be evaluated using only the
  * attributes of the left or right side of sub query when applicable.
  *
  * Check https://cwiki.apache.org/confluence/display/Hive/OuterJoinBehavior for more details
@@ -821,7 +821,7 @@ object SimplifyCasts extends Rule[LogicalPlan] {
 /**
  * Removes nodes that are not necessary.
  */
-object RemoveDispensable extends Rule[LogicalPlan] {
+object RemoveDispensableExpressions extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
     case UnaryPositive(child) => child
     case PromotePrecision(child) => child
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
index 0128c220baaca..fba4c5ca77d64 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala
@@ -734,7 +734,7 @@ class FilterPushdownSuite extends PlanTest {
     comparePlans(optimized, correctAnswer)
   }
 
-  test("aggregate: don't push down filters which is nondeterministic") {
+  test("aggregate: don't push down filters that are nondeterministic") {
     val originalQuery = testRelation
       .select('a, 'b)
       .groupBy('a)('a + Rand(10) as 'aa, count('b) as 'c, Rand(11).as("rnd"))

From 4376b5bea8171e4e73b3dbabbfdf84fa1afd140b Mon Sep 17 00:00:00 2001
From: muxator <muxator@users.noreply.github.com>
Date: Thu, 26 Nov 2015 18:52:20 -0800
Subject: [PATCH 477/518] doc typo: "classificaion" -> "classification"

Author: muxator <muxator@users.noreply.github.com>

Closes #10008 from muxator/patch-1.
---
 docs/mllib-linear-methods.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/mllib-linear-methods.md b/docs/mllib-linear-methods.md
index 0c76e6e999465..132f8c354aa9c 100644
--- a/docs/mllib-linear-methods.md
+++ b/docs/mllib-linear-methods.md
@@ -122,7 +122,7 @@ Under the hood, linear methods use convex optimization methods to optimize the o
 [Classification](http://en.wikipedia.org/wiki/Statistical_classification) aims to divide items into
 categories.
 The most common classification type is
-[binary classificaion](http://en.wikipedia.org/wiki/Binary_classification), where there are two
+[binary classification](http://en.wikipedia.org/wiki/Binary_classification), where there are two
 categories, usually named positive and negative.
 If there are more than two categories, it is called
 [multiclass classification](http://en.wikipedia.org/wiki/Multiclass_classification).

From 0c1e72e7f79231e537299b57a1ab7cd843171923 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Thu, 26 Nov 2015 18:56:22 -0800
Subject: [PATCH 478/518] [SPARK-11996][CORE] Make the executor thread dump
 work again

In the previous implementation, the driver needs to know the executor listening address to send the thread dump request. However, in Netty RPC, the executor doesn't listen to any port, so the executor thread dump feature is broken.

This patch makes the driver use the endpointRef stored in BlockManagerMasterEndpoint to send the thread dump request to fix it.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #9976 from zsxwing/executor-thread-dump.
---
 .../scala/org/apache/spark/SparkContext.scala | 10 ++---
 .../org/apache/spark/executor/Executor.scala  |  5 ---
 .../spark/executor/ExecutorEndpoint.scala     | 43 -------------------
 .../spark/storage/BlockManagerMaster.scala    |  4 +-
 .../storage/BlockManagerMasterEndpoint.scala  | 12 +++---
 .../spark/storage/BlockManagerMessages.scala  |  7 ++-
 .../storage/BlockManagerSlaveEndpoint.scala   |  7 ++-
 project/MimaExcludes.scala                    |  8 ++++
 8 files changed, 29 insertions(+), 67 deletions(-)
 delete mode 100644 core/src/main/scala/org/apache/spark/executor/ExecutorEndpoint.scala

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 2c10779f2b893..b030d3c71dc20 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -48,20 +48,20 @@ import org.apache.mesos.MesosNativeLibrary
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.deploy.{LocalSparkCluster, SparkHadoopUtil}
-import org.apache.spark.executor.{ExecutorEndpoint, TriggerThreadDump}
 import org.apache.spark.input.{StreamInputFormat, PortableDataStream, WholeTextFileInputFormat,
   FixedLengthBinaryInputFormat}
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.metrics.MetricsSystem
 import org.apache.spark.partial.{ApproximateEvaluator, PartialResult}
 import org.apache.spark.rdd._
-import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef}
+import org.apache.spark.rpc.RpcEndpointRef
 import org.apache.spark.scheduler._
 import org.apache.spark.scheduler.cluster.{CoarseGrainedSchedulerBackend,
   SparkDeploySchedulerBackend, SimrSchedulerBackend}
 import org.apache.spark.scheduler.cluster.mesos.{CoarseMesosSchedulerBackend, MesosSchedulerBackend}
 import org.apache.spark.scheduler.local.LocalBackend
 import org.apache.spark.storage._
+import org.apache.spark.storage.BlockManagerMessages.TriggerThreadDump
 import org.apache.spark.ui.{SparkUI, ConsoleProgressBar}
 import org.apache.spark.ui.jobs.JobProgressListener
 import org.apache.spark.util._
@@ -619,11 +619,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
       if (executorId == SparkContext.DRIVER_IDENTIFIER) {
         Some(Utils.getThreadDump())
       } else {
-        val (host, port) = env.blockManager.master.getRpcHostPortForExecutor(executorId).get
-        val endpointRef = env.rpcEnv.setupEndpointRef(
-          SparkEnv.executorActorSystemName,
-          RpcAddress(host, port),
-          ExecutorEndpoint.EXECUTOR_ENDPOINT_NAME)
+        val endpointRef = env.blockManager.master.getExecutorEndpointRef(executorId).get
         Some(endpointRef.askWithRetry[Array[ThreadStackTrace]](TriggerThreadDump))
       }
     } catch {
diff --git a/core/src/main/scala/org/apache/spark/executor/Executor.scala b/core/src/main/scala/org/apache/spark/executor/Executor.scala
index 9e88d488c0379..6154f06e3ac11 100644
--- a/core/src/main/scala/org/apache/spark/executor/Executor.scala
+++ b/core/src/main/scala/org/apache/spark/executor/Executor.scala
@@ -85,10 +85,6 @@ private[spark] class Executor(
     env.blockManager.initialize(conf.getAppId)
   }
 
-  // Create an RpcEndpoint for receiving RPCs from the driver
-  private val executorEndpoint = env.rpcEnv.setupEndpoint(
-    ExecutorEndpoint.EXECUTOR_ENDPOINT_NAME, new ExecutorEndpoint(env.rpcEnv, executorId))
-
   // Whether to load classes in user jars before those in Spark jars
   private val userClassPathFirst = conf.getBoolean("spark.executor.userClassPathFirst", false)
 
@@ -136,7 +132,6 @@ private[spark] class Executor(
 
   def stop(): Unit = {
     env.metricsSystem.report()
-    env.rpcEnv.stop(executorEndpoint)
     heartbeater.shutdown()
     heartbeater.awaitTermination(10, TimeUnit.SECONDS)
     threadPool.shutdown()
diff --git a/core/src/main/scala/org/apache/spark/executor/ExecutorEndpoint.scala b/core/src/main/scala/org/apache/spark/executor/ExecutorEndpoint.scala
deleted file mode 100644
index cf362f8464735..0000000000000
--- a/core/src/main/scala/org/apache/spark/executor/ExecutorEndpoint.scala
+++ /dev/null
@@ -1,43 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.executor
-
-import org.apache.spark.rpc.{RpcEnv, RpcCallContext, RpcEndpoint}
-import org.apache.spark.util.Utils
-
-/**
- * Driver -> Executor message to trigger a thread dump.
- */
-private[spark] case object TriggerThreadDump
-
-/**
- * [[RpcEndpoint]] that runs inside of executors to enable driver -> executor RPC.
- */
-private[spark]
-class ExecutorEndpoint(override val rpcEnv: RpcEnv, executorId: String) extends RpcEndpoint {
-
-  override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
-    case TriggerThreadDump =>
-      context.reply(Utils.getThreadDump())
-  }
-
-}
-
-object ExecutorEndpoint {
-  val EXECUTOR_ENDPOINT_NAME = "ExecutorEndpoint"
-}
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala
index f45bff34d4dbc..440c4c18aadd0 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala
@@ -87,8 +87,8 @@ class BlockManagerMaster(
     driverEndpoint.askWithRetry[Seq[BlockManagerId]](GetPeers(blockManagerId))
   }
 
-  def getRpcHostPortForExecutor(executorId: String): Option[(String, Int)] = {
-    driverEndpoint.askWithRetry[Option[(String, Int)]](GetRpcHostPortForExecutor(executorId))
+  def getExecutorEndpointRef(executorId: String): Option[RpcEndpointRef] = {
+    driverEndpoint.askWithRetry[Option[RpcEndpointRef]](GetExecutorEndpointRef(executorId))
   }
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala
index 7db6035553ae6..41892b4ffce5b 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala
@@ -19,7 +19,6 @@ package org.apache.spark.storage
 
 import java.util.{HashMap => JHashMap}
 
-import scala.collection.immutable.HashSet
 import scala.collection.mutable
 import scala.collection.JavaConverters._
 import scala.concurrent.{ExecutionContext, Future}
@@ -75,8 +74,8 @@ class BlockManagerMasterEndpoint(
     case GetPeers(blockManagerId) =>
       context.reply(getPeers(blockManagerId))
 
-    case GetRpcHostPortForExecutor(executorId) =>
-      context.reply(getRpcHostPortForExecutor(executorId))
+    case GetExecutorEndpointRef(executorId) =>
+      context.reply(getExecutorEndpointRef(executorId))
 
     case GetMemoryStatus =>
       context.reply(memoryStatus)
@@ -388,15 +387,14 @@ class BlockManagerMasterEndpoint(
   }
 
   /**
-   * Returns the hostname and port of an executor, based on the [[RpcEnv]] address of its
-   * [[BlockManagerSlaveEndpoint]].
+   * Returns an [[RpcEndpointRef]] of the [[BlockManagerSlaveEndpoint]] for sending RPC messages.
    */
-  private def getRpcHostPortForExecutor(executorId: String): Option[(String, Int)] = {
+  private def getExecutorEndpointRef(executorId: String): Option[RpcEndpointRef] = {
     for (
       blockManagerId <- blockManagerIdByExecutor.get(executorId);
       info <- blockManagerInfo.get(blockManagerId)
     ) yield {
-      (info.slaveEndpoint.address.host, info.slaveEndpoint.address.port)
+      info.slaveEndpoint
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala
index 376e9eb48843d..f392a4a0cd9be 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala
@@ -42,6 +42,11 @@ private[spark] object BlockManagerMessages {
   case class RemoveBroadcast(broadcastId: Long, removeFromDriver: Boolean = true)
     extends ToBlockManagerSlave
 
+  /**
+   * Driver -> Executor message to trigger a thread dump.
+   */
+  case object TriggerThreadDump extends ToBlockManagerSlave
+
   //////////////////////////////////////////////////////////////////////////////////
   // Messages from slaves to the master.
   //////////////////////////////////////////////////////////////////////////////////
@@ -90,7 +95,7 @@ private[spark] object BlockManagerMessages {
 
   case class GetPeers(blockManagerId: BlockManagerId) extends ToBlockManagerMaster
 
-  case class GetRpcHostPortForExecutor(executorId: String) extends ToBlockManagerMaster
+  case class GetExecutorEndpointRef(executorId: String) extends ToBlockManagerMaster
 
   case class RemoveExecutor(execId: String) extends ToBlockManagerMaster
 
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerSlaveEndpoint.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerSlaveEndpoint.scala
index e749631bf6f19..9eca902f7454e 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockManagerSlaveEndpoint.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerSlaveEndpoint.scala
@@ -19,10 +19,10 @@ package org.apache.spark.storage
 
 import scala.concurrent.{ExecutionContext, Future}
 
-import org.apache.spark.rpc.{ThreadSafeRpcEndpoint, RpcEnv, RpcCallContext, RpcEndpoint}
-import org.apache.spark.util.ThreadUtils
 import org.apache.spark.{Logging, MapOutputTracker, SparkEnv}
+import org.apache.spark.rpc.{RpcCallContext, RpcEnv, ThreadSafeRpcEndpoint}
 import org.apache.spark.storage.BlockManagerMessages._
+import org.apache.spark.util.{ThreadUtils, Utils}
 
 /**
  * An RpcEndpoint to take commands from the master to execute options. For example,
@@ -70,6 +70,9 @@ class BlockManagerSlaveEndpoint(
 
     case GetMatchingBlockIds(filter, _) =>
       context.reply(blockManager.getMatchingBlockIds(filter))
+
+    case TriggerThreadDump =>
+      context.reply(Utils.getThreadDump())
   }
 
   private def doAsync[T](actionMessage: String, context: RpcCallContext)(body: => T) {
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 54a9ad956d119..566bfe8efb7a4 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -147,6 +147,14 @@ object MimaExcludes {
         // SPARK-4557 Changed foreachRDD to use VoidFunction
         ProblemFilters.exclude[MissingMethodProblem](
           "org.apache.spark.streaming.api.java.JavaDStreamLike.foreachRDD")
+      ) ++ Seq(
+        // SPARK-11996 Make the executor thread dump work again
+        ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.executor.ExecutorEndpoint"),
+        ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.executor.ExecutorEndpoint$"),
+        ProblemFilters.exclude[MissingClassProblem](
+          "org.apache.spark.storage.BlockManagerMessages$GetRpcHostPortForExecutor"),
+        ProblemFilters.exclude[MissingClassProblem](
+          "org.apache.spark.storage.BlockManagerMessages$GetRpcHostPortForExecutor$")
       )
     case v if v.startsWith("1.5") =>
       Seq(

From 6f6bb0e893c8370cbab4d63a56d74e00cb7f3cf6 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 26 Nov 2015 19:00:36 -0800
Subject: [PATCH 479/518] [SPARK-12011][SQL] Stddev/Variance etc should support
 columnName as arguments

Spark SQL aggregate function:
```Java
stddev
stddev_pop
stddev_samp
variance
var_pop
var_samp
skewness
kurtosis
collect_list
collect_set
```
should support ```columnName``` as arguments like other aggregate function(max/min/count/sum).

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9994 from yanboliang/SPARK-12011.
---
 .../org/apache/spark/sql/functions.scala      | 86 +++++++++++++++++++
 .../spark/sql/DataFrameAggregateSuite.scala   |  3 +
 2 files changed, 89 insertions(+)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 276c5dfc8b062..e79defbbbdeea 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -214,6 +214,16 @@ object functions extends LegacyFunctions {
    */
   def collect_list(e: Column): Column = callUDF("collect_list", e)
 
+  /**
+   * Aggregate function: returns a list of objects with duplicates.
+   *
+   * For now this is an alias for the collect_list Hive UDAF.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def collect_list(columnName: String): Column = collect_list(Column(columnName))
+
   /**
    * Aggregate function: returns a set of objects with duplicate elements eliminated.
    *
@@ -224,6 +234,16 @@ object functions extends LegacyFunctions {
    */
   def collect_set(e: Column): Column = callUDF("collect_set", e)
 
+  /**
+   * Aggregate function: returns a set of objects with duplicate elements eliminated.
+   *
+   * For now this is an alias for the collect_set Hive UDAF.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def collect_set(columnName: String): Column = collect_set(Column(columnName))
+
   /**
    * Aggregate function: returns the Pearson Correlation Coefficient for two columns.
    *
@@ -312,6 +332,14 @@ object functions extends LegacyFunctions {
    */
   def kurtosis(e: Column): Column = withAggregateFunction { Kurtosis(e.expr) }
 
+  /**
+   * Aggregate function: returns the kurtosis of the values in a group.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def kurtosis(columnName: String): Column = kurtosis(Column(columnName))
+
   /**
    * Aggregate function: returns the last value in a group.
    *
@@ -386,6 +414,14 @@ object functions extends LegacyFunctions {
    */
   def skewness(e: Column): Column = withAggregateFunction { Skewness(e.expr) }
 
+  /**
+   * Aggregate function: returns the skewness of the values in a group.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def skewness(columnName: String): Column = skewness(Column(columnName))
+
   /**
    * Aggregate function: alias for [[stddev_samp]].
    *
@@ -394,6 +430,14 @@ object functions extends LegacyFunctions {
    */
   def stddev(e: Column): Column = withAggregateFunction { StddevSamp(e.expr) }
 
+  /**
+   * Aggregate function: alias for [[stddev_samp]].
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def stddev(columnName: String): Column = stddev(Column(columnName))
+
   /**
    * Aggregate function: returns the sample standard deviation of
    * the expression in a group.
@@ -403,6 +447,15 @@ object functions extends LegacyFunctions {
    */
   def stddev_samp(e: Column): Column = withAggregateFunction { StddevSamp(e.expr) }
 
+  /**
+   * Aggregate function: returns the sample standard deviation of
+   * the expression in a group.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def stddev_samp(columnName: String): Column = stddev_samp(Column(columnName))
+
   /**
    * Aggregate function: returns the population standard deviation of
    * the expression in a group.
@@ -412,6 +465,15 @@ object functions extends LegacyFunctions {
    */
   def stddev_pop(e: Column): Column = withAggregateFunction { StddevPop(e.expr) }
 
+  /**
+   * Aggregate function: returns the population standard deviation of
+   * the expression in a group.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def stddev_pop(columnName: String): Column = stddev_pop(Column(columnName))
+
   /**
    * Aggregate function: returns the sum of all values in the expression.
    *
@@ -452,6 +514,14 @@ object functions extends LegacyFunctions {
    */
   def variance(e: Column): Column = withAggregateFunction { VarianceSamp(e.expr) }
 
+  /**
+   * Aggregate function: alias for [[var_samp]].
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def variance(columnName: String): Column = variance(Column(columnName))
+
   /**
    * Aggregate function: returns the unbiased variance of the values in a group.
    *
@@ -460,6 +530,14 @@ object functions extends LegacyFunctions {
    */
   def var_samp(e: Column): Column = withAggregateFunction { VarianceSamp(e.expr) }
 
+  /**
+   * Aggregate function: returns the unbiased variance of the values in a group.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def var_samp(columnName: String): Column = var_samp(Column(columnName))
+
   /**
    * Aggregate function: returns the population variance of the values in a group.
    *
@@ -468,6 +546,14 @@ object functions extends LegacyFunctions {
    */
   def var_pop(e: Column): Column = withAggregateFunction { VariancePop(e.expr) }
 
+  /**
+   * Aggregate function: returns the population variance of the values in a group.
+   *
+   * @group agg_funcs
+   * @since 1.6.0
+   */
+  def var_pop(columnName: String): Column = var_pop(Column(columnName))
+
   //////////////////////////////////////////////////////////////////////////////////////////////
   // Window functions
   //////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index 9c42f65bb6f52..b5c636d0de1d1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -261,6 +261,9 @@ class DataFrameAggregateSuite extends QueryTest with SharedSQLContext {
     checkAnswer(
       testData2.agg(stddev('a), stddev_pop('a), stddev_samp('a)),
       Row(testData2ADev, math.sqrt(4 / 6.0), testData2ADev))
+    checkAnswer(
+      testData2.agg(stddev("a"), stddev_pop("a"), stddev_samp("a")),
+      Row(testData2ADev, math.sqrt(4 / 6.0), testData2ADev))
   }
 
   test("zero stddev") {

From b63938a8b04a30feb6b2255c4d4e530a74855afc Mon Sep 17 00:00:00 2001
From: mariusvniekerk <marius.v.niekerk@gmail.com>
Date: Thu, 26 Nov 2015 19:13:16 -0800
Subject: [PATCH 480/518] [SPARK-11881][SQL] Fix for postgresql fetchsize > 0

Reference: https://jdbc.postgresql.org/documentation/head/query.html#query-with-cursor
In order for PostgreSQL to honor the fetchSize non-zero setting, its Connection.autoCommit needs to be set to false. Otherwise, it will just quietly ignore the fetchSize setting.

This adds a new side-effecting dialect specific beforeFetch method that will fire before a select query is ran.

Author: mariusvniekerk <marius.v.niekerk@gmail.com>

Closes #9861 from mariusvniekerk/SPARK-11881.
---
 .../execution/datasources/jdbc/JDBCRDD.scala    | 12 ++++++++++++
 .../apache/spark/sql/jdbc/JdbcDialects.scala    | 11 +++++++++++
 .../apache/spark/sql/jdbc/PostgresDialect.scala | 17 ++++++++++++++++-
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
index 89c850ce238d7..f9b72597dd2a9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
@@ -224,6 +224,7 @@ private[sql] object JDBCRDD extends Logging {
       quotedColumns,
       filters,
       parts,
+      url,
       properties)
   }
 }
@@ -241,6 +242,7 @@ private[sql] class JDBCRDD(
     columns: Array[String],
     filters: Array[Filter],
     partitions: Array[Partition],
+    url: String,
     properties: Properties)
   extends RDD[InternalRow](sc, Nil) {
 
@@ -361,6 +363,9 @@ private[sql] class JDBCRDD(
     context.addTaskCompletionListener{ context => close() }
     val part = thePart.asInstanceOf[JDBCPartition]
     val conn = getConnection()
+    val dialect = JdbcDialects.get(url)
+    import scala.collection.JavaConverters._
+    dialect.beforeFetch(conn, properties.asScala.toMap)
 
     // H2's JDBC driver does not support the setSchema() method.  We pass a
     // fully-qualified table name in the SELECT statement.  I don't know how to
@@ -489,6 +494,13 @@ private[sql] class JDBCRDD(
       }
       try {
         if (null != conn) {
+          if (!conn.getAutoCommit && !conn.isClosed) {
+            try {
+              conn.commit()
+            } catch {
+              case e: Throwable => logWarning("Exception committing transaction", e)
+            }
+          }
           conn.close()
         }
         logInfo("closed connection")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
index b3b2cb6178c52..13db141f27db6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.jdbc
 
+import java.sql.Connection
+
 import org.apache.spark.sql.types._
 import org.apache.spark.annotation.DeveloperApi
 
@@ -97,6 +99,15 @@ abstract class JdbcDialect extends Serializable {
     s"SELECT * FROM $table WHERE 1=0"
   }
 
+  /**
+    * Override connection specific properties to run before a select is made.  This is in place to
+    * allow dialects that need special treatment to optimize behavior.
+    * @param connection The connection object
+    * @param properties The connection properties.  This is passed through from the relation.
+    */
+  def beforeFetch(connection: Connection, properties: Map[String, String]): Unit = {
+  }
+
 }
 
 /**
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
index ed3faa1268635..3cf80f576e92c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.sql.jdbc
 
-import java.sql.Types
+import java.sql.{Connection, Types}
 
 import org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils
 import org.apache.spark.sql.types._
@@ -70,4 +70,19 @@ private object PostgresDialect extends JdbcDialect {
   override def getTableExistsQuery(table: String): String = {
     s"SELECT 1 FROM $table LIMIT 1"
   }
+
+  override def beforeFetch(connection: Connection, properties: Map[String, String]): Unit = {
+    super.beforeFetch(connection, properties)
+
+    // According to the postgres jdbc documentation we need to be in autocommit=false if we actually
+    // want to have fetchsize be non 0 (all the rows).  This allows us to not have to cache all the
+    // rows inside the driver when fetching.
+    //
+    // See: https://jdbc.postgresql.org/documentation/head/query.html#query-with-cursor
+    //
+    if (properties.getOrElse("fetchsize", "0").toInt > 0) {
+      connection.setAutoCommit(false)
+    }
+
+  }
 }

From d8220885c492141dfc61e8ffb92934f2339fe8d3 Mon Sep 17 00:00:00 2001
From: Jeff Zhang <zjffdu@apache.org>
Date: Thu, 26 Nov 2015 19:15:22 -0800
Subject: [PATCH 481/518] [SPARK-11917][PYSPARK] Add SQLContext#dropTempTable
 to PySpark

Author: Jeff Zhang <zjffdu@apache.org>

Closes #9903 from zjffdu/SPARK-11917.
---
 python/pyspark/sql/context.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/python/pyspark/sql/context.py b/python/pyspark/sql/context.py
index a49c1b58d0180..b05aa2f5c4cd7 100644
--- a/python/pyspark/sql/context.py
+++ b/python/pyspark/sql/context.py
@@ -445,6 +445,15 @@ def registerDataFrameAsTable(self, df, tableName):
         else:
             raise ValueError("Can only register DataFrame as table")
 
+    @since(1.6)
+    def dropTempTable(self, tableName):
+        """ Remove the temp table from catalog.
+
+        >>> sqlContext.registerDataFrameAsTable(df, "table1")
+        >>> sqlContext.dropTempTable("table1")
+        """
+        self._ssql_ctx.dropTempTable(tableName)
+
     def parquetFile(self, *paths):
         """Loads a Parquet file, returning the result as a :class:`DataFrame`.
 

From 4d4cbc034bef559f47f8b74cecd8196dc8a85348 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@oc0558782468.ibm.com>
Date: Thu, 26 Nov 2015 19:17:46 -0800
Subject: [PATCH 482/518] [SPARK-11778][SQL] add regression test

Fix regression test for SPARK-11778.
 marmbrus
Could you please take a look?
Thank you very much!!

Author: Huaxin Gao <huaxing@oc0558782468.ibm.com>

Closes #9890 from huaxingao/spark-11778-regression-test.
---
 .../hive/HiveDataFrameAnalyticsSuite.scala    | 10 ------
 .../spark/sql/hive/HiveDataFrameSuite.scala   | 32 +++++++++++++++++++
 2 files changed, 32 insertions(+), 10 deletions(-)
 create mode 100644 sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameSuite.scala

diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
index f19a74d4b3724..9864acf765265 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameAnalyticsSuite.scala
@@ -34,14 +34,10 @@ class HiveDataFrameAnalyticsSuite extends QueryTest with TestHiveSingleton with
   override def beforeAll() {
     testData = Seq((1, 2), (2, 2), (3, 4)).toDF("a", "b")
     hiveContext.registerDataFrameAsTable(testData, "mytable")
-    hiveContext.sql("create schema usrdb")
-    hiveContext.sql("create table usrdb.test(c1 int)")
   }
 
   override def afterAll(): Unit = {
     hiveContext.dropTempTable("mytable")
-    hiveContext.sql("drop table usrdb.test")
-    hiveContext.sql("drop schema usrdb")
   }
 
   test("rollup") {
@@ -78,10 +74,4 @@ class HiveDataFrameAnalyticsSuite extends QueryTest with TestHiveSingleton with
       sql("select a, b, sum(b) from mytable group by a, b with cube").collect()
     )
   }
-
-  // There was a bug in DataFrameFrameReader.table and it has problem for table with schema name,
-  // Before fix, it throw Exceptionorg.apache.spark.sql.catalyst.analysis.NoSuchTableException
-  test("table name with schema") {
-    hiveContext.read.table("usrdb.test")
-  }
 }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameSuite.scala
new file mode 100644
index 0000000000000..7fdc5d71937ff
--- /dev/null
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveDataFrameSuite.scala
@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.hive
+
+import org.apache.spark.sql.hive.test.TestHiveSingleton
+import org.apache.spark.sql.QueryTest
+
+class HiveDataFrameSuite extends QueryTest with TestHiveSingleton {
+  test("table name with schema") {
+    // regression test for SPARK-11778
+    hiveContext.sql("create schema usrdb")
+    hiveContext.sql("create table usrdb.test(c int)")
+    hiveContext.read.table("usrdb.test")
+    hiveContext.sql("drop table usrdb.test")
+    hiveContext.sql("drop schema usrdb")
+  }
+}

From 5eaed4e45c6c57e995ac7438016fad545716e596 Mon Sep 17 00:00:00 2001
From: Jeremy Derr <jcderr@radius.com>
Date: Thu, 26 Nov 2015 19:25:13 -0800
Subject: [PATCH 483/518] [SPARK-11991] fixes

If `--private-ips` is required but not provided, spark_ec2.py may behave inappropriately, including attempting to ssh to localhost in attempts to verify ssh connectivity to the cluster.

This fixes that behavior by raising a `UsageError` exception if `get_dns_name` is unable to determine a hostname as a result.

Author: Jeremy Derr <jcderr@radius.com>

Closes #9975 from jcderr/SPARK-11991/ec_spark.py_hostname_check.
---
 ec2/spark_ec2.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/ec2/spark_ec2.py b/ec2/spark_ec2.py
index 9fd652a3df4c4..84a950c9f6529 100755
--- a/ec2/spark_ec2.py
+++ b/ec2/spark_ec2.py
@@ -1242,6 +1242,10 @@ def get_ip_address(instance, private_ips=False):
 def get_dns_name(instance, private_ips=False):
     dns = instance.public_dns_name if not private_ips else \
         instance.private_ip_address
+    if not dns:
+        raise UsageError("Failed to determine hostname of {0}.\n"
+                         "Please check that you provided --private-ips if "
+                         "necessary".format(instance))
     return dns
 
 

From 10e315c28c933b967674ae51e1b2f24160c2e8a5 Mon Sep 17 00:00:00 2001
From: Reynold Xin <rxin@databricks.com>
Date: Thu, 26 Nov 2015 19:36:43 -0800
Subject: [PATCH 484/518] Fix style violation for b63938a8b04

---
 .../apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
index f9b72597dd2a9..57a8a044a37cd 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala
@@ -20,6 +20,8 @@ package org.apache.spark.sql.execution.datasources.jdbc
 import java.sql.{Connection, DriverManager, ResultSet, ResultSetMetaData, SQLException}
 import java.util.Properties
 
+import scala.util.control.NonFatal
+
 import org.apache.commons.lang3.StringUtils
 
 import org.apache.spark.rdd.RDD
@@ -498,7 +500,7 @@ private[sql] class JDBCRDD(
             try {
               conn.commit()
             } catch {
-              case e: Throwable => logWarning("Exception committing transaction", e)
+              case NonFatal(e) => logWarning("Exception committing transaction", e)
             }
           }
           conn.close()

From a374e20b5492c775f20d32e8fbddadbd8098a655 Mon Sep 17 00:00:00 2001
From: Dilip Biswal <dbiswal@us.ibm.com>
Date: Thu, 26 Nov 2015 21:04:40 -0800
Subject: [PATCH 485/518] [SPARK-11997] [SQL] NPE when save a DataFrame as
 parquet and partitioned by long column

Check for partition column null-ability while building the partition spec.

Author: Dilip Biswal <dbiswal@us.ibm.com>

Closes #10001 from dilipbiswal/spark-11997.
---
 .../org/apache/spark/sql/sources/interfaces.scala   |  2 +-
 .../datasources/parquet/ParquetQuerySuite.scala     | 13 +++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
index f9465157c936d..9ace25dc7d21b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala
@@ -607,7 +607,7 @@ abstract class HadoopFsRelation private[sql](
         def castPartitionValuesToUserSchema(row: InternalRow) = {
           InternalRow((0 until row.numFields).map { i =>
             Cast(
-              Literal.create(row.getString(i), StringType),
+              Literal.create(row.getUTF8String(i), StringType),
               userProvidedSchema.fields(i).dataType).eval()
           }: _*)
         }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
index 70fae32b7e7a1..f777e973052d3 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala
@@ -252,6 +252,19 @@ class ParquetQuerySuite extends QueryTest with ParquetTest with SharedSQLContext
     }
   }
 
+  test("SPARK-11997 parquet with null partition values") {
+    withTempPath { dir =>
+      val path = dir.getCanonicalPath
+      sqlContext.range(1, 3)
+        .selectExpr("if(id % 2 = 0, null, id) AS n", "id")
+        .write.partitionBy("n").parquet(path)
+
+      checkAnswer(
+        sqlContext.read.parquet(path).filter("n is null"),
+        Row(2, null))
+    }
+  }
+
   // This test case is ignored because of parquet-mr bug PARQUET-370
   ignore("SPARK-10301 requested schema clipping - schemas with disjoint sets of fields") {
     withTempPath { dir =>

From ba02f6cb5a40511cefa511d410be93c035d43f23 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Fri, 27 Nov 2015 11:48:01 -0800
Subject: [PATCH 486/518] [SPARK-12025][SPARKR] Rename some window rank
 function names for SparkR

Change ```cumeDist -> cume_dist, denseRank -> dense_rank, percentRank -> percent_rank, rowNumber -> row_number``` at SparkR side.
There are two reasons that we should make this change:
* We should follow the [naming convention rule of R](http://www.inside-r.org/node/230645)
* Spark DataFrame has deprecated the old convention (such as ```cumeDist```) and will remove it in Spark 2.0.

It's better to fix this issue before 1.6 release, otherwise we will make breaking API change.
cc shivaram sun-rui

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #10016 from yanboliang/SPARK-12025.
---
 R/pkg/NAMESPACE                  |  8 ++---
 R/pkg/R/functions.R              | 54 ++++++++++++++++----------------
 R/pkg/R/generics.R               | 16 +++++-----
 R/pkg/inst/tests/test_sparkSQL.R |  4 +--
 4 files changed, 41 insertions(+), 41 deletions(-)

diff --git a/R/pkg/NAMESPACE b/R/pkg/NAMESPACE
index 260c9edce62e0..5d04dd6acaab8 100644
--- a/R/pkg/NAMESPACE
+++ b/R/pkg/NAMESPACE
@@ -123,14 +123,14 @@ exportMethods("%in%",
               "count",
               "countDistinct",
               "crc32",
-              "cumeDist",
+              "cume_dist",
               "date_add",
               "date_format",
               "date_sub",
               "datediff",
               "dayofmonth",
               "dayofyear",
-              "denseRank",
+              "dense_rank",
               "desc",
               "endsWith",
               "exp",
@@ -188,7 +188,7 @@ exportMethods("%in%",
               "next_day",
               "ntile",
               "otherwise",
-              "percentRank",
+              "percent_rank",
               "pmod",
               "quarter",
               "rand",
@@ -200,7 +200,7 @@ exportMethods("%in%",
               "rint",
               "rlike",
               "round",
-              "rowNumber",
+              "row_number",
               "rpad",
               "rtrim",
               "second",
diff --git a/R/pkg/R/functions.R b/R/pkg/R/functions.R
index 25a1f22101494..e98e7a0117ca0 100644
--- a/R/pkg/R/functions.R
+++ b/R/pkg/R/functions.R
@@ -2146,47 +2146,47 @@ setMethod("ifelse",
 
 ###################### Window functions######################
 
-#' cumeDist
+#' cume_dist
 #'
 #' Window function: returns the cumulative distribution of values within a window partition,
 #' i.e. the fraction of rows that are below the current row.
 #'
 #'   N = total number of rows in the partition
-#'   cumeDist(x) = number of values before (and including) x / N
+#'   cume_dist(x) = number of values before (and including) x / N
 #'
 #' This is equivalent to the CUME_DIST function in SQL.
 #'
-#' @rdname cumeDist
-#' @name cumeDist
+#' @rdname cume_dist
+#' @name cume_dist
 #' @family window_funcs
 #' @export
-#' @examples \dontrun{cumeDist()}
-setMethod("cumeDist",
+#' @examples \dontrun{cume_dist()}
+setMethod("cume_dist",
           signature(x = "missing"),
           function() {
-            jc <- callJStatic("org.apache.spark.sql.functions", "cumeDist")
+            jc <- callJStatic("org.apache.spark.sql.functions", "cume_dist")
             column(jc)
           })
 
-#' denseRank
+#' dense_rank
 #'
 #' Window function: returns the rank of rows within a window partition, without any gaps.
-#' The difference between rank and denseRank is that denseRank leaves no gaps in ranking
-#' sequence when there are ties. That is, if you were ranking a competition using denseRank
+#' The difference between rank and dense_rank is that dense_rank leaves no gaps in ranking
+#' sequence when there are ties. That is, if you were ranking a competition using dense_rank
 #' and had three people tie for second place, you would say that all three were in second
 #' place and that the next person came in third.
 #'
 #' This is equivalent to the DENSE_RANK function in SQL.
 #'
-#' @rdname denseRank
-#' @name denseRank
+#' @rdname dense_rank
+#' @name dense_rank
 #' @family window_funcs
 #' @export
-#' @examples \dontrun{denseRank()}
-setMethod("denseRank",
+#' @examples \dontrun{dense_rank()}
+setMethod("dense_rank",
           signature(x = "missing"),
           function() {
-            jc <- callJStatic("org.apache.spark.sql.functions", "denseRank")
+            jc <- callJStatic("org.apache.spark.sql.functions", "dense_rank")
             column(jc)
           })
 
@@ -2264,7 +2264,7 @@ setMethod("ntile",
             column(jc)
           })
 
-#' percentRank
+#' percent_rank
 #'
 #' Window function: returns the relative rank (i.e. percentile) of rows within a window partition.
 #'
@@ -2274,15 +2274,15 @@ setMethod("ntile",
 #'
 #' This is equivalent to the PERCENT_RANK function in SQL.
 #'
-#' @rdname percentRank
-#' @name percentRank
+#' @rdname percent_rank
+#' @name percent_rank
 #' @family window_funcs
 #' @export
-#' @examples \dontrun{percentRank()}
-setMethod("percentRank",
+#' @examples \dontrun{percent_rank()}
+setMethod("percent_rank",
           signature(x = "missing"),
           function() {
-            jc <- callJStatic("org.apache.spark.sql.functions", "percentRank")
+            jc <- callJStatic("org.apache.spark.sql.functions", "percent_rank")
             column(jc)
           })
 
@@ -2316,21 +2316,21 @@ setMethod("rank",
             base::rank(x, ...)
           })
 
-#' rowNumber
+#' row_number
 #'
 #' Window function: returns a sequential number starting at 1 within a window partition.
 #'
 #' This is equivalent to the ROW_NUMBER function in SQL.
 #'
-#' @rdname rowNumber
-#' @name rowNumber
+#' @rdname row_number
+#' @name row_number
 #' @family window_funcs
 #' @export
-#' @examples \dontrun{rowNumber()}
-setMethod("rowNumber",
+#' @examples \dontrun{row_number()}
+setMethod("row_number",
           signature(x = "missing"),
           function() {
-            jc <- callJStatic("org.apache.spark.sql.functions", "rowNumber")
+            jc <- callJStatic("org.apache.spark.sql.functions", "row_number")
             column(jc)
           })
 
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 1b3f10ea04643..0c305441e043e 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -700,9 +700,9 @@ setGeneric("countDistinct", function(x, ...) { standardGeneric("countDistinct")
 #' @export
 setGeneric("crc32", function(x) { standardGeneric("crc32") })
 
-#' @rdname cumeDist
+#' @rdname cume_dist
 #' @export
-setGeneric("cumeDist", function(x) { standardGeneric("cumeDist") })
+setGeneric("cume_dist", function(x) { standardGeneric("cume_dist") })
 
 #' @rdname datediff
 #' @export
@@ -728,9 +728,9 @@ setGeneric("dayofmonth", function(x) { standardGeneric("dayofmonth") })
 #' @export
 setGeneric("dayofyear", function(x) { standardGeneric("dayofyear") })
 
-#' @rdname denseRank
+#' @rdname dense_rank
 #' @export
-setGeneric("denseRank", function(x) { standardGeneric("denseRank") })
+setGeneric("dense_rank", function(x) { standardGeneric("dense_rank") })
 
 #' @rdname explode
 #' @export
@@ -872,9 +872,9 @@ setGeneric("ntile", function(x) { standardGeneric("ntile") })
 #' @export
 setGeneric("n_distinct", function(x, ...) { standardGeneric("n_distinct") })
 
-#' @rdname percentRank
+#' @rdname percent_rank
 #' @export
-setGeneric("percentRank", function(x) { standardGeneric("percentRank") })
+setGeneric("percent_rank", function(x) { standardGeneric("percent_rank") })
 
 #' @rdname pmod
 #' @export
@@ -913,9 +913,9 @@ setGeneric("reverse", function(x) { standardGeneric("reverse") })
 #' @export
 setGeneric("rint", function(x, ...) { standardGeneric("rint") })
 
-#' @rdname rowNumber
+#' @rdname row_number
 #' @export
-setGeneric("rowNumber", function(x) { standardGeneric("rowNumber") })
+setGeneric("row_number", function(x) { standardGeneric("row_number") })
 
 #' @rdname rpad
 #' @export
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 3f4f319fe745d..0fbe0658265b2 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -861,8 +861,8 @@ test_that("column functions", {
   c11 <- to_date(c) + trim(c) + unbase64(c) + unhex(c) + upper(c)
   c12 <- variance(c)
   c13 <- lead("col", 1) + lead(c, 1) + lag("col", 1) + lag(c, 1)
-  c14 <- cumeDist() + ntile(1)
-  c15 <- denseRank() + percentRank() + rank() + rowNumber()
+  c14 <- cume_dist() + ntile(1)
+  c15 <- dense_rank() + percent_rank() + rank() + row_number()
 
   # Test if base::rank() is exposed
   expect_equal(class(rank())[[1]], "Column")

From f57e6c9effdb9e282fc8ae66dc30fe053fed5272 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Fri, 27 Nov 2015 11:50:18 -0800
Subject: [PATCH 487/518] [SPARK-12021][STREAMING][TESTS] Fix the potential
 dead-lock in StreamingListenerSuite

In StreamingListenerSuite."don't call ssc.stop in listener", after the main thread calls `ssc.stop()`,  `StreamingContextStoppingCollector` may call  `ssc.stop()` in the listener bus thread, which is a dead-lock. This PR updated `StreamingContextStoppingCollector` to only call `ssc.stop()` in the first batch to avoid the dead-lock.

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #10011 from zsxwing/fix-test-deadlock.
---
 .../streaming/StreamingListenerSuite.scala    | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/streaming/src/test/scala/org/apache/spark/streaming/StreamingListenerSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/StreamingListenerSuite.scala
index df4575ab25aad..04cd5bdc26be2 100644
--- a/streaming/src/test/scala/org/apache/spark/streaming/StreamingListenerSuite.scala
+++ b/streaming/src/test/scala/org/apache/spark/streaming/StreamingListenerSuite.scala
@@ -222,7 +222,11 @@ class StreamingListenerSuite extends TestSuiteBase with Matchers {
     val batchCounter = new BatchCounter(_ssc)
     _ssc.start()
     // Make sure running at least one batch
-    batchCounter.waitUntilBatchesCompleted(expectedNumCompletedBatches = 1, timeout = 10000)
+    if (!batchCounter.waitUntilBatchesCompleted(expectedNumCompletedBatches = 1, timeout = 10000)) {
+      fail("The first batch cannot complete in 10 seconds")
+    }
+    // When reaching here, we can make sure `StreamingContextStoppingCollector` won't call
+    // `ssc.stop()`, so it's safe to call `_ssc.stop()` now.
     _ssc.stop()
     assert(contextStoppingCollector.sparkExSeen)
   }
@@ -345,12 +349,21 @@ class FailureReasonsCollector extends StreamingListener {
  */
 class StreamingContextStoppingCollector(val ssc: StreamingContext) extends StreamingListener {
   @volatile var sparkExSeen = false
+
+  private var isFirstBatch = true
+
   override def onBatchCompleted(batchCompleted: StreamingListenerBatchCompleted) {
-    try {
-      ssc.stop()
-    } catch {
-      case se: SparkException =>
-        sparkExSeen = true
+    if (isFirstBatch) {
+      // We should only call `ssc.stop()` in the first batch. Otherwise, it's possible that the main
+      // thread is calling `ssc.stop()`, while StreamingContextStoppingCollector is also calling
+      // `ssc.stop()` in the listener thread, which becomes a dead-lock.
+      isFirstBatch = false
+      try {
+        ssc.stop()
+      } catch {
+        case se: SparkException =>
+          sparkExSeen = true
+      }
     }
   }
 }

From b9921524d970f9413039967c1f17ae2e736982f0 Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Fri, 27 Nov 2015 15:11:13 -0800
Subject: [PATCH 488/518] [SPARK-12020][TESTS][TEST-HADOOP2.0] PR builder
 cannot trigger hadoop 2.0 test

https://issues.apache.org/jira/browse/SPARK-12020

Author: Yin Huai <yhuai@databricks.com>

Closes #10010 from yhuai/SPARK-12020.
---
 dev/run-tests-jenkins.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dev/run-tests-jenkins.py b/dev/run-tests-jenkins.py
index 623004310e189..4f390ef1eaa32 100755
--- a/dev/run-tests-jenkins.py
+++ b/dev/run-tests-jenkins.py
@@ -164,7 +164,7 @@ def main():
     # Switch the Hadoop profile based on the PR title:
     if "test-hadoop1.0" in ghprb_pull_title:
         os.environ["AMPLAB_JENKINS_BUILD_PROFILE"] = "hadoop1.0"
-    if "test-hadoop2.2" in ghprb_pull_title:
+    if "test-hadoop2.0" in ghprb_pull_title:
         os.environ["AMPLAB_JENKINS_BUILD_PROFILE"] = "hadoop2.0"
     if "test-hadoop2.2" in ghprb_pull_title:
         os.environ["AMPLAB_JENKINS_BUILD_PROFILE"] = "hadoop2.2"

From 149cd692ee2e127d79386fd8e584f4f70a2906ba Mon Sep 17 00:00:00 2001
From: gatorsmile <gatorsmile@gmail.com>
Date: Fri, 27 Nov 2015 22:44:08 -0800
Subject: [PATCH 489/518] [SPARK-12028] [SQL] get_json_object returns an
 incorrect result when the value is null literals

When calling `get_json_object` for the following two cases, both results are `"null"`:

```scala
    val tuple: Seq[(String, String)] = ("5", """{"f1": null}""") :: Nil
    val df: DataFrame = tuple.toDF("key", "jstring")
    val res = df.select(functions.get_json_object($"jstring", "$.f1")).collect()
```
```scala
    val tuple2: Seq[(String, String)] = ("5", """{"f1": "null"}""") :: Nil
    val df2: DataFrame = tuple2.toDF("key", "jstring")
    val res3 = df2.select(functions.get_json_object($"jstring", "$.f1")).collect()
```

Fixed the problem and also added a test case.

Author: gatorsmile <gatorsmile@gmail.com>

Closes #10018 from gatorsmile/get_json_object.
---
 .../expressions/jsonExpressions.scala         |  7 +++++--
 .../apache/spark/sql/JsonFunctionsSuite.scala | 20 +++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
index 8cd73236a7876..4991b9cb54e5e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala
@@ -298,8 +298,11 @@ case class GetJsonObject(json: Expression, path: Expression)
 
       case (FIELD_NAME, Named(name) :: xs) if p.getCurrentName == name =>
         // exact field match
-        p.nextToken()
-        evaluatePath(p, g, style, xs)
+        if (p.nextToken() != JsonToken.VALUE_NULL) {
+          evaluatePath(p, g, style, xs)
+        } else {
+          false
+        }
 
       case (FIELD_NAME, Wildcard :: xs) =>
         // wildcard field match
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala
index 14fd56fc8c222..1f384edf321b0 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala
@@ -39,6 +39,26 @@ class JsonFunctionsSuite extends QueryTest with SharedSQLContext {
     ("6", "[invalid JSON string]") ::
     Nil
 
+  test("function get_json_object - null") {
+    val df: DataFrame = tuples.toDF("key", "jstring")
+    val expected =
+      Row("1", "value1", "value2", "3", null, "5.23") ::
+        Row("2", "value12", "2", "value3", "4.01", null) ::
+        Row("3", "value13", "2", "value33", "value44", "5.01") ::
+        Row("4", null, null, null, null, null) ::
+        Row("5", "", null, null, null, null) ::
+        Row("6", null, null, null, null, null) ::
+        Nil
+
+    checkAnswer(
+      df.select($"key", functions.get_json_object($"jstring", "$.f1"),
+        functions.get_json_object($"jstring", "$.f2"),
+        functions.get_json_object($"jstring", "$.f3"),
+        functions.get_json_object($"jstring", "$.f4"),
+        functions.get_json_object($"jstring", "$.f5")),
+      expected)
+  }
+
   test("json_tuple select") {
     val df: DataFrame = tuples.toDF("key", "jstring")
     val expected =

From 28e46ab46368ea3833c8e805163893bbb6f2a265 Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Sat, 28 Nov 2015 21:02:05 -0800
Subject: [PATCH 490/518] [SPARK-12029][SPARKR] Improve column functions
 signature, param check, tests, fix doc and add examples

shivaram sun-rui

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #10019 from felixcheung/rfunctionsdoc.
---
 R/pkg/R/functions.R              | 121 +++++++++++++++++++++++--------
 R/pkg/inst/tests/test_sparkSQL.R |   9 ++-
 2 files changed, 96 insertions(+), 34 deletions(-)

diff --git a/R/pkg/R/functions.R b/R/pkg/R/functions.R
index e98e7a0117ca0..b30331c61c9a7 100644
--- a/R/pkg/R/functions.R
+++ b/R/pkg/R/functions.R
@@ -878,7 +878,7 @@ setMethod("rtrim",
 #'}
 setMethod("sd",
           signature(x = "Column"),
-          function(x, na.rm = FALSE) {
+          function(x) {
             # In R, sample standard deviation is calculated with the sd() function.
             stddev_samp(x)
           })
@@ -1250,7 +1250,7 @@ setMethod("upper",
 #'}
 setMethod("var",
           signature(x = "Column"),
-          function(x, y = NULL, na.rm = FALSE, use) {
+          function(x) {
             # In R, sample variance is calculated with the var() function.
             var_samp(x)
           })
@@ -1467,6 +1467,7 @@ setMethod("pmod", signature(y = "Column"),
 #' @name approxCountDistinct
 #' @return the approximate number of distinct items in a group.
 #' @export
+#' @examples \dontrun{approxCountDistinct(df$c, 0.02)}
 setMethod("approxCountDistinct",
           signature(x = "Column"),
           function(x, rsd = 0.05) {
@@ -1481,14 +1482,16 @@ setMethod("approxCountDistinct",
 #' @name countDistinct
 #' @return the number of distinct items in a group.
 #' @export
+#' @examples \dontrun{countDistinct(df$c)}
 setMethod("countDistinct",
           signature(x = "Column"),
           function(x, ...) {
-            jcol <- lapply(list(...), function (x) {
+            jcols <- lapply(list(...), function (x) {
+              stopifnot(class(x) == "Column")
               x@jc
             })
             jc <- callJStatic("org.apache.spark.sql.functions", "countDistinct", x@jc,
-                              jcol)
+                              jcols)
             column(jc)
           })
 
@@ -1501,10 +1504,14 @@ setMethod("countDistinct",
 #' @rdname concat
 #' @name concat
 #' @export
+#' @examples \dontrun{concat(df$strings, df$strings2)}
 setMethod("concat",
           signature(x = "Column"),
           function(x, ...) {
-            jcols <- lapply(list(x, ...), function(x) { x@jc })
+            jcols <- lapply(list(x, ...), function (x) {
+              stopifnot(class(x) == "Column")
+              x@jc
+            })
             jc <- callJStatic("org.apache.spark.sql.functions", "concat", jcols)
             column(jc)
           })
@@ -1518,11 +1525,15 @@ setMethod("concat",
 #' @rdname greatest
 #' @name greatest
 #' @export
+#' @examples \dontrun{greatest(df$c, df$d)}
 setMethod("greatest",
           signature(x = "Column"),
           function(x, ...) {
             stopifnot(length(list(...)) > 0)
-            jcols <- lapply(list(x, ...), function(x) { x@jc })
+            jcols <- lapply(list(x, ...), function (x) {
+              stopifnot(class(x) == "Column")
+              x@jc
+            })
             jc <- callJStatic("org.apache.spark.sql.functions", "greatest", jcols)
             column(jc)
           })
@@ -1530,17 +1541,21 @@ setMethod("greatest",
 #' least
 #'
 #' Returns the least value of the list of column names, skipping null values.
-#' This function takes at least 2 parameters. It will return null iff all parameters are null.
+#' This function takes at least 2 parameters. It will return null if all parameters are null.
 #'
 #' @family normal_funcs
 #' @rdname least
 #' @name least
 #' @export
+#' @examples \dontrun{least(df$c, df$d)}
 setMethod("least",
           signature(x = "Column"),
           function(x, ...) {
             stopifnot(length(list(...)) > 0)
-            jcols <- lapply(list(x, ...), function(x) { x@jc })
+            jcols <- lapply(list(x, ...), function (x) {
+              stopifnot(class(x) == "Column")
+              x@jc
+            })
             jc <- callJStatic("org.apache.spark.sql.functions", "least", jcols)
             column(jc)
           })
@@ -1549,11 +1564,10 @@ setMethod("least",
 #'
 #' Computes the ceiling of the given value.
 #'
-#' @family math_funcs
 #' @rdname ceil
-#' @name ceil
-#' @aliases ceil
+#' @name ceiling
 #' @export
+#' @examples \dontrun{ceiling(df$c)}
 setMethod("ceiling",
           signature(x = "Column"),
           function(x) {
@@ -1564,11 +1578,10 @@ setMethod("ceiling",
 #'
 #' Computes the signum of the given value.
 #'
-#' @family math_funcs
 #' @rdname signum
-#' @name signum
-#' @aliases signum
+#' @name sign
 #' @export
+#' @examples \dontrun{sign(df$c)}
 setMethod("sign", signature(x = "Column"),
           function(x) {
             signum(x)
@@ -1578,11 +1591,10 @@ setMethod("sign", signature(x = "Column"),
 #'
 #' Aggregate function: returns the number of distinct items in a group.
 #'
-#' @family agg_funcs
 #' @rdname countDistinct
-#' @name countDistinct
-#' @aliases countDistinct
+#' @name n_distinct
 #' @export
+#' @examples \dontrun{n_distinct(df$c)}
 setMethod("n_distinct", signature(x = "Column"),
           function(x, ...) {
             countDistinct(x, ...)
@@ -1592,11 +1604,10 @@ setMethod("n_distinct", signature(x = "Column"),
 #'
 #' Aggregate function: returns the number of items in a group.
 #'
-#' @family agg_funcs
 #' @rdname count
-#' @name count
-#' @aliases count
+#' @name n
 #' @export
+#' @examples \dontrun{n(df$c)}
 setMethod("n", signature(x = "Column"),
           function(x) {
             count(x)
@@ -1617,6 +1628,7 @@ setMethod("n", signature(x = "Column"),
 #' @rdname date_format
 #' @name date_format
 #' @export
+#' @examples \dontrun{date_format(df$t, 'MM/dd/yyy')}
 setMethod("date_format", signature(y = "Column", x = "character"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "date_format", y@jc, x)
@@ -1631,6 +1643,7 @@ setMethod("date_format", signature(y = "Column", x = "character"),
 #' @rdname from_utc_timestamp
 #' @name from_utc_timestamp
 #' @export
+#' @examples \dontrun{from_utc_timestamp(df$t, 'PST')}
 setMethod("from_utc_timestamp", signature(y = "Column", x = "character"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "from_utc_timestamp", y@jc, x)
@@ -1649,6 +1662,7 @@ setMethod("from_utc_timestamp", signature(y = "Column", x = "character"),
 #' @rdname instr
 #' @name instr
 #' @export
+#' @examples \dontrun{instr(df$c, 'b')}
 setMethod("instr", signature(y = "Column", x = "character"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "instr", y@jc, x)
@@ -1663,13 +1677,18 @@ setMethod("instr", signature(y = "Column", x = "character"),
 #' For example, \code{next_day('2015-07-27', "Sunday")} returns 2015-08-02 because that is the first
 #' Sunday after 2015-07-27.
 #'
-#' Day of the week parameter is case insensitive, and accepts:
+#' Day of the week parameter is case insensitive, and accepts first three or two characters:
 #' "Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun".
 #'
 #' @family datetime_funcs
 #' @rdname next_day
 #' @name next_day
 #' @export
+#' @examples
+#'\dontrun{
+#'next_day(df$d, 'Sun')
+#'next_day(df$d, 'Sunday')
+#'}
 setMethod("next_day", signature(y = "Column", x = "character"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "next_day", y@jc, x)
@@ -1684,6 +1703,7 @@ setMethod("next_day", signature(y = "Column", x = "character"),
 #' @rdname to_utc_timestamp
 #' @name to_utc_timestamp
 #' @export
+#' @examples \dontrun{to_utc_timestamp(df$t, 'PST')}
 setMethod("to_utc_timestamp", signature(y = "Column", x = "character"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "to_utc_timestamp", y@jc, x)
@@ -1697,8 +1717,8 @@ setMethod("to_utc_timestamp", signature(y = "Column", x = "character"),
 #' @name add_months
 #' @family datetime_funcs
 #' @rdname add_months
-#' @name add_months
 #' @export
+#' @examples \dontrun{add_months(df$d, 1)}
 setMethod("add_months", signature(y = "Column", x = "numeric"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "add_months", y@jc, as.integer(x))
@@ -1713,6 +1733,7 @@ setMethod("add_months", signature(y = "Column", x = "numeric"),
 #' @rdname date_add
 #' @name date_add
 #' @export
+#' @examples \dontrun{date_add(df$d, 1)}
 setMethod("date_add", signature(y = "Column", x = "numeric"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "date_add", y@jc, as.integer(x))
@@ -1727,6 +1748,7 @@ setMethod("date_add", signature(y = "Column", x = "numeric"),
 #' @rdname date_sub
 #' @name date_sub
 #' @export
+#' @examples \dontrun{date_sub(df$d, 1)}
 setMethod("date_sub", signature(y = "Column", x = "numeric"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "date_sub", y@jc, as.integer(x))
@@ -1735,16 +1757,19 @@ setMethod("date_sub", signature(y = "Column", x = "numeric"),
 
 #' format_number
 #'
-#' Formats numeric column x to a format like '#,###,###.##', rounded to d decimal places,
+#' Formats numeric column y to a format like '#,###,###.##', rounded to x decimal places,
 #' and returns the result as a string column.
 #'
-#' If d is 0, the result has no decimal point or fractional part.
-#' If d < 0, the result will be null.'
+#' If x is 0, the result has no decimal point or fractional part.
+#' If x < 0, the result will be null.
 #'
+#' @param y column to format
+#' @param x number of decimal place to format to
 #' @family string_funcs
 #' @rdname format_number
 #' @name format_number
 #' @export
+#' @examples \dontrun{format_number(df$n, 4)}
 setMethod("format_number", signature(y = "Column", x = "numeric"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions",
@@ -1764,6 +1789,7 @@ setMethod("format_number", signature(y = "Column", x = "numeric"),
 #' @rdname sha2
 #' @name sha2
 #' @export
+#' @examples \dontrun{sha2(df$c, 256)}
 setMethod("sha2", signature(y = "Column", x = "numeric"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "sha2", y@jc, as.integer(x))
@@ -1779,6 +1805,7 @@ setMethod("sha2", signature(y = "Column", x = "numeric"),
 #' @rdname shiftLeft
 #' @name shiftLeft
 #' @export
+#' @examples \dontrun{shiftLeft(df$c, 1)}
 setMethod("shiftLeft", signature(y = "Column", x = "numeric"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions",
@@ -1796,6 +1823,7 @@ setMethod("shiftLeft", signature(y = "Column", x = "numeric"),
 #' @rdname shiftRight
 #' @name shiftRight
 #' @export
+#' @examples \dontrun{shiftRight(df$c, 1)}
 setMethod("shiftRight", signature(y = "Column", x = "numeric"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions",
@@ -1813,6 +1841,7 @@ setMethod("shiftRight", signature(y = "Column", x = "numeric"),
 #' @rdname shiftRightUnsigned
 #' @name shiftRightUnsigned
 #' @export
+#' @examples \dontrun{shiftRightUnsigned(df$c, 1)}
 setMethod("shiftRightUnsigned", signature(y = "Column", x = "numeric"),
           function(y, x) {
             jc <- callJStatic("org.apache.spark.sql.functions",
@@ -1830,6 +1859,7 @@ setMethod("shiftRightUnsigned", signature(y = "Column", x = "numeric"),
 #' @rdname concat_ws
 #' @name concat_ws
 #' @export
+#' @examples \dontrun{concat_ws('-', df$s, df$d)}
 setMethod("concat_ws", signature(sep = "character", x = "Column"),
           function(sep, x, ...) {
             jcols <- lapply(list(x, ...), function(x) { x@jc })
@@ -1845,6 +1875,7 @@ setMethod("concat_ws", signature(sep = "character", x = "Column"),
 #' @rdname conv
 #' @name conv
 #' @export
+#' @examples \dontrun{conv(df$n, 2, 16)}
 setMethod("conv", signature(x = "Column", fromBase = "numeric", toBase = "numeric"),
           function(x, fromBase, toBase) {
             fromBase <- as.integer(fromBase)
@@ -1864,6 +1895,7 @@ setMethod("conv", signature(x = "Column", fromBase = "numeric", toBase = "numeri
 #' @rdname expr
 #' @name expr
 #' @export
+#' @examples \dontrun{expr('length(name)')}
 setMethod("expr", signature(x = "character"),
           function(x) {
             jc <- callJStatic("org.apache.spark.sql.functions", "expr", x)
@@ -1878,6 +1910,7 @@ setMethod("expr", signature(x = "character"),
 #' @rdname format_string
 #' @name format_string
 #' @export
+#' @examples \dontrun{format_string('%d %s', df$a, df$b)}
 setMethod("format_string", signature(format = "character", x = "Column"),
           function(format, x, ...) {
             jcols <- lapply(list(x, ...), function(arg) { arg@jc })
@@ -1897,6 +1930,11 @@ setMethod("format_string", signature(format = "character", x = "Column"),
 #' @rdname from_unixtime
 #' @name from_unixtime
 #' @export
+#' @examples
+#'\dontrun{
+#'from_unixtime(df$t)
+#'from_unixtime(df$t, 'yyyy/MM/dd HH')
+#'}
 setMethod("from_unixtime", signature(x = "Column"),
           function(x, format = "yyyy-MM-dd HH:mm:ss") {
             jc <- callJStatic("org.apache.spark.sql.functions",
@@ -1915,6 +1953,7 @@ setMethod("from_unixtime", signature(x = "Column"),
 #' @rdname locate
 #' @name locate
 #' @export
+#' @examples \dontrun{locate('b', df$c, 1)}
 setMethod("locate", signature(substr = "character", str = "Column"),
           function(substr, str, pos = 0) {
             jc <- callJStatic("org.apache.spark.sql.functions",
@@ -1931,6 +1970,7 @@ setMethod("locate", signature(substr = "character", str = "Column"),
 #' @rdname lpad
 #' @name lpad
 #' @export
+#' @examples \dontrun{lpad(df$c, 6, '#')}
 setMethod("lpad", signature(x = "Column", len = "numeric", pad = "character"),
           function(x, len, pad) {
             jc <- callJStatic("org.apache.spark.sql.functions",
@@ -1947,12 +1987,13 @@ setMethod("lpad", signature(x = "Column", len = "numeric", pad = "character"),
 #' @rdname rand
 #' @name rand
 #' @export
+#' @examples \dontrun{rand()}
 setMethod("rand", signature(seed = "missing"),
           function(seed) {
             jc <- callJStatic("org.apache.spark.sql.functions", "rand")
             column(jc)
           })
-#' @family normal_funcs
+
 #' @rdname rand
 #' @name rand
 #' @export
@@ -1970,12 +2011,13 @@ setMethod("rand", signature(seed = "numeric"),
 #' @rdname randn
 #' @name randn
 #' @export
+#' @examples \dontrun{randn()}
 setMethod("randn", signature(seed = "missing"),
           function(seed) {
             jc <- callJStatic("org.apache.spark.sql.functions", "randn")
             column(jc)
           })
-#' @family normal_funcs
+
 #' @rdname randn
 #' @name randn
 #' @export
@@ -1993,6 +2035,7 @@ setMethod("randn", signature(seed = "numeric"),
 #' @rdname regexp_extract
 #' @name regexp_extract
 #' @export
+#' @examples \dontrun{regexp_extract(df$c, '(\d+)-(\d+)', 1)}
 setMethod("regexp_extract",
           signature(x = "Column", pattern = "character", idx = "numeric"),
           function(x, pattern, idx) {
@@ -2010,6 +2053,7 @@ setMethod("regexp_extract",
 #' @rdname regexp_replace
 #' @name regexp_replace
 #' @export
+#' @examples \dontrun{regexp_replace(df$c, '(\\d+)', '--')}
 setMethod("regexp_replace",
           signature(x = "Column", pattern = "character", replacement = "character"),
           function(x, pattern, replacement) {
@@ -2027,6 +2071,7 @@ setMethod("regexp_replace",
 #' @rdname rpad
 #' @name rpad
 #' @export
+#' @examples \dontrun{rpad(df$c, 6, '#')}
 setMethod("rpad", signature(x = "Column", len = "numeric", pad = "character"),
           function(x, len, pad) {
             jc <- callJStatic("org.apache.spark.sql.functions",
@@ -2040,12 +2085,17 @@ setMethod("rpad", signature(x = "Column", len = "numeric", pad = "character"),
 #' Returns the substring from string str before count occurrences of the delimiter delim.
 #' If count is positive, everything the left of the final delimiter (counting from left) is
 #' returned. If count is negative, every to the right of the final delimiter (counting from the
-#' right) is returned. substring <- index performs a case-sensitive match when searching for delim.
+#' right) is returned. substring_index performs a case-sensitive match when searching for delim.
 #'
 #' @family string_funcs
 #' @rdname substring_index
 #' @name substring_index
 #' @export
+#' @examples
+#'\dontrun{
+#'substring_index(df$c, '.', 2)
+#'substring_index(df$c, '.', -1)
+#'}
 setMethod("substring_index",
           signature(x = "Column", delim = "character", count = "numeric"),
           function(x, delim, count) {
@@ -2066,6 +2116,7 @@ setMethod("substring_index",
 #' @rdname translate
 #' @name translate
 #' @export
+#' @examples \dontrun{translate(df$c, 'rnlt', '123')}
 setMethod("translate",
           signature(x = "Column", matchingString = "character", replaceString = "character"),
           function(x, matchingString, replaceString) {
@@ -2082,12 +2133,18 @@ setMethod("translate",
 #' @rdname unix_timestamp
 #' @name unix_timestamp
 #' @export
+#' @examples
+#'\dontrun{
+#'unix_timestamp()
+#'unix_timestamp(df$t)
+#'unix_timestamp(df$t, 'yyyy-MM-dd HH')
+#'}
 setMethod("unix_timestamp", signature(x = "missing", format = "missing"),
           function(x, format) {
             jc <- callJStatic("org.apache.spark.sql.functions", "unix_timestamp")
             column(jc)
           })
-#' @family datetime_funcs
+
 #' @rdname unix_timestamp
 #' @name unix_timestamp
 #' @export
@@ -2096,7 +2153,7 @@ setMethod("unix_timestamp", signature(x = "Column", format = "missing"),
             jc <- callJStatic("org.apache.spark.sql.functions", "unix_timestamp", x@jc)
             column(jc)
           })
-#' @family datetime_funcs
+
 #' @rdname unix_timestamp
 #' @name unix_timestamp
 #' @export
@@ -2113,7 +2170,9 @@ setMethod("unix_timestamp", signature(x = "Column", format = "character"),
 #' @family normal_funcs
 #' @rdname when
 #' @name when
+#' @seealso \link{ifelse}
 #' @export
+#' @examples \dontrun{when(df$age == 2, df$age + 1)}
 setMethod("when", signature(condition = "Column", value = "ANY"),
           function(condition, value) {
               condition <- condition@jc
@@ -2130,7 +2189,9 @@ setMethod("when", signature(condition = "Column", value = "ANY"),
 #' @family normal_funcs
 #' @rdname ifelse
 #' @name ifelse
+#' @seealso \link{when}
 #' @export
+#' @examples \dontrun{ifelse(df$a > 1 & df$b > 2, 0, 1)}
 setMethod("ifelse",
           signature(test = "Column", yes = "ANY", no = "ANY"),
           function(test, yes, no) {
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 0fbe0658265b2..899fc3b977385 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -880,14 +880,15 @@ test_that("column functions", {
   expect_equal(collect(df3)[[2, 1]], FALSE)
   expect_equal(collect(df3)[[3, 1]], TRUE)
 
-  expect_equal(collect(select(df, sum(df$age)))[1, 1], 49)
+  df4 <- select(df, countDistinct(df$age, df$name))
+  expect_equal(collect(df4)[[1, 1]], 2)
 
+  expect_equal(collect(select(df, sum(df$age)))[1, 1], 49)
   expect_true(abs(collect(select(df, stddev(df$age)))[1, 1] - 7.778175) < 1e-6)
-
   expect_equal(collect(select(df, var_pop(df$age)))[1, 1], 30.25)
 
-  df4 <- createDataFrame(sqlContext, list(list(a = "010101")))
-  expect_equal(collect(select(df4, conv(df4$a, 2, 16)))[1, 1], "15")
+  df5 <- createDataFrame(sqlContext, list(list(a = "010101")))
+  expect_equal(collect(select(df5, conv(df5$a, 2, 16)))[1, 1], "15")
 
   # Test array_contains() and sort_array()
   df <- createDataFrame(sqlContext, list(list(list(1L, 2L, 3L)), list(list(6L, 5L, 4L))))

From c793d2d9a1ccc203fc103eb0636958fe8d71f471 Mon Sep 17 00:00:00 2001
From: felixcheung <felixcheung_m@hotmail.com>
Date: Sat, 28 Nov 2015 21:16:21 -0800
Subject: [PATCH 491/518] [SPARK-9319][SPARKR] Add support for setting column
 names, types

Add support for for colnames, colnames<-, coltypes<-
Also added tests for names, names<- which have no test previously.

I merged with PR 8984 (coltypes). Clicked the wrong thing, crewed up the PR. Recreated it here. Was #9218

shivaram sun-rui

Author: felixcheung <felixcheung_m@hotmail.com>

Closes #9654 from felixcheung/colnamescoltypes.
---
 R/pkg/NAMESPACE                  |   6 +-
 R/pkg/R/DataFrame.R              | 166 ++++++++++++++++++++++---------
 R/pkg/R/generics.R               |  20 +++-
 R/pkg/R/types.R                  |   8 ++
 R/pkg/inst/tests/test_sparkSQL.R |  40 +++++++-
 5 files changed, 185 insertions(+), 55 deletions(-)

diff --git a/R/pkg/NAMESPACE b/R/pkg/NAMESPACE
index 5d04dd6acaab8..43e5e0119e7fe 100644
--- a/R/pkg/NAMESPACE
+++ b/R/pkg/NAMESPACE
@@ -27,7 +27,10 @@ exportMethods("arrange",
               "attach",
               "cache",
               "collect",
+              "colnames",
+              "colnames<-",
               "coltypes",
+              "coltypes<-",
               "columns",
               "count",
               "cov",
@@ -56,6 +59,7 @@ exportMethods("arrange",
               "mutate",
               "na.omit",
               "names",
+              "names<-",
               "ncol",
               "nrow",
               "orderBy",
@@ -276,4 +280,4 @@ export("structField",
        "structType",
        "structType.jobj",
        "structType.structField",
-       "print.structType")
\ No newline at end of file
+       "print.structType")
diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index 8a13e7a36766d..f89e2682d9e29 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -254,6 +254,7 @@ setMethod("dtypes",
 #' @family DataFrame functions
 #' @rdname columns
 #' @name columns
+
 #' @export
 #' @examples
 #'\dontrun{
@@ -262,6 +263,7 @@ setMethod("dtypes",
 #' path <- "path/to/file.json"
 #' df <- jsonFile(sqlContext, path)
 #' columns(df)
+#' colnames(df)
 #'}
 setMethod("columns",
           signature(x = "DataFrame"),
@@ -290,6 +292,121 @@ setMethod("names<-",
             }
           })
 
+#' @rdname columns
+#' @name colnames
+setMethod("colnames",
+          signature(x = "DataFrame"),
+          function(x) {
+            columns(x)
+          })
+
+#' @rdname columns
+#' @name colnames<-
+setMethod("colnames<-",
+          signature(x = "DataFrame", value = "character"),
+          function(x, value) {
+            sdf <- callJMethod(x@sdf, "toDF", as.list(value))
+            dataFrame(sdf)
+          })
+
+#' coltypes
+#'
+#' Get column types of a DataFrame
+#'
+#' @param x A SparkSQL DataFrame
+#' @return value A character vector with the column types of the given DataFrame
+#' @rdname coltypes
+#' @name coltypes
+#' @family DataFrame functions
+#' @export
+#' @examples
+#'\dontrun{
+#' irisDF <- createDataFrame(sqlContext, iris)
+#' coltypes(irisDF)
+#'}
+setMethod("coltypes",
+          signature(x = "DataFrame"),
+          function(x) {
+            # Get the data types of the DataFrame by invoking dtypes() function
+            types <- sapply(dtypes(x), function(x) {x[[2]]})
+
+            # Map Spark data types into R's data types using DATA_TYPES environment
+            rTypes <- sapply(types, USE.NAMES=F, FUN=function(x) {
+              # Check for primitive types
+              type <- PRIMITIVE_TYPES[[x]]
+
+              if (is.null(type)) {
+                # Check for complex types
+                for (t in names(COMPLEX_TYPES)) {
+                  if (substring(x, 1, nchar(t)) == t) {
+                    type <- COMPLEX_TYPES[[t]]
+                    break
+                  }
+                }
+
+                if (is.null(type)) {
+                  stop(paste("Unsupported data type: ", x))
+                }
+              }
+              type
+            })
+
+            # Find which types don't have mapping to R
+            naIndices <- which(is.na(rTypes))
+
+            # Assign the original scala data types to the unmatched ones
+            rTypes[naIndices] <- types[naIndices]
+
+            rTypes
+          })
+
+#' coltypes
+#'
+#' Set the column types of a DataFrame.
+#'
+#' @param x A SparkSQL DataFrame
+#' @param value A character vector with the target column types for the given
+#'    DataFrame. Column types can be one of integer, numeric/double, character, logical, or NA
+#'    to keep that column as-is.
+#' @rdname coltypes
+#' @name coltypes<-
+#' @export
+#' @examples
+#'\dontrun{
+#' sc <- sparkR.init()
+#' sqlContext <- sparkRSQL.init(sc)
+#' path <- "path/to/file.json"
+#' df <- jsonFile(sqlContext, path)
+#' coltypes(df) <- c("character", "integer")
+#' coltypes(df) <- c(NA, "numeric")
+#'}
+setMethod("coltypes<-",
+          signature(x = "DataFrame", value = "character"),
+          function(x, value) {
+            cols <- columns(x)
+            ncols <- length(cols)
+            if (length(value) == 0) {
+              stop("Cannot set types of an empty DataFrame with no Column")
+            }
+            if (length(value) != ncols) {
+              stop("Length of type vector should match the number of columns for DataFrame")
+            }
+            newCols <- lapply(seq_len(ncols), function(i) {
+              col <- getColumn(x, cols[i])
+              if (!is.na(value[i])) {
+                stype <- rToSQLTypes[[value[i]]]
+                if (is.null(stype)) {
+                  stop("Only atomic type is supported for column types")
+                }
+                cast(col, stype)
+              } else {
+                col
+              }
+            })
+            nx <- select(x, newCols)
+            dataFrame(nx@sdf)
+          })
+
 #' Register Temporary Table
 #'
 #' Registers a DataFrame as a Temporary Table in the SQLContext
@@ -2102,52 +2219,3 @@ setMethod("with",
             newEnv <- assignNewEnv(data)
             eval(substitute(expr), envir = newEnv, enclos = newEnv)
           })
-
-#' Returns the column types of a DataFrame.
-#'
-#' @name coltypes
-#' @title Get column types of a DataFrame
-#' @family dataframe_funcs
-#' @param x (DataFrame)
-#' @return value (character) A character vector with the column types of the given DataFrame
-#' @rdname coltypes
-#' @examples \dontrun{
-#' irisDF <- createDataFrame(sqlContext, iris)
-#' coltypes(irisDF)
-#' }
-setMethod("coltypes",
-          signature(x = "DataFrame"),
-          function(x) {
-            # Get the data types of the DataFrame by invoking dtypes() function
-            types <- sapply(dtypes(x), function(x) {x[[2]]})
-
-            # Map Spark data types into R's data types using DATA_TYPES environment
-            rTypes <- sapply(types, USE.NAMES=F, FUN=function(x) {
-
-              # Check for primitive types
-              type <- PRIMITIVE_TYPES[[x]]
-
-              if (is.null(type)) {
-                # Check for complex types
-                for (t in names(COMPLEX_TYPES)) {
-                  if (substring(x, 1, nchar(t)) == t) {
-                    type <- COMPLEX_TYPES[[t]]
-                    break
-                  }
-                }
-
-                if (is.null(type)) {
-                  stop(paste("Unsupported data type: ", x))
-                }
-              }
-              type
-            })
-
-            # Find which types don't have mapping to R
-            naIndices <- which(is.na(rTypes))
-
-            # Assign the original scala data types to the unmatched ones
-            rTypes[naIndices] <- types[naIndices]
-
-            rTypes
-          })
diff --git a/R/pkg/R/generics.R b/R/pkg/R/generics.R
index 0c305441e043e..711ce38f9e104 100644
--- a/R/pkg/R/generics.R
+++ b/R/pkg/R/generics.R
@@ -385,6 +385,22 @@ setGeneric("agg", function (x, ...) { standardGeneric("agg") })
 #' @export
 setGeneric("arrange", function(x, col, ...) { standardGeneric("arrange") })
 
+#' @rdname columns
+#' @export
+setGeneric("colnames", function(x, do.NULL = TRUE, prefix = "col") { standardGeneric("colnames") })
+
+#' @rdname columns
+#' @export
+setGeneric("colnames<-", function(x, value) { standardGeneric("colnames<-") })
+
+#' @rdname coltypes
+#' @export
+setGeneric("coltypes", function(x) { standardGeneric("coltypes") })
+
+#' @rdname coltypes
+#' @export
+setGeneric("coltypes<-", function(x, value) { standardGeneric("coltypes<-") })
+
 #' @rdname schema
 #' @export
 setGeneric("columns", function(x) {standardGeneric("columns") })
@@ -1081,7 +1097,3 @@ setGeneric("attach")
 #' @rdname with
 #' @export
 setGeneric("with")
-
-#' @rdname coltypes
-#' @export
-setGeneric("coltypes", function(x) { standardGeneric("coltypes") })
diff --git a/R/pkg/R/types.R b/R/pkg/R/types.R
index 1828c23ab0f6d..dae4fe858bdbc 100644
--- a/R/pkg/R/types.R
+++ b/R/pkg/R/types.R
@@ -41,3 +41,11 @@ COMPLEX_TYPES <- list(
 
 # The full list of data types.
 DATA_TYPES <- as.environment(c(as.list(PRIMITIVE_TYPES), COMPLEX_TYPES))
+
+# An environment for mapping R to Scala, names are R types and values are Scala types.
+rToSQLTypes <- as.environment(list(
+  "integer"   = "integer", # in R, integer is 32bit
+  "numeric"   = "double",  # in R, numeric == double which is 64bit
+  "double"    = "double",
+  "character" = "string",
+  "logical"   = "boolean"))
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index 899fc3b977385..d3b2f20bf81c5 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -622,6 +622,26 @@ test_that("schema(), dtypes(), columns(), names() return the correct values/form
   expect_equal(testNames[2], "name")
 })
 
+test_that("names() colnames() set the column names", {
+  df <- jsonFile(sqlContext, jsonPath)
+  names(df) <- c("col1", "col2")
+  expect_equal(colnames(df)[2], "col2")
+
+  colnames(df) <- c("col3", "col4")
+  expect_equal(names(df)[1], "col3")
+
+  # Test base::colnames base::names
+  m2 <- cbind(1, 1:4)
+  expect_equal(colnames(m2, do.NULL = FALSE), c("col1", "col2"))
+  colnames(m2) <- c("x","Y")
+  expect_equal(colnames(m2), c("x", "Y"))
+
+  z <- list(a = 1, b = "c", c = 1:3)
+  expect_equal(names(z)[3], "c")
+  names(z)[3] <- "c2"
+  expect_equal(names(z)[3], "c2")
+})
+
 test_that("head() and first() return the correct data", {
   df <- jsonFile(sqlContext, jsonPath)
   testHead <- head(df)
@@ -1617,7 +1637,7 @@ test_that("with() on a DataFrame", {
   expect_equal(nrow(sum2), 35)
 })
 
-test_that("Method coltypes() to get R's data types of a DataFrame", {
+test_that("Method coltypes() to get and set R's data types of a DataFrame", {
   expect_equal(coltypes(irisDF), c(rep("numeric", 4), "character"))
 
   data <- data.frame(c1=c(1,2,3),
@@ -1636,6 +1656,24 @@ test_that("Method coltypes() to get R's data types of a DataFrame", {
   x <- createDataFrame(sqlContext, list(list(as.environment(
     list("a"="b", "c"="d", "e"="f")))))
   expect_equal(coltypes(x), "map<string,string>")
+
+  df <- selectExpr(jsonFile(sqlContext, jsonPath), "name", "(age * 1.21) as age")
+  expect_equal(dtypes(df), list(c("name", "string"), c("age", "decimal(24,2)")))
+
+  df1 <- select(df, cast(df$age, "integer"))
+  coltypes(df) <- c("character", "integer")
+  expect_equal(dtypes(df), list(c("name", "string"), c("age", "int")))
+  value <- collect(df[, 2])[[3, 1]]
+  expect_equal(value, collect(df1)[[3, 1]])
+  expect_equal(value, 22)
+
+  coltypes(df) <- c(NA, "numeric")
+  expect_equal(dtypes(df), list(c("name", "string"), c("age", "double")))
+
+  expect_error(coltypes(df) <- c("character"),
+               "Length of type vector should match the number of columns for DataFrame")
+  expect_error(coltypes(df) <- c("environment", "list"),
+               "Only atomic type is supported for column types")
 })
 
 unlink(parquetPath)

From cc7a1bc9370b163f51230e5ca4be612d133a5086 Mon Sep 17 00:00:00 2001
From: Sun Rui <rui.sun@intel.com>
Date: Sun, 29 Nov 2015 11:08:26 -0800
Subject: [PATCH 492/518] [SPARK-11781][SPARKR] SparkR has problem in inferring
 type of raw type.

Author: Sun Rui <rui.sun@intel.com>

Closes #9769 from sun-rui/SPARK-11781.
---
 R/pkg/R/DataFrame.R              | 34 ++++++++++++++++-------------
 R/pkg/R/SQLContext.R             |  2 +-
 R/pkg/R/types.R                  | 37 ++++++++++++++++++--------------
 R/pkg/inst/tests/test_sparkSQL.R |  6 ++++++
 4 files changed, 47 insertions(+), 32 deletions(-)

diff --git a/R/pkg/R/DataFrame.R b/R/pkg/R/DataFrame.R
index f89e2682d9e29..a82ded9c51fac 100644
--- a/R/pkg/R/DataFrame.R
+++ b/R/pkg/R/DataFrame.R
@@ -793,8 +793,8 @@ setMethod("dim",
 setMethod("collect",
           signature(x = "DataFrame"),
           function(x, stringsAsFactors = FALSE) {
-            names <- columns(x)
-            ncol <- length(names)
+            dtypes <- dtypes(x)
+            ncol <- length(dtypes)
             if (ncol <= 0) {
               # empty data.frame with 0 columns and 0 rows
               data.frame()
@@ -817,25 +817,29 @@ setMethod("collect",
                 # data of complex type can be held. But getting a cell from a column
                 # of list type returns a list instead of a vector. So for columns of
                 # non-complex type, append them as vector.
+                #
+                # For columns of complex type, be careful to access them.
+                # Get a column of complex type returns a list.
+                # Get a cell from a column of complex type returns a list instead of a vector.
                 col <- listCols[[colIndex]]
+                colName <- dtypes[[colIndex]][[1]]
                 if (length(col) <= 0) {
-                  df[[names[colIndex]]] <- col
+                  df[[colName]] <- col
                 } else {
-                  # TODO: more robust check on column of primitive types
-                  vec <- do.call(c, col)
-                  if (class(vec) != "list") {
-                    df[[names[colIndex]]] <- vec
+                  colType <- dtypes[[colIndex]][[2]]
+                  # Note that "binary" columns behave like complex types.
+                  if (!is.null(PRIMITIVE_TYPES[[colType]]) && colType != "binary") {
+                    vec <- do.call(c, col)
+                    stopifnot(class(vec) != "list")
+                    df[[colName]] <- vec
                   } else {
-                    # For columns of complex type, be careful to access them.
-                    # Get a column of complex type returns a list.
-                    # Get a cell from a column of complex type returns a list instead of a vector.
-                    df[[names[colIndex]]] <- col
-                 }
+                    df[[colName]] <- col
+                  }
+                }
               }
+              df
             }
-            df
-          }
-        })
+          })
 
 #' Limit
 #'
diff --git a/R/pkg/R/SQLContext.R b/R/pkg/R/SQLContext.R
index a62b25fde926d..85541c8e22447 100644
--- a/R/pkg/R/SQLContext.R
+++ b/R/pkg/R/SQLContext.R
@@ -63,7 +63,7 @@ infer_type <- function(x) {
     })
     type <- Reduce(paste0, type)
     type <- paste0("struct<", substr(type, 1, nchar(type) - 1), ">")
-  } else if (length(x) > 1) {
+  } else if (length(x) > 1 && type != "binary") {
     paste0("array<", infer_type(x[[1]]), ">")
   } else {
     type
diff --git a/R/pkg/R/types.R b/R/pkg/R/types.R
index dae4fe858bdbc..1f06af7e904fe 100644
--- a/R/pkg/R/types.R
+++ b/R/pkg/R/types.R
@@ -19,25 +19,30 @@
 # values are equivalent R types. This is stored in an environment to allow for
 # more efficient look up (environments use hashmaps).
 PRIMITIVE_TYPES <- as.environment(list(
-  "byte"="integer",
-  "tinyint"="integer",
-  "smallint"="integer",
-  "integer"="integer",
-  "bigint"="numeric",
-  "float"="numeric",
-  "double"="numeric",
-  "decimal"="numeric",
-  "string"="character",
-  "binary"="raw",
-  "boolean"="logical",
-  "timestamp"="POSIXct",
-  "date"="Date"))
+  "tinyint" = "integer",
+  "smallint" = "integer",
+  "int" = "integer",
+  "bigint" = "numeric",
+  "float" = "numeric",
+  "double" = "numeric",
+  "decimal" = "numeric",
+  "string" = "character",
+  "binary" = "raw",
+  "boolean" = "logical",
+  "timestamp" = "POSIXct",
+  "date" = "Date",
+  # following types are not SQL types returned by dtypes(). They are listed here for usage
+  # by checkType() in schema.R.
+  # TODO: refactor checkType() in schema.R.
+  "byte" = "integer",
+  "integer" = "integer"
+  ))
 
 # The complex data types. These do not have any direct mapping to R's types.
 COMPLEX_TYPES <- list(
-  "map"=NA,
-  "array"=NA,
-  "struct"=NA)
+  "map" = NA,
+  "array" = NA,
+  "struct" = NA)
 
 # The full list of data types.
 DATA_TYPES <- as.environment(c(as.list(PRIMITIVE_TYPES), COMPLEX_TYPES))
diff --git a/R/pkg/inst/tests/test_sparkSQL.R b/R/pkg/inst/tests/test_sparkSQL.R
index d3b2f20bf81c5..92ec82096c6df 100644
--- a/R/pkg/inst/tests/test_sparkSQL.R
+++ b/R/pkg/inst/tests/test_sparkSQL.R
@@ -72,6 +72,8 @@ test_that("infer types and check types", {
   expect_equal(infer_type(e), "map<string,integer>")
 
   expect_error(checkType("map<integer,integer>"), "Key type in a map must be string or character")
+
+  expect_equal(infer_type(as.raw(c(1, 2, 3))), "binary")
 })
 
 test_that("structType and structField", {
@@ -250,6 +252,10 @@ test_that("create DataFrame from list or data.frame", {
 
   mtcarsdf <- createDataFrame(sqlContext, mtcars)
   expect_equivalent(collect(mtcarsdf), mtcars)
+
+  bytes <- as.raw(c(1, 2, 3))
+  df <- createDataFrame(sqlContext, list(list(bytes)))
+  expect_equal(collect(df)[[1]][[1]], bytes)
 })
 
 test_that("create DataFrame with different data types", {

From 3d28081e53698ed77e93c04299957c02bcaba9bf Mon Sep 17 00:00:00 2001
From: Herman van Hovell <hvanhovell@questtec.nl>
Date: Sun, 29 Nov 2015 14:13:11 -0800
Subject: [PATCH 493/518] [SPARK-12024][SQL] More efficient multi-column
 counting.

In https://github.com/apache/spark/pull/9409 we enabled multi-column counting. The approach taken in that PR introduces a bit of overhead by first creating a row only to check if all of the columns are non-null.

This PR fixes that technical debt. Count now takes multiple columns as its input. In order to make this work I have also added support for multiple columns in the single distinct code path.

cc yhuai

Author: Herman van Hovell <hvanhovell@questtec.nl>

Closes #10015 from hvanhovell/SPARK-12024.
---
 .../expressions/aggregate/Count.scala         | 21 ++--------
 .../expressions/conditionalExpressions.scala  | 27 -------------
 .../sql/catalyst/optimizer/Optimizer.scala    | 14 ++++---
 .../ConditionalExpressionSuite.scala          | 14 -------
 .../spark/sql/execution/aggregate/utils.scala | 39 +++++++++----------
 .../spark/sql/expressions/WindowSpec.scala    |  4 +-
 6 files changed, 33 insertions(+), 86 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
index 09a1da9200df0..441f52ab5ca58 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala
@@ -21,8 +21,7 @@ import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.types._
 
-case class Count(child: Expression) extends DeclarativeAggregate {
-  override def children: Seq[Expression] = child :: Nil
+case class Count(children: Seq[Expression]) extends DeclarativeAggregate {
 
   override def nullable: Boolean = false
 
@@ -30,7 +29,7 @@ case class Count(child: Expression) extends DeclarativeAggregate {
   override def dataType: DataType = LongType
 
   // Expected input data type.
-  override def inputTypes: Seq[AbstractDataType] = Seq(AnyDataType)
+  override def inputTypes: Seq[AbstractDataType] = Seq.fill(children.size)(AnyDataType)
 
   private lazy val count = AttributeReference("count", LongType)()
 
@@ -41,7 +40,7 @@ case class Count(child: Expression) extends DeclarativeAggregate {
   )
 
   override lazy val updateExpressions = Seq(
-    /* count = */ If(IsNull(child), count, count + 1L)
+    /* count = */ If(children.map(IsNull).reduce(Or), count, count + 1L)
   )
 
   override lazy val mergeExpressions = Seq(
@@ -54,17 +53,5 @@ case class Count(child: Expression) extends DeclarativeAggregate {
 }
 
 object Count {
-  def apply(children: Seq[Expression]): Count = {
-    // This is used to deal with COUNT DISTINCT. When we have multiple
-    // children (COUNT(DISTINCT col1, col2, ...)), we wrap them in a STRUCT (i.e. a Row).
-    // Also, the semantic of COUNT(DISTINCT col1, col2, ...) is that if there is any
-    // null in the arguments, we will not count that row. So, we use DropAnyNull at here
-    // to return a null when any field of the created STRUCT is null.
-    val child = if (children.size > 1) {
-      DropAnyNull(CreateStruct(children))
-    } else {
-      children.head
-    }
-    Count(child)
-  }
+  def apply(child: Expression): Count = Count(child :: Nil)
 }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
index 694a2a7c54a90..40b1eec63e551 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala
@@ -426,30 +426,3 @@ case class Greatest(children: Seq[Expression]) extends Expression {
   }
 }
 
-/** Operator that drops a row when it contains any nulls. */
-case class DropAnyNull(child: Expression) extends UnaryExpression with ExpectsInputTypes {
-  override def nullable: Boolean = true
-  override def dataType: DataType = child.dataType
-  override def inputTypes: Seq[AbstractDataType] = Seq(StructType)
-
-  protected override def nullSafeEval(input: Any): InternalRow = {
-    val row = input.asInstanceOf[InternalRow]
-    if (row.anyNull) {
-      null
-    } else {
-      row
-    }
-  }
-
-  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
-    nullSafeCodeGen(ctx, ev, eval => {
-      s"""
-        if ($eval.anyNull()) {
-          ${ev.isNull} = true;
-        } else {
-          ${ev.value} = $eval;
-        }
-      """
-    })
-  }
-}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 2901d8f2efddf..06d14fcf8b9c2 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -362,9 +362,14 @@ object LikeSimplification extends Rule[LogicalPlan] {
  * Null value propagation from bottom to top of the expression tree.
  */
 object NullPropagation extends Rule[LogicalPlan] {
+  def nonNullLiteral(e: Expression): Boolean = e match {
+    case Literal(null, _) => false
+    case _ => true
+  }
+
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
     case q: LogicalPlan => q transformExpressionsUp {
-      case e @ AggregateExpression(Count(Literal(null, _)), _, _) =>
+      case e @ AggregateExpression(Count(exprs), _, _) if !exprs.exists(nonNullLiteral) =>
         Cast(Literal(0L), e.dataType)
       case e @ IsNull(c) if !c.nullable => Literal.create(false, BooleanType)
       case e @ IsNotNull(c) if !c.nullable => Literal.create(true, BooleanType)
@@ -377,16 +382,13 @@ object NullPropagation extends Rule[LogicalPlan] {
         Literal.create(null, e.dataType)
       case e @ EqualNullSafe(Literal(null, _), r) => IsNull(r)
       case e @ EqualNullSafe(l, Literal(null, _)) => IsNull(l)
-      case e @ AggregateExpression(Count(expr), mode, false) if !expr.nullable =>
+      case e @ AggregateExpression(Count(exprs), mode, false) if !exprs.exists(_.nullable) =>
         // This rule should be only triggered when isDistinct field is false.
         AggregateExpression(Count(Literal(1)), mode, isDistinct = false)
 
       // For Coalesce, remove null literals.
       case e @ Coalesce(children) =>
-        val newChildren = children.filter {
-          case Literal(null, _) => false
-          case _ => true
-        }
+        val newChildren = children.filter(nonNullLiteral)
         if (newChildren.length == 0) {
           Literal.create(null, e.dataType)
         } else if (newChildren.length == 1) {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConditionalExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConditionalExpressionSuite.scala
index c1e3c17b87102..0df673bb9fa02 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConditionalExpressionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ConditionalExpressionSuite.scala
@@ -231,18 +231,4 @@ class ConditionalExpressionSuite extends SparkFunSuite with ExpressionEvalHelper
       checkConsistencyBetweenInterpretedAndCodegen(Greatest, dt, 2)
     }
   }
-
-  test("function dropAnyNull") {
-    val drop = DropAnyNull(CreateStruct(Seq('a.string.at(0), 'b.string.at(1))))
-    val a = create_row("a", "q")
-    val nullStr: String = null
-    checkEvaluation(drop, a, a)
-    checkEvaluation(drop, null, create_row("b", nullStr))
-    checkEvaluation(drop, null, create_row(nullStr, nullStr))
-
-    val row = 'r.struct(
-      StructField("a", StringType, false),
-      StructField("b", StringType, true)).at(0)
-    checkEvaluation(DropAnyNull(row), null, create_row(null))
-  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
index a70e41436c7aa..76b938cdb694e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/utils.scala
@@ -146,20 +146,16 @@ object Utils {
       aggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes))
 
     // functionsWithDistinct is guaranteed to be non-empty. Even though it may contain more than one
-    // DISTINCT aggregate function, all of those functions will have the same column expression.
+    // DISTINCT aggregate function, all of those functions will have the same column expressions.
     // For example, it would be valid for functionsWithDistinct to be
     // [COUNT(DISTINCT foo), MAX(DISTINCT foo)], but [COUNT(DISTINCT bar), COUNT(DISTINCT foo)] is
     // disallowed because those two distinct aggregates have different column expressions.
-    val distinctColumnExpression: Expression = {
-      val allDistinctColumnExpressions = functionsWithDistinct.head.aggregateFunction.children
-      assert(allDistinctColumnExpressions.length == 1)
-      allDistinctColumnExpressions.head
-    }
-    val namedDistinctColumnExpression: NamedExpression = distinctColumnExpression match {
+    val distinctColumnExpressions = functionsWithDistinct.head.aggregateFunction.children
+    val namedDistinctColumnExpressions = distinctColumnExpressions.map {
       case ne: NamedExpression => ne
       case other => Alias(other, other.toString)()
     }
-    val distinctColumnAttribute: Attribute = namedDistinctColumnExpression.toAttribute
+    val distinctColumnAttributes = namedDistinctColumnExpressions.map(_.toAttribute)
     val groupingAttributes = groupingExpressions.map(_.toAttribute)
 
     // 1. Create an Aggregate Operator for partial aggregations.
@@ -170,10 +166,11 @@ object Utils {
       // We will group by the original grouping expression, plus an additional expression for the
       // DISTINCT column. For example, for AVG(DISTINCT value) GROUP BY key, the grouping
       // expressions will be [key, value].
-      val partialAggregateGroupingExpressions = groupingExpressions :+ namedDistinctColumnExpression
+      val partialAggregateGroupingExpressions =
+        groupingExpressions ++ namedDistinctColumnExpressions
       val partialAggregateResult =
         groupingAttributes ++
-          Seq(distinctColumnAttribute) ++
+          distinctColumnAttributes ++
           partialAggregateExpressions.flatMap(_.aggregateFunction.inputAggBufferAttributes)
       if (usesTungstenAggregate) {
         TungstenAggregate(
@@ -208,28 +205,28 @@ object Utils {
         partialMergeAggregateExpressions.flatMap(_.aggregateFunction.aggBufferAttributes)
       val partialMergeAggregateResult =
         groupingAttributes ++
-          Seq(distinctColumnAttribute) ++
+          distinctColumnAttributes ++
           partialMergeAggregateExpressions.flatMap(_.aggregateFunction.inputAggBufferAttributes)
       if (usesTungstenAggregate) {
         TungstenAggregate(
           requiredChildDistributionExpressions = Some(groupingAttributes),
-          groupingExpressions = groupingAttributes :+ distinctColumnAttribute,
+          groupingExpressions = groupingAttributes ++ distinctColumnAttributes,
           nonCompleteAggregateExpressions = partialMergeAggregateExpressions,
           nonCompleteAggregateAttributes = partialMergeAggregateAttributes,
           completeAggregateExpressions = Nil,
           completeAggregateAttributes = Nil,
-          initialInputBufferOffset = (groupingAttributes :+ distinctColumnAttribute).length,
+          initialInputBufferOffset = (groupingAttributes ++ distinctColumnAttributes).length,
           resultExpressions = partialMergeAggregateResult,
           child = partialAggregate)
       } else {
         SortBasedAggregate(
           requiredChildDistributionExpressions = Some(groupingAttributes),
-          groupingExpressions = groupingAttributes :+ distinctColumnAttribute,
+          groupingExpressions = groupingAttributes ++ distinctColumnAttributes,
           nonCompleteAggregateExpressions = partialMergeAggregateExpressions,
           nonCompleteAggregateAttributes = partialMergeAggregateAttributes,
           completeAggregateExpressions = Nil,
           completeAggregateAttributes = Nil,
-          initialInputBufferOffset = (groupingAttributes :+ distinctColumnAttribute).length,
+          initialInputBufferOffset = (groupingAttributes ++ distinctColumnAttributes).length,
           resultExpressions = partialMergeAggregateResult,
           child = partialAggregate)
       }
@@ -244,14 +241,16 @@ object Utils {
         expr => aggregateFunctionToAttribute(expr.aggregateFunction, expr.isDistinct)
       }
 
+      val distinctColumnAttributeLookup =
+        distinctColumnExpressions.zip(distinctColumnAttributes).toMap
       val (completeAggregateExpressions, completeAggregateAttributes) = functionsWithDistinct.map {
         // Children of an AggregateFunction with DISTINCT keyword has already
         // been evaluated. At here, we need to replace original children
         // to AttributeReferences.
         case agg @ AggregateExpression(aggregateFunction, mode, true) =>
-          val rewrittenAggregateFunction = aggregateFunction.transformDown {
-            case expr if expr == distinctColumnExpression => distinctColumnAttribute
-          }.asInstanceOf[AggregateFunction]
+          val rewrittenAggregateFunction = aggregateFunction
+            .transformDown(distinctColumnAttributeLookup)
+            .asInstanceOf[AggregateFunction]
           // We rewrite the aggregate function to a non-distinct aggregation because
           // its input will have distinct arguments.
           // We just keep the isDistinct setting to true, so when users look at the query plan,
@@ -270,7 +269,7 @@ object Utils {
           nonCompleteAggregateAttributes = finalAggregateAttributes,
           completeAggregateExpressions = completeAggregateExpressions,
           completeAggregateAttributes = completeAggregateAttributes,
-          initialInputBufferOffset = (groupingAttributes :+ distinctColumnAttribute).length,
+          initialInputBufferOffset = (groupingAttributes ++ distinctColumnAttributes).length,
           resultExpressions = resultExpressions,
           child = partialMergeAggregate)
       } else {
@@ -281,7 +280,7 @@ object Utils {
           nonCompleteAggregateAttributes = finalAggregateAttributes,
           completeAggregateExpressions = completeAggregateExpressions,
           completeAggregateAttributes = completeAggregateAttributes,
-          initialInputBufferOffset = (groupingAttributes :+ distinctColumnAttribute).length,
+          initialInputBufferOffset = (groupingAttributes ++ distinctColumnAttributes).length,
           resultExpressions = resultExpressions,
           child = partialMergeAggregate)
       }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/expressions/WindowSpec.scala b/sql/core/src/main/scala/org/apache/spark/sql/expressions/WindowSpec.scala
index fc873c04f88f0..893e800a61438 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/expressions/WindowSpec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/expressions/WindowSpec.scala
@@ -152,8 +152,8 @@ class WindowSpec private[sql](
           case Sum(child) => WindowExpression(
             UnresolvedWindowFunction("sum", child :: Nil),
             WindowSpecDefinition(partitionSpec, orderSpec, frame))
-          case Count(child) => WindowExpression(
-            UnresolvedWindowFunction("count", child :: Nil),
+          case Count(children) => WindowExpression(
+            UnresolvedWindowFunction("count", children),
             WindowSpecDefinition(partitionSpec, orderSpec, frame))
           case First(child, ignoreNulls) => WindowExpression(
             // TODO this is a hack for Hive UDAF first_value

From 0ddfe7868948e302858a2b03b50762eaefbeb53e Mon Sep 17 00:00:00 2001
From: Yin Huai <yhuai@databricks.com>
Date: Sun, 29 Nov 2015 19:02:15 -0800
Subject: [PATCH 494/518] [SPARK-12039] [SQL] Ignore HiveSparkSubmitSuite's
 "SPARK-9757 Persist Parquet relation with decimal column".

https://issues.apache.org/jira/browse/SPARK-12039

Since it is pretty flaky in hadoop 1 tests, we can disable it while we are investigating the cause.

Author: Yin Huai <yhuai@databricks.com>

Closes #10035 from yhuai/SPARK-12039-ignore.
---
 .../scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala
index 24a3afee148c5..92962193311d4 100644
--- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala
+++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveSparkSubmitSuite.scala
@@ -101,7 +101,7 @@ class HiveSparkSubmitSuite
     runSparkSubmit(args)
   }
 
-  test("SPARK-9757 Persist Parquet relation with decimal column") {
+  ignore("SPARK-9757 Persist Parquet relation with decimal column") {
     val unusedJar = TestUtils.createJarWithClasses(Seq.empty)
     val args = Seq(
       "--class", SPARK_9757.getClass.getName.stripSuffix("$"),

From e0749442051d6e29dae4f4cdcb2937c0b015f98f Mon Sep 17 00:00:00 2001
From: toddwan <tawan0109@outlook.com>
Date: Mon, 30 Nov 2015 09:26:29 +0000
Subject: [PATCH 495/518] [SPARK-11859][MESOS] SparkContext accepts invalid
 Master URLs in the form zk://host:port for a multi-master Mesos cluster using
 ZooKeeper

* According to below doc and validation logic in [SparkSubmit.scala](https://github.com/apache/spark/blob/master/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala#L231), master URL for a mesos cluster should always start with `mesos://`

http://spark.apache.org/docs/latest/running-on-mesos.html
`The Master URLs for Mesos are in the form mesos://host:5050 for a single-master Mesos cluster, or mesos://zk://host:2181 for a multi-master Mesos cluster using ZooKeeper.`

* However, [SparkContext.scala](https://github.com/apache/spark/blob/master/core/src/main/scala/org/apache/spark/SparkContext.scala#L2749) fails the validation and can receive master URL in the form `zk://host:port`

* For the master URLs in the form `zk:host:port`, the valid form should be `mesos://zk://host:port`

* This PR restrict the validation in `SparkContext.scala`, and now only mesos master URLs prefixed with `mesos://` can be accepted.

* This PR also updated corresponding unit test.

Author: toddwan <tawan0109@outlook.com>

Closes #9886 from toddwan/S11859.
---
 .../scala/org/apache/spark/SparkContext.scala    | 16 ++++++++++------
 .../SparkContextSchedulerCreationSuite.scala     |  5 +++++
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index b030d3c71dc20..8a62b71c3fa68 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -2708,15 +2708,14 @@ object SparkContext extends Logging {
         scheduler.initialize(backend)
         (backend, scheduler)
 
-      case mesosUrl @ MESOS_REGEX(_) =>
+      case MESOS_REGEX(mesosUrl) =>
         MesosNativeLibrary.load()
         val scheduler = new TaskSchedulerImpl(sc)
         val coarseGrained = sc.conf.getBoolean("spark.mesos.coarse", defaultValue = true)
-        val url = mesosUrl.stripPrefix("mesos://") // strip scheme from raw Mesos URLs
         val backend = if (coarseGrained) {
-          new CoarseMesosSchedulerBackend(scheduler, sc, url, sc.env.securityManager)
+          new CoarseMesosSchedulerBackend(scheduler, sc, mesosUrl, sc.env.securityManager)
         } else {
-          new MesosSchedulerBackend(scheduler, sc, url)
+          new MesosSchedulerBackend(scheduler, sc, mesosUrl)
         }
         scheduler.initialize(backend)
         (backend, scheduler)
@@ -2727,6 +2726,11 @@ object SparkContext extends Logging {
         scheduler.initialize(backend)
         (backend, scheduler)
 
+      case zkUrl if zkUrl.startsWith("zk://") =>
+        logWarning("Master URL for a multi-master Mesos cluster managed by ZooKeeper should be " +
+          "in the form mesos://zk://host:port. Current Master URL will stop working in Spark 2.0.")
+        createTaskScheduler(sc, "mesos://" + zkUrl)
+
       case _ =>
         throw new SparkException("Could not parse Master URL: '" + master + "'")
     }
@@ -2745,8 +2749,8 @@ private object SparkMasterRegex {
   val LOCAL_CLUSTER_REGEX = """local-cluster\[\s*([0-9]+)\s*,\s*([0-9]+)\s*,\s*([0-9]+)\s*]""".r
   // Regular expression for connecting to Spark deploy clusters
   val SPARK_REGEX = """spark://(.*)""".r
-  // Regular expression for connection to Mesos cluster by mesos:// or zk:// url
-  val MESOS_REGEX = """(mesos|zk)://.*""".r
+  // Regular expression for connection to Mesos cluster by mesos:// or mesos://zk:// url
+  val MESOS_REGEX = """mesos://(.*)""".r
   // Regular expression for connection to Simr cluster
   val SIMR_REGEX = """simr://(.*)""".r
 }
diff --git a/core/src/test/scala/org/apache/spark/SparkContextSchedulerCreationSuite.scala b/core/src/test/scala/org/apache/spark/SparkContextSchedulerCreationSuite.scala
index e5a14a69ef05f..d18e0782c0392 100644
--- a/core/src/test/scala/org/apache/spark/SparkContextSchedulerCreationSuite.scala
+++ b/core/src/test/scala/org/apache/spark/SparkContextSchedulerCreationSuite.scala
@@ -175,6 +175,11 @@ class SparkContextSchedulerCreationSuite
   }
 
   test("mesos with zookeeper") {
+    testMesos("mesos://zk://localhost:1234,localhost:2345",
+      classOf[MesosSchedulerBackend], coarse = false)
+  }
+
+  test("mesos with zookeeper and Master URL starting with zk://") {
     testMesos("zk://localhost:1234,localhost:2345", classOf[MesosSchedulerBackend], coarse = false)
   }
 }

From 953e8e6dcb32cd88005834e9c3720740e201826c Mon Sep 17 00:00:00 2001
From: Prashant Sharma <scrapcodes@gmail.com>
Date: Mon, 30 Nov 2015 09:30:58 +0000
Subject: [PATCH 496/518] [MINOR][BUILD] Changed the comment to reflect the
 plugin project is there to support SBT pom reader only.

Author: Prashant Sharma <scrapcodes@gmail.com>

Closes #10012 from ScrapCodes/minor-build-comment.
---
 project/project/SparkPluginBuild.scala | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/project/project/SparkPluginBuild.scala b/project/project/SparkPluginBuild.scala
index 471d00bd8223f..cbb88dc7dd1dd 100644
--- a/project/project/SparkPluginBuild.scala
+++ b/project/project/SparkPluginBuild.scala
@@ -19,9 +19,8 @@ import sbt._
 import sbt.Keys._
 
 /**
- * This plugin project is there to define new scala style rules for spark. This is
- * a plugin project so that this gets compiled first and is put on the classpath and
- * becomes available for scalastyle sbt plugin.
+ * This plugin project is there because we use our custom fork of sbt-pom-reader plugin. This is
+ * a plugin project so that this gets compiled first and is available on the classpath for SBT build.
  */
 object SparkPluginDef extends Build {
   lazy val root = Project("plugins", file(".")) dependsOn(sbtPomReader)

From 26c3581f17f475fab2f3b5301b8f253ff2fa6438 Mon Sep 17 00:00:00 2001
From: Wieland Hoffmann <themineo@gmail.com>
Date: Mon, 30 Nov 2015 09:32:48 +0000
Subject: [PATCH 497/518] [DOC] Explicitly state that top maintains the order
 of elements

Top is implemented in terms of takeOrdered, which already maintains the
order, so top should, too.

Author: Wieland Hoffmann <themineo@gmail.com>

Closes #10013 from mineo/top-order.
---
 .../main/scala/org/apache/spark/api/java/JavaRDDLike.scala    | 4 ++--
 core/src/main/scala/org/apache/spark/rdd/RDD.scala            | 3 ++-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala b/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
index 871be0b1f39ea..1e9d4f1803a81 100644
--- a/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
+++ b/core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala
@@ -556,7 +556,7 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
 
   /**
    * Returns the top k (largest) elements from this RDD as defined by
-   * the specified Comparator[T].
+   * the specified Comparator[T] and maintains the order.
    * @param num k, the number of top elements to return
    * @param comp the comparator that defines the order
    * @return an array of top elements
@@ -567,7 +567,7 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
 
   /**
    * Returns the top k (largest) elements from this RDD using the
-   * natural ordering for T.
+   * natural ordering for T and maintains the order.
    * @param num k, the number of top elements to return
    * @return an array of top elements
    */
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index 2aeb5eeaad32c..8b3731d935788 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -1327,7 +1327,8 @@ abstract class RDD[T: ClassTag](
 
   /**
    * Returns the top k (largest) elements from this RDD as defined by the specified
-   * implicit Ordering[T]. This does the opposite of [[takeOrdered]]. For example:
+   * implicit Ordering[T] and maintains the ordering. This does the opposite of
+   * [[takeOrdered]]. For example:
    * {{{
    *   sc.parallelize(Seq(10, 4, 2, 12, 3)).top(1)
    *   // returns Array(12)

From bf0e85a70a54a2d7fd6804b6bd00c63c20e2bb00 Mon Sep 17 00:00:00 2001
From: Prashant Sharma <scrapcodes@gmail.com>
Date: Mon, 30 Nov 2015 10:11:27 +0000
Subject: [PATCH 498/518] [SPARK-12023][BUILD] Fix warnings while packaging
 spark with maven.

this is a trivial fix, discussed [here](http://stackoverflow.com/questions/28500401/maven-assembly-plugin-warning-the-assembly-descriptor-contains-a-filesystem-roo/).

Author: Prashant Sharma <scrapcodes@gmail.com>

Closes #10014 from ScrapCodes/assembly-warning.
---
 assembly/src/main/assembly/assembly.xml      | 8 ++++----
 external/mqtt/src/main/assembly/assembly.xml | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/assembly/src/main/assembly/assembly.xml b/assembly/src/main/assembly/assembly.xml
index 711156337b7c3..009d4b92f406c 100644
--- a/assembly/src/main/assembly/assembly.xml
+++ b/assembly/src/main/assembly/assembly.xml
@@ -32,7 +32,7 @@
       <directory>
         ${project.parent.basedir}/core/src/main/resources/org/apache/spark/ui/static/
       </directory>
-      <outputDirectory>/ui-resources/org/apache/spark/ui/static</outputDirectory>
+      <outputDirectory>ui-resources/org/apache/spark/ui/static</outputDirectory>
       <includes>
         <include>**/*</include>
       </includes>
@@ -41,7 +41,7 @@
       <directory>
         ${project.parent.basedir}/sbin/
       </directory>
-      <outputDirectory>/sbin</outputDirectory>
+      <outputDirectory>sbin</outputDirectory>
       <includes>
         <include>**/*</include>
       </includes>
@@ -50,7 +50,7 @@
       <directory>
         ${project.parent.basedir}/bin/
       </directory>
-      <outputDirectory>/bin</outputDirectory>
+      <outputDirectory>bin</outputDirectory>
       <includes>
         <include>**/*</include>
       </includes>
@@ -59,7 +59,7 @@
       <directory>
         ${project.parent.basedir}/assembly/target/${spark.jar.dir}
       </directory>
-      <outputDirectory>/</outputDirectory>
+      <outputDirectory></outputDirectory>
       <includes>
         <include>${spark.jar.basename}</include>
       </includes>
diff --git a/external/mqtt/src/main/assembly/assembly.xml b/external/mqtt/src/main/assembly/assembly.xml
index ecab5b360eb3e..c110b01b34e10 100644
--- a/external/mqtt/src/main/assembly/assembly.xml
+++ b/external/mqtt/src/main/assembly/assembly.xml
@@ -24,7 +24,7 @@
   <fileSets>
     <fileSet>
       <directory>${project.build.directory}/scala-${scala.binary.version}/test-classes</directory>
-      <outputDirectory>/</outputDirectory>
+      <outputDirectory></outputDirectory>
     </fileSet>
   </fileSets>
 

From 2db4662fe2d72749c06ad5f11f641a388343f77c Mon Sep 17 00:00:00 2001
From: CK50 <christian.kurz@oracle.com>
Date: Mon, 30 Nov 2015 20:08:49 +0800
Subject: [PATCH 499/518] [SPARK-11989][SQL] Only use commit in JDBC data
 source if the underlying database supports transactions

Fixes [SPARK-11989](https://issues.apache.org/jira/browse/SPARK-11989)

Author: CK50 <christian.kurz@oracle.com>
Author: Christian Kurz <christian.kurz@oracle.com>

Closes #9973 from CK50/branch-1.6_non-transactional.

(cherry picked from commit a589736a1b237ef2f3bd59fbaeefe143ddcc8f4e)
Signed-off-by: Reynold Xin <rxin@databricks.com>
---
 .../datasources/jdbc/JdbcUtils.scala          | 22 ++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
index 7375a5c09123f..252f1cfd5d9c5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
@@ -21,6 +21,7 @@ import java.sql.{Connection, PreparedStatement}
 import java.util.Properties
 
 import scala.util.Try
+import scala.util.control.NonFatal
 
 import org.apache.spark.Logging
 import org.apache.spark.sql.jdbc.{JdbcDialect, JdbcType, JdbcDialects}
@@ -125,8 +126,19 @@ object JdbcUtils extends Logging {
       dialect: JdbcDialect): Iterator[Byte] = {
     val conn = getConnection()
     var committed = false
+    val supportsTransactions = try {
+      conn.getMetaData().supportsDataManipulationTransactionsOnly() ||
+      conn.getMetaData().supportsDataDefinitionAndDataManipulationTransactions()
+    } catch {
+      case NonFatal(e) =>
+        logWarning("Exception while detecting transaction support", e)
+        true
+    }
+
     try {
-      conn.setAutoCommit(false) // Everything in the same db transaction.
+      if (supportsTransactions) {
+        conn.setAutoCommit(false) // Everything in the same db transaction.
+      }
       val stmt = insertStatement(conn, table, rddSchema)
       try {
         var rowCount = 0
@@ -175,14 +187,18 @@ object JdbcUtils extends Logging {
       } finally {
         stmt.close()
       }
-      conn.commit()
+      if (supportsTransactions) {
+        conn.commit()
+      }
       committed = true
     } finally {
       if (!committed) {
         // The stage must fail.  We got here through an exception path, so
         // let the exception through unless rollback() or close() want to
         // tell the user about another problem.
-        conn.rollback()
+        if (supportsTransactions) {
+          conn.rollback()
+        }
         conn.close()
       } else {
         // The stage must succeed.  We cannot propagate any exception close() might throw.

From 17275fa99c670537c52843df405279a52b5c9594 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Mon, 30 Nov 2015 10:32:13 -0800
Subject: [PATCH 500/518] [SPARK-11700] [SQL] Remove thread local SQLContext 
 in SparkPlan

In 1.6, we introduce a public API to have a SQLContext for current thread, SparkPlan should use that.

Author: Davies Liu <davies@databricks.com>

Closes #9990 from davies/leak_context.
---
 .../scala/org/apache/spark/sql/SQLContext.scala    | 10 +++++-----
 .../spark/sql/execution/QueryExecution.scala       |  3 +--
 .../org/apache/spark/sql/execution/SparkPlan.scala | 14 ++++----------
 .../apache/spark/sql/MultiSQLContextsSuite.scala   |  2 +-
 .../sql/execution/ExchangeCoordinatorSuite.scala   |  2 +-
 .../sql/execution/RowFormatConvertersSuite.scala   |  4 ++--
 6 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index 1c2ac5f6f11bf..8d2783952532a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -26,7 +26,6 @@ import scala.collection.immutable
 import scala.reflect.runtime.universe.TypeTag
 import scala.util.control.NonFatal
 
-import org.apache.spark.{SparkException, SparkContext}
 import org.apache.spark.annotation.{DeveloperApi, Experimental}
 import org.apache.spark.api.java.{JavaRDD, JavaSparkContext}
 import org.apache.spark.rdd.RDD
@@ -45,9 +44,10 @@ import org.apache.spark.sql.execution.datasources._
 import org.apache.spark.sql.execution.ui.{SQLListener, SQLTab}
 import org.apache.spark.sql.sources.BaseRelation
 import org.apache.spark.sql.types._
-import org.apache.spark.sql.{execution => sparkexecution}
 import org.apache.spark.sql.util.ExecutionListenerManager
+import org.apache.spark.sql.{execution => sparkexecution}
 import org.apache.spark.util.Utils
+import org.apache.spark.{SparkContext, SparkException}
 
 /**
  * The entry point for working with structured data (rows and columns) in Spark.  Allows the
@@ -401,7 +401,7 @@ class SQLContext private[sql](
    */
   @Experimental
   def createDataFrame[A <: Product : TypeTag](rdd: RDD[A]): DataFrame = {
-    SparkPlan.currentContext.set(self)
+    SQLContext.setActive(self)
     val schema = ScalaReflection.schemaFor[A].dataType.asInstanceOf[StructType]
     val attributeSeq = schema.toAttributes
     val rowRDD = RDDConversions.productToRowRdd(rdd, schema.map(_.dataType))
@@ -417,7 +417,7 @@ class SQLContext private[sql](
    */
   @Experimental
   def createDataFrame[A <: Product : TypeTag](data: Seq[A]): DataFrame = {
-    SparkPlan.currentContext.set(self)
+    SQLContext.setActive(self)
     val schema = ScalaReflection.schemaFor[A].dataType.asInstanceOf[StructType]
     val attributeSeq = schema.toAttributes
     DataFrame(self, LocalRelation.fromProduct(attributeSeq, data))
@@ -1334,7 +1334,7 @@ object SQLContext {
     activeContext.remove()
   }
 
-  private[sql] def getActiveContextOption(): Option[SQLContext] = {
+  private[sql] def getActive(): Option[SQLContext] = {
     Option(activeContext.get())
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index 5da5aea17e25b..107570f9dbcc8 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
@@ -42,9 +42,8 @@ class QueryExecution(val sqlContext: SQLContext, val logical: LogicalPlan) {
 
   lazy val optimizedPlan: LogicalPlan = sqlContext.optimizer.execute(withCachedData)
 
-  // TODO: Don't just pick the first one...
   lazy val sparkPlan: SparkPlan = {
-    SparkPlan.currentContext.set(sqlContext)
+    SQLContext.setActive(sqlContext)
     sqlContext.planner.plan(optimizedPlan).next()
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
index 534a3bcb8364d..507641ff8263e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
@@ -23,21 +23,15 @@ import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark.Logging
 import org.apache.spark.rdd.{RDD, RDDOperationScope}
-import org.apache.spark.sql.SQLContext
-import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.CatalystTypeConverters
+import org.apache.spark.sql.{Row, SQLContext}
+import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow}
 import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.physical._
 import org.apache.spark.sql.execution.metric.{LongSQLMetric, SQLMetric}
 import org.apache.spark.sql.types.DataType
 
-object SparkPlan {
-  protected[sql] val currentContext = new ThreadLocal[SQLContext]()
-}
-
 /**
  * The base class for physical operators.
  */
@@ -49,7 +43,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
    * populated by the query planning infrastructure.
    */
   @transient
-  protected[spark] final val sqlContext = SparkPlan.currentContext.get()
+  protected[spark] final val sqlContext = SQLContext.getActive().get
 
   protected def sparkContext = sqlContext.sparkContext
 
@@ -69,7 +63,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
 
   /** Overridden make copy also propogates sqlContext to copied plan. */
   override def makeCopy(newArgs: Array[AnyRef]): SparkPlan = {
-    SparkPlan.currentContext.set(sqlContext)
+    SQLContext.setActive(sqlContext)
     super.makeCopy(newArgs)
   }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/MultiSQLContextsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/MultiSQLContextsSuite.scala
index 34c5c68fd1c18..162c0b56c6e11 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/MultiSQLContextsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/MultiSQLContextsSuite.scala
@@ -27,7 +27,7 @@ class MultiSQLContextsSuite extends SparkFunSuite with BeforeAndAfterAll {
   private var sparkConf: SparkConf = _
 
   override protected def beforeAll(): Unit = {
-    originalActiveSQLContext = SQLContext.getActiveContextOption()
+    originalActiveSQLContext = SQLContext.getActive()
     originalInstantiatedSQLContext = SQLContext.getInstantiatedContextOption()
 
     SQLContext.clearActive()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala
index b96d50a70b85c..180050bdac00f 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeCoordinatorSuite.scala
@@ -30,7 +30,7 @@ class ExchangeCoordinatorSuite extends SparkFunSuite with BeforeAndAfterAll {
   private var originalInstantiatedSQLContext: Option[SQLContext] = _
 
   override protected def beforeAll(): Unit = {
-    originalActiveSQLContext = SQLContext.getActiveContextOption()
+    originalActiveSQLContext = SQLContext.getActive()
     originalInstantiatedSQLContext = SQLContext.getInstantiatedContextOption()
 
     SQLContext.clearActive()
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/RowFormatConvertersSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/RowFormatConvertersSuite.scala
index 6876ab0f02b10..13d68a103a225 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/RowFormatConvertersSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/RowFormatConvertersSuite.scala
@@ -18,7 +18,7 @@
 package org.apache.spark.sql.execution
 
 import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.Row
+import org.apache.spark.sql.{SQLContext, Row}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Attribute, Literal, IsNull}
 import org.apache.spark.sql.catalyst.util.GenericArrayData
@@ -94,7 +94,7 @@ class RowFormatConvertersSuite extends SparkPlanTest with SharedSQLContext {
   }
 
   test("SPARK-9683: copy UTF8String when convert unsafe array/map to safe") {
-    SparkPlan.currentContext.set(sqlContext)
+    SQLContext.setActive(sqlContext)
     val schema = ArrayType(StringType)
     val rows = (1 to 100).map { i =>
       InternalRow(new GenericArrayData(Array[Any](UTF8String.fromString(i.toString))))

From 8df584b0200402d8b2ce0a8de24f7a760ced8655 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies@databricks.com>
Date: Mon, 30 Nov 2015 11:54:18 -0800
Subject: [PATCH 501/518] [SPARK-11982] [SQL] improve performance of cartesian
 product

This PR improve the performance of CartesianProduct by caching the result of right plan.

After this patch, the query time of TPC-DS Q65 go down to 4 seconds from 28 minutes (420X faster).

cc nongli

Author: Davies Liu <davies@databricks.com>

Closes #9969 from davies/improve_cartesian.
---
 .../unsafe/sort/UnsafeExternalSorter.java     | 63 +++++++++++++++
 .../unsafe/sort/UnsafeInMemorySorter.java     |  7 ++
 .../execution/joins/CartesianProduct.scala    | 76 +++++++++++++++++--
 .../execution/metric/SQLMetricsSuite.scala    |  2 +-
 4 files changed, 139 insertions(+), 9 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
index 2e40312674737..5a97f4f11340c 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeExternalSorter.java
@@ -21,6 +21,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.util.LinkedList;
+import java.util.Queue;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
@@ -521,4 +522,66 @@ public long getKeyPrefix() {
       return upstream.getKeyPrefix();
     }
   }
+
+  /**
+   * Returns a iterator, which will return the rows in the order as inserted.
+   *
+   * It is the caller's responsibility to call `cleanupResources()`
+   * after consuming this iterator.
+   */
+  public UnsafeSorterIterator getIterator() throws IOException {
+    if (spillWriters.isEmpty()) {
+      assert(inMemSorter != null);
+      return inMemSorter.getIterator();
+    } else {
+      LinkedList<UnsafeSorterIterator> queue = new LinkedList<>();
+      for (UnsafeSorterSpillWriter spillWriter : spillWriters) {
+        queue.add(spillWriter.getReader(blockManager));
+      }
+      if (inMemSorter != null) {
+        queue.add(inMemSorter.getIterator());
+      }
+      return new ChainedIterator(queue);
+    }
+  }
+
+  /**
+   * Chain multiple UnsafeSorterIterator together as single one.
+   */
+  class ChainedIterator extends UnsafeSorterIterator {
+
+    private final Queue<UnsafeSorterIterator> iterators;
+    private UnsafeSorterIterator current;
+
+    public ChainedIterator(Queue<UnsafeSorterIterator> iterators) {
+      assert iterators.size() > 0;
+      this.iterators = iterators;
+      this.current = iterators.remove();
+    }
+
+    @Override
+    public boolean hasNext() {
+      while (!current.hasNext() && !iterators.isEmpty()) {
+        current = iterators.remove();
+      }
+      return current.hasNext();
+    }
+
+    @Override
+    public void loadNext() throws IOException {
+      current.loadNext();
+    }
+
+    @Override
+    public Object getBaseObject() { return current.getBaseObject(); }
+
+    @Override
+    public long getBaseOffset() { return current.getBaseOffset(); }
+
+    @Override
+    public int getRecordLength() { return current.getRecordLength(); }
+
+    @Override
+    public long getKeyPrefix() { return current.getKeyPrefix(); }
+  }
 }
diff --git a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
index dce1f15a2963c..c91e88f31bf9b 100644
--- a/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/util/collection/unsafe/sort/UnsafeInMemorySorter.java
@@ -226,4 +226,11 @@ public SortedIterator getSortedIterator() {
     sorter.sort(array, 0, pos / 2, sortComparator);
     return new SortedIterator(pos / 2);
   }
+
+  /**
+   * Returns an iterator over record pointers in original order (inserted).
+   */
+  public SortedIterator getIterator() {
+    return new SortedIterator(pos / 2);
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/CartesianProduct.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/CartesianProduct.scala
index f467519b802a7..fa2bc7672131c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/CartesianProduct.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/CartesianProduct.scala
@@ -17,16 +17,75 @@
 
 package org.apache.spark.sql.execution.joins
 
-import org.apache.spark.rdd.RDD
+import org.apache.spark._
+import org.apache.spark.rdd.{CartesianPartition, CartesianRDD, RDD}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Attribute, JoinedRow}
-import org.apache.spark.sql.execution.{BinaryNode, SparkPlan}
+import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeRowJoiner
+import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeRow}
 import org.apache.spark.sql.execution.metric.SQLMetrics
+import org.apache.spark.sql.execution.{BinaryNode, SparkPlan}
+import org.apache.spark.util.CompletionIterator
+import org.apache.spark.util.collection.unsafe.sort.UnsafeExternalSorter
+
+
+/**
+  * An optimized CartesianRDD for UnsafeRow, which will cache the rows from second child RDD,
+  * will be much faster than building the right partition for every row in left RDD, it also
+  * materialize the right RDD (in case of the right RDD is nondeterministic).
+  */
+private[spark]
+class UnsafeCartesianRDD(left : RDD[UnsafeRow], right : RDD[UnsafeRow], numFieldsOfRight: Int)
+  extends CartesianRDD[UnsafeRow, UnsafeRow](left.sparkContext, left, right) {
+
+  override def compute(split: Partition, context: TaskContext): Iterator[(UnsafeRow, UnsafeRow)] = {
+    // We will not sort the rows, so prefixComparator and recordComparator are null.
+    val sorter = UnsafeExternalSorter.create(
+      context.taskMemoryManager(),
+      SparkEnv.get.blockManager,
+      context,
+      null,
+      null,
+      1024,
+      SparkEnv.get.memoryManager.pageSizeBytes)
+
+    val partition = split.asInstanceOf[CartesianPartition]
+    for (y <- rdd2.iterator(partition.s2, context)) {
+      sorter.insertRecord(y.getBaseObject, y.getBaseOffset, y.getSizeInBytes, 0)
+    }
+
+    // Create an iterator from sorter and wrapper it as Iterator[UnsafeRow]
+    def createIter(): Iterator[UnsafeRow] = {
+      val iter = sorter.getIterator
+      val unsafeRow = new UnsafeRow
+      new Iterator[UnsafeRow] {
+        override def hasNext: Boolean = {
+          iter.hasNext
+        }
+        override def next(): UnsafeRow = {
+          iter.loadNext()
+          unsafeRow.pointTo(iter.getBaseObject, iter.getBaseOffset, numFieldsOfRight,
+            iter.getRecordLength)
+          unsafeRow
+        }
+      }
+    }
+
+    val resultIter =
+      for (x <- rdd1.iterator(partition.s1, context);
+           y <- createIter()) yield (x, y)
+    CompletionIterator[(UnsafeRow, UnsafeRow), Iterator[(UnsafeRow, UnsafeRow)]](
+      resultIter, sorter.cleanupResources)
+  }
+}
 
 
 case class CartesianProduct(left: SparkPlan, right: SparkPlan) extends BinaryNode {
   override def output: Seq[Attribute] = left.output ++ right.output
 
+  override def canProcessSafeRows: Boolean = false
+  override def canProcessUnsafeRows: Boolean = true
+  override def outputsUnsafeRows: Boolean = true
+
   override private[sql] lazy val metrics = Map(
     "numLeftRows" -> SQLMetrics.createLongMetric(sparkContext, "number of left rows"),
     "numRightRows" -> SQLMetrics.createLongMetric(sparkContext, "number of right rows"),
@@ -39,18 +98,19 @@ case class CartesianProduct(left: SparkPlan, right: SparkPlan) extends BinaryNod
 
     val leftResults = left.execute().map { row =>
       numLeftRows += 1
-      row.copy()
+      row.asInstanceOf[UnsafeRow]
     }
     val rightResults = right.execute().map { row =>
       numRightRows += 1
-      row.copy()
+      row.asInstanceOf[UnsafeRow]
     }
 
-    leftResults.cartesian(rightResults).mapPartitionsInternal { iter =>
-      val joinedRow = new JoinedRow
+    val pair = new UnsafeCartesianRDD(leftResults, rightResults, right.output.size)
+    pair.mapPartitionsInternal { iter =>
+      val joiner = GenerateUnsafeRowJoiner.create(left.schema, right.schema)
       iter.map { r =>
         numOutputRows += 1
-        joinedRow(r._1, r._2)
+        joiner.join(r._1, r._2)
       }
     }
   }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index ebfa1eaf3e5bc..4f2cad19bfb6b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -317,7 +317,7 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
       testSparkPlanMetrics(df, 1, Map(
         1L -> ("CartesianProduct", Map(
           "number of left rows" -> 12L, // left needs to be scanned twice
-          "number of right rows" -> 12L, // right is read 6 times
+          "number of right rows" -> 4L, // right is read twice
           "number of output rows" -> 12L)))
       )
     }

From f2fbfa444f6e8d27953ec2d1c0b3abd603c963f9 Mon Sep 17 00:00:00 2001
From: BenFradet <benjamin.fradet@gmail.com>
Date: Mon, 30 Nov 2015 13:02:08 -0800
Subject: [PATCH 502/518] [MINOR][DOCS] fixed list display in ml-ensembles

The list in ml-ensembles.md wasn't properly formatted and, as a result, was looking like this:
![old](http://i.imgur.com/2ZhELLR.png)

This PR aims to make it look like this:
![new](http://i.imgur.com/0Xriwd2.png)

Author: BenFradet <benjamin.fradet@gmail.com>

Closes #10025 from BenFradet/ml-ensembles-doc.
---
 docs/ml-ensembles.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/ml-ensembles.md b/docs/ml-ensembles.md
index f6c3c30d5334f..14fef76f260ff 100644
--- a/docs/ml-ensembles.md
+++ b/docs/ml-ensembles.md
@@ -20,6 +20,7 @@ Both use [MLlib decision trees](ml-decision-tree.html) as their base models.
 Users can find more information about ensemble algorithms in the [MLlib Ensemble guide](mllib-ensembles.html).  In this section, we demonstrate the Pipelines API for ensembles.
 
 The main differences between this API and the [original MLlib ensembles API](mllib-ensembles.html) are:
+
 * support for ML Pipelines
 * separation of classification vs. regression
 * use of DataFrame metadata to distinguish continuous and categorical features

From 2c5dee0fb8e4d1734ea3a0f22e0b5bfd2f6dba46 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Mon, 30 Nov 2015 13:41:52 -0800
Subject: [PATCH 503/518] Revert "[SPARK-11206] Support SQL UI on the history
 server"

This reverts commit cc243a079b1c039d6e7f0b410d1654d94a090e14 / PR #9297

I'm reverting this because it broke SQLListenerMemoryLeakSuite in the master Maven builds.

See #9991 for a discussion of why this broke the tests.
---
 .rat-excludes                                 |   1 -
 .../org/apache/spark/JavaSparkListener.java   |   3 -
 .../apache/spark/SparkFirehoseListener.java   |   4 -
 .../scheduler/EventLoggingListener.scala      |   4 -
 .../spark/scheduler/SparkListener.scala       |  24 +--
 .../spark/scheduler/SparkListenerBus.scala    |   1 -
 .../scala/org/apache/spark/ui/SparkUI.scala   |  16 +-
 .../org/apache/spark/util/JsonProtocol.scala  |  11 +-
 ...park.scheduler.SparkHistoryListenerFactory |   1 -
 .../org/apache/spark/sql/SQLContext.scala     |  18 +--
 .../spark/sql/execution/SQLExecution.scala    |  24 ++-
 .../spark/sql/execution/SparkPlanInfo.scala   |  46 ------
 .../sql/execution/metric/SQLMetricInfo.scala  |  30 ----
 .../sql/execution/metric/SQLMetrics.scala     |  56 +++----
 .../sql/execution/ui/ExecutionPage.scala      |   4 +-
 .../spark/sql/execution/ui/SQLListener.scala  | 139 ++++++------------
 .../spark/sql/execution/ui/SQLTab.scala       |  12 +-
 .../sql/execution/ui/SparkPlanGraph.scala     |  20 +--
 .../execution/metric/SQLMetricsSuite.scala    |   4 +-
 .../sql/execution/ui/SQLListenerSuite.scala   |  43 +++---
 .../spark/sql/test/SharedSQLContext.scala     |   1 -
 21 files changed, 135 insertions(+), 327 deletions(-)
 delete mode 100644 sql/core/src/main/resources/META-INF/services/org.apache.spark.scheduler.SparkHistoryListenerFactory
 delete mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
 delete mode 100644 sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetricInfo.scala

diff --git a/.rat-excludes b/.rat-excludes
index 7262c960ed6bb..08fba6d351d6a 100644
--- a/.rat-excludes
+++ b/.rat-excludes
@@ -82,5 +82,4 @@ INDEX
 gen-java.*
 .*avpr
 org.apache.spark.sql.sources.DataSourceRegister
-org.apache.spark.scheduler.SparkHistoryListenerFactory
 .*parquet
diff --git a/core/src/main/java/org/apache/spark/JavaSparkListener.java b/core/src/main/java/org/apache/spark/JavaSparkListener.java
index 23bc9a2e81727..fa9acf0a15b88 100644
--- a/core/src/main/java/org/apache/spark/JavaSparkListener.java
+++ b/core/src/main/java/org/apache/spark/JavaSparkListener.java
@@ -82,7 +82,4 @@ public void onExecutorRemoved(SparkListenerExecutorRemoved executorRemoved) { }
   @Override
   public void onBlockUpdated(SparkListenerBlockUpdated blockUpdated) { }
 
-  @Override
-  public void onOtherEvent(SparkListenerEvent event) { }
-
 }
diff --git a/core/src/main/java/org/apache/spark/SparkFirehoseListener.java b/core/src/main/java/org/apache/spark/SparkFirehoseListener.java
index e6b24afd88ad4..1214d05ba6063 100644
--- a/core/src/main/java/org/apache/spark/SparkFirehoseListener.java
+++ b/core/src/main/java/org/apache/spark/SparkFirehoseListener.java
@@ -118,8 +118,4 @@ public void onBlockUpdated(SparkListenerBlockUpdated blockUpdated) {
         onEvent(blockUpdated);
     }
 
-    @Override
-    public void onOtherEvent(SparkListenerEvent event) {
-        onEvent(event);
-    }
 }
diff --git a/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala b/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala
index eaa07acc5132e..000a021a528cf 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala
@@ -207,10 +207,6 @@ private[spark] class EventLoggingListener(
   // No-op because logging every update would be overkill
   override def onExecutorMetricsUpdate(event: SparkListenerExecutorMetricsUpdate): Unit = { }
 
-  override def onOtherEvent(event: SparkListenerEvent): Unit = {
-    logEvent(event, flushLogger = true)
-  }
-
   /**
    * Stop logging events. The event log file will be renamed so that it loses the
    * ".inprogress" suffix.
diff --git a/core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala b/core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala
index 075a7f13172de..896f1743332f1 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/SparkListener.scala
@@ -22,19 +22,15 @@ import java.util.Properties
 import scala.collection.Map
 import scala.collection.mutable
 
-import com.fasterxml.jackson.annotation.JsonTypeInfo
-
-import org.apache.spark.{Logging, SparkConf, TaskEndReason}
+import org.apache.spark.{Logging, TaskEndReason}
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.scheduler.cluster.ExecutorInfo
 import org.apache.spark.storage.{BlockManagerId, BlockUpdatedInfo}
 import org.apache.spark.util.{Distribution, Utils}
-import org.apache.spark.ui.SparkUI
 
 @DeveloperApi
-@JsonTypeInfo(use = JsonTypeInfo.Id.CLASS, include = JsonTypeInfo.As.PROPERTY, property = "Event")
-trait SparkListenerEvent
+sealed trait SparkListenerEvent
 
 @DeveloperApi
 case class SparkListenerStageSubmitted(stageInfo: StageInfo, properties: Properties = null)
@@ -134,17 +130,6 @@ case class SparkListenerApplicationEnd(time: Long) extends SparkListenerEvent
  */
 private[spark] case class SparkListenerLogStart(sparkVersion: String) extends SparkListenerEvent
 
-/**
- * Interface for creating history listeners defined in other modules like SQL, which are used to
- * rebuild the history UI.
- */
-private[spark] trait SparkHistoryListenerFactory {
-  /**
-   * Create listeners used to rebuild the history UI.
-   */
-  def createListeners(conf: SparkConf, sparkUI: SparkUI): Seq[SparkListener]
-}
-
 /**
  * :: DeveloperApi ::
  * Interface for listening to events from the Spark scheduler. Note that this is an internal
@@ -238,11 +223,6 @@ trait SparkListener {
    * Called when the driver receives a block update info.
    */
   def onBlockUpdated(blockUpdated: SparkListenerBlockUpdated) { }
-
-  /**
-   * Called when other events like SQL-specific events are posted.
-   */
-  def onOtherEvent(event: SparkListenerEvent) { }
 }
 
 /**
diff --git a/core/src/main/scala/org/apache/spark/scheduler/SparkListenerBus.scala b/core/src/main/scala/org/apache/spark/scheduler/SparkListenerBus.scala
index 95722a07144ec..04afde33f5aad 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/SparkListenerBus.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/SparkListenerBus.scala
@@ -61,7 +61,6 @@ private[spark] trait SparkListenerBus extends ListenerBus[SparkListener, SparkLi
       case blockUpdated: SparkListenerBlockUpdated =>
         listener.onBlockUpdated(blockUpdated)
       case logStart: SparkListenerLogStart => // ignore event log metadata
-      case _ => listener.onOtherEvent(event)
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/ui/SparkUI.scala b/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
index 8da6884a38535..4608bce202ec8 100644
--- a/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
+++ b/core/src/main/scala/org/apache/spark/ui/SparkUI.scala
@@ -17,13 +17,10 @@
 
 package org.apache.spark.ui
 
-import java.util.{Date, ServiceLoader}
-
-import scala.collection.JavaConverters._
+import java.util.Date
 
 import org.apache.spark.status.api.v1.{ApiRootResource, ApplicationAttemptInfo, ApplicationInfo,
   UIRoot}
-import org.apache.spark.util.Utils
 import org.apache.spark.{Logging, SecurityManager, SparkConf, SparkContext}
 import org.apache.spark.scheduler._
 import org.apache.spark.storage.StorageStatusListener
@@ -157,16 +154,7 @@ private[spark] object SparkUI {
       appName: String,
       basePath: String,
       startTime: Long): SparkUI = {
-    val sparkUI = create(
-      None, conf, listenerBus, securityManager, appName, basePath, startTime = startTime)
-
-    val listenerFactories = ServiceLoader.load(classOf[SparkHistoryListenerFactory],
-      Utils.getContextOrSparkClassLoader).asScala
-    listenerFactories.foreach { listenerFactory =>
-      val listeners = listenerFactory.createListeners(conf, sparkUI)
-      listeners.foreach(listenerBus.addListener)
-    }
-    sparkUI
+    create(None, conf, listenerBus, securityManager, appName, basePath, startTime = startTime)
   }
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
index 7f5d713ec6505..c9beeb25e05af 100644
--- a/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
+++ b/core/src/main/scala/org/apache/spark/util/JsonProtocol.scala
@@ -19,21 +19,19 @@ package org.apache.spark.util
 
 import java.util.{Properties, UUID}
 
+import org.apache.spark.scheduler.cluster.ExecutorInfo
+
 import scala.collection.JavaConverters._
 import scala.collection.Map
 
-import com.fasterxml.jackson.databind.ObjectMapper
-import com.fasterxml.jackson.module.scala.DefaultScalaModule
 import org.json4s.DefaultFormats
 import org.json4s.JsonDSL._
 import org.json4s.JsonAST._
-import org.json4s.jackson.JsonMethods._
 
 import org.apache.spark._
 import org.apache.spark.executor._
 import org.apache.spark.rdd.RDDOperationScope
 import org.apache.spark.scheduler._
-import org.apache.spark.scheduler.cluster.ExecutorInfo
 import org.apache.spark.storage._
 
 /**
@@ -56,8 +54,6 @@ private[spark] object JsonProtocol {
 
   private implicit val format = DefaultFormats
 
-  private val mapper = new ObjectMapper().registerModule(DefaultScalaModule)
-
   /** ------------------------------------------------- *
    * JSON serialization methods for SparkListenerEvents |
    * -------------------------------------------------- */
@@ -100,7 +96,6 @@ private[spark] object JsonProtocol {
         executorMetricsUpdateToJson(metricsUpdate)
       case blockUpdated: SparkListenerBlockUpdated =>
         throw new MatchError(blockUpdated)  // TODO(ekl) implement this
-      case _ => parse(mapper.writeValueAsString(event))
     }
   }
 
@@ -516,8 +511,6 @@ private[spark] object JsonProtocol {
       case `executorRemoved` => executorRemovedFromJson(json)
       case `logStart` => logStartFromJson(json)
       case `metricsUpdate` => executorMetricsUpdateFromJson(json)
-      case other => mapper.readValue(compact(render(json)), Utils.classForName(other))
-        .asInstanceOf[SparkListenerEvent]
     }
   }
 
diff --git a/sql/core/src/main/resources/META-INF/services/org.apache.spark.scheduler.SparkHistoryListenerFactory b/sql/core/src/main/resources/META-INF/services/org.apache.spark.scheduler.SparkHistoryListenerFactory
deleted file mode 100644
index 507100be90967..0000000000000
--- a/sql/core/src/main/resources/META-INF/services/org.apache.spark.scheduler.SparkHistoryListenerFactory
+++ /dev/null
@@ -1 +0,0 @@
-org.apache.spark.sql.execution.ui.SQLHistoryListenerFactory
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
index 8d2783952532a..9cc65de19180a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala
@@ -1263,8 +1263,6 @@ object SQLContext {
    */
   @transient private val instantiatedContext = new AtomicReference[SQLContext]()
 
-  @transient private val sqlListener = new AtomicReference[SQLListener]()
-
   /**
    * Get the singleton SQLContext if it exists or create a new one using the given SparkContext.
    *
@@ -1309,10 +1307,6 @@ object SQLContext {
     Option(instantiatedContext.get())
   }
 
-  private[sql] def clearSqlListener(): Unit = {
-    sqlListener.set(null)
-  }
-
   /**
    * Changes the SQLContext that will be returned in this thread and its children when
    * SQLContext.getOrCreate() is called. This can be used to ensure that a given thread receives
@@ -1361,13 +1355,9 @@ object SQLContext {
    * Create a SQLListener then add it into SparkContext, and create an SQLTab if there is SparkUI.
    */
   private[sql] def createListenerAndUI(sc: SparkContext): SQLListener = {
-    if (sqlListener.get() == null) {
-      val listener = new SQLListener(sc.conf)
-      if (sqlListener.compareAndSet(null, listener)) {
-        sc.addSparkListener(listener)
-        sc.ui.foreach(new SQLTab(listener, _))
-      }
-    }
-    sqlListener.get()
+    val listener = new SQLListener(sc.conf)
+    sc.addSparkListener(listener)
+    sc.ui.foreach(new SQLTab(listener, _))
+    listener
   }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
index 34971986261c2..1422e15549c94 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
@@ -21,8 +21,7 @@ import java.util.concurrent.atomic.AtomicLong
 
 import org.apache.spark.SparkContext
 import org.apache.spark.sql.SQLContext
-import org.apache.spark.sql.execution.ui.{SparkListenerSQLExecutionStart,
-  SparkListenerSQLExecutionEnd}
+import org.apache.spark.sql.execution.ui.SparkPlanGraph
 import org.apache.spark.util.Utils
 
 private[sql] object SQLExecution {
@@ -46,14 +45,25 @@ private[sql] object SQLExecution {
       sc.setLocalProperty(EXECUTION_ID_KEY, executionId.toString)
       val r = try {
         val callSite = Utils.getCallSite()
-        sqlContext.sparkContext.listenerBus.post(SparkListenerSQLExecutionStart(
-          executionId, callSite.shortForm, callSite.longForm, queryExecution.toString,
-          SparkPlanInfo.fromSparkPlan(queryExecution.executedPlan), System.currentTimeMillis()))
+        sqlContext.listener.onExecutionStart(
+          executionId,
+          callSite.shortForm,
+          callSite.longForm,
+          queryExecution.toString,
+          SparkPlanGraph(queryExecution.executedPlan),
+          System.currentTimeMillis())
         try {
           body
         } finally {
-          sqlContext.sparkContext.listenerBus.post(SparkListenerSQLExecutionEnd(
-            executionId, System.currentTimeMillis()))
+          // Ideally, we need to make sure onExecutionEnd happens after onJobStart and onJobEnd.
+          // However, onJobStart and onJobEnd run in the listener thread. Because we cannot add new
+          // SQL event types to SparkListener since it's a public API, we cannot guarantee that.
+          //
+          // SQLListener should handle the case that onExecutionEnd happens before onJobEnd.
+          //
+          // The worst case is onExecutionEnd may happen before onJobStart when the listener thread
+          // is very busy. If so, we cannot track the jobs for the execution. It seems acceptable.
+          sqlContext.listener.onExecutionEnd(executionId, System.currentTimeMillis())
         }
       } finally {
         sc.setLocalProperty(EXECUTION_ID_KEY, null)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
deleted file mode 100644
index 486ce34064e43..0000000000000
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala
+++ /dev/null
@@ -1,46 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.execution
-
-import org.apache.spark.annotation.DeveloperApi
-import org.apache.spark.sql.execution.metric.SQLMetricInfo
-import org.apache.spark.util.Utils
-
-/**
- * :: DeveloperApi ::
- * Stores information about a SQL SparkPlan.
- */
-@DeveloperApi
-class SparkPlanInfo(
-    val nodeName: String,
-    val simpleString: String,
-    val children: Seq[SparkPlanInfo],
-    val metrics: Seq[SQLMetricInfo])
-
-private[sql] object SparkPlanInfo {
-
-  def fromSparkPlan(plan: SparkPlan): SparkPlanInfo = {
-    val metrics = plan.metrics.toSeq.map { case (key, metric) =>
-      new SQLMetricInfo(metric.name.getOrElse(key), metric.id,
-        Utils.getFormattedClassName(metric.param))
-    }
-    val children = plan.children.map(fromSparkPlan)
-
-    new SparkPlanInfo(plan.nodeName, plan.simpleString, children, metrics)
-  }
-}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetricInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetricInfo.scala
deleted file mode 100644
index 2708219ad3485..0000000000000
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetricInfo.scala
+++ /dev/null
@@ -1,30 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.sql.execution.metric
-
-import org.apache.spark.annotation.DeveloperApi
-
-/**
- * :: DeveloperApi ::
- * Stores information about a SQL Metric.
- */
-@DeveloperApi
-class SQLMetricInfo(
-    val name: String,
-    val accumulatorId: Long,
-    val metricParam: String)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala
index 6c0f6f8a52dc5..1c253e3942e95 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLMetrics.scala
@@ -104,39 +104,21 @@ private class LongSQLMetricParam(val stringValue: Seq[Long] => String, initialVa
   override def zero: LongSQLMetricValue = new LongSQLMetricValue(initialValue)
 }
 
-private object LongSQLMetricParam extends LongSQLMetricParam(_.sum.toString, 0L)
-
-private object StaticsLongSQLMetricParam extends LongSQLMetricParam(
-  (values: Seq[Long]) => {
-    // This is a workaround for SPARK-11013.
-    // We use -1 as initial value of the accumulator, if the accumulator is valid, we will update
-    // it at the end of task and the value will be at least 0.
-    val validValues = values.filter(_ >= 0)
-    val Seq(sum, min, med, max) = {
-      val metric = if (validValues.length == 0) {
-        Seq.fill(4)(0L)
-      } else {
-        val sorted = validValues.sorted
-        Seq(sorted.sum, sorted(0), sorted(validValues.length / 2), sorted(validValues.length - 1))
-      }
-      metric.map(Utils.bytesToString)
-    }
-    s"\n$sum ($min, $med, $max)"
-  }, -1L)
-
 private[sql] object SQLMetrics {
 
   private def createLongMetric(
       sc: SparkContext,
       name: String,
-      param: LongSQLMetricParam): LongSQLMetric = {
+      stringValue: Seq[Long] => String,
+      initialValue: Long): LongSQLMetric = {
+    val param = new LongSQLMetricParam(stringValue, initialValue)
     val acc = new LongSQLMetric(name, param)
     sc.cleaner.foreach(_.registerAccumulatorForCleanup(acc))
     acc
   }
 
   def createLongMetric(sc: SparkContext, name: String): LongSQLMetric = {
-    createLongMetric(sc, name, LongSQLMetricParam)
+    createLongMetric(sc, name, _.sum.toString, 0L)
   }
 
   /**
@@ -144,25 +126,31 @@ private[sql] object SQLMetrics {
    * spill size, etc.
    */
   def createSizeMetric(sc: SparkContext, name: String): LongSQLMetric = {
+    val stringValue = (values: Seq[Long]) => {
+      // This is a workaround for SPARK-11013.
+      // We use -1 as initial value of the accumulator, if the accumulator is valid, we will update
+      // it at the end of task and the value will be at least 0.
+      val validValues = values.filter(_ >= 0)
+      val Seq(sum, min, med, max) = {
+        val metric = if (validValues.length == 0) {
+          Seq.fill(4)(0L)
+        } else {
+          val sorted = validValues.sorted
+          Seq(sorted.sum, sorted(0), sorted(validValues.length / 2), sorted(validValues.length - 1))
+        }
+        metric.map(Utils.bytesToString)
+      }
+      s"\n$sum ($min, $med, $max)"
+    }
     // The final result of this metric in physical operator UI may looks like:
     // data size total (min, med, max):
     // 100GB (100MB, 1GB, 10GB)
-    createLongMetric(sc, s"$name total (min, med, max)", StaticsLongSQLMetricParam)
-  }
-
-  def getMetricParam(metricParamName: String): SQLMetricParam[SQLMetricValue[Any], Any] = {
-    val longSQLMetricParam = Utils.getFormattedClassName(LongSQLMetricParam)
-    val staticsSQLMetricParam = Utils.getFormattedClassName(StaticsLongSQLMetricParam)
-    val metricParam = metricParamName match {
-      case `longSQLMetricParam` => LongSQLMetricParam
-      case `staticsSQLMetricParam` => StaticsLongSQLMetricParam
-    }
-    metricParam.asInstanceOf[SQLMetricParam[SQLMetricValue[Any], Any]]
+    createLongMetric(sc, s"$name total (min, med, max)", stringValue, -1L)
   }
 
   /**
    * A metric that its value will be ignored. Use this one when we need a metric parameter but don't
    * care about the value.
    */
-  val nullLongMetric = new LongSQLMetric("null", LongSQLMetricParam)
+  val nullLongMetric = new LongSQLMetric("null", new LongSQLMetricParam(_.sum.toString, 0L))
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala
index c74ad40406992..e74d6fb396e1c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/ExecutionPage.scala
@@ -19,7 +19,9 @@ package org.apache.spark.sql.execution.ui
 
 import javax.servlet.http.HttpServletRequest
 
-import scala.xml.Node
+import scala.xml.{Node, Unparsed}
+
+import org.apache.commons.lang3.StringEscapeUtils
 
 import org.apache.spark.Logging
 import org.apache.spark.ui.{UIUtils, WebUIPage}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala
index e19a1e3e5851f..5a072de400b6a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLListener.scala
@@ -19,34 +19,11 @@ package org.apache.spark.sql.execution.ui
 
 import scala.collection.mutable
 
-import org.apache.spark.annotation.DeveloperApi
+import org.apache.spark.executor.TaskMetrics
 import org.apache.spark.scheduler._
 import org.apache.spark.sql.execution.SQLExecution
-import org.apache.spark.sql.execution.SparkPlanInfo
-import org.apache.spark.sql.execution.metric.{LongSQLMetricValue, SQLMetricValue, SQLMetricParam}
+import org.apache.spark.sql.execution.metric.{SQLMetricParam, SQLMetricValue}
 import org.apache.spark.{JobExecutionStatus, Logging, SparkConf}
-import org.apache.spark.ui.SparkUI
-
-@DeveloperApi
-case class SparkListenerSQLExecutionStart(
-    executionId: Long,
-    description: String,
-    details: String,
-    physicalPlanDescription: String,
-    sparkPlanInfo: SparkPlanInfo,
-    time: Long)
-  extends SparkListenerEvent
-
-@DeveloperApi
-case class SparkListenerSQLExecutionEnd(executionId: Long, time: Long)
-  extends SparkListenerEvent
-
-private[sql] class SQLHistoryListenerFactory extends SparkHistoryListenerFactory {
-
-  override def createListeners(conf: SparkConf, sparkUI: SparkUI): Seq[SparkListener] = {
-    List(new SQLHistoryListener(conf, sparkUI))
-  }
-}
 
 private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Logging {
 
@@ -141,8 +118,7 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
   override def onExecutorMetricsUpdate(
       executorMetricsUpdate: SparkListenerExecutorMetricsUpdate): Unit = synchronized {
     for ((taskId, stageId, stageAttemptID, metrics) <- executorMetricsUpdate.taskMetrics) {
-      updateTaskAccumulatorValues(taskId, stageId, stageAttemptID, metrics.accumulatorUpdates(),
-        finishTask = false)
+      updateTaskAccumulatorValues(taskId, stageId, stageAttemptID, metrics, finishTask = false)
     }
   }
 
@@ -164,7 +140,7 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
       taskEnd.taskInfo.taskId,
       taskEnd.stageId,
       taskEnd.stageAttemptId,
-      taskEnd.taskMetrics.accumulatorUpdates(),
+      taskEnd.taskMetrics,
       finishTask = true)
   }
 
@@ -172,12 +148,15 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
    * Update the accumulator values of a task with the latest metrics for this task. This is called
    * every time we receive an executor heartbeat or when a task finishes.
    */
-  protected def updateTaskAccumulatorValues(
+  private def updateTaskAccumulatorValues(
       taskId: Long,
       stageId: Int,
       stageAttemptID: Int,
-      accumulatorUpdates: Map[Long, Any],
+      metrics: TaskMetrics,
       finishTask: Boolean): Unit = {
+    if (metrics == null) {
+      return
+    }
 
     _stageIdToStageMetrics.get(stageId) match {
       case Some(stageMetrics) =>
@@ -195,9 +174,9 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
             case Some(taskMetrics) =>
               if (finishTask) {
                 taskMetrics.finished = true
-                taskMetrics.accumulatorUpdates = accumulatorUpdates
+                taskMetrics.accumulatorUpdates = metrics.accumulatorUpdates()
               } else if (!taskMetrics.finished) {
-                taskMetrics.accumulatorUpdates = accumulatorUpdates
+                taskMetrics.accumulatorUpdates = metrics.accumulatorUpdates()
               } else {
                 // If a task is finished, we should not override with accumulator updates from
                 // heartbeat reports
@@ -206,7 +185,7 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
               // TODO Now just set attemptId to 0. Should fix here when we can get the attempt
               // id from SparkListenerExecutorMetricsUpdate
               stageMetrics.taskIdToMetricUpdates(taskId) = new SQLTaskMetrics(
-                  attemptId = 0, finished = finishTask, accumulatorUpdates)
+                  attemptId = 0, finished = finishTask, metrics.accumulatorUpdates())
           }
         }
       case None =>
@@ -214,40 +193,38 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
     }
   }
 
-  override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
-    case SparkListenerSQLExecutionStart(executionId, description, details,
-      physicalPlanDescription, sparkPlanInfo, time) =>
-      val physicalPlanGraph = SparkPlanGraph(sparkPlanInfo)
-      val sqlPlanMetrics = physicalPlanGraph.nodes.flatMap { node =>
-        node.metrics.map(metric => metric.accumulatorId -> metric)
-      }
-      val executionUIData = new SQLExecutionUIData(
-        executionId,
-        description,
-        details,
-        physicalPlanDescription,
-        physicalPlanGraph,
-        sqlPlanMetrics.toMap,
-        time)
-      synchronized {
-        activeExecutions(executionId) = executionUIData
-        _executionIdToData(executionId) = executionUIData
-      }
-    case SparkListenerSQLExecutionEnd(executionId, time) => synchronized {
-      _executionIdToData.get(executionId).foreach { executionUIData =>
-        executionUIData.completionTime = Some(time)
-        if (!executionUIData.hasRunningJobs) {
-          // onExecutionEnd happens after all "onJobEnd"s
-          // So we should update the execution lists.
-          markExecutionFinished(executionId)
-        } else {
-          // There are some running jobs, onExecutionEnd happens before some "onJobEnd"s.
-          // Then we don't if the execution is successful, so let the last onJobEnd updates the
-          // execution lists.
-        }
+  def onExecutionStart(
+      executionId: Long,
+      description: String,
+      details: String,
+      physicalPlanDescription: String,
+      physicalPlanGraph: SparkPlanGraph,
+      time: Long): Unit = {
+    val sqlPlanMetrics = physicalPlanGraph.nodes.flatMap { node =>
+      node.metrics.map(metric => metric.accumulatorId -> metric)
+    }
+
+    val executionUIData = new SQLExecutionUIData(executionId, description, details,
+      physicalPlanDescription, physicalPlanGraph, sqlPlanMetrics.toMap, time)
+    synchronized {
+      activeExecutions(executionId) = executionUIData
+      _executionIdToData(executionId) = executionUIData
+    }
+  }
+
+  def onExecutionEnd(executionId: Long, time: Long): Unit = synchronized {
+    _executionIdToData.get(executionId).foreach { executionUIData =>
+      executionUIData.completionTime = Some(time)
+      if (!executionUIData.hasRunningJobs) {
+        // onExecutionEnd happens after all "onJobEnd"s
+        // So we should update the execution lists.
+        markExecutionFinished(executionId)
+      } else {
+        // There are some running jobs, onExecutionEnd happens before some "onJobEnd"s.
+        // Then we don't if the execution is successful, so let the last onJobEnd updates the
+        // execution lists.
       }
     }
-    case _ => // Ignore
   }
 
   private def markExecutionFinished(executionId: Long): Unit = {
@@ -312,38 +289,6 @@ private[sql] class SQLListener(conf: SparkConf) extends SparkListener with Loggi
 
 }
 
-private[spark] class SQLHistoryListener(conf: SparkConf, sparkUI: SparkUI)
-  extends SQLListener(conf) {
-
-  private var sqlTabAttached = false
-
-  override def onExecutorMetricsUpdate(
-      executorMetricsUpdate: SparkListenerExecutorMetricsUpdate): Unit = synchronized {
-    // Do nothing
-  }
-
-  override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = synchronized {
-    updateTaskAccumulatorValues(
-      taskEnd.taskInfo.taskId,
-      taskEnd.stageId,
-      taskEnd.stageAttemptId,
-      taskEnd.taskInfo.accumulables.map { acc =>
-        (acc.id, new LongSQLMetricValue(acc.update.getOrElse("0").toLong))
-      }.toMap,
-      finishTask = true)
-  }
-
-  override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
-    case _: SparkListenerSQLExecutionStart =>
-      if (!sqlTabAttached) {
-        new SQLTab(this, sparkUI)
-        sqlTabAttached = true
-      }
-      super.onOtherEvent(event)
-    case _ => super.onOtherEvent(event)
-  }
-}
-
 /**
  * Represent all necessary data for an execution that will be used in Web UI.
  */
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLTab.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLTab.scala
index 4f50b2ecdc8f8..9c27944d42fc6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLTab.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLTab.scala
@@ -17,11 +17,13 @@
 
 package org.apache.spark.sql.execution.ui
 
+import java.util.concurrent.atomic.AtomicInteger
+
 import org.apache.spark.Logging
 import org.apache.spark.ui.{SparkUI, SparkUITab}
 
 private[sql] class SQLTab(val listener: SQLListener, sparkUI: SparkUI)
-  extends SparkUITab(sparkUI, "SQL") with Logging {
+  extends SparkUITab(sparkUI, SQLTab.nextTabName) with Logging {
 
   val parent = sparkUI
 
@@ -33,5 +35,13 @@ private[sql] class SQLTab(val listener: SQLListener, sparkUI: SparkUI)
 }
 
 private[sql] object SQLTab {
+
   private val STATIC_RESOURCE_DIR = "org/apache/spark/sql/execution/ui/static"
+
+  private val nextTabId = new AtomicInteger(0)
+
+  private def nextTabName: String = {
+    val nextId = nextTabId.getAndIncrement()
+    if (nextId == 0) "SQL" else s"SQL$nextId"
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SparkPlanGraph.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SparkPlanGraph.scala
index 7af0ff09c5c6d..f1fce5478a3fe 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SparkPlanGraph.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SparkPlanGraph.scala
@@ -21,8 +21,8 @@ import java.util.concurrent.atomic.AtomicLong
 
 import scala.collection.mutable
 
-import org.apache.spark.sql.execution.SparkPlanInfo
-import org.apache.spark.sql.execution.metric.SQLMetrics
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.metric.{SQLMetricParam, SQLMetricValue}
 
 /**
  * A graph used for storing information of an executionPlan of DataFrame.
@@ -48,27 +48,27 @@ private[sql] object SparkPlanGraph {
   /**
    * Build a SparkPlanGraph from the root of a SparkPlan tree.
    */
-  def apply(planInfo: SparkPlanInfo): SparkPlanGraph = {
+  def apply(plan: SparkPlan): SparkPlanGraph = {
     val nodeIdGenerator = new AtomicLong(0)
     val nodes = mutable.ArrayBuffer[SparkPlanGraphNode]()
     val edges = mutable.ArrayBuffer[SparkPlanGraphEdge]()
-    buildSparkPlanGraphNode(planInfo, nodeIdGenerator, nodes, edges)
+    buildSparkPlanGraphNode(plan, nodeIdGenerator, nodes, edges)
     new SparkPlanGraph(nodes, edges)
   }
 
   private def buildSparkPlanGraphNode(
-      planInfo: SparkPlanInfo,
+      plan: SparkPlan,
       nodeIdGenerator: AtomicLong,
       nodes: mutable.ArrayBuffer[SparkPlanGraphNode],
       edges: mutable.ArrayBuffer[SparkPlanGraphEdge]): SparkPlanGraphNode = {
-    val metrics = planInfo.metrics.map { metric =>
-      SQLPlanMetric(metric.name, metric.accumulatorId,
-        SQLMetrics.getMetricParam(metric.metricParam))
+    val metrics = plan.metrics.toSeq.map { case (key, metric) =>
+      SQLPlanMetric(metric.name.getOrElse(key), metric.id,
+        metric.param.asInstanceOf[SQLMetricParam[SQLMetricValue[Any], Any]])
     }
     val node = SparkPlanGraphNode(
-      nodeIdGenerator.getAndIncrement(), planInfo.nodeName, planInfo.simpleString, metrics)
+      nodeIdGenerator.getAndIncrement(), plan.nodeName, plan.simpleString, metrics)
     nodes += node
-    val childrenNodes = planInfo.children.map(
+    val childrenNodes = plan.children.map(
       child => buildSparkPlanGraphNode(child, nodeIdGenerator, nodes, edges))
     for (child <- childrenNodes) {
       edges += SparkPlanGraphEdge(child.id, node.id)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
index 4f2cad19bfb6b..82867ab4967bb 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLMetricsSuite.scala
@@ -26,7 +26,6 @@ import org.apache.xbean.asm5.Opcodes._
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql._
-import org.apache.spark.sql.execution.SparkPlanInfo
 import org.apache.spark.sql.execution.ui.SparkPlanGraph
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SharedSQLContext
@@ -83,8 +82,7 @@ class SQLMetricsSuite extends SparkFunSuite with SharedSQLContext {
     if (jobs.size == expectedNumOfJobs) {
       // If we can track all jobs, check the metric values
       val metricValues = sqlContext.listener.getExecutionMetrics(executionId)
-      val actualMetrics = SparkPlanGraph(SparkPlanInfo.fromSparkPlan(
-        df.queryExecution.executedPlan)).nodes.filter { node =>
+      val actualMetrics = SparkPlanGraph(df.queryExecution.executedPlan).nodes.filter { node =>
         expectedMetrics.contains(node.id)
       }.map { node =>
         val nodeMetrics = node.metrics.map { metric =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
index f93d081d0c30e..c15aac775096c 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SQLListenerSuite.scala
@@ -21,10 +21,10 @@ import java.util.Properties
 
 import org.apache.spark.{SparkException, SparkContext, SparkConf, SparkFunSuite}
 import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.sql.execution.metric.LongSQLMetricValue
 import org.apache.spark.scheduler._
 import org.apache.spark.sql.{DataFrame, SQLContext}
-import org.apache.spark.sql.execution.{SparkPlanInfo, SQLExecution}
-import org.apache.spark.sql.execution.metric.LongSQLMetricValue
+import org.apache.spark.sql.execution.SQLExecution
 import org.apache.spark.sql.test.SharedSQLContext
 
 class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
@@ -82,8 +82,7 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
     val executionId = 0
     val df = createTestDataFrame
     val accumulatorIds =
-      SparkPlanGraph(SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan))
-        .nodes.flatMap(_.metrics.map(_.accumulatorId))
+      SparkPlanGraph(df.queryExecution.executedPlan).nodes.flatMap(_.metrics.map(_.accumulatorId))
     // Assume all accumulators are long
     var accumulatorValue = 0L
     val accumulatorUpdates = accumulatorIds.map { id =>
@@ -91,13 +90,13 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
       (id, accumulatorValue)
     }.toMap
 
-    listener.onOtherEvent(SparkListenerSQLExecutionStart(
+    listener.onExecutionStart(
       executionId,
       "test",
       "test",
       df.queryExecution.toString,
-      SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      SparkPlanGraph(df.queryExecution.executedPlan),
+      System.currentTimeMillis())
 
     val executionUIData = listener.executionIdToData(0)
 
@@ -207,8 +206,7 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
       time = System.currentTimeMillis(),
       JobSucceeded
     ))
-    listener.onOtherEvent(SparkListenerSQLExecutionEnd(
-      executionId, System.currentTimeMillis()))
+    listener.onExecutionEnd(executionId, System.currentTimeMillis())
 
     assert(executionUIData.runningJobs.isEmpty)
     assert(executionUIData.succeededJobs === Seq(0))
@@ -221,20 +219,19 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
     val listener = new SQLListener(sqlContext.sparkContext.conf)
     val executionId = 0
     val df = createTestDataFrame
-    listener.onOtherEvent(SparkListenerSQLExecutionStart(
+    listener.onExecutionStart(
       executionId,
       "test",
       "test",
       df.queryExecution.toString,
-      SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      SparkPlanGraph(df.queryExecution.executedPlan),
+      System.currentTimeMillis())
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
       stageInfos = Nil,
       createProperties(executionId)))
-    listener.onOtherEvent(SparkListenerSQLExecutionEnd(
-      executionId, System.currentTimeMillis()))
+    listener.onExecutionEnd(executionId, System.currentTimeMillis())
     listener.onJobEnd(SparkListenerJobEnd(
       jobId = 0,
       time = System.currentTimeMillis(),
@@ -251,13 +248,13 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
     val listener = new SQLListener(sqlContext.sparkContext.conf)
     val executionId = 0
     val df = createTestDataFrame
-    listener.onOtherEvent(SparkListenerSQLExecutionStart(
+    listener.onExecutionStart(
       executionId,
       "test",
       "test",
       df.queryExecution.toString,
-      SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      SparkPlanGraph(df.queryExecution.executedPlan),
+      System.currentTimeMillis())
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
@@ -274,8 +271,7 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
       time = System.currentTimeMillis(),
       stageInfos = Nil,
       createProperties(executionId)))
-    listener.onOtherEvent(SparkListenerSQLExecutionEnd(
-      executionId, System.currentTimeMillis()))
+    listener.onExecutionEnd(executionId, System.currentTimeMillis())
     listener.onJobEnd(SparkListenerJobEnd(
       jobId = 1,
       time = System.currentTimeMillis(),
@@ -292,20 +288,19 @@ class SQLListenerSuite extends SparkFunSuite with SharedSQLContext {
     val listener = new SQLListener(sqlContext.sparkContext.conf)
     val executionId = 0
     val df = createTestDataFrame
-    listener.onOtherEvent(SparkListenerSQLExecutionStart(
+    listener.onExecutionStart(
       executionId,
       "test",
       "test",
       df.queryExecution.toString,
-      SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan),
-      System.currentTimeMillis()))
+      SparkPlanGraph(df.queryExecution.executedPlan),
+      System.currentTimeMillis())
     listener.onJobStart(SparkListenerJobStart(
       jobId = 0,
       time = System.currentTimeMillis(),
       stageInfos = Seq.empty,
       createProperties(executionId)))
-    listener.onOtherEvent(SparkListenerSQLExecutionEnd(
-      executionId, System.currentTimeMillis()))
+    listener.onExecutionEnd(executionId, System.currentTimeMillis())
     listener.onJobEnd(SparkListenerJobEnd(
       jobId = 0,
       time = System.currentTimeMillis(),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
index e7b376548787c..963d10eed62ed 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/test/SharedSQLContext.scala
@@ -42,7 +42,6 @@ trait SharedSQLContext extends SQLTestUtils {
    * Initialize the [[TestSQLContext]].
    */
   protected override def beforeAll(): Unit = {
-    SQLContext.clearSqlListener()
     if (_ctx == null) {
       _ctx = new TestSQLContext
     }

From a8ceec5e8c1572dd3d74783c06c78b7ca0b8a7ce Mon Sep 17 00:00:00 2001
From: Teng Qiu <teng.qiu@gmail.com>
Date: Tue, 1 Dec 2015 07:27:32 +0900
Subject: [PATCH 504/518] [SPARK-12053][CORE] EventLoggingListener.getLogPath
 needs 4 parameters

```EventLoggingListener.getLogPath``` needs 4 input arguments:
https://github.com/apache/spark/blob/v1.6.0-preview2/core/src/main/scala/org/apache/spark/scheduler/EventLoggingListener.scala#L276-L280

the 3rd parameter should be appAttemptId, 4th parameter is codec...

Author: Teng Qiu <teng.qiu@gmail.com>

Closes #10044 from chutium/SPARK-12053.
---
 core/src/main/scala/org/apache/spark/deploy/master/Master.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/master/Master.scala b/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
index 9952c97dbdffc..1355e1ad1b523 100644
--- a/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/master/Master.scala
@@ -934,7 +934,7 @@ private[deploy] class Master(
         }
 
       val eventLogFilePrefix = EventLoggingListener.getLogPath(
-          eventLogDir, app.id, app.desc.eventLogCodec)
+          eventLogDir, app.id, appAttemptId = None, compressionCodecName = app.desc.eventLogCodec)
       val fs = Utils.getHadoopFileSystem(eventLogDir, hadoopConf)
       val inProgressExists = fs.exists(new Path(eventLogFilePrefix +
           EventLoggingListener.IN_PROGRESS))

From e232720a65dfb9ae6135cbb7674e35eddd88d625 Mon Sep 17 00:00:00 2001
From: Yuhao Yang <hhbyyh@gmail.com>
Date: Mon, 30 Nov 2015 14:56:51 -0800
Subject: [PATCH 505/518] [SPARK-11689][ML] Add user guide and example code for
 LDA under spark.ml

jira: https://issues.apache.org/jira/browse/SPARK-11689

Add simple user guide for LDA under spark.ml and example code under examples/. Use include_example to include example code in the user guide markdown. Check SPARK-11606 for instructions.

Original PR is reverted due to document build error. https://github.com/apache/spark/pull/9722

mengxr feynmanliang yinxusen  Sorry for the troubling.

Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #9974 from hhbyyh/ldaMLExample.
---
 docs/ml-clustering.md                         | 31 ++++++
 docs/ml-guide.md                              |  3 +-
 docs/mllib-guide.md                           |  1 +
 .../spark/examples/ml/JavaLDAExample.java     | 97 +++++++++++++++++++
 .../apache/spark/examples/ml/LDAExample.scala | 77 +++++++++++++++
 5 files changed, 208 insertions(+), 1 deletion(-)
 create mode 100644 docs/ml-clustering.md
 create mode 100644 examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
 create mode 100644 examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala

diff --git a/docs/ml-clustering.md b/docs/ml-clustering.md
new file mode 100644
index 0000000000000..cfefb5dfbde9e
--- /dev/null
+++ b/docs/ml-clustering.md
@@ -0,0 +1,31 @@
+---
+layout: global
+title: Clustering - ML
+displayTitle: <a href="ml-guide.html">ML</a> - Clustering
+---
+
+In this section, we introduce the pipeline API for [clustering in mllib](mllib-clustering.html).
+
+## Latent Dirichlet allocation (LDA)
+
+`LDA` is implemented as an `Estimator` that supports both `EMLDAOptimizer` and `OnlineLDAOptimizer`,
+and generates a `LDAModel` as the base models. Expert users may cast a `LDAModel` generated by
+`EMLDAOptimizer` to a `DistributedLDAModel` if needed.
+
+<div class="codetabs">
+
+<div data-lang="scala" markdown="1">
+
+Refer to the [Scala API docs](api/scala/index.html#org.apache.spark.ml.clustering.LDA) for more details.
+
+{% include_example scala/org/apache/spark/examples/ml/LDAExample.scala %}
+</div>
+
+<div data-lang="java" markdown="1">
+
+Refer to the [Java API docs](api/java/org/apache/spark/ml/clustering/LDA.html) for more details.
+
+{% include_example java/org/apache/spark/examples/ml/JavaLDAExample.java %}
+</div>
+
+</div>
\ No newline at end of file
diff --git a/docs/ml-guide.md b/docs/ml-guide.md
index be18a05361a17..6f35b30c3d4df 100644
--- a/docs/ml-guide.md
+++ b/docs/ml-guide.md
@@ -40,6 +40,7 @@ Also, some algorithms have additional capabilities in the `spark.ml` API; e.g.,
 provide class probabilities, and linear models provide model summaries.
 
 * [Feature extraction, transformation, and selection](ml-features.html)
+* [Clustering](ml-clustering.html)
 * [Decision Trees for classification and regression](ml-decision-tree.html)
 * [Ensembles](ml-ensembles.html)
 * [Linear methods with elastic net regularization](ml-linear-methods.html)
@@ -950,4 +951,4 @@ model.transform(test)
 {% endhighlight %}
 </div>
 
-</div>
+</div>
\ No newline at end of file
diff --git a/docs/mllib-guide.md b/docs/mllib-guide.md
index 91e50ccfecec4..54e35fcbb15af 100644
--- a/docs/mllib-guide.md
+++ b/docs/mllib-guide.md
@@ -69,6 +69,7 @@ We list major functionality from both below, with links to detailed guides.
 concepts. It also contains sections on using algorithms within the Pipelines API, for example:
 
 * [Feature extraction, transformation, and selection](ml-features.html)
+* [Clustering](ml-clustering.html)
 * [Decision trees for classification and regression](ml-decision-tree.html)
 * [Ensembles](ml-ensembles.html)
 * [Linear methods with elastic net regularization](ml-linear-methods.html)
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
new file mode 100644
index 0000000000000..3a5d3237c85f6
--- /dev/null
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaLDAExample.java
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml;
+// $example on$
+import java.util.regex.Pattern;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function;
+import org.apache.spark.ml.clustering.LDA;
+import org.apache.spark.ml.clustering.LDAModel;
+import org.apache.spark.mllib.linalg.Vector;
+import org.apache.spark.mllib.linalg.VectorUDT;
+import org.apache.spark.mllib.linalg.Vectors;
+import org.apache.spark.sql.DataFrame;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.sql.catalyst.expressions.GenericRow;
+import org.apache.spark.sql.types.Metadata;
+import org.apache.spark.sql.types.StructField;
+import org.apache.spark.sql.types.StructType;
+// $example off$
+
+/**
+ * An example demonstrating LDA
+ * Run with
+ * <pre>
+ * bin/run-example ml.JavaLDAExample
+ * </pre>
+ */
+public class JavaLDAExample {
+
+  // $example on$
+  private static class ParseVector implements Function<String, Row> {
+    private static final Pattern separator = Pattern.compile(" ");
+
+    @Override
+    public Row call(String line) {
+      String[] tok = separator.split(line);
+      double[] point = new double[tok.length];
+      for (int i = 0; i < tok.length; ++i) {
+        point[i] = Double.parseDouble(tok[i]);
+      }
+      Vector[] points = {Vectors.dense(point)};
+      return new GenericRow(points);
+    }
+  }
+
+  public static void main(String[] args) {
+
+    String inputFile = "data/mllib/sample_lda_data.txt";
+
+    // Parses the arguments
+    SparkConf conf = new SparkConf().setAppName("JavaLDAExample");
+    JavaSparkContext jsc = new JavaSparkContext(conf);
+    SQLContext sqlContext = new SQLContext(jsc);
+
+    // Loads data
+    JavaRDD<Row> points = jsc.textFile(inputFile).map(new ParseVector());
+    StructField[] fields = {new StructField("features", new VectorUDT(), false, Metadata.empty())};
+    StructType schema = new StructType(fields);
+    DataFrame dataset = sqlContext.createDataFrame(points, schema);
+
+    // Trains a LDA model
+    LDA lda = new LDA()
+      .setK(10)
+      .setMaxIter(10);
+    LDAModel model = lda.fit(dataset);
+
+    System.out.println(model.logLikelihood(dataset));
+    System.out.println(model.logPerplexity(dataset));
+
+    // Shows the result
+    DataFrame topics = model.describeTopics(3);
+    topics.show(false);
+    model.transform(dataset).show(false);
+
+    jsc.stop();
+  }
+  // $example off$
+}
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala
new file mode 100644
index 0000000000000..419ce3d87a6ac
--- /dev/null
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/LDAExample.scala
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.examples.ml
+
+// scalastyle:off println
+import org.apache.spark.{SparkContext, SparkConf}
+import org.apache.spark.mllib.linalg.{VectorUDT, Vectors}
+// $example on$
+import org.apache.spark.ml.clustering.LDA
+import org.apache.spark.sql.{Row, SQLContext}
+import org.apache.spark.sql.types.{StructField, StructType}
+// $example off$
+
+/**
+ * An example demonstrating a LDA of ML pipeline.
+ * Run with
+ * {{{
+ * bin/run-example ml.LDAExample
+ * }}}
+ */
+object LDAExample {
+
+  final val FEATURES_COL = "features"
+
+  def main(args: Array[String]): Unit = {
+
+    val input = "data/mllib/sample_lda_data.txt"
+    // Creates a Spark context and a SQL context
+    val conf = new SparkConf().setAppName(s"${this.getClass.getSimpleName}")
+    val sc = new SparkContext(conf)
+    val sqlContext = new SQLContext(sc)
+
+    // $example on$
+    // Loads data
+    val rowRDD = sc.textFile(input).filter(_.nonEmpty)
+      .map(_.split(" ").map(_.toDouble)).map(Vectors.dense).map(Row(_))
+    val schema = StructType(Array(StructField(FEATURES_COL, new VectorUDT, false)))
+    val dataset = sqlContext.createDataFrame(rowRDD, schema)
+
+    // Trains a LDA model
+    val lda = new LDA()
+      .setK(10)
+      .setMaxIter(10)
+      .setFeaturesCol(FEATURES_COL)
+    val model = lda.fit(dataset)
+    val transformed = model.transform(dataset)
+
+    val ll = model.logLikelihood(dataset)
+    val lp = model.logPerplexity(dataset)
+
+    // describeTopics
+    val topics = model.describeTopics(3)
+
+    // Shows the result
+    topics.show(false)
+    transformed.show(false)
+
+    // $example off$
+    sc.stop()
+  }
+}
+// scalastyle:on println

From de64b65f7cf2ac58c1abc310ba547637fdbb8557 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Mon, 30 Nov 2015 15:01:08 -0800
Subject: [PATCH 506/518] [SPARK-11975][ML] Remove duplicate mllib example
 (DT/RF/GBT in Java/Python)

Remove duplicate mllib example (DT/RF/GBT in Java/Python).
Since we have tutorial code for DT/RF/GBT classification/regression in Scala/Java/Python and example applications for DT/RF/GBT in Scala, so we mark these as duplicated and remove them.
mengxr

Author: Yanbo Liang <ybliang8@gmail.com>

Closes #9954 from yanboliang/SPARK-11975.
---
 .../examples/mllib/JavaDecisionTree.java      | 116 --------------
 .../mllib/JavaGradientBoostedTreesRunner.java | 126 ---------------
 .../mllib/JavaRandomForestExample.java        | 139 -----------------
 .../main/python/mllib/decision_tree_runner.py | 144 ------------------
 .../python/mllib/gradient_boosted_trees.py    |  77 ----------
 .../python/mllib/random_forest_example.py     |  90 -----------
 6 files changed, 692 deletions(-)
 delete mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTree.java
 delete mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostedTreesRunner.java
 delete mode 100644 examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestExample.java
 delete mode 100755 examples/src/main/python/mllib/decision_tree_runner.py
 delete mode 100644 examples/src/main/python/mllib/gradient_boosted_trees.py
 delete mode 100755 examples/src/main/python/mllib/random_forest_example.py

diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTree.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTree.java
deleted file mode 100644
index 1f82e3f4cb18e..0000000000000
--- a/examples/src/main/java/org/apache/spark/examples/mllib/JavaDecisionTree.java
+++ /dev/null
@@ -1,116 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.examples.mllib;
-
-import java.util.HashMap;
-
-import scala.Tuple2;
-
-import org.apache.spark.api.java.function.Function2;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.DecisionTree;
-import org.apache.spark.mllib.tree.model.DecisionTreeModel;
-import org.apache.spark.mllib.util.MLUtils;
-import org.apache.spark.SparkConf;
-
-/**
- * Classification and regression using decision trees.
- */
-public final class JavaDecisionTree {
-
-  public static void main(String[] args) {
-    String datapath = "data/mllib/sample_libsvm_data.txt";
-    if (args.length == 1) {
-      datapath = args[0];
-    } else if (args.length > 1) {
-      System.err.println("Usage: JavaDecisionTree <libsvm format data file>");
-      System.exit(1);
-    }
-    SparkConf sparkConf = new SparkConf().setAppName("JavaDecisionTree");
-    JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD().cache();
-
-    // Compute the number of classes from the data.
-    Integer numClasses = data.map(new Function<LabeledPoint, Double>() {
-      @Override public Double call(LabeledPoint p) {
-        return p.label();
-      }
-    }).countByValue().size();
-
-    // Set parameters.
-    //  Empty categoricalFeaturesInfo indicates all features are continuous.
-    HashMap<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-    String impurity = "gini";
-    Integer maxDepth = 5;
-    Integer maxBins = 32;
-
-    // Train a DecisionTree model for classification.
-    final DecisionTreeModel model = DecisionTree.trainClassifier(data, numClasses,
-      categoricalFeaturesInfo, impurity, maxDepth, maxBins);
-
-    // Evaluate model on training instances and compute training error
-    JavaPairRDD<Double, Double> predictionAndLabel =
-      data.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-        @Override public Tuple2<Double, Double> call(LabeledPoint p) {
-          return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-        }
-      });
-    Double trainErr =
-      1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
-        @Override public Boolean call(Tuple2<Double, Double> pl) {
-          return !pl._1().equals(pl._2());
-        }
-      }).count() / data.count();
-    System.out.println("Training error: " + trainErr);
-    System.out.println("Learned classification tree model:\n" + model);
-
-    // Train a DecisionTree model for regression.
-    impurity = "variance";
-    final DecisionTreeModel regressionModel = DecisionTree.trainRegressor(data,
-        categoricalFeaturesInfo, impurity, maxDepth, maxBins);
-
-    // Evaluate model on training instances and compute training error
-    JavaPairRDD<Double, Double> regressorPredictionAndLabel =
-      data.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-        @Override public Tuple2<Double, Double> call(LabeledPoint p) {
-          return new Tuple2<Double, Double>(regressionModel.predict(p.features()), p.label());
-        }
-      });
-    Double trainMSE =
-      regressorPredictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
-        @Override public Double call(Tuple2<Double, Double> pl) {
-          Double diff = pl._1() - pl._2();
-          return diff * diff;
-        }
-      }).reduce(new Function2<Double, Double, Double>() {
-        @Override public Double call(Double a, Double b) {
-          return a + b;
-        }
-      }) / data.count();
-    System.out.println("Training Mean Squared Error: " + trainMSE);
-    System.out.println("Learned regression tree model:\n" + regressionModel);
-
-    sc.stop();
-  }
-}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostedTreesRunner.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostedTreesRunner.java
deleted file mode 100644
index a1844d5d07ad4..0000000000000
--- a/examples/src/main/java/org/apache/spark/examples/mllib/JavaGradientBoostedTreesRunner.java
+++ /dev/null
@@ -1,126 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.examples.mllib;
-
-import scala.Tuple2;
-
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.Function2;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.GradientBoostedTrees;
-import org.apache.spark.mllib.tree.configuration.BoostingStrategy;
-import org.apache.spark.mllib.tree.model.GradientBoostedTreesModel;
-import org.apache.spark.mllib.util.MLUtils;
-
-/**
- * Classification and regression using gradient-boosted decision trees.
- */
-public final class JavaGradientBoostedTreesRunner {
-
-  private static void usage() {
-    System.err.println("Usage: JavaGradientBoostedTreesRunner <libsvm format data file>" +
-        " <Classification/Regression>");
-    System.exit(-1);
-  }
-
-  public static void main(String[] args) {
-    String datapath = "data/mllib/sample_libsvm_data.txt";
-    String algo = "Classification";
-    if (args.length >= 1) {
-      datapath = args[0];
-    }
-    if (args.length >= 2) {
-      algo = args[1];
-    }
-    if (args.length > 2) {
-      usage();
-    }
-    SparkConf sparkConf = new SparkConf().setAppName("JavaGradientBoostedTreesRunner");
-    JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD().cache();
-
-    // Set parameters.
-    //  Note: All features are treated as continuous.
-    BoostingStrategy boostingStrategy = BoostingStrategy.defaultParams(algo);
-    boostingStrategy.setNumIterations(10);
-    boostingStrategy.treeStrategy().setMaxDepth(5);
-
-    if (algo.equals("Classification")) {
-      // Compute the number of classes from the data.
-      Integer numClasses = data.map(new Function<LabeledPoint, Double>() {
-        @Override public Double call(LabeledPoint p) {
-          return p.label();
-        }
-      }).countByValue().size();
-      boostingStrategy.treeStrategy().setNumClasses(numClasses);
-
-      // Train a GradientBoosting model for classification.
-      final GradientBoostedTreesModel model = GradientBoostedTrees.train(data, boostingStrategy);
-
-      // Evaluate model on training instances and compute training error
-      JavaPairRDD<Double, Double> predictionAndLabel =
-          data.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-            @Override public Tuple2<Double, Double> call(LabeledPoint p) {
-              return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-            }
-          });
-      Double trainErr =
-          1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
-            @Override public Boolean call(Tuple2<Double, Double> pl) {
-              return !pl._1().equals(pl._2());
-            }
-          }).count() / data.count();
-      System.out.println("Training error: " + trainErr);
-      System.out.println("Learned classification tree model:\n" + model);
-    } else if (algo.equals("Regression")) {
-      // Train a GradientBoosting model for classification.
-      final GradientBoostedTreesModel model = GradientBoostedTrees.train(data, boostingStrategy);
-
-      // Evaluate model on training instances and compute training error
-      JavaPairRDD<Double, Double> predictionAndLabel =
-          data.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-            @Override public Tuple2<Double, Double> call(LabeledPoint p) {
-              return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-            }
-          });
-      Double trainMSE =
-          predictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
-            @Override public Double call(Tuple2<Double, Double> pl) {
-              Double diff = pl._1() - pl._2();
-              return diff * diff;
-            }
-          }).reduce(new Function2<Double, Double, Double>() {
-            @Override public Double call(Double a, Double b) {
-              return a + b;
-            }
-          }) / data.count();
-      System.out.println("Training Mean Squared Error: " + trainMSE);
-      System.out.println("Learned regression tree model:\n" + model);
-    } else {
-      usage();
-    }
-
-    sc.stop();
-  }
-}
diff --git a/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestExample.java b/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestExample.java
deleted file mode 100644
index 89a4e092a5af7..0000000000000
--- a/examples/src/main/java/org/apache/spark/examples/mllib/JavaRandomForestExample.java
+++ /dev/null
@@ -1,139 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements.  See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License.  You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.apache.spark.examples.mllib;
-
-import scala.Tuple2;
-
-import java.util.HashMap;
-
-import org.apache.spark.SparkConf;
-import org.apache.spark.api.java.JavaPairRDD;
-import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.Function;
-import org.apache.spark.api.java.function.Function2;
-import org.apache.spark.api.java.function.PairFunction;
-import org.apache.spark.mllib.regression.LabeledPoint;
-import org.apache.spark.mllib.tree.RandomForest;
-import org.apache.spark.mllib.tree.model.RandomForestModel;
-import org.apache.spark.mllib.util.MLUtils;
-
-public final class JavaRandomForestExample {
-
-  /**
-   * Note: This example illustrates binary classification.
-   * For information on multiclass classification, please refer to the JavaDecisionTree.java
-   * example.
-   */
-  private static void testClassification(JavaRDD<LabeledPoint> trainingData,
-                                         JavaRDD<LabeledPoint> testData) {
-    // Train a RandomForest model.
-    //  Empty categoricalFeaturesInfo indicates all features are continuous.
-    Integer numClasses = 2;
-    HashMap<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-    Integer numTrees = 3; // Use more in practice.
-    String featureSubsetStrategy = "auto"; // Let the algorithm choose.
-    String impurity = "gini";
-    Integer maxDepth = 4;
-    Integer maxBins = 32;
-    Integer seed = 12345;
-
-    final RandomForestModel model = RandomForest.trainClassifier(trainingData, numClasses,
-        categoricalFeaturesInfo, numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins,
-        seed);
-
-    // Evaluate model on test instances and compute test error
-    JavaPairRDD<Double, Double> predictionAndLabel =
-        testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-          @Override
-          public Tuple2<Double, Double> call(LabeledPoint p) {
-            return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-          }
-        });
-    Double testErr =
-        1.0 * predictionAndLabel.filter(new Function<Tuple2<Double, Double>, Boolean>() {
-          @Override
-          public Boolean call(Tuple2<Double, Double> pl) {
-            return !pl._1().equals(pl._2());
-          }
-        }).count() / testData.count();
-    System.out.println("Test Error: " + testErr);
-    System.out.println("Learned classification forest model:\n" + model.toDebugString());
-  }
-
-  private static void testRegression(JavaRDD<LabeledPoint> trainingData,
-                                     JavaRDD<LabeledPoint> testData) {
-    // Train a RandomForest model.
-    //  Empty categoricalFeaturesInfo indicates all features are continuous.
-    HashMap<Integer, Integer> categoricalFeaturesInfo = new HashMap<Integer, Integer>();
-    Integer numTrees = 3; // Use more in practice.
-    String featureSubsetStrategy = "auto"; // Let the algorithm choose.
-    String impurity = "variance";
-    Integer maxDepth = 4;
-    Integer maxBins = 32;
-    Integer seed = 12345;
-
-    final RandomForestModel model = RandomForest.trainRegressor(trainingData,
-        categoricalFeaturesInfo, numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins,
-        seed);
-
-    // Evaluate model on test instances and compute test error
-    JavaPairRDD<Double, Double> predictionAndLabel =
-        testData.mapToPair(new PairFunction<LabeledPoint, Double, Double>() {
-          @Override
-          public Tuple2<Double, Double> call(LabeledPoint p) {
-            return new Tuple2<Double, Double>(model.predict(p.features()), p.label());
-          }
-        });
-    Double testMSE =
-        predictionAndLabel.map(new Function<Tuple2<Double, Double>, Double>() {
-          @Override
-          public Double call(Tuple2<Double, Double> pl) {
-            Double diff = pl._1() - pl._2();
-            return diff * diff;
-          }
-        }).reduce(new Function2<Double, Double, Double>() {
-          @Override
-          public Double call(Double a, Double b) {
-            return a + b;
-          }
-        }) / testData.count();
-    System.out.println("Test Mean Squared Error: " + testMSE);
-    System.out.println("Learned regression forest model:\n" + model.toDebugString());
-  }
-
-  public static void main(String[] args) {
-    SparkConf sparkConf = new SparkConf().setAppName("JavaRandomForestExample");
-    JavaSparkContext sc = new JavaSparkContext(sparkConf);
-
-    // Load and parse the data file.
-    String datapath = "data/mllib/sample_libsvm_data.txt";
-    JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(sc.sc(), datapath).toJavaRDD();
-    // Split the data into training and test sets (30% held out for testing)
-    JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
-    JavaRDD<LabeledPoint> trainingData = splits[0];
-    JavaRDD<LabeledPoint> testData = splits[1];
-
-    System.out.println("\nRunning example of classification using RandomForest\n");
-    testClassification(trainingData, testData);
-
-    System.out.println("\nRunning example of regression using RandomForest\n");
-    testRegression(trainingData, testData);
-    sc.stop();
-  }
-}
diff --git a/examples/src/main/python/mllib/decision_tree_runner.py b/examples/src/main/python/mllib/decision_tree_runner.py
deleted file mode 100755
index 513ed8fd51450..0000000000000
--- a/examples/src/main/python/mllib/decision_tree_runner.py
+++ /dev/null
@@ -1,144 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-"""
-Decision tree classification and regression using MLlib.
-
-This example requires NumPy (http://www.numpy.org/).
-"""
-from __future__ import print_function
-
-import numpy
-import os
-import sys
-
-from operator import add
-
-from pyspark import SparkContext
-from pyspark.mllib.regression import LabeledPoint
-from pyspark.mllib.tree import DecisionTree
-from pyspark.mllib.util import MLUtils
-
-
-def getAccuracy(dtModel, data):
-    """
-    Return accuracy of DecisionTreeModel on the given RDD[LabeledPoint].
-    """
-    seqOp = (lambda acc, x: acc + (x[0] == x[1]))
-    predictions = dtModel.predict(data.map(lambda x: x.features))
-    truth = data.map(lambda p: p.label)
-    trainCorrect = predictions.zip(truth).aggregate(0, seqOp, add)
-    if data.count() == 0:
-        return 0
-    return trainCorrect / (0.0 + data.count())
-
-
-def getMSE(dtModel, data):
-    """
-    Return mean squared error (MSE) of DecisionTreeModel on the given
-    RDD[LabeledPoint].
-    """
-    seqOp = (lambda acc, x: acc + numpy.square(x[0] - x[1]))
-    predictions = dtModel.predict(data.map(lambda x: x.features))
-    truth = data.map(lambda p: p.label)
-    trainMSE = predictions.zip(truth).aggregate(0, seqOp, add)
-    if data.count() == 0:
-        return 0
-    return trainMSE / (0.0 + data.count())
-
-
-def reindexClassLabels(data):
-    """
-    Re-index class labels in a dataset to the range {0,...,numClasses-1}.
-    If all labels in that range already appear at least once,
-     then the returned RDD is the same one (without a mapping).
-    Note: If a label simply does not appear in the data,
-          the index will not include it.
-          Be aware of this when reindexing subsampled data.
-    :param data: RDD of LabeledPoint where labels are integer values
-                 denoting labels for a classification problem.
-    :return: Pair (reindexedData, origToNewLabels) where
-             reindexedData is an RDD of LabeledPoint with labels in
-              the range {0,...,numClasses-1}, and
-             origToNewLabels is a dictionary mapping original labels
-              to new labels.
-    """
-    # classCounts: class --> # examples in class
-    classCounts = data.map(lambda x: x.label).countByValue()
-    numExamples = sum(classCounts.values())
-    sortedClasses = sorted(classCounts.keys())
-    numClasses = len(classCounts)
-    # origToNewLabels: class --> index in 0,...,numClasses-1
-    if (numClasses < 2):
-        print("Dataset for classification should have at least 2 classes."
-              " The given dataset had only %d classes." % numClasses, file=sys.stderr)
-        exit(1)
-    origToNewLabels = dict([(sortedClasses[i], i) for i in range(0, numClasses)])
-
-    print("numClasses = %d" % numClasses)
-    print("Per-class example fractions, counts:")
-    print("Class\tFrac\tCount")
-    for c in sortedClasses:
-        frac = classCounts[c] / (numExamples + 0.0)
-        print("%g\t%g\t%d" % (c, frac, classCounts[c]))
-
-    if (sortedClasses[0] == 0 and sortedClasses[-1] == numClasses - 1):
-        return (data, origToNewLabels)
-    else:
-        reindexedData = \
-            data.map(lambda x: LabeledPoint(origToNewLabels[x.label], x.features))
-        return (reindexedData, origToNewLabels)
-
-
-def usage():
-    print("Usage: decision_tree_runner [libsvm format data filepath]", file=sys.stderr)
-    exit(1)
-
-
-if __name__ == "__main__":
-    if len(sys.argv) > 2:
-        usage()
-    sc = SparkContext(appName="PythonDT")
-
-    # Load data.
-    dataPath = 'data/mllib/sample_libsvm_data.txt'
-    if len(sys.argv) == 2:
-        dataPath = sys.argv[1]
-    if not os.path.isfile(dataPath):
-        sc.stop()
-        usage()
-    points = MLUtils.loadLibSVMFile(sc, dataPath)
-
-    # Re-index class labels if needed.
-    (reindexedData, origToNewLabels) = reindexClassLabels(points)
-    numClasses = len(origToNewLabels)
-
-    # Train a classifier.
-    categoricalFeaturesInfo = {}  # no categorical features
-    model = DecisionTree.trainClassifier(reindexedData, numClasses=numClasses,
-                                         categoricalFeaturesInfo=categoricalFeaturesInfo)
-    # Print learned tree and stats.
-    print("Trained DecisionTree for classification:")
-    print("  Model numNodes: %d" % model.numNodes())
-    print("  Model depth: %d" % model.depth())
-    print("  Training accuracy: %g" % getAccuracy(model, reindexedData))
-    if model.numNodes() < 20:
-        print(model.toDebugString())
-    else:
-        print(model)
-
-    sc.stop()
diff --git a/examples/src/main/python/mllib/gradient_boosted_trees.py b/examples/src/main/python/mllib/gradient_boosted_trees.py
deleted file mode 100644
index 781bd61c9d2b5..0000000000000
--- a/examples/src/main/python/mllib/gradient_boosted_trees.py
+++ /dev/null
@@ -1,77 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-"""
-Gradient boosted Trees classification and regression using MLlib.
-"""
-from __future__ import print_function
-
-import sys
-
-from pyspark.context import SparkContext
-from pyspark.mllib.tree import GradientBoostedTrees
-from pyspark.mllib.util import MLUtils
-
-
-def testClassification(trainingData, testData):
-    # Train a GradientBoostedTrees model.
-    #  Empty categoricalFeaturesInfo indicates all features are continuous.
-    model = GradientBoostedTrees.trainClassifier(trainingData, categoricalFeaturesInfo={},
-                                                 numIterations=30, maxDepth=4)
-    # Evaluate model on test instances and compute test error
-    predictions = model.predict(testData.map(lambda x: x.features))
-    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-    testErr = labelsAndPredictions.filter(lambda v_p: v_p[0] != v_p[1]).count() \
-        / float(testData.count())
-    print('Test Error = ' + str(testErr))
-    print('Learned classification ensemble model:')
-    print(model.toDebugString())
-
-
-def testRegression(trainingData, testData):
-    # Train a GradientBoostedTrees model.
-    #  Empty categoricalFeaturesInfo indicates all features are continuous.
-    model = GradientBoostedTrees.trainRegressor(trainingData, categoricalFeaturesInfo={},
-                                                numIterations=30, maxDepth=4)
-    # Evaluate model on test instances and compute test error
-    predictions = model.predict(testData.map(lambda x: x.features))
-    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-    testMSE = labelsAndPredictions.map(lambda vp: (vp[0] - vp[1]) * (vp[0] - vp[1])).sum() \
-        / float(testData.count())
-    print('Test Mean Squared Error = ' + str(testMSE))
-    print('Learned regression ensemble model:')
-    print(model.toDebugString())
-
-
-if __name__ == "__main__":
-    if len(sys.argv) > 1:
-        print("Usage: gradient_boosted_trees", file=sys.stderr)
-        exit(1)
-    sc = SparkContext(appName="PythonGradientBoostedTrees")
-
-    # Load and parse the data file into an RDD of LabeledPoint.
-    data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
-    # Split the data into training and test sets (30% held out for testing)
-    (trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-    print('\nRunning example of classification using GradientBoostedTrees\n')
-    testClassification(trainingData, testData)
-
-    print('\nRunning example of regression using GradientBoostedTrees\n')
-    testRegression(trainingData, testData)
-
-    sc.stop()
diff --git a/examples/src/main/python/mllib/random_forest_example.py b/examples/src/main/python/mllib/random_forest_example.py
deleted file mode 100755
index 4cfdad868c66e..0000000000000
--- a/examples/src/main/python/mllib/random_forest_example.py
+++ /dev/null
@@ -1,90 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-"""
-Random Forest classification and regression using MLlib.
-
-Note: This example illustrates binary classification.
-      For information on multiclass classification, please refer to the decision_tree_runner.py
-      example.
-"""
-from __future__ import print_function
-
-import sys
-
-from pyspark.context import SparkContext
-from pyspark.mllib.tree import RandomForest
-from pyspark.mllib.util import MLUtils
-
-
-def testClassification(trainingData, testData):
-    # Train a RandomForest model.
-    #  Empty categoricalFeaturesInfo indicates all features are continuous.
-    #  Note: Use larger numTrees in practice.
-    #  Setting featureSubsetStrategy="auto" lets the algorithm choose.
-    model = RandomForest.trainClassifier(trainingData, numClasses=2,
-                                         categoricalFeaturesInfo={},
-                                         numTrees=3, featureSubsetStrategy="auto",
-                                         impurity='gini', maxDepth=4, maxBins=32)
-
-    # Evaluate model on test instances and compute test error
-    predictions = model.predict(testData.map(lambda x: x.features))
-    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-    testErr = labelsAndPredictions.filter(lambda v_p: v_p[0] != v_p[1]).count()\
-        / float(testData.count())
-    print('Test Error = ' + str(testErr))
-    print('Learned classification forest model:')
-    print(model.toDebugString())
-
-
-def testRegression(trainingData, testData):
-    # Train a RandomForest model.
-    #  Empty categoricalFeaturesInfo indicates all features are continuous.
-    #  Note: Use larger numTrees in practice.
-    #  Setting featureSubsetStrategy="auto" lets the algorithm choose.
-    model = RandomForest.trainRegressor(trainingData, categoricalFeaturesInfo={},
-                                        numTrees=3, featureSubsetStrategy="auto",
-                                        impurity='variance', maxDepth=4, maxBins=32)
-
-    # Evaluate model on test instances and compute test error
-    predictions = model.predict(testData.map(lambda x: x.features))
-    labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)
-    testMSE = labelsAndPredictions.map(lambda v_p1: (v_p1[0] - v_p1[1]) * (v_p1[0] - v_p1[1]))\
-        .sum() / float(testData.count())
-    print('Test Mean Squared Error = ' + str(testMSE))
-    print('Learned regression forest model:')
-    print(model.toDebugString())
-
-
-if __name__ == "__main__":
-    if len(sys.argv) > 1:
-        print("Usage: random_forest_example", file=sys.stderr)
-        exit(1)
-    sc = SparkContext(appName="PythonRandomForestExample")
-
-    # Load and parse the data file into an RDD of LabeledPoint.
-    data = MLUtils.loadLibSVMFile(sc, 'data/mllib/sample_libsvm_data.txt')
-    # Split the data into training and test sets (30% held out for testing)
-    (trainingData, testData) = data.randomSplit([0.7, 0.3])
-
-    print('\nRunning example of classification using RandomForest\n')
-    testClassification(trainingData, testData)
-
-    print('\nRunning example of regression using RandomForest\n')
-    testRegression(trainingData, testData)
-
-    sc.stop()

From 55358889309cf2d856b72e72e0f3081dfdf61cfa Mon Sep 17 00:00:00 2001
From: Feynman Liang <feynman.liang@gmail.com>
Date: Mon, 30 Nov 2015 15:38:44 -0800
Subject: [PATCH 507/518] [SPARK-11960][MLLIB][DOC] User guide for streaming
 tests

CC jkbradley mengxr josepablocam

Author: Feynman Liang <feynman.liang@gmail.com>

Closes #10005 from feynmanliang/streaming-test-user-guide.
---
 docs/mllib-guide.md                           |  1 +
 docs/mllib-statistics.md                      | 25 +++++++++++++++++++
 .../examples/mllib/StreamingTestExample.scala |  2 ++
 3 files changed, 28 insertions(+)

diff --git a/docs/mllib-guide.md b/docs/mllib-guide.md
index 54e35fcbb15af..43772adcf26e1 100644
--- a/docs/mllib-guide.md
+++ b/docs/mllib-guide.md
@@ -34,6 +34,7 @@ We list major functionality from both below, with links to detailed guides.
   * [correlations](mllib-statistics.html#correlations)
   * [stratified sampling](mllib-statistics.html#stratified-sampling)
   * [hypothesis testing](mllib-statistics.html#hypothesis-testing)
+  * [streaming significance testing](mllib-statistics.html#streaming-significance-testing)
   * [random data generation](mllib-statistics.html#random-data-generation)
 * [Classification and regression](mllib-classification-regression.html)
   * [linear models (SVMs, logistic regression, linear regression)](mllib-linear-methods.html)
diff --git a/docs/mllib-statistics.md b/docs/mllib-statistics.md
index ade5b0768aefe..de209f68e19ca 100644
--- a/docs/mllib-statistics.md
+++ b/docs/mllib-statistics.md
@@ -521,6 +521,31 @@ print(testResult) # summary of the test including the p-value, test statistic,
 </div>
 </div>
 
+### Streaming Significance Testing
+MLlib provides online implementations of some tests to support use cases
+like A/B testing. These tests may be performed on a Spark Streaming
+`DStream[(Boolean,Double)]` where the first element of each tuple
+indicates control group (`false`) or treatment group (`true`) and the
+second element is the value of an observation.
+
+Streaming significance testing supports the following parameters:
+
+* `peacePeriod` - The number of initial data points from the stream to
+ignore, used to mitigate novelty effects.
+* `windowSize` - The number of past batches to perform hypothesis
+testing over. Setting to `0` will perform cumulative processing using
+all prior batches.
+
+
+<div class="codetabs">
+<div data-lang="scala" markdown="1">
+[`StreamingTest`](api/scala/index.html#org.apache.spark.mllib.stat.test.StreamingTest)
+provides streaming hypothesis testing.
+
+{% include_example scala/org/apache/spark/examples/mllib/StreamingTestExample.scala %}
+</div>
+</div>
+
 
 ## Random data generation
 
diff --git a/examples/src/main/scala/org/apache/spark/examples/mllib/StreamingTestExample.scala b/examples/src/main/scala/org/apache/spark/examples/mllib/StreamingTestExample.scala
index ab29f90254d34..b6677c6476639 100644
--- a/examples/src/main/scala/org/apache/spark/examples/mllib/StreamingTestExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/mllib/StreamingTestExample.scala
@@ -64,6 +64,7 @@ object StreamingTestExample {
       dir.toString
     })
 
+    // $example on$
     val data = ssc.textFileStream(dataDir).map(line => line.split(",") match {
       case Array(label, value) => (label.toBoolean, value.toDouble)
     })
@@ -75,6 +76,7 @@ object StreamingTestExample {
 
     val out = streamingTest.registerStream(data)
     out.print()
+    // $example off$
 
     // Stop processing if test becomes significant or we time out
     var timeoutCounter = numBatchesTimeout

From ecc00ec3faf09b0e21bc4b468cb45e45d05cf482 Mon Sep 17 00:00:00 2001
From: Davies Liu <davies.liu@gmail.com>
Date: Mon, 30 Nov 2015 15:42:10 -0800
Subject: [PATCH 508/518] fix Maven build

---
 .../main/scala/org/apache/spark/sql/execution/SparkPlan.scala   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
index 507641ff8263e..a78177751c9dc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlan.scala
@@ -43,7 +43,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
    * populated by the query planning infrastructure.
    */
   @transient
-  protected[spark] final val sqlContext = SQLContext.getActive().get
+  protected[spark] final val sqlContext = SQLContext.getActive().getOrElse(null)
 
   protected def sparkContext = sqlContext.sparkContext
 

From edb26e7f4e1164645971c9a139eb29ddec8acc5d Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Mon, 30 Nov 2015 16:31:59 -0800
Subject: [PATCH 509/518] [SPARK-12058][HOTFIX] Disable KinesisStreamTests
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

KinesisStreamTests in test.py is broken because of #9403. See https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/46896/testReport/(root)/KinesisStreamTests/test_kinesis_stream/

Because Streaming Python didn’t work when merging https://github.com/apache/spark/pull/9403, the PR build didn’t report the Python test failure actually.

This PR just disabled the test to unblock #10039

Author: Shixiong Zhu <shixiong@databricks.com>

Closes #10047 from zsxwing/disable-python-kinesis-test.
---
 python/pyspark/streaming/tests.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/pyspark/streaming/tests.py b/python/pyspark/streaming/tests.py
index d380d697bc51c..a647e6bf39581 100644
--- a/python/pyspark/streaming/tests.py
+++ b/python/pyspark/streaming/tests.py
@@ -1409,6 +1409,7 @@ def test_kinesis_stream_api(self):
             InitialPositionInStream.LATEST, 2, StorageLevel.MEMORY_AND_DISK_2,
             "awsAccessKey", "awsSecretKey")
 
+    @unittest.skip("Enable it when we fix SPAKR-12058")
     def test_kinesis_stream(self):
         if not are_kinesis_tests_enabled:
             sys.stderr.write(

From d3ca8cfac286ae19f8bedc736877ea9d0a0a072c Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Mon, 30 Nov 2015 16:37:27 -0800
Subject: [PATCH 510/518] [SPARK-12000] Fix API doc generation issues

This pull request fixes multiple issues with API doc generation.

- Modify the Jekyll plugin so that the entire doc build fails if API docs cannot be generated. This will make it easy to detect when the doc build breaks, since this will now trigger Jenkins failures.
- Change how we handle the `-target` compiler option flag in order to fix `javadoc` generation.
- Incorporate doc changes from thunterdb (in #10048).

Closes #10048.

Author: Josh Rosen <joshrosen@databricks.com>
Author: Timothy Hunter <timhunter@databricks.com>

Closes #10049 from JoshRosen/fix-doc-build.
---
 docs/_plugins/copy_api_dirs.rb                        |  6 +++---
 .../apache/spark/network/client/StreamCallback.java   |  4 ++--
 .../org/apache/spark/network/server/RpcHandler.java   |  2 +-
 project/SparkBuild.scala                              | 11 ++++++++---
 4 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/docs/_plugins/copy_api_dirs.rb b/docs/_plugins/copy_api_dirs.rb
index 01718d98dffe0..f2f3e2e653149 100644
--- a/docs/_plugins/copy_api_dirs.rb
+++ b/docs/_plugins/copy_api_dirs.rb
@@ -27,7 +27,7 @@
     cd("..")
 
     puts "Running 'build/sbt -Pkinesis-asl clean compile unidoc' from " + pwd + "; this may take a few minutes..."
-    puts `build/sbt -Pkinesis-asl clean compile unidoc`
+    system("build/sbt -Pkinesis-asl clean compile unidoc") || raise("Unidoc generation failed")
 
     puts "Moving back into docs dir."
     cd("docs")
@@ -117,7 +117,7 @@
 
   puts "Moving to python/docs directory and building sphinx."
   cd("../python/docs")
-  puts `make html`
+  system(make html) || raise("Python doc generation failed")
 
   puts "Moving back into home dir."
   cd("../../")
@@ -131,7 +131,7 @@
   # Build SparkR API docs
   puts "Moving to R directory and building roxygen docs."
   cd("R")
-  puts `./create-docs.sh`
+  system("./create-docs.sh") || raise("R doc generation failed")
 
   puts "Moving back into home dir."
   cd("../")
diff --git a/network/common/src/main/java/org/apache/spark/network/client/StreamCallback.java b/network/common/src/main/java/org/apache/spark/network/client/StreamCallback.java
index 093fada320cc3..51d34cac6e636 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/StreamCallback.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/StreamCallback.java
@@ -21,8 +21,8 @@
 import java.nio.ByteBuffer;
 
 /**
- * Callback for streaming data. Stream data will be offered to the {@link onData(ByteBuffer)}
- * method as it arrives. Once all the stream data is received, {@link onComplete()} will be
+ * Callback for streaming data. Stream data will be offered to the {@link onData(String, ByteBuffer)}
+ * method as it arrives. Once all the stream data is received, {@link onComplete(String)} will be
  * called.
  * <p>
  * The network library guarantees that a single thread will call these methods at a time, but
diff --git a/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java b/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
index 65109ddfe13b9..1a11f7b3820c6 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
@@ -55,7 +55,7 @@ public abstract void receive(
 
   /**
    * Receives an RPC message that does not expect a reply. The default implementation will
-   * call "{@link receive(TransportClient, byte[], RpcResponseCallback}" and log a warning if
+   * call "{@link receive(TransportClient, byte[], RpcResponseCallback)}" and log a warning if
    * any of the callback methods are called.
    *
    * @param client A channel client which enables the handler to make requests back to the sender
diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala
index f575f0012d59e..63290d8a666e6 100644
--- a/project/SparkBuild.scala
+++ b/project/SparkBuild.scala
@@ -160,7 +160,12 @@ object SparkBuild extends PomBuild {
 
     javacOptions in Compile ++= Seq(
       "-encoding", "UTF-8",
-      "-source", javacJVMVersion.value,
+      "-source", javacJVMVersion.value
+    ),
+    // This -target option cannot be set in the Compile configuration scope since `javadoc` doesn't
+    // play nicely with it; see https://github.com/sbt/sbt/issues/355#issuecomment-3817629 for
+    // additional discussion and explanation.
+    javacOptions in (Compile, compile) ++= Seq(
       "-target", javacJVMVersion.value
     ),
 
@@ -547,9 +552,9 @@ object Unidoc {
     publish := {},
 
     unidocProjectFilter in(ScalaUnidoc, unidoc) :=
-      inAnyProject -- inProjects(OldDeps.project, repl, examples, tools, streamingFlumeSink, yarn),
+      inAnyProject -- inProjects(OldDeps.project, repl, examples, tools, streamingFlumeSink, yarn, testTags),
     unidocProjectFilter in(JavaUnidoc, unidoc) :=
-      inAnyProject -- inProjects(OldDeps.project, repl, bagel, examples, tools, streamingFlumeSink, yarn),
+      inAnyProject -- inProjects(OldDeps.project, repl, bagel, examples, tools, streamingFlumeSink, yarn, testTags),
 
     // Skip actual catalyst, but include the subproject.
     // Catalyst is not public API and contains quasiquotes which break scaladoc.

From e6dc89a33951e9197a77dbcacf022c27469ae41e Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Mon, 30 Nov 2015 17:18:44 -0800
Subject: [PATCH 511/518] [SPARK-12035] Add more debug information in
 include_example tag of Jekyll

https://issues.apache.org/jira/browse/SPARK-12035

When we debuging lots of example code files, like in https://github.com/apache/spark/pull/10002, it's hard to know which file causes errors due to limited information in `include_example.rb`. With their filenames, we can locate bugs easily.

Author: Xusen Yin <yinxusen@gmail.com>

Closes #10026 from yinxusen/SPARK-12035.
---
 docs/_plugins/include_example.rb | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/docs/_plugins/include_example.rb b/docs/_plugins/include_example.rb
index 564c86680f68e..f7485826a762d 100644
--- a/docs/_plugins/include_example.rb
+++ b/docs/_plugins/include_example.rb
@@ -75,10 +75,10 @@ def select_lines(code)
         .select { |l, i| l.include? "$example off$" }
         .map { |l, i| i }
 
-      raise "Start indices amount is not equal to end indices amount, please check the code." \
+      raise "Start indices amount is not equal to end indices amount, see #{@file}." \
         unless startIndices.size == endIndices.size
 
-      raise "No code is selected by include_example, please check the code." \
+      raise "No code is selected by include_example, see #{@file}." \
         if startIndices.size == 0
 
       # Select and join code blocks together, with a space line between each of two continuous
@@ -86,8 +86,10 @@ def select_lines(code)
       lastIndex = -1
       result = ""
       startIndices.zip(endIndices).each do |start, endline|
-        raise "Overlapping between two example code blocks are not allowed." if start <= lastIndex
-        raise "$example on$ should not be in the same line with $example off$." if start == endline
+        raise "Overlapping between two example code blocks are not allowed, see #{@file}." \
+            if start <= lastIndex
+        raise "$example on$ should not be in the same line with $example off$, see #{@file}." \
+            if start == endline
         lastIndex = endline
         range = Range.new(start + 1, endline - 1)
         result += trim_codeblock(lines[range]).join

From 0a46e4377216a1f7de478f220c3b3042a77789e2 Mon Sep 17 00:00:00 2001
From: CodingCat <zhunansjtu@gmail.com>
Date: Mon, 30 Nov 2015 17:19:26 -0800
Subject: [PATCH 512/518] [SPARK-12037][CORE] initialize heartbeatReceiverRef
 before calling startDriverHeartbeat

https://issues.apache.org/jira/browse/SPARK-12037

a simple fix by changing the order of the statements

Author: CodingCat <zhunansjtu@gmail.com>

Closes #10032 from CodingCat/SPARK-12037.
---
 .../main/scala/org/apache/spark/executor/Executor.scala    | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/executor/Executor.scala b/core/src/main/scala/org/apache/spark/executor/Executor.scala
index 6154f06e3ac11..7b68dfe5ad06e 100644
--- a/core/src/main/scala/org/apache/spark/executor/Executor.scala
+++ b/core/src/main/scala/org/apache/spark/executor/Executor.scala
@@ -109,6 +109,10 @@ private[spark] class Executor(
   // Executor for the heartbeat task.
   private val heartbeater = ThreadUtils.newDaemonSingleThreadScheduledExecutor("driver-heartbeater")
 
+  // must be initialized before running startDriverHeartbeat()
+  private val heartbeatReceiverRef =
+    RpcUtils.makeDriverRef(HeartbeatReceiver.ENDPOINT_NAME, conf, env.rpcEnv)
+
   startDriverHeartbeater()
 
   def launchTask(
@@ -411,9 +415,6 @@ private[spark] class Executor(
     }
   }
 
-  private val heartbeatReceiverRef =
-    RpcUtils.makeDriverRef(HeartbeatReceiver.ENDPOINT_NAME, conf, env.rpcEnv)
-
   /** Reports heartbeat and metrics for active tasks to the driver. */
   private def reportHeartBeat(): Unit = {
     // list of (task id, metrics) to send back to the driver

From 9bf2120672ae0f620a217ccd96bef189ab75e0d6 Mon Sep 17 00:00:00 2001
From: Marcelo Vanzin <vanzin@cloudera.com>
Date: Mon, 30 Nov 2015 17:22:05 -0800
Subject: [PATCH 513/518] [SPARK-12007][NETWORK] Avoid copies in the network
 lib's RPC layer.

This change seems large, but most of it is just replacing `byte[]`
with `ByteBuffer` and `new byte[]` with `ByteBuffer.allocate()`,
since it changes the network library's API.

The following are parts of the code that actually have meaningful
changes:

- The Message implementations were changed to inherit from a new
  AbstractMessage that can optionally hold a reference to a body
  (in the form of a ManagedBuffer); this is similar to how
  ResponseWithBody worked before, except now it's not restricted
  to just responses.

- The TransportFrameDecoder was pretty much rewritten to avoid
  copies as much as possible; it doesn't rely on CompositeByteBuf
  to accumulate incoming data anymore, since CompositeByteBuf
  has issues when slices are retained. The code now is able to
  create frames without having to resort to copying bytes except
  for a few bytes (containing the frame length) in very rare cases.

- Some minor changes in the SASL layer to convert things back to
  `byte[]` since the JDK SASL API operates on those.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #9987 from vanzin/SPARK-12007.
---
 .../mesos/MesosExternalShuffleService.scala   |   3 +-
 .../network/netty/NettyBlockRpcServer.scala   |   8 +-
 .../netty/NettyBlockTransferService.scala     |   6 +-
 .../apache/spark/rpc/netty/NettyRpcEnv.scala  |  16 +-
 .../org/apache/spark/rpc/netty/Outbox.scala   |   9 +-
 .../rpc/netty/NettyRpcHandlerSuite.scala      |   3 +-
 .../network/client/RpcResponseCallback.java   |   4 +-
 .../spark/network/client/TransportClient.java |  16 +-
 .../client/TransportResponseHandler.java      |  16 +-
 .../network/protocol/AbstractMessage.java     |  54 +++++++
 ...Body.java => AbstractResponseMessage.java} |  16 +-
 .../network/protocol/ChunkFetchFailure.java   |   2 +-
 .../network/protocol/ChunkFetchRequest.java   |   2 +-
 .../network/protocol/ChunkFetchSuccess.java   |   8 +-
 .../spark/network/protocol/Message.java       |  11 +-
 .../network/protocol/MessageEncoder.java      |  29 ++--
 .../spark/network/protocol/OneWayMessage.java |  33 ++--
 .../spark/network/protocol/RpcFailure.java    |   2 +-
 .../spark/network/protocol/RpcRequest.java    |  34 +++--
 .../spark/network/protocol/RpcResponse.java   |  39 +++--
 .../spark/network/protocol/StreamFailure.java |   2 +-
 .../spark/network/protocol/StreamRequest.java |   2 +-
 .../network/protocol/StreamResponse.java      |  19 +--
 .../network/sasl/SaslClientBootstrap.java     |  12 +-
 .../spark/network/sasl/SaslMessage.java       |  31 ++--
 .../spark/network/sasl/SaslRpcHandler.java    |  26 +++-
 .../spark/network/server/MessageHandler.java  |   2 +-
 .../spark/network/server/NoOpRpcHandler.java  |   8 +-
 .../spark/network/server/RpcHandler.java      |   8 +-
 .../server/TransportChannelHandler.java       |   2 +-
 .../server/TransportRequestHandler.java       |  15 +-
 .../apache/spark/network/util/JavaUtils.java  |  48 ++++--
 .../network/util/TransportFrameDecoder.java   | 142 ++++++++++++------
 .../network/ChunkFetchIntegrationSuite.java   |   5 +-
 .../apache/spark/network/ProtocolSuite.java   |  10 +-
 .../RequestTimeoutIntegrationSuite.java       |  51 ++++---
 .../spark/network/RpcIntegrationSuite.java    |  26 ++--
 .../org/apache/spark/network/StreamSuite.java |   5 +-
 .../TransportResponseHandlerSuite.java        |  24 +--
 .../spark/network/sasl/SparkSaslSuite.java    |  43 ++++--
 .../util/TransportFrameDecoderSuite.java      |  23 ++-
 .../shuffle/ExternalShuffleBlockHandler.java  |   9 +-
 .../shuffle/ExternalShuffleClient.java        |   3 +-
 .../shuffle/OneForOneBlockFetcher.java        |   7 +-
 .../mesos/MesosExternalShuffleClient.java     |   5 +-
 .../protocol/BlockTransferMessage.java        |   8 +-
 .../network/sasl/SaslIntegrationSuite.java    |  18 ++-
 .../shuffle/BlockTransferMessagesSuite.java   |   2 +-
 .../ExternalShuffleBlockHandlerSuite.java     |  21 +--
 .../shuffle/OneForOneBlockFetcherSuite.java   |   8 +-
 50 files changed, 589 insertions(+), 307 deletions(-)
 create mode 100644 network/common/src/main/java/org/apache/spark/network/protocol/AbstractMessage.java
 rename network/common/src/main/java/org/apache/spark/network/protocol/{ResponseWithBody.java => AbstractResponseMessage.java} (63%)

diff --git a/core/src/main/scala/org/apache/spark/deploy/mesos/MesosExternalShuffleService.scala b/core/src/main/scala/org/apache/spark/deploy/mesos/MesosExternalShuffleService.scala
index 12337a940a414..8ffcfc0878a42 100644
--- a/core/src/main/scala/org/apache/spark/deploy/mesos/MesosExternalShuffleService.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/mesos/MesosExternalShuffleService.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.deploy.mesos
 
 import java.net.SocketAddress
+import java.nio.ByteBuffer
 
 import scala.collection.mutable
 
@@ -56,7 +57,7 @@ private[mesos] class MesosExternalShuffleBlockHandler(transportConf: TransportCo
           }
         }
         connectedApps(address) = appId
-        callback.onSuccess(new Array[Byte](0))
+        callback.onSuccess(ByteBuffer.allocate(0))
       case _ => super.handleMessage(message, client, callback)
     }
   }
diff --git a/core/src/main/scala/org/apache/spark/network/netty/NettyBlockRpcServer.scala b/core/src/main/scala/org/apache/spark/network/netty/NettyBlockRpcServer.scala
index 76968249fb625..df8c21fb837ed 100644
--- a/core/src/main/scala/org/apache/spark/network/netty/NettyBlockRpcServer.scala
+++ b/core/src/main/scala/org/apache/spark/network/netty/NettyBlockRpcServer.scala
@@ -47,9 +47,9 @@ class NettyBlockRpcServer(
 
   override def receive(
       client: TransportClient,
-      messageBytes: Array[Byte],
+      rpcMessage: ByteBuffer,
       responseContext: RpcResponseCallback): Unit = {
-    val message = BlockTransferMessage.Decoder.fromByteArray(messageBytes)
+    val message = BlockTransferMessage.Decoder.fromByteBuffer(rpcMessage)
     logTrace(s"Received request: $message")
 
     message match {
@@ -58,7 +58,7 @@ class NettyBlockRpcServer(
           openBlocks.blockIds.map(BlockId.apply).map(blockManager.getBlockData)
         val streamId = streamManager.registerStream(appId, blocks.iterator.asJava)
         logTrace(s"Registered streamId $streamId with ${blocks.size} buffers")
-        responseContext.onSuccess(new StreamHandle(streamId, blocks.size).toByteArray)
+        responseContext.onSuccess(new StreamHandle(streamId, blocks.size).toByteBuffer)
 
       case uploadBlock: UploadBlock =>
         // StorageLevel is serialized as bytes using our JavaSerializer.
@@ -66,7 +66,7 @@ class NettyBlockRpcServer(
           serializer.newInstance().deserialize(ByteBuffer.wrap(uploadBlock.metadata))
         val data = new NioManagedBuffer(ByteBuffer.wrap(uploadBlock.blockData))
         blockManager.putBlockData(BlockId(uploadBlock.blockId), data, level)
-        responseContext.onSuccess(new Array[Byte](0))
+        responseContext.onSuccess(ByteBuffer.allocate(0))
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala b/core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala
index b0694e3c6c8af..82c16e855b0c0 100644
--- a/core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala
+++ b/core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala
@@ -17,6 +17,8 @@
 
 package org.apache.spark.network.netty
 
+import java.nio.ByteBuffer
+
 import scala.collection.JavaConverters._
 import scala.concurrent.{Future, Promise}
 
@@ -133,9 +135,9 @@ class NettyBlockTransferService(conf: SparkConf, securityManager: SecurityManage
       data
     }
 
-    client.sendRpc(new UploadBlock(appId, execId, blockId.toString, levelBytes, array).toByteArray,
+    client.sendRpc(new UploadBlock(appId, execId, blockId.toString, levelBytes, array).toByteBuffer,
       new RpcResponseCallback {
-        override def onSuccess(response: Array[Byte]): Unit = {
+        override def onSuccess(response: ByteBuffer): Unit = {
           logTrace(s"Successfully uploaded block $blockId")
           result.success((): Unit)
         }
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
index c7d74fa1d9195..68c5f44145b0d 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/NettyRpcEnv.scala
@@ -241,16 +241,14 @@ private[netty] class NettyRpcEnv(
     promise.future.mapTo[T].recover(timeout.addMessageIfTimeout)(ThreadUtils.sameThread)
   }
 
-  private[netty] def serialize(content: Any): Array[Byte] = {
-    val buffer = javaSerializerInstance.serialize(content)
-    java.util.Arrays.copyOfRange(
-      buffer.array(), buffer.arrayOffset + buffer.position, buffer.arrayOffset + buffer.limit)
+  private[netty] def serialize(content: Any): ByteBuffer = {
+    javaSerializerInstance.serialize(content)
   }
 
-  private[netty] def deserialize[T: ClassTag](client: TransportClient, bytes: Array[Byte]): T = {
+  private[netty] def deserialize[T: ClassTag](client: TransportClient, bytes: ByteBuffer): T = {
     NettyRpcEnv.currentClient.withValue(client) {
       deserialize { () =>
-        javaSerializerInstance.deserialize[T](ByteBuffer.wrap(bytes))
+        javaSerializerInstance.deserialize[T](bytes)
       }
     }
   }
@@ -557,7 +555,7 @@ private[netty] class NettyRpcHandler(
 
   override def receive(
       client: TransportClient,
-      message: Array[Byte],
+      message: ByteBuffer,
       callback: RpcResponseCallback): Unit = {
     val messageToDispatch = internalReceive(client, message)
     dispatcher.postRemoteMessage(messageToDispatch, callback)
@@ -565,12 +563,12 @@ private[netty] class NettyRpcHandler(
 
   override def receive(
       client: TransportClient,
-      message: Array[Byte]): Unit = {
+      message: ByteBuffer): Unit = {
     val messageToDispatch = internalReceive(client, message)
     dispatcher.postOneWayMessage(messageToDispatch)
   }
 
-  private def internalReceive(client: TransportClient, message: Array[Byte]): RequestMessage = {
+  private def internalReceive(client: TransportClient, message: ByteBuffer): RequestMessage = {
     val addr = client.getChannel().remoteAddress().asInstanceOf[InetSocketAddress]
     assert(addr != null)
     val clientAddr = RpcAddress(addr.getHostName, addr.getPort)
diff --git a/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala b/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
index 36fdd00bbc4c2..2316ebe347bb7 100644
--- a/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
+++ b/core/src/main/scala/org/apache/spark/rpc/netty/Outbox.scala
@@ -17,6 +17,7 @@
 
 package org.apache.spark.rpc.netty
 
+import java.nio.ByteBuffer
 import java.util.concurrent.Callable
 import javax.annotation.concurrent.GuardedBy
 
@@ -34,7 +35,7 @@ private[netty] sealed trait OutboxMessage {
 
 }
 
-private[netty] case class OneWayOutboxMessage(content: Array[Byte]) extends OutboxMessage
+private[netty] case class OneWayOutboxMessage(content: ByteBuffer) extends OutboxMessage
   with Logging {
 
   override def sendWith(client: TransportClient): Unit = {
@@ -48,9 +49,9 @@ private[netty] case class OneWayOutboxMessage(content: Array[Byte]) extends Outb
 }
 
 private[netty] case class RpcOutboxMessage(
-    content: Array[Byte],
+    content: ByteBuffer,
     _onFailure: (Throwable) => Unit,
-    _onSuccess: (TransportClient, Array[Byte]) => Unit)
+    _onSuccess: (TransportClient, ByteBuffer) => Unit)
   extends OutboxMessage with RpcResponseCallback {
 
   private var client: TransportClient = _
@@ -70,7 +71,7 @@ private[netty] case class RpcOutboxMessage(
     _onFailure(e)
   }
 
-  override def onSuccess(response: Array[Byte]): Unit = {
+  override def onSuccess(response: ByteBuffer): Unit = {
     _onSuccess(client, response)
   }
 
diff --git a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
index 323184cdd9b6e..ebd6f700710bd 100644
--- a/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rpc/netty/NettyRpcHandlerSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.rpc.netty
 
 import java.net.InetSocketAddress
+import java.nio.ByteBuffer
 
 import io.netty.channel.Channel
 import org.mockito.Mockito._
@@ -32,7 +33,7 @@ class NettyRpcHandlerSuite extends SparkFunSuite {
 
   val env = mock(classOf[NettyRpcEnv])
   val sm = mock(classOf[StreamManager])
-  when(env.deserialize(any(classOf[TransportClient]), any(classOf[Array[Byte]]))(any()))
+  when(env.deserialize(any(classOf[TransportClient]), any(classOf[ByteBuffer]))(any()))
     .thenReturn(RequestMessage(RpcAddress("localhost", 12345), null, null))
 
   test("receive") {
diff --git a/network/common/src/main/java/org/apache/spark/network/client/RpcResponseCallback.java b/network/common/src/main/java/org/apache/spark/network/client/RpcResponseCallback.java
index 6ec960d795420..47e93f9846fa6 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/RpcResponseCallback.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/RpcResponseCallback.java
@@ -17,13 +17,15 @@
 
 package org.apache.spark.network.client;
 
+import java.nio.ByteBuffer;
+
 /**
  * Callback for the result of a single RPC. This will be invoked once with either success or
  * failure.
  */
 public interface RpcResponseCallback {
   /** Successful serialized result from server. */
-  void onSuccess(byte[] response);
+  void onSuccess(ByteBuffer response);
 
   /** Exception either propagated from server or raised on client side. */
   void onFailure(Throwable e);
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java b/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
index 8a58e7b24585b..c49ca4d5ee925 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportClient.java
@@ -20,6 +20,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 import java.net.SocketAddress;
+import java.nio.ByteBuffer;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -36,6 +37,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.spark.network.buffer.NioManagedBuffer;
 import org.apache.spark.network.protocol.ChunkFetchRequest;
 import org.apache.spark.network.protocol.OneWayMessage;
 import org.apache.spark.network.protocol.RpcRequest;
@@ -212,7 +214,7 @@ public void operationComplete(ChannelFuture future) throws Exception {
    * @param callback Callback to handle the RPC's reply.
    * @return The RPC's id.
    */
-  public long sendRpc(byte[] message, final RpcResponseCallback callback) {
+  public long sendRpc(ByteBuffer message, final RpcResponseCallback callback) {
     final String serverAddr = NettyUtils.getRemoteAddress(channel);
     final long startTime = System.currentTimeMillis();
     logger.trace("Sending RPC to {}", serverAddr);
@@ -220,7 +222,7 @@ public long sendRpc(byte[] message, final RpcResponseCallback callback) {
     final long requestId = Math.abs(UUID.randomUUID().getLeastSignificantBits());
     handler.addRpcRequest(requestId, callback);
 
-    channel.writeAndFlush(new RpcRequest(requestId, message)).addListener(
+    channel.writeAndFlush(new RpcRequest(requestId, new NioManagedBuffer(message))).addListener(
       new ChannelFutureListener() {
         @Override
         public void operationComplete(ChannelFuture future) throws Exception {
@@ -249,12 +251,12 @@ public void operationComplete(ChannelFuture future) throws Exception {
    * Synchronously sends an opaque message to the RpcHandler on the server-side, waiting for up to
    * a specified timeout for a response.
    */
-  public byte[] sendRpcSync(byte[] message, long timeoutMs) {
-    final SettableFuture<byte[]> result = SettableFuture.create();
+  public ByteBuffer sendRpcSync(ByteBuffer message, long timeoutMs) {
+    final SettableFuture<ByteBuffer> result = SettableFuture.create();
 
     sendRpc(message, new RpcResponseCallback() {
       @Override
-      public void onSuccess(byte[] response) {
+      public void onSuccess(ByteBuffer response) {
         result.set(response);
       }
 
@@ -279,8 +281,8 @@ public void onFailure(Throwable e) {
    *
    * @param message The message to send.
    */
-  public void send(byte[] message) {
-    channel.writeAndFlush(new OneWayMessage(message));
+  public void send(ByteBuffer message) {
+    channel.writeAndFlush(new OneWayMessage(new NioManagedBuffer(message)));
   }
 
   /**
diff --git a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
index 4c15045363b84..23a8dba593442 100644
--- a/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java
@@ -136,7 +136,7 @@ public void exceptionCaught(Throwable cause) {
   }
 
   @Override
-  public void handle(ResponseMessage message) {
+  public void handle(ResponseMessage message) throws Exception {
     String remoteAddress = NettyUtils.getRemoteAddress(channel);
     if (message instanceof ChunkFetchSuccess) {
       ChunkFetchSuccess resp = (ChunkFetchSuccess) message;
@@ -144,11 +144,11 @@ public void handle(ResponseMessage message) {
       if (listener == null) {
         logger.warn("Ignoring response for block {} from {} since it is not outstanding",
           resp.streamChunkId, remoteAddress);
-        resp.body.release();
+        resp.body().release();
       } else {
         outstandingFetches.remove(resp.streamChunkId);
-        listener.onSuccess(resp.streamChunkId.chunkIndex, resp.body);
-        resp.body.release();
+        listener.onSuccess(resp.streamChunkId.chunkIndex, resp.body());
+        resp.body().release();
       }
     } else if (message instanceof ChunkFetchFailure) {
       ChunkFetchFailure resp = (ChunkFetchFailure) message;
@@ -166,10 +166,14 @@ public void handle(ResponseMessage message) {
       RpcResponseCallback listener = outstandingRpcs.get(resp.requestId);
       if (listener == null) {
         logger.warn("Ignoring response for RPC {} from {} ({} bytes) since it is not outstanding",
-          resp.requestId, remoteAddress, resp.response.length);
+          resp.requestId, remoteAddress, resp.body().size());
       } else {
         outstandingRpcs.remove(resp.requestId);
-        listener.onSuccess(resp.response);
+        try {
+          listener.onSuccess(resp.body().nioByteBuffer());
+        } finally {
+          resp.body().release();
+        }
       }
     } else if (message instanceof RpcFailure) {
       RpcFailure resp = (RpcFailure) message;
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/AbstractMessage.java b/network/common/src/main/java/org/apache/spark/network/protocol/AbstractMessage.java
new file mode 100644
index 0000000000000..2924218c2f08b
--- /dev/null
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/AbstractMessage.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.network.protocol;
+
+import com.google.common.base.Objects;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
+
+/**
+ * Abstract class for messages which optionally contain a body kept in a separate buffer.
+ */
+public abstract class AbstractMessage implements Message {
+  private final ManagedBuffer body;
+  private final boolean isBodyInFrame;
+
+  protected AbstractMessage() {
+    this(null, false);
+  }
+
+  protected AbstractMessage(ManagedBuffer body, boolean isBodyInFrame) {
+    this.body = body;
+    this.isBodyInFrame = isBodyInFrame;
+  }
+
+  @Override
+  public ManagedBuffer body() {
+    return body;
+  }
+
+  @Override
+  public boolean isBodyInFrame() {
+    return isBodyInFrame;
+  }
+
+  protected boolean equals(AbstractMessage other) {
+    return isBodyInFrame == other.isBodyInFrame && Objects.equal(body, other.body);
+  }
+
+}
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/ResponseWithBody.java b/network/common/src/main/java/org/apache/spark/network/protocol/AbstractResponseMessage.java
similarity index 63%
rename from network/common/src/main/java/org/apache/spark/network/protocol/ResponseWithBody.java
rename to network/common/src/main/java/org/apache/spark/network/protocol/AbstractResponseMessage.java
index 67be77e39f711..c362c92fc4f52 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/ResponseWithBody.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/AbstractResponseMessage.java
@@ -17,23 +17,15 @@
 
 package org.apache.spark.network.protocol;
 
-import com.google.common.base.Objects;
-import io.netty.buffer.ByteBuf;
-
 import org.apache.spark.network.buffer.ManagedBuffer;
-import org.apache.spark.network.buffer.NettyManagedBuffer;
 
 /**
- * Abstract class for response messages that contain a large data portion kept in a separate
- * buffer. These messages are treated especially by MessageEncoder.
+ * Abstract class for response messages.
  */
-public abstract class ResponseWithBody implements ResponseMessage {
-  public final ManagedBuffer body;
-  public final boolean isBodyInFrame;
+public abstract class AbstractResponseMessage extends AbstractMessage implements ResponseMessage {
 
-  protected ResponseWithBody(ManagedBuffer body, boolean isBodyInFrame) {
-    this.body = body;
-    this.isBodyInFrame = isBodyInFrame;
+  protected AbstractResponseMessage(ManagedBuffer body, boolean isBodyInFrame) {
+    super(body, isBodyInFrame);
   }
 
   public abstract ResponseMessage createFailureResponse(String error);
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchFailure.java b/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchFailure.java
index f0363830b61ac..7b28a9a969486 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchFailure.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchFailure.java
@@ -23,7 +23,7 @@
 /**
  * Response to {@link ChunkFetchRequest} when there is an error fetching the chunk.
  */
-public final class ChunkFetchFailure implements ResponseMessage {
+public final class ChunkFetchFailure extends AbstractMessage implements ResponseMessage {
   public final StreamChunkId streamChunkId;
   public final String errorString;
 
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchRequest.java b/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchRequest.java
index 5a173af54f618..26d063feb5fe3 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchRequest.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchRequest.java
@@ -24,7 +24,7 @@
  * Request to fetch a sequence of a single chunk of a stream. This will correspond to a single
  * {@link org.apache.spark.network.protocol.ResponseMessage} (either success or failure).
  */
-public final class ChunkFetchRequest implements RequestMessage {
+public final class ChunkFetchRequest extends AbstractMessage implements RequestMessage {
   public final StreamChunkId streamChunkId;
 
   public ChunkFetchRequest(StreamChunkId streamChunkId) {
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchSuccess.java b/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchSuccess.java
index e6a7e9a8b4145..94c2ac9b20e43 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchSuccess.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/ChunkFetchSuccess.java
@@ -30,7 +30,7 @@
  * may be written by Netty in a more efficient manner (i.e., zero-copy write).
  * Similarly, the client-side decoding will reuse the Netty ByteBuf as the buffer.
  */
-public final class ChunkFetchSuccess extends ResponseWithBody {
+public final class ChunkFetchSuccess extends AbstractResponseMessage {
   public final StreamChunkId streamChunkId;
 
   public ChunkFetchSuccess(StreamChunkId streamChunkId, ManagedBuffer buffer) {
@@ -67,14 +67,14 @@ public static ChunkFetchSuccess decode(ByteBuf buf) {
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(streamChunkId, body);
+    return Objects.hashCode(streamChunkId, body());
   }
 
   @Override
   public boolean equals(Object other) {
     if (other instanceof ChunkFetchSuccess) {
       ChunkFetchSuccess o = (ChunkFetchSuccess) other;
-      return streamChunkId.equals(o.streamChunkId) && body.equals(o.body);
+      return streamChunkId.equals(o.streamChunkId) && super.equals(o);
     }
     return false;
   }
@@ -83,7 +83,7 @@ public boolean equals(Object other) {
   public String toString() {
     return Objects.toStringHelper(this)
       .add("streamChunkId", streamChunkId)
-      .add("buffer", body)
+      .add("buffer", body())
       .toString();
   }
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/Message.java b/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
index 39afd03db60ee..66f5b8b3a59c8 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/Message.java
@@ -19,17 +19,25 @@
 
 import io.netty.buffer.ByteBuf;
 
+import org.apache.spark.network.buffer.ManagedBuffer;
+
 /** An on-the-wire transmittable message. */
 public interface Message extends Encodable {
   /** Used to identify this request type. */
   Type type();
 
+  /** An optional body for the message. */
+  ManagedBuffer body();
+
+  /** Whether to include the body of the message in the same frame as the message. */
+  boolean isBodyInFrame();
+
   /** Preceding every serialized Message is its type, which allows us to deserialize it. */
   public static enum Type implements Encodable {
     ChunkFetchRequest(0), ChunkFetchSuccess(1), ChunkFetchFailure(2),
     RpcRequest(3), RpcResponse(4), RpcFailure(5),
     StreamRequest(6), StreamResponse(7), StreamFailure(8),
-    OneWayMessage(9);
+    OneWayMessage(9), User(-1);
 
     private final byte id;
 
@@ -57,6 +65,7 @@ public static Type decode(ByteBuf buf) {
         case 7: return StreamResponse;
         case 8: return StreamFailure;
         case 9: return OneWayMessage;
+        case -1: throw new IllegalArgumentException("User type messages cannot be decoded.");
         default: throw new IllegalArgumentException("Unknown message type: " + id);
       }
     }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java b/network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java
index 6cce97c807dc0..abca22347b783 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java
@@ -42,25 +42,28 @@ public final class MessageEncoder extends MessageToMessageEncoder<Message> {
    * data to 'out', in order to enable zero-copy transfer.
    */
   @Override
-  public void encode(ChannelHandlerContext ctx, Message in, List<Object> out) {
+  public void encode(ChannelHandlerContext ctx, Message in, List<Object> out) throws Exception {
     Object body = null;
     long bodyLength = 0;
     boolean isBodyInFrame = false;
 
-    // Detect ResponseWithBody messages and get the data buffer out of them.
-    // The body is used in order to enable zero-copy transfer for the payload.
-    if (in instanceof ResponseWithBody) {
-      ResponseWithBody resp = (ResponseWithBody) in;
+    // If the message has a body, take it out to enable zero-copy transfer for the payload.
+    if (in.body() != null) {
       try {
-        bodyLength = resp.body.size();
-        body = resp.body.convertToNetty();
-        isBodyInFrame = resp.isBodyInFrame;
+        bodyLength = in.body().size();
+        body = in.body().convertToNetty();
+        isBodyInFrame = in.isBodyInFrame();
       } catch (Exception e) {
-        // Re-encode this message as a failure response.
-        String error = e.getMessage() != null ? e.getMessage() : "null";
-        logger.error(String.format("Error processing %s for client %s",
-          resp, ctx.channel().remoteAddress()), e);
-        encode(ctx, resp.createFailureResponse(error), out);
+        if (in instanceof AbstractResponseMessage) {
+          AbstractResponseMessage resp = (AbstractResponseMessage) in;
+          // Re-encode this message as a failure response.
+          String error = e.getMessage() != null ? e.getMessage() : "null";
+          logger.error(String.format("Error processing %s for client %s",
+            in, ctx.channel().remoteAddress()), e);
+          encode(ctx, resp.createFailureResponse(error), out);
+        } else {
+          throw e;
+        }
         return;
       }
     }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/OneWayMessage.java b/network/common/src/main/java/org/apache/spark/network/protocol/OneWayMessage.java
index 95a0270be3da9..efe0470f35875 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/OneWayMessage.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/OneWayMessage.java
@@ -17,21 +17,21 @@
 
 package org.apache.spark.network.protocol;
 
-import java.util.Arrays;
-
 import com.google.common.base.Objects;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
 
 /**
  * A RPC that does not expect a reply, which is handled by a remote
  * {@link org.apache.spark.network.server.RpcHandler}.
  */
-public final class OneWayMessage implements RequestMessage {
-  /** Serialized message to send to remote RpcHandler. */
-  public final byte[] message;
+public final class OneWayMessage extends AbstractMessage implements RequestMessage {
 
-  public OneWayMessage(byte[] message) {
-    this.message = message;
+  public OneWayMessage(ManagedBuffer body) {
+    super(body, true);
   }
 
   @Override
@@ -39,29 +39,34 @@ public OneWayMessage(byte[] message) {
 
   @Override
   public int encodedLength() {
-    return Encoders.ByteArrays.encodedLength(message);
+    // The integer (a.k.a. the body size) is not really used, since that information is already
+    // encoded in the frame length. But this maintains backwards compatibility with versions of
+    // RpcRequest that use Encoders.ByteArrays.
+    return 4;
   }
 
   @Override
   public void encode(ByteBuf buf) {
-    Encoders.ByteArrays.encode(buf, message);
+    // See comment in encodedLength().
+    buf.writeInt((int) body().size());
   }
 
   public static OneWayMessage decode(ByteBuf buf) {
-    byte[] message = Encoders.ByteArrays.decode(buf);
-    return new OneWayMessage(message);
+    // See comment in encodedLength().
+    buf.readInt();
+    return new OneWayMessage(new NettyManagedBuffer(buf.retain()));
   }
 
   @Override
   public int hashCode() {
-    return Arrays.hashCode(message);
+    return Objects.hashCode(body());
   }
 
   @Override
   public boolean equals(Object other) {
     if (other instanceof OneWayMessage) {
       OneWayMessage o = (OneWayMessage) other;
-      return Arrays.equals(message, o.message);
+      return super.equals(o);
     }
     return false;
   }
@@ -69,7 +74,7 @@ public boolean equals(Object other) {
   @Override
   public String toString() {
     return Objects.toStringHelper(this)
-      .add("message", message)
+      .add("body", body())
       .toString();
   }
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/RpcFailure.java b/network/common/src/main/java/org/apache/spark/network/protocol/RpcFailure.java
index 2dfc7876ba328..a76624ef5dc96 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/RpcFailure.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/RpcFailure.java
@@ -21,7 +21,7 @@
 import io.netty.buffer.ByteBuf;
 
 /** Response to {@link RpcRequest} for a failed RPC. */
-public final class RpcFailure implements ResponseMessage {
+public final class RpcFailure extends AbstractMessage implements ResponseMessage {
   public final long requestId;
   public final String errorString;
 
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/RpcRequest.java b/network/common/src/main/java/org/apache/spark/network/protocol/RpcRequest.java
index 745039db742fa..96213794a8015 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/RpcRequest.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/RpcRequest.java
@@ -17,26 +17,25 @@
 
 package org.apache.spark.network.protocol;
 
-import java.util.Arrays;
-
 import com.google.common.base.Objects;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
 
 /**
  * A generic RPC which is handled by a remote {@link org.apache.spark.network.server.RpcHandler}.
  * This will correspond to a single
  * {@link org.apache.spark.network.protocol.ResponseMessage} (either success or failure).
  */
-public final class RpcRequest implements RequestMessage {
+public final class RpcRequest extends AbstractMessage implements RequestMessage {
   /** Used to link an RPC request with its response. */
   public final long requestId;
 
-  /** Serialized message to send to remote RpcHandler. */
-  public final byte[] message;
-
-  public RpcRequest(long requestId, byte[] message) {
+  public RpcRequest(long requestId, ManagedBuffer message) {
+    super(message, true);
     this.requestId = requestId;
-    this.message = message;
   }
 
   @Override
@@ -44,31 +43,36 @@ public RpcRequest(long requestId, byte[] message) {
 
   @Override
   public int encodedLength() {
-    return 8 + Encoders.ByteArrays.encodedLength(message);
+    // The integer (a.k.a. the body size) is not really used, since that information is already
+    // encoded in the frame length. But this maintains backwards compatibility with versions of
+    // RpcRequest that use Encoders.ByteArrays.
+    return 8 + 4;
   }
 
   @Override
   public void encode(ByteBuf buf) {
     buf.writeLong(requestId);
-    Encoders.ByteArrays.encode(buf, message);
+    // See comment in encodedLength().
+    buf.writeInt((int) body().size());
   }
 
   public static RpcRequest decode(ByteBuf buf) {
     long requestId = buf.readLong();
-    byte[] message = Encoders.ByteArrays.decode(buf);
-    return new RpcRequest(requestId, message);
+    // See comment in encodedLength().
+    buf.readInt();
+    return new RpcRequest(requestId, new NettyManagedBuffer(buf.retain()));
   }
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(requestId, Arrays.hashCode(message));
+    return Objects.hashCode(requestId, body());
   }
 
   @Override
   public boolean equals(Object other) {
     if (other instanceof RpcRequest) {
       RpcRequest o = (RpcRequest) other;
-      return requestId == o.requestId && Arrays.equals(message, o.message);
+      return requestId == o.requestId && super.equals(o);
     }
     return false;
   }
@@ -77,7 +81,7 @@ public boolean equals(Object other) {
   public String toString() {
     return Objects.toStringHelper(this)
       .add("requestId", requestId)
-      .add("message", message)
+      .add("body", body())
       .toString();
   }
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/RpcResponse.java b/network/common/src/main/java/org/apache/spark/network/protocol/RpcResponse.java
index 1671cd444f039..bae866e14a1e1 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/RpcResponse.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/RpcResponse.java
@@ -17,49 +17,62 @@
 
 package org.apache.spark.network.protocol;
 
-import java.util.Arrays;
-
 import com.google.common.base.Objects;
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
 
 /** Response to {@link RpcRequest} for a successful RPC. */
-public final class RpcResponse implements ResponseMessage {
+public final class RpcResponse extends AbstractResponseMessage {
   public final long requestId;
-  public final byte[] response;
 
-  public RpcResponse(long requestId, byte[] response) {
+  public RpcResponse(long requestId, ManagedBuffer message) {
+    super(message, true);
     this.requestId = requestId;
-    this.response = response;
   }
 
   @Override
   public Type type() { return Type.RpcResponse; }
 
   @Override
-  public int encodedLength() { return 8 + Encoders.ByteArrays.encodedLength(response); }
+  public int encodedLength() {
+    // The integer (a.k.a. the body size) is not really used, since that information is already
+    // encoded in the frame length. But this maintains backwards compatibility with versions of
+    // RpcRequest that use Encoders.ByteArrays.
+    return 8 + 4;
+  }
 
   @Override
   public void encode(ByteBuf buf) {
     buf.writeLong(requestId);
-    Encoders.ByteArrays.encode(buf, response);
+    // See comment in encodedLength().
+    buf.writeInt((int) body().size());
+  }
+
+  @Override
+  public ResponseMessage createFailureResponse(String error) {
+    return new RpcFailure(requestId, error);
   }
 
   public static RpcResponse decode(ByteBuf buf) {
     long requestId = buf.readLong();
-    byte[] response = Encoders.ByteArrays.decode(buf);
-    return new RpcResponse(requestId, response);
+    // See comment in encodedLength().
+    buf.readInt();
+    return new RpcResponse(requestId, new NettyManagedBuffer(buf.retain()));
   }
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(requestId, Arrays.hashCode(response));
+    return Objects.hashCode(requestId, body());
   }
 
   @Override
   public boolean equals(Object other) {
     if (other instanceof RpcResponse) {
       RpcResponse o = (RpcResponse) other;
-      return requestId == o.requestId && Arrays.equals(response, o.response);
+      return requestId == o.requestId && super.equals(o);
     }
     return false;
   }
@@ -68,7 +81,7 @@ public boolean equals(Object other) {
   public String toString() {
     return Objects.toStringHelper(this)
       .add("requestId", requestId)
-      .add("response", response)
+      .add("body", body())
       .toString();
   }
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/StreamFailure.java b/network/common/src/main/java/org/apache/spark/network/protocol/StreamFailure.java
index e3dade2ebf905..26747ee55b4de 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/StreamFailure.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/StreamFailure.java
@@ -26,7 +26,7 @@
 /**
  * Message indicating an error when transferring a stream.
  */
-public final class StreamFailure implements ResponseMessage {
+public final class StreamFailure extends AbstractMessage implements ResponseMessage {
   public final String streamId;
   public final String error;
 
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/StreamRequest.java b/network/common/src/main/java/org/apache/spark/network/protocol/StreamRequest.java
index 821e8f53884d7..35af5a84ba6bd 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/StreamRequest.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/StreamRequest.java
@@ -29,7 +29,7 @@
  * The stream ID is an arbitrary string that needs to be negotiated between the two endpoints before
  * the data can be streamed.
  */
-public final class StreamRequest implements RequestMessage {
+public final class StreamRequest extends AbstractMessage implements RequestMessage {
    public final String streamId;
 
    public StreamRequest(String streamId) {
diff --git a/network/common/src/main/java/org/apache/spark/network/protocol/StreamResponse.java b/network/common/src/main/java/org/apache/spark/network/protocol/StreamResponse.java
index ac5ab9a323a11..51b899930f721 100644
--- a/network/common/src/main/java/org/apache/spark/network/protocol/StreamResponse.java
+++ b/network/common/src/main/java/org/apache/spark/network/protocol/StreamResponse.java
@@ -30,15 +30,15 @@
  * sender. The receiver is expected to set a temporary channel handler that will consume the
  * number of bytes this message says the stream has.
  */
-public final class StreamResponse extends ResponseWithBody {
-   public final String streamId;
-   public final long byteCount;
+public final class StreamResponse extends AbstractResponseMessage {
+  public final String streamId;
+  public final long byteCount;
 
-   public StreamResponse(String streamId, long byteCount, ManagedBuffer buffer) {
-     super(buffer, false);
-     this.streamId = streamId;
-     this.byteCount = byteCount;
-   }
+  public StreamResponse(String streamId, long byteCount, ManagedBuffer buffer) {
+    super(buffer, false);
+    this.streamId = streamId;
+    this.byteCount = byteCount;
+  }
 
   @Override
   public Type type() { return Type.StreamResponse; }
@@ -68,7 +68,7 @@ public static StreamResponse decode(ByteBuf buf) {
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(byteCount, streamId);
+    return Objects.hashCode(byteCount, streamId, body());
   }
 
   @Override
@@ -85,6 +85,7 @@ public String toString() {
     return Objects.toStringHelper(this)
       .add("streamId", streamId)
       .add("byteCount", byteCount)
+      .add("body", body())
       .toString();
   }
 
diff --git a/network/common/src/main/java/org/apache/spark/network/sasl/SaslClientBootstrap.java b/network/common/src/main/java/org/apache/spark/network/sasl/SaslClientBootstrap.java
index 69923769d44b4..68381037d6891 100644
--- a/network/common/src/main/java/org/apache/spark/network/sasl/SaslClientBootstrap.java
+++ b/network/common/src/main/java/org/apache/spark/network/sasl/SaslClientBootstrap.java
@@ -17,6 +17,8 @@
 
 package org.apache.spark.network.sasl;
 
+import java.io.IOException;
+import java.nio.ByteBuffer;
 import javax.security.sasl.Sasl;
 import javax.security.sasl.SaslException;
 
@@ -28,6 +30,7 @@
 
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.client.TransportClientBootstrap;
+import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.TransportConf;
 
 /**
@@ -70,11 +73,12 @@ public void doBootstrap(TransportClient client, Channel channel) {
 
       while (!saslClient.isComplete()) {
         SaslMessage msg = new SaslMessage(appId, payload);
-        ByteBuf buf = Unpooled.buffer(msg.encodedLength());
+        ByteBuf buf = Unpooled.buffer(msg.encodedLength() + (int) msg.body().size());
         msg.encode(buf);
+        buf.writeBytes(msg.body().nioByteBuffer());
 
-        byte[] response = client.sendRpcSync(buf.array(), conf.saslRTTimeoutMs());
-        payload = saslClient.response(response);
+        ByteBuffer response = client.sendRpcSync(buf.nioBuffer(), conf.saslRTTimeoutMs());
+        payload = saslClient.response(JavaUtils.bufferToArray(response));
       }
 
       client.setClientId(appId);
@@ -88,6 +92,8 @@ public void doBootstrap(TransportClient client, Channel channel) {
         saslClient = null;
         logger.debug("Channel {} configured for SASL encryption.", client);
       }
+    } catch (IOException ioe) {
+      throw new RuntimeException(ioe);
     } finally {
       if (saslClient != null) {
         try {
diff --git a/network/common/src/main/java/org/apache/spark/network/sasl/SaslMessage.java b/network/common/src/main/java/org/apache/spark/network/sasl/SaslMessage.java
index cad76ab7aa54e..e52b526f09c77 100644
--- a/network/common/src/main/java/org/apache/spark/network/sasl/SaslMessage.java
+++ b/network/common/src/main/java/org/apache/spark/network/sasl/SaslMessage.java
@@ -18,38 +18,50 @@
 package org.apache.spark.network.sasl;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
 
-import org.apache.spark.network.protocol.Encodable;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
 import org.apache.spark.network.protocol.Encoders;
+import org.apache.spark.network.protocol.AbstractMessage;
 
 /**
  * Encodes a Sasl-related message which is attempting to authenticate using some credentials tagged
  * with the given appId. This appId allows a single SaslRpcHandler to multiplex different
  * applications which may be using different sets of credentials.
  */
-class SaslMessage implements Encodable {
+class SaslMessage extends AbstractMessage {
 
   /** Serialization tag used to catch incorrect payloads. */
   private static final byte TAG_BYTE = (byte) 0xEA;
 
   public final String appId;
-  public final byte[] payload;
 
-  public SaslMessage(String appId, byte[] payload) {
+  public SaslMessage(String appId, byte[] message) {
+    this(appId, Unpooled.wrappedBuffer(message));
+  }
+
+  public SaslMessage(String appId, ByteBuf message) {
+    super(new NettyManagedBuffer(message), true);
     this.appId = appId;
-    this.payload = payload;
   }
 
+  @Override
+  public Type type() { return Type.User; }
+
   @Override
   public int encodedLength() {
-    return 1 + Encoders.Strings.encodedLength(appId) + Encoders.ByteArrays.encodedLength(payload);
+    // The integer (a.k.a. the body size) is not really used, since that information is already
+    // encoded in the frame length. But this maintains backwards compatibility with versions of
+    // RpcRequest that use Encoders.ByteArrays.
+    return 1 + Encoders.Strings.encodedLength(appId) + 4;
   }
 
   @Override
   public void encode(ByteBuf buf) {
     buf.writeByte(TAG_BYTE);
     Encoders.Strings.encode(buf, appId);
-    Encoders.ByteArrays.encode(buf, payload);
+    // See comment in encodedLength().
+    buf.writeInt((int) body().size());
   }
 
   public static SaslMessage decode(ByteBuf buf) {
@@ -59,7 +71,8 @@ public static SaslMessage decode(ByteBuf buf) {
     }
 
     String appId = Encoders.Strings.decode(buf);
-    byte[] payload = Encoders.ByteArrays.decode(buf);
-    return new SaslMessage(appId, payload);
+    // See comment in encodedLength().
+    buf.readInt();
+    return new SaslMessage(appId, buf.retain());
   }
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java b/network/common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
index 830db94b890c5..c215bd9d15045 100644
--- a/network/common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/sasl/SaslRpcHandler.java
@@ -17,8 +17,11 @@
 
 package org.apache.spark.network.sasl;
 
+import java.io.IOException;
+import java.nio.ByteBuffer;
 import javax.security.sasl.Sasl;
 
+import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
 import org.slf4j.Logger;
@@ -28,6 +31,7 @@
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.server.RpcHandler;
 import org.apache.spark.network.server.StreamManager;
+import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.TransportConf;
 
 /**
@@ -70,14 +74,20 @@ class SaslRpcHandler extends RpcHandler {
   }
 
   @Override
-  public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+  public void receive(TransportClient client, ByteBuffer message, RpcResponseCallback callback) {
     if (isComplete) {
       // Authentication complete, delegate to base handler.
       delegate.receive(client, message, callback);
       return;
     }
 
-    SaslMessage saslMessage = SaslMessage.decode(Unpooled.wrappedBuffer(message));
+    ByteBuf nettyBuf = Unpooled.wrappedBuffer(message);
+    SaslMessage saslMessage;
+    try {
+      saslMessage = SaslMessage.decode(nettyBuf);
+    } finally {
+      nettyBuf.release();
+    }
 
     if (saslServer == null) {
       // First message in the handshake, setup the necessary state.
@@ -86,8 +96,14 @@ public void receive(TransportClient client, byte[] message, RpcResponseCallback
         conf.saslServerAlwaysEncrypt());
     }
 
-    byte[] response = saslServer.response(saslMessage.payload);
-    callback.onSuccess(response);
+    byte[] response;
+    try {
+      response = saslServer.response(JavaUtils.bufferToArray(
+        saslMessage.body().nioByteBuffer()));
+    } catch (IOException ioe) {
+      throw new RuntimeException(ioe);
+    }
+    callback.onSuccess(ByteBuffer.wrap(response));
 
     // Setup encryption after the SASL response is sent, otherwise the client can't parse the
     // response. It's ok to change the channel pipeline here since we are processing an incoming
@@ -109,7 +125,7 @@ public void receive(TransportClient client, byte[] message, RpcResponseCallback
   }
 
   @Override
-  public void receive(TransportClient client, byte[] message) {
+  public void receive(TransportClient client, ByteBuffer message) {
     delegate.receive(client, message);
   }
 
diff --git a/network/common/src/main/java/org/apache/spark/network/server/MessageHandler.java b/network/common/src/main/java/org/apache/spark/network/server/MessageHandler.java
index b80c15106ecbd..3843406b27403 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/MessageHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/MessageHandler.java
@@ -26,7 +26,7 @@
  */
 public abstract class MessageHandler<T extends Message> {
   /** Handles the receipt of a single message. */
-  public abstract void handle(T message);
+  public abstract void handle(T message) throws Exception;
 
   /** Invoked when an exception was caught on the Channel. */
   public abstract void exceptionCaught(Throwable cause);
diff --git a/network/common/src/main/java/org/apache/spark/network/server/NoOpRpcHandler.java b/network/common/src/main/java/org/apache/spark/network/server/NoOpRpcHandler.java
index 1502b7489e864..6ed61da5c7eff 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/NoOpRpcHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/NoOpRpcHandler.java
@@ -1,5 +1,3 @@
-package org.apache.spark.network.server;
-
 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
@@ -17,6 +15,10 @@
  * limitations under the License.
  */
 
+package org.apache.spark.network.server;
+
+import java.nio.ByteBuffer;
+
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
 
@@ -29,7 +31,7 @@ public NoOpRpcHandler() {
   }
 
   @Override
-  public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+  public void receive(TransportClient client, ByteBuffer message, RpcResponseCallback callback) {
     throw new UnsupportedOperationException("Cannot handle messages");
   }
 
diff --git a/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java b/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
index 1a11f7b3820c6..ee1c683699478 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/RpcHandler.java
@@ -17,6 +17,8 @@
 
 package org.apache.spark.network.server;
 
+import java.nio.ByteBuffer;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -44,7 +46,7 @@ public abstract class RpcHandler {
    */
   public abstract void receive(
       TransportClient client,
-      byte[] message,
+      ByteBuffer message,
       RpcResponseCallback callback);
 
   /**
@@ -62,7 +64,7 @@ public abstract void receive(
    *               of this RPC. This will always be the exact same object for a particular channel.
    * @param message The serialized bytes of the RPC.
    */
-  public void receive(TransportClient client, byte[] message) {
+  public void receive(TransportClient client, ByteBuffer message) {
     receive(client, message, ONE_WAY_CALLBACK);
   }
 
@@ -79,7 +81,7 @@ private static class OneWayRpcCallback implements RpcResponseCallback {
     private final Logger logger = LoggerFactory.getLogger(OneWayRpcCallback.class);
 
     @Override
-    public void onSuccess(byte[] response) {
+    public void onSuccess(ByteBuffer response) {
       logger.warn("Response provided for one-way RPC.");
     }
 
diff --git a/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java b/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
index 3164e00679035..09435bcbab35e 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/TransportChannelHandler.java
@@ -99,7 +99,7 @@ public void channelUnregistered(ChannelHandlerContext ctx) throws Exception {
   }
 
   @Override
-  public void channelRead0(ChannelHandlerContext ctx, Message request) {
+  public void channelRead0(ChannelHandlerContext ctx, Message request) throws Exception {
     if (request instanceof RequestMessage) {
       requestHandler.handle((RequestMessage) request);
     } else {
diff --git a/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java b/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
index db18ea77d1073..c864d7ce16bd3 100644
--- a/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
+++ b/network/common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java
@@ -17,6 +17,8 @@
 
 package org.apache.spark.network.server;
 
+import java.nio.ByteBuffer;
+
 import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
 import io.netty.channel.Channel;
@@ -26,6 +28,7 @@
 import org.slf4j.LoggerFactory;
 
 import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NioManagedBuffer;
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.protocol.ChunkFetchRequest;
@@ -143,10 +146,10 @@ private void processStreamRequest(final StreamRequest req) {
 
   private void processRpcRequest(final RpcRequest req) {
     try {
-      rpcHandler.receive(reverseClient, req.message, new RpcResponseCallback() {
+      rpcHandler.receive(reverseClient, req.body().nioByteBuffer(), new RpcResponseCallback() {
         @Override
-        public void onSuccess(byte[] response) {
-          respond(new RpcResponse(req.requestId, response));
+        public void onSuccess(ByteBuffer response) {
+          respond(new RpcResponse(req.requestId, new NioManagedBuffer(response)));
         }
 
         @Override
@@ -157,14 +160,18 @@ public void onFailure(Throwable e) {
     } catch (Exception e) {
       logger.error("Error while invoking RpcHandler#receive() on RPC id " + req.requestId, e);
       respond(new RpcFailure(req.requestId, Throwables.getStackTraceAsString(e)));
+    } finally {
+      req.body().release();
     }
   }
 
   private void processOneWayMessage(OneWayMessage req) {
     try {
-      rpcHandler.receive(reverseClient, req.message);
+      rpcHandler.receive(reverseClient, req.body().nioByteBuffer());
     } catch (Exception e) {
       logger.error("Error while invoking RpcHandler#receive() for one-way message.", e);
+    } finally {
+      req.body().release();
     }
   }
 
diff --git a/network/common/src/main/java/org/apache/spark/network/util/JavaUtils.java b/network/common/src/main/java/org/apache/spark/network/util/JavaUtils.java
index 7d27439cfde7a..b3d8e0cd7cdcd 100644
--- a/network/common/src/main/java/org/apache/spark/network/util/JavaUtils.java
+++ b/network/common/src/main/java/org/apache/spark/network/util/JavaUtils.java
@@ -132,7 +132,7 @@ private static boolean isSymlink(File file) throws IOException {
     return !fileInCanonicalDir.getCanonicalFile().equals(fileInCanonicalDir.getAbsoluteFile());
   }
 
-  private static final ImmutableMap<String, TimeUnit> timeSuffixes = 
+  private static final ImmutableMap<String, TimeUnit> timeSuffixes =
     ImmutableMap.<String, TimeUnit>builder()
       .put("us", TimeUnit.MICROSECONDS)
       .put("ms", TimeUnit.MILLISECONDS)
@@ -164,32 +164,32 @@ private static boolean isSymlink(File file) throws IOException {
    */
   private static long parseTimeString(String str, TimeUnit unit) {
     String lower = str.toLowerCase().trim();
-    
+
     try {
       Matcher m = Pattern.compile("(-?[0-9]+)([a-z]+)?").matcher(lower);
       if (!m.matches()) {
         throw new NumberFormatException("Failed to parse time string: " + str);
       }
-      
+
       long val = Long.parseLong(m.group(1));
       String suffix = m.group(2);
-      
+
       // Check for invalid suffixes
       if (suffix != null && !timeSuffixes.containsKey(suffix)) {
         throw new NumberFormatException("Invalid suffix: \"" + suffix + "\"");
       }
-      
+
       // If suffix is valid use that, otherwise none was provided and use the default passed
       return unit.convert(val, suffix != null ? timeSuffixes.get(suffix) : unit);
     } catch (NumberFormatException e) {
       String timeError = "Time must be specified as seconds (s), " +
               "milliseconds (ms), microseconds (us), minutes (m or min), hour (h), or day (d). " +
               "E.g. 50s, 100ms, or 250us.";
-      
+
       throw new NumberFormatException(timeError + "\n" + e.getMessage());
     }
   }
-  
+
   /**
    * Convert a time parameter such as (50s, 100ms, or 250us) to milliseconds for internal use. If
    * no suffix is provided, the passed number is assumed to be in ms.
@@ -205,10 +205,10 @@ public static long timeStringAsMs(String str) {
   public static long timeStringAsSec(String str) {
     return parseTimeString(str, TimeUnit.SECONDS);
   }
-  
+
   /**
    * Convert a passed byte string (e.g. 50b, 100kb, or 250mb) to a ByteUnit for
-   * internal use. If no suffix is provided a direct conversion of the provided default is 
+   * internal use. If no suffix is provided a direct conversion of the provided default is
    * attempted.
    */
   private static long parseByteString(String str, ByteUnit unit) {
@@ -217,7 +217,7 @@ private static long parseByteString(String str, ByteUnit unit) {
     try {
       Matcher m = Pattern.compile("([0-9]+)([a-z]+)?").matcher(lower);
       Matcher fractionMatcher = Pattern.compile("([0-9]+\\.[0-9]+)([a-z]+)?").matcher(lower);
-      
+
       if (m.matches()) {
         long val = Long.parseLong(m.group(1));
         String suffix = m.group(2);
@@ -228,14 +228,14 @@ private static long parseByteString(String str, ByteUnit unit) {
         }
 
         // If suffix is valid use that, otherwise none was provided and use the default passed
-        return unit.convertFrom(val, suffix != null ? byteSuffixes.get(suffix) : unit);  
+        return unit.convertFrom(val, suffix != null ? byteSuffixes.get(suffix) : unit);
       } else if (fractionMatcher.matches()) {
-        throw new NumberFormatException("Fractional values are not supported. Input was: " 
+        throw new NumberFormatException("Fractional values are not supported. Input was: "
           + fractionMatcher.group(1));
       } else {
-        throw new NumberFormatException("Failed to parse byte string: " + str);  
+        throw new NumberFormatException("Failed to parse byte string: " + str);
       }
-      
+
     } catch (NumberFormatException e) {
       String timeError = "Size must be specified as bytes (b), " +
         "kibibytes (k), mebibytes (m), gibibytes (g), tebibytes (t), or pebibytes(p). " +
@@ -248,7 +248,7 @@ private static long parseByteString(String str, ByteUnit unit) {
   /**
    * Convert a passed byte string (e.g. 50b, 100k, or 250m) to bytes for
    * internal use.
-   * 
+   *
    * If no suffix is provided, the passed number is assumed to be in bytes.
    */
   public static long byteStringAsBytes(String str) {
@@ -264,7 +264,7 @@ public static long byteStringAsBytes(String str) {
   public static long byteStringAsKb(String str) {
     return parseByteString(str, ByteUnit.KiB);
   }
-  
+
   /**
    * Convert a passed byte string (e.g. 50b, 100k, or 250m) to mebibytes for
    * internal use.
@@ -284,4 +284,20 @@ public static long byteStringAsMb(String str) {
   public static long byteStringAsGb(String str) {
     return parseByteString(str, ByteUnit.GiB);
   }
+
+  /**
+   * Returns a byte array with the buffer's contents, trying to avoid copying the data if
+   * possible.
+   */
+  public static byte[] bufferToArray(ByteBuffer buffer) {
+    if (buffer.hasArray() && buffer.arrayOffset() == 0 &&
+        buffer.array().length == buffer.remaining()) {
+      return buffer.array();
+    } else {
+      byte[] bytes = new byte[buffer.remaining()];
+      buffer.get(bytes);
+      return bytes;
+    }
+  }
+
 }
diff --git a/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java b/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
index 5889562dd9705..a466c729154aa 100644
--- a/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
+++ b/network/common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java
@@ -17,9 +17,13 @@
 
 package org.apache.spark.network.util;
 
+import java.util.Iterator;
+import java.util.LinkedList;
+
 import com.google.common.base.Preconditions;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelInboundHandlerAdapter;
 
@@ -44,84 +48,138 @@ public class TransportFrameDecoder extends ChannelInboundHandlerAdapter {
   public static final String HANDLER_NAME = "frameDecoder";
   private static final int LENGTH_SIZE = 8;
   private static final int MAX_FRAME_SIZE = Integer.MAX_VALUE;
+  private static final int UNKNOWN_FRAME_SIZE = -1;
+
+  private final LinkedList<ByteBuf> buffers = new LinkedList<>();
+  private final ByteBuf frameLenBuf = Unpooled.buffer(LENGTH_SIZE, LENGTH_SIZE);
 
-  private CompositeByteBuf buffer;
+  private long totalSize = 0;
+  private long nextFrameSize = UNKNOWN_FRAME_SIZE;
   private volatile Interceptor interceptor;
 
   @Override
   public void channelRead(ChannelHandlerContext ctx, Object data) throws Exception {
     ByteBuf in = (ByteBuf) data;
+    buffers.add(in);
+    totalSize += in.readableBytes();
+
+    while (!buffers.isEmpty()) {
+      // First, feed the interceptor, and if it's still, active, try again.
+      if (interceptor != null) {
+        ByteBuf first = buffers.getFirst();
+        int available = first.readableBytes();
+        if (feedInterceptor(first)) {
+          assert !first.isReadable() : "Interceptor still active but buffer has data.";
+        }
 
-    if (buffer == null) {
-      buffer = in.alloc().compositeBuffer();
-    }
-
-    buffer.addComponent(in).writerIndex(buffer.writerIndex() + in.readableBytes());
-
-    while (buffer.isReadable()) {
-      discardReadBytes();
-      if (!feedInterceptor()) {
+        int read = available - first.readableBytes();
+        if (read == available) {
+          buffers.removeFirst().release();
+        }
+        totalSize -= read;
+      } else {
+        // Interceptor is not active, so try to decode one frame.
         ByteBuf frame = decodeNext();
         if (frame == null) {
           break;
         }
-
         ctx.fireChannelRead(frame);
       }
     }
-
-    discardReadBytes();
   }
 
-  private void discardReadBytes() {
-    // If the buffer's been retained by downstream code, then make a copy of the remaining
-    // bytes into a new buffer. Otherwise, just discard stale components.
-    if (buffer.refCnt() > 1) {
-      CompositeByteBuf newBuffer = buffer.alloc().compositeBuffer();
+  private long decodeFrameSize() {
+    if (nextFrameSize != UNKNOWN_FRAME_SIZE || totalSize < LENGTH_SIZE) {
+      return nextFrameSize;
+    }
 
-      if (buffer.readableBytes() > 0) {
-        ByteBuf spillBuf = buffer.alloc().buffer(buffer.readableBytes());
-        spillBuf.writeBytes(buffer);
-        newBuffer.addComponent(spillBuf).writerIndex(spillBuf.readableBytes());
+    // We know there's enough data. If the first buffer contains all the data, great. Otherwise,
+    // hold the bytes for the frame length in a composite buffer until we have enough data to read
+    // the frame size. Normally, it should be rare to need more than one buffer to read the frame
+    // size.
+    ByteBuf first = buffers.getFirst();
+    if (first.readableBytes() >= LENGTH_SIZE) {
+      nextFrameSize = first.readLong() - LENGTH_SIZE;
+      totalSize -= LENGTH_SIZE;
+      if (!first.isReadable()) {
+        buffers.removeFirst().release();
       }
+      return nextFrameSize;
+    }
 
-      buffer.release();
-      buffer = newBuffer;
-    } else {
-      buffer.discardReadComponents();
+    while (frameLenBuf.readableBytes() < LENGTH_SIZE) {
+      ByteBuf next = buffers.getFirst();
+      int toRead = Math.min(next.readableBytes(), LENGTH_SIZE - frameLenBuf.readableBytes());
+      frameLenBuf.writeBytes(next, toRead);
+      if (!next.isReadable()) {
+        buffers.removeFirst().release();
+      }
     }
+
+    nextFrameSize = frameLenBuf.readLong() - LENGTH_SIZE;
+    totalSize -= LENGTH_SIZE;
+    frameLenBuf.clear();
+    return nextFrameSize;
   }
 
   private ByteBuf decodeNext() throws Exception {
-    if (buffer.readableBytes() < LENGTH_SIZE) {
+    long frameSize = decodeFrameSize();
+    if (frameSize == UNKNOWN_FRAME_SIZE || totalSize < frameSize) {
       return null;
     }
 
-    int frameLen = (int) buffer.readLong() - LENGTH_SIZE;
-    if (buffer.readableBytes() < frameLen) {
-      buffer.readerIndex(buffer.readerIndex() - LENGTH_SIZE);
-      return null;
+    // Reset size for next frame.
+    nextFrameSize = UNKNOWN_FRAME_SIZE;
+
+    Preconditions.checkArgument(frameSize < MAX_FRAME_SIZE, "Too large frame: %s", frameSize);
+    Preconditions.checkArgument(frameSize > 0, "Frame length should be positive: %s", frameSize);
+
+    // If the first buffer holds the entire frame, return it.
+    int remaining = (int) frameSize;
+    if (buffers.getFirst().readableBytes() >= remaining) {
+      return nextBufferForFrame(remaining);
     }
 
-    Preconditions.checkArgument(frameLen < MAX_FRAME_SIZE, "Too large frame: %s", frameLen);
-    Preconditions.checkArgument(frameLen > 0, "Frame length should be positive: %s", frameLen);
+    // Otherwise, create a composite buffer.
+    CompositeByteBuf frame = buffers.getFirst().alloc().compositeBuffer();
+    while (remaining > 0) {
+      ByteBuf next = nextBufferForFrame(remaining);
+      remaining -= next.readableBytes();
+      frame.addComponent(next).writerIndex(frame.writerIndex() + next.readableBytes());
+    }
+    assert remaining == 0;
+    return frame;
+  }
+
+  /**
+   * Takes the first buffer in the internal list, and either adjust it to fit in the frame
+   * (by taking a slice out of it) or remove it from the internal list.
+   */
+  private ByteBuf nextBufferForFrame(int bytesToRead) {
+    ByteBuf buf = buffers.getFirst();
+    ByteBuf frame;
+
+    if (buf.readableBytes() > bytesToRead) {
+      frame = buf.retain().readSlice(bytesToRead);
+      totalSize -= bytesToRead;
+    } else {
+      frame = buf;
+      buffers.removeFirst();
+      totalSize -= frame.readableBytes();
+    }
 
-    ByteBuf frame = buffer.readSlice(frameLen);
-    frame.retain();
     return frame;
   }
 
   @Override
   public void channelInactive(ChannelHandlerContext ctx) throws Exception {
-    if (buffer != null) {
-      if (buffer.isReadable()) {
-        feedInterceptor();
-      }
-      buffer.release();
+    for (ByteBuf b : buffers) {
+      b.release();
     }
     if (interceptor != null) {
       interceptor.channelInactive();
     }
+    frameLenBuf.release();
     super.channelInactive(ctx);
   }
 
@@ -141,8 +199,8 @@ public void setInterceptor(Interceptor interceptor) {
   /**
    * @return Whether the interceptor is still active after processing the data.
    */
-  private boolean feedInterceptor() throws Exception {
-    if (interceptor != null && !interceptor.handle(buffer)) {
+  private boolean feedInterceptor(ByteBuf buf) throws Exception {
+    if (interceptor != null && !interceptor.handle(buf)) {
       interceptor = null;
     }
     return interceptor != null;
diff --git a/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java b/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
index 50a324e293386..70c849d60e0a6 100644
--- a/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java
@@ -107,7 +107,10 @@ public ManagedBuffer getChunk(long streamId, int chunkIndex) {
     };
     RpcHandler handler = new RpcHandler() {
       @Override
-      public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+      public void receive(
+          TransportClient client,
+          ByteBuffer message,
+          RpcResponseCallback callback) {
         throw new UnsupportedOperationException();
       }
 
diff --git a/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java b/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
index 1aa20900ffe74..6c8dd742f4b64 100644
--- a/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/ProtocolSuite.java
@@ -82,10 +82,10 @@ private void testClientToServer(Message msg) {
   @Test
   public void requests() {
     testClientToServer(new ChunkFetchRequest(new StreamChunkId(1, 2)));
-    testClientToServer(new RpcRequest(12345, new byte[0]));
-    testClientToServer(new RpcRequest(12345, new byte[100]));
+    testClientToServer(new RpcRequest(12345, new TestManagedBuffer(0)));
+    testClientToServer(new RpcRequest(12345, new TestManagedBuffer(10)));
     testClientToServer(new StreamRequest("abcde"));
-    testClientToServer(new OneWayMessage(new byte[100]));
+    testClientToServer(new OneWayMessage(new TestManagedBuffer(10)));
   }
 
   @Test
@@ -94,8 +94,8 @@ public void responses() {
     testServerToClient(new ChunkFetchSuccess(new StreamChunkId(1, 2), new TestManagedBuffer(0)));
     testServerToClient(new ChunkFetchFailure(new StreamChunkId(1, 2), "this is an error"));
     testServerToClient(new ChunkFetchFailure(new StreamChunkId(1, 2), ""));
-    testServerToClient(new RpcResponse(12345, new byte[0]));
-    testServerToClient(new RpcResponse(12345, new byte[1000]));
+    testServerToClient(new RpcResponse(12345, new TestManagedBuffer(0)));
+    testServerToClient(new RpcResponse(12345, new TestManagedBuffer(100)));
     testServerToClient(new RpcFailure(0, "this is an error"));
     testServerToClient(new RpcFailure(0, ""));
     // Note: buffer size must be "0" since StreamResponse's buffer is written differently to the
diff --git a/network/common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java b/network/common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java
index 42955ef69235a..f9b5bf96d6215 100644
--- a/network/common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java
@@ -31,6 +31,7 @@
 import org.apache.spark.network.util.MapConfigProvider;
 import org.apache.spark.network.util.TransportConf;
 import org.junit.*;
+import static org.junit.Assert.*;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -84,13 +85,16 @@ public void tearDown() {
   @Test
   public void timeoutInactiveRequests() throws Exception {
     final Semaphore semaphore = new Semaphore(1);
-    final byte[] response = new byte[16];
+    final int responseSize = 16;
     RpcHandler handler = new RpcHandler() {
       @Override
-      public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+      public void receive(
+          TransportClient client,
+          ByteBuffer message,
+          RpcResponseCallback callback) {
         try {
           semaphore.tryAcquire(FOREVER, TimeUnit.MILLISECONDS);
-          callback.onSuccess(response);
+          callback.onSuccess(ByteBuffer.allocate(responseSize));
         } catch (InterruptedException e) {
           // do nothing
         }
@@ -110,15 +114,15 @@ public StreamManager getStreamManager() {
     // First completes quickly (semaphore starts at 1).
     TestCallback callback0 = new TestCallback();
     synchronized (callback0) {
-      client.sendRpc(new byte[0], callback0);
+      client.sendRpc(ByteBuffer.allocate(0), callback0);
       callback0.wait(FOREVER);
-      assert (callback0.success.length == response.length);
+      assertEquals(responseSize, callback0.successLength);
     }
 
     // Second times out after 2 seconds, with slack. Must be IOException.
     TestCallback callback1 = new TestCallback();
     synchronized (callback1) {
-      client.sendRpc(new byte[0], callback1);
+      client.sendRpc(ByteBuffer.allocate(0), callback1);
       callback1.wait(4 * 1000);
       assert (callback1.failure != null);
       assert (callback1.failure instanceof IOException);
@@ -131,13 +135,16 @@ public StreamManager getStreamManager() {
   @Test
   public void timeoutCleanlyClosesClient() throws Exception {
     final Semaphore semaphore = new Semaphore(0);
-    final byte[] response = new byte[16];
+    final int responseSize = 16;
     RpcHandler handler = new RpcHandler() {
       @Override
-      public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+      public void receive(
+          TransportClient client,
+          ByteBuffer message,
+          RpcResponseCallback callback) {
         try {
           semaphore.tryAcquire(FOREVER, TimeUnit.MILLISECONDS);
-          callback.onSuccess(response);
+          callback.onSuccess(ByteBuffer.allocate(responseSize));
         } catch (InterruptedException e) {
           // do nothing
         }
@@ -158,7 +165,7 @@ public StreamManager getStreamManager() {
       clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
     TestCallback callback0 = new TestCallback();
     synchronized (callback0) {
-      client0.sendRpc(new byte[0], callback0);
+      client0.sendRpc(ByteBuffer.allocate(0), callback0);
       callback0.wait(FOREVER);
       assert (callback0.failure instanceof IOException);
       assert (!client0.isActive());
@@ -170,10 +177,10 @@ public StreamManager getStreamManager() {
       clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
     TestCallback callback1 = new TestCallback();
     synchronized (callback1) {
-      client1.sendRpc(new byte[0], callback1);
+      client1.sendRpc(ByteBuffer.allocate(0), callback1);
       callback1.wait(FOREVER);
-      assert (callback1.success.length == response.length);
-      assert (callback1.failure == null);
+      assertEquals(responseSize, callback1.successLength);
+      assertNull(callback1.failure);
     }
   }
 
@@ -191,7 +198,10 @@ public ManagedBuffer getChunk(long streamId, int chunkIndex) {
     };
     RpcHandler handler = new RpcHandler() {
       @Override
-      public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+      public void receive(
+          TransportClient client,
+          ByteBuffer message,
+          RpcResponseCallback callback) {
         throw new UnsupportedOperationException();
       }
 
@@ -218,9 +228,10 @@ public StreamManager getStreamManager() {
 
     synchronized (callback0) {
       // not complete yet, but should complete soon
-      assert (callback0.success == null && callback0.failure == null);
+      assertEquals(-1, callback0.successLength);
+      assertNull(callback0.failure);
       callback0.wait(2 * 1000);
-      assert (callback0.failure instanceof IOException);
+      assertTrue(callback0.failure instanceof IOException);
     }
 
     synchronized (callback1) {
@@ -235,13 +246,13 @@ public StreamManager getStreamManager() {
    */
   class TestCallback implements RpcResponseCallback, ChunkReceivedCallback {
 
-    byte[] success;
+    int successLength = -1;
     Throwable failure;
 
     @Override
-    public void onSuccess(byte[] response) {
+    public void onSuccess(ByteBuffer response) {
       synchronized(this) {
-        success = response;
+        successLength = response.remaining();
         this.notifyAll();
       }
     }
@@ -258,7 +269,7 @@ public void onFailure(Throwable e) {
     public void onSuccess(int chunkIndex, ManagedBuffer buffer) {
       synchronized(this) {
         try {
-          success = buffer.nioByteBuffer().array();
+          successLength = buffer.nioByteBuffer().remaining();
           this.notifyAll();
         } catch (IOException e) {
           // weird
diff --git a/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java b/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
index 88fa2258bb794..9e9be98c140b7 100644
--- a/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/RpcIntegrationSuite.java
@@ -17,6 +17,7 @@
 
 package org.apache.spark.network;
 
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
@@ -26,7 +27,6 @@
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 
-import com.google.common.base.Charsets;
 import com.google.common.collect.Sets;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -41,6 +41,7 @@
 import org.apache.spark.network.server.RpcHandler;
 import org.apache.spark.network.server.StreamManager;
 import org.apache.spark.network.server.TransportServer;
+import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.SystemPropertyConfigProvider;
 import org.apache.spark.network.util.TransportConf;
 
@@ -55,11 +56,14 @@ public static void setUp() throws Exception {
     TransportConf conf = new TransportConf("shuffle", new SystemPropertyConfigProvider());
     rpcHandler = new RpcHandler() {
       @Override
-      public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
-        String msg = new String(message, Charsets.UTF_8);
+      public void receive(
+          TransportClient client,
+          ByteBuffer message,
+          RpcResponseCallback callback) {
+        String msg = JavaUtils.bytesToString(message);
         String[] parts = msg.split("/");
         if (parts[0].equals("hello")) {
-          callback.onSuccess(("Hello, " + parts[1] + "!").getBytes(Charsets.UTF_8));
+          callback.onSuccess(JavaUtils.stringToBytes("Hello, " + parts[1] + "!"));
         } else if (parts[0].equals("return error")) {
           callback.onFailure(new RuntimeException("Returned: " + parts[1]));
         } else if (parts[0].equals("throw error")) {
@@ -68,9 +72,8 @@ public void receive(TransportClient client, byte[] message, RpcResponseCallback
       }
 
       @Override
-      public void receive(TransportClient client, byte[] message) {
-        String msg = new String(message, Charsets.UTF_8);
-        oneWayMsgs.add(msg);
+      public void receive(TransportClient client, ByteBuffer message) {
+        oneWayMsgs.add(JavaUtils.bytesToString(message));
       }
 
       @Override
@@ -103,8 +106,9 @@ private RpcResult sendRPC(String ... commands) throws Exception {
 
     RpcResponseCallback callback = new RpcResponseCallback() {
       @Override
-      public void onSuccess(byte[] message) {
-        res.successMessages.add(new String(message, Charsets.UTF_8));
+      public void onSuccess(ByteBuffer message) {
+        String response = JavaUtils.bytesToString(message);
+        res.successMessages.add(response);
         sem.release();
       }
 
@@ -116,7 +120,7 @@ public void onFailure(Throwable e) {
     };
 
     for (String command : commands) {
-      client.sendRpc(command.getBytes(Charsets.UTF_8), callback);
+      client.sendRpc(JavaUtils.stringToBytes(command), callback);
     }
 
     if (!sem.tryAcquire(commands.length, 5, TimeUnit.SECONDS)) {
@@ -173,7 +177,7 @@ public void sendOneWayMessage() throws Exception {
     final String message = "no reply";
     TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
     try {
-      client.send(message.getBytes(Charsets.UTF_8));
+      client.send(JavaUtils.stringToBytes(message));
       assertEquals(0, client.getHandler().numOutstandingRequests());
 
       // Make sure the message arrives.
diff --git a/network/common/src/test/java/org/apache/spark/network/StreamSuite.java b/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
index 538f3efe8d6f2..9c49556927f0b 100644
--- a/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/StreamSuite.java
@@ -116,7 +116,10 @@ public ManagedBuffer openStream(String streamId) {
     };
     RpcHandler handler = new RpcHandler() {
       @Override
-      public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+      public void receive(
+          TransportClient client,
+          ByteBuffer message,
+          RpcResponseCallback callback) {
         throw new UnsupportedOperationException();
       }
 
diff --git a/network/common/src/test/java/org/apache/spark/network/TransportResponseHandlerSuite.java b/network/common/src/test/java/org/apache/spark/network/TransportResponseHandlerSuite.java
index 30144f4a9fc7a..128f7cba74350 100644
--- a/network/common/src/test/java/org/apache/spark/network/TransportResponseHandlerSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/TransportResponseHandlerSuite.java
@@ -17,6 +17,8 @@
 
 package org.apache.spark.network;
 
+import java.nio.ByteBuffer;
+
 import io.netty.channel.Channel;
 import io.netty.channel.local.LocalChannel;
 import org.junit.Test;
@@ -27,6 +29,7 @@
 import static org.mockito.Mockito.*;
 
 import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NioManagedBuffer;
 import org.apache.spark.network.client.ChunkReceivedCallback;
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.StreamCallback;
@@ -42,7 +45,7 @@
 
 public class TransportResponseHandlerSuite {
   @Test
-  public void handleSuccessfulFetch() {
+  public void handleSuccessfulFetch() throws Exception {
     StreamChunkId streamChunkId = new StreamChunkId(1, 0);
 
     TransportResponseHandler handler = new TransportResponseHandler(new LocalChannel());
@@ -56,7 +59,7 @@ public void handleSuccessfulFetch() {
   }
 
   @Test
-  public void handleFailedFetch() {
+  public void handleFailedFetch() throws Exception {
     StreamChunkId streamChunkId = new StreamChunkId(1, 0);
     TransportResponseHandler handler = new TransportResponseHandler(new LocalChannel());
     ChunkReceivedCallback callback = mock(ChunkReceivedCallback.class);
@@ -69,7 +72,7 @@ public void handleFailedFetch() {
   }
 
   @Test
-  public void clearAllOutstandingRequests() {
+  public void clearAllOutstandingRequests() throws Exception {
     TransportResponseHandler handler = new TransportResponseHandler(new LocalChannel());
     ChunkReceivedCallback callback = mock(ChunkReceivedCallback.class);
     handler.addFetchRequest(new StreamChunkId(1, 0), callback);
@@ -88,23 +91,24 @@ public void clearAllOutstandingRequests() {
   }
 
   @Test
-  public void handleSuccessfulRPC() {
+  public void handleSuccessfulRPC() throws Exception {
     TransportResponseHandler handler = new TransportResponseHandler(new LocalChannel());
     RpcResponseCallback callback = mock(RpcResponseCallback.class);
     handler.addRpcRequest(12345, callback);
     assertEquals(1, handler.numOutstandingRequests());
 
-    handler.handle(new RpcResponse(54321, new byte[7])); // should be ignored
+    // This response should be ignored.
+    handler.handle(new RpcResponse(54321, new NioManagedBuffer(ByteBuffer.allocate(7))));
     assertEquals(1, handler.numOutstandingRequests());
 
-    byte[] arr = new byte[10];
-    handler.handle(new RpcResponse(12345, arr));
-    verify(callback, times(1)).onSuccess(eq(arr));
+    ByteBuffer resp = ByteBuffer.allocate(10);
+    handler.handle(new RpcResponse(12345, new NioManagedBuffer(resp)));
+    verify(callback, times(1)).onSuccess(eq(ByteBuffer.allocate(10)));
     assertEquals(0, handler.numOutstandingRequests());
   }
 
   @Test
-  public void handleFailedRPC() {
+  public void handleFailedRPC() throws Exception {
     TransportResponseHandler handler = new TransportResponseHandler(new LocalChannel());
     RpcResponseCallback callback = mock(RpcResponseCallback.class);
     handler.addRpcRequest(12345, callback);
@@ -119,7 +123,7 @@ public void handleFailedRPC() {
   }
 
   @Test
-  public void testActiveStreams() {
+  public void testActiveStreams() throws Exception {
     Channel c = new LocalChannel();
     c.pipeline().addLast(TransportFrameDecoder.HANDLER_NAME, new TransportFrameDecoder());
     TransportResponseHandler handler = new TransportResponseHandler(c);
diff --git a/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java b/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
index a6f180bc40c9a..751516b9d82a1 100644
--- a/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java
@@ -22,7 +22,7 @@
 
 import java.io.File;
 import java.lang.reflect.Method;
-import java.nio.charset.StandardCharsets;
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Random;
@@ -57,6 +57,7 @@
 import org.apache.spark.network.server.TransportServer;
 import org.apache.spark.network.server.TransportServerBootstrap;
 import org.apache.spark.network.util.ByteArrayWritableChannel;
+import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.SystemPropertyConfigProvider;
 import org.apache.spark.network.util.TransportConf;
 
@@ -123,39 +124,53 @@ public void testNonMatching() {
   }
 
   @Test
-  public void testSaslAuthentication() throws Exception {
+  public void testSaslAuthentication() throws Throwable {
     testBasicSasl(false);
   }
 
   @Test
-  public void testSaslEncryption() throws Exception {
+  public void testSaslEncryption() throws Throwable {
     testBasicSasl(true);
   }
 
-  private void testBasicSasl(boolean encrypt) throws Exception {
+  private void testBasicSasl(boolean encrypt) throws Throwable {
     RpcHandler rpcHandler = mock(RpcHandler.class);
     doAnswer(new Answer<Void>() {
         @Override
         public Void answer(InvocationOnMock invocation) {
-          byte[] message = (byte[]) invocation.getArguments()[1];
+          ByteBuffer message = (ByteBuffer) invocation.getArguments()[1];
           RpcResponseCallback cb = (RpcResponseCallback) invocation.getArguments()[2];
-          assertEquals("Ping", new String(message, StandardCharsets.UTF_8));
-          cb.onSuccess("Pong".getBytes(StandardCharsets.UTF_8));
+          assertEquals("Ping", JavaUtils.bytesToString(message));
+          cb.onSuccess(JavaUtils.stringToBytes("Pong"));
           return null;
         }
       })
       .when(rpcHandler)
-      .receive(any(TransportClient.class), any(byte[].class), any(RpcResponseCallback.class));
+      .receive(any(TransportClient.class), any(ByteBuffer.class), any(RpcResponseCallback.class));
 
     SaslTestCtx ctx = new SaslTestCtx(rpcHandler, encrypt, false);
     try {
-      byte[] response = ctx.client.sendRpcSync("Ping".getBytes(StandardCharsets.UTF_8),
-                                               TimeUnit.SECONDS.toMillis(10));
-      assertEquals("Pong", new String(response, StandardCharsets.UTF_8));
+      ByteBuffer response = ctx.client.sendRpcSync(JavaUtils.stringToBytes("Ping"),
+        TimeUnit.SECONDS.toMillis(10));
+      assertEquals("Pong", JavaUtils.bytesToString(response));
     } finally {
       ctx.close();
       // There should be 2 terminated events; one for the client, one for the server.
-      verify(rpcHandler, times(2)).connectionTerminated(any(TransportClient.class));
+      Throwable error = null;
+      long deadline = System.nanoTime() + TimeUnit.NANOSECONDS.convert(10, TimeUnit.SECONDS);
+      while (deadline > System.nanoTime()) {
+        try {
+          verify(rpcHandler, times(2)).connectionTerminated(any(TransportClient.class));
+          error = null;
+          break;
+        } catch (Throwable t) {
+          error = t;
+          TimeUnit.MILLISECONDS.sleep(10);
+        }
+      }
+      if (error != null) {
+        throw error;
+      }
     }
   }
 
@@ -325,8 +340,8 @@ public void testDataEncryptionIsActuallyEnabled() throws Exception {
     SaslTestCtx ctx = null;
     try {
       ctx = new SaslTestCtx(mock(RpcHandler.class), true, true);
-      ctx.client.sendRpcSync("Ping".getBytes(StandardCharsets.UTF_8),
-                             TimeUnit.SECONDS.toMillis(10));
+      ctx.client.sendRpcSync(JavaUtils.stringToBytes("Ping"),
+        TimeUnit.SECONDS.toMillis(10));
       fail("Should have failed to send RPC to server.");
     } catch (Exception e) {
       assertFalse(e.getCause() instanceof TimeoutException);
diff --git a/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java b/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java
index 19475c21ffce9..d4de4a941d480 100644
--- a/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java
+++ b/network/common/src/test/java/org/apache/spark/network/util/TransportFrameDecoderSuite.java
@@ -118,6 +118,27 @@ public Void answer(InvocationOnMock in) {
     }
   }
 
+  @Test
+  public void testSplitLengthField() throws Exception {
+    byte[] frame = new byte[1024 * (RND.nextInt(31) + 1)];
+    ByteBuf buf = Unpooled.buffer(frame.length + 8);
+    buf.writeLong(frame.length + 8);
+    buf.writeBytes(frame);
+
+    TransportFrameDecoder decoder = new TransportFrameDecoder();
+    ChannelHandlerContext ctx = mockChannelHandlerContext();
+    try {
+      decoder.channelRead(ctx, buf.readSlice(RND.nextInt(7)).retain());
+      verify(ctx, never()).fireChannelRead(any(ByteBuf.class));
+      decoder.channelRead(ctx, buf);
+      verify(ctx).fireChannelRead(any(ByteBuf.class));
+      assertEquals(0, buf.refCnt());
+    } finally {
+      decoder.channelInactive(ctx);
+      release(buf);
+    }
+  }
+
   @Test(expected = IllegalArgumentException.class)
   public void testNegativeFrameSize() throws Exception {
     testInvalidFrame(-1);
@@ -183,7 +204,7 @@ private void testInvalidFrame(long size) throws Exception {
     try {
       decoder.channelRead(ctx, frame);
     } finally {
-      frame.release();
+      release(frame);
     }
   }
 
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandler.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandler.java
index 3ddf5c3c39189..f22187a01db02 100644
--- a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandler.java
+++ b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandler.java
@@ -19,6 +19,7 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.List;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -66,8 +67,8 @@ public ExternalShuffleBlockHandler(
   }
 
   @Override
-  public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
-    BlockTransferMessage msgObj = BlockTransferMessage.Decoder.fromByteArray(message);
+  public void receive(TransportClient client, ByteBuffer message, RpcResponseCallback callback) {
+    BlockTransferMessage msgObj = BlockTransferMessage.Decoder.fromByteBuffer(message);
     handleMessage(msgObj, client, callback);
   }
 
@@ -85,13 +86,13 @@ protected void handleMessage(
       }
       long streamId = streamManager.registerStream(client.getClientId(), blocks.iterator());
       logger.trace("Registered streamId {} with {} buffers", streamId, msg.blockIds.length);
-      callback.onSuccess(new StreamHandle(streamId, msg.blockIds.length).toByteArray());
+      callback.onSuccess(new StreamHandle(streamId, msg.blockIds.length).toByteBuffer());
 
     } else if (msgObj instanceof RegisterExecutor) {
       RegisterExecutor msg = (RegisterExecutor) msgObj;
       checkAuth(client, msg.appId);
       blockManager.registerExecutor(msg.appId, msg.execId, msg.executorInfo);
-      callback.onSuccess(new byte[0]);
+      callback.onSuccess(ByteBuffer.wrap(new byte[0]));
 
     } else {
       throw new UnsupportedOperationException("Unexpected message: " + msgObj);
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java
index ef3a9dcc8711f..58ca87d9d3b13 100644
--- a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java
+++ b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleClient.java
@@ -18,6 +18,7 @@
 package org.apache.spark.network.shuffle;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.List;
 
 import com.google.common.base.Preconditions;
@@ -139,7 +140,7 @@ public void registerWithShuffleServer(
     checkInit();
     TransportClient client = clientFactory.createUnmanagedClient(host, port);
     try {
-      byte[] registerMessage = new RegisterExecutor(appId, execId, executorInfo).toByteArray();
+      ByteBuffer registerMessage = new RegisterExecutor(appId, execId, executorInfo).toByteBuffer();
       client.sendRpcSync(registerMessage, 5000 /* timeoutMs */);
     } finally {
       client.close();
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/OneForOneBlockFetcher.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/OneForOneBlockFetcher.java
index e653f5cb147ee..1b2ddbf1ed917 100644
--- a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/OneForOneBlockFetcher.java
+++ b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/OneForOneBlockFetcher.java
@@ -17,6 +17,7 @@
 
 package org.apache.spark.network.shuffle;
 
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 
 import org.slf4j.Logger;
@@ -89,11 +90,11 @@ public void start() {
       throw new IllegalArgumentException("Zero-sized blockIds array");
     }
 
-    client.sendRpc(openMessage.toByteArray(), new RpcResponseCallback() {
+    client.sendRpc(openMessage.toByteBuffer(), new RpcResponseCallback() {
       @Override
-      public void onSuccess(byte[] response) {
+      public void onSuccess(ByteBuffer response) {
         try {
-          streamHandle = (StreamHandle) BlockTransferMessage.Decoder.fromByteArray(response);
+          streamHandle = (StreamHandle) BlockTransferMessage.Decoder.fromByteBuffer(response);
           logger.trace("Successfully opened blocks {}, preparing to fetch chunks.", streamHandle);
 
           // Immediately request all chunks -- we expect that the total size of the request is
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/mesos/MesosExternalShuffleClient.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/mesos/MesosExternalShuffleClient.java
index 7543b6be4f2a1..675820308bd4c 100644
--- a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/mesos/MesosExternalShuffleClient.java
+++ b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/mesos/MesosExternalShuffleClient.java
@@ -18,6 +18,7 @@
 package org.apache.spark.network.shuffle.mesos;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -54,11 +55,11 @@ public MesosExternalShuffleClient(
 
   public void registerDriverWithShuffleService(String host, int port) throws IOException {
     checkInit();
-    byte[] registerDriver = new RegisterDriver(appId).toByteArray();
+    ByteBuffer registerDriver = new RegisterDriver(appId).toByteBuffer();
     TransportClient client = clientFactory.createClient(host, port);
     client.sendRpc(registerDriver, new RpcResponseCallback() {
       @Override
-      public void onSuccess(byte[] response) {
+      public void onSuccess(ByteBuffer response) {
         logger.info("Successfully registered app " + appId + " with external shuffle service.");
       }
 
diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java
index fcb52363e632c..7fbe3384b4d4f 100644
--- a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java
+++ b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java
@@ -17,6 +17,8 @@
 
 package org.apache.spark.network.shuffle.protocol;
 
+import java.nio.ByteBuffer;
+
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 
@@ -53,7 +55,7 @@ private Type(int id) {
   // NB: Java does not support static methods in interfaces, so we must put this in a static class.
   public static class Decoder {
     /** Deserializes the 'type' byte followed by the message itself. */
-    public static BlockTransferMessage fromByteArray(byte[] msg) {
+    public static BlockTransferMessage fromByteBuffer(ByteBuffer msg) {
       ByteBuf buf = Unpooled.wrappedBuffer(msg);
       byte type = buf.readByte();
       switch (type) {
@@ -68,12 +70,12 @@ public static BlockTransferMessage fromByteArray(byte[] msg) {
   }
 
   /** Serializes the 'type' byte followed by the message itself. */
-  public byte[] toByteArray() {
+  public ByteBuffer toByteBuffer() {
     // Allow room for encoded message, plus the type byte
     ByteBuf buf = Unpooled.buffer(encodedLength() + 1);
     buf.writeByte(type().id);
     encode(buf);
     assert buf.writableBytes() == 0 : "Writable bytes remain: " + buf.writableBytes();
-    return buf.array();
+    return buf.nioBuffer();
   }
 }
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java
index 1c2fa4d0d462c..19c870aebb023 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java
@@ -18,6 +18,7 @@
 package org.apache.spark.network.sasl;
 
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -52,6 +53,7 @@
 import org.apache.spark.network.shuffle.protocol.OpenBlocks;
 import org.apache.spark.network.shuffle.protocol.RegisterExecutor;
 import org.apache.spark.network.shuffle.protocol.StreamHandle;
+import org.apache.spark.network.util.JavaUtils;
 import org.apache.spark.network.util.SystemPropertyConfigProvider;
 import org.apache.spark.network.util.TransportConf;
 
@@ -107,8 +109,8 @@ public void testGoodClient() throws IOException {
 
     TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
     String msg = "Hello, World!";
-    byte[] resp = client.sendRpcSync(msg.getBytes(), TIMEOUT_MS);
-    assertEquals(msg, new String(resp)); // our rpc handler should just return the given msg
+    ByteBuffer resp = client.sendRpcSync(JavaUtils.stringToBytes(msg), TIMEOUT_MS);
+    assertEquals(msg, JavaUtils.bytesToString(resp));
   }
 
   @Test
@@ -136,7 +138,7 @@ public void testNoSaslClient() throws IOException {
 
     TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
     try {
-      client.sendRpcSync(new byte[13], TIMEOUT_MS);
+      client.sendRpcSync(ByteBuffer.allocate(13), TIMEOUT_MS);
       fail("Should have failed");
     } catch (Exception e) {
       assertTrue(e.getMessage(), e.getMessage().contains("Expected SaslMessage"));
@@ -144,7 +146,7 @@ public void testNoSaslClient() throws IOException {
 
     try {
       // Guessing the right tag byte doesn't magically get you in...
-      client.sendRpcSync(new byte[] { (byte) 0xEA }, TIMEOUT_MS);
+      client.sendRpcSync(ByteBuffer.wrap(new byte[] { (byte) 0xEA }), TIMEOUT_MS);
       fail("Should have failed");
     } catch (Exception e) {
       assertTrue(e.getMessage(), e.getMessage().contains("java.lang.IndexOutOfBoundsException"));
@@ -222,13 +224,13 @@ public synchronized void onBlockFetchFailure(String blockId, Throwable t) {
         new String[] { System.getProperty("java.io.tmpdir") }, 1,
         "org.apache.spark.shuffle.sort.SortShuffleManager");
       RegisterExecutor regmsg = new RegisterExecutor("app-1", "0", executorInfo);
-      client1.sendRpcSync(regmsg.toByteArray(), TIMEOUT_MS);
+      client1.sendRpcSync(regmsg.toByteBuffer(), TIMEOUT_MS);
 
       // Make a successful request to fetch blocks, which creates a new stream. But do not actually
       // fetch any blocks, to keep the stream open.
       OpenBlocks openMessage = new OpenBlocks("app-1", "0", blockIds);
-      byte[] response = client1.sendRpcSync(openMessage.toByteArray(), TIMEOUT_MS);
-      StreamHandle stream = (StreamHandle) BlockTransferMessage.Decoder.fromByteArray(response);
+      ByteBuffer response = client1.sendRpcSync(openMessage.toByteBuffer(), TIMEOUT_MS);
+      StreamHandle stream = (StreamHandle) BlockTransferMessage.Decoder.fromByteBuffer(response);
       long streamId = stream.streamId;
 
       // Create a second client, authenticated with a different app ID, and try to read from
@@ -275,7 +277,7 @@ public synchronized void onFailure(int chunkIndex, Throwable t) {
   /** RPC handler which simply responds with the message it received. */
   public static class TestRpcHandler extends RpcHandler {
     @Override
-    public void receive(TransportClient client, byte[] message, RpcResponseCallback callback) {
+    public void receive(TransportClient client, ByteBuffer message, RpcResponseCallback callback) {
       callback.onSuccess(message);
     }
 
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/BlockTransferMessagesSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/BlockTransferMessagesSuite.java
index d65de9ca550a3..86c8609e7070b 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/BlockTransferMessagesSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/BlockTransferMessagesSuite.java
@@ -36,7 +36,7 @@ public void serializeOpenShuffleBlocks() {
   }
 
   private void checkSerializeDeserialize(BlockTransferMessage msg) {
-    BlockTransferMessage msg2 = BlockTransferMessage.Decoder.fromByteArray(msg.toByteArray());
+    BlockTransferMessage msg2 = BlockTransferMessage.Decoder.fromByteBuffer(msg.toByteBuffer());
     assertEquals(msg, msg2);
     assertEquals(msg.hashCode(), msg2.hashCode());
     assertEquals(msg.toString(), msg2.toString());
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandlerSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandlerSuite.java
index e61390cf57061..9379412155e88 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandlerSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandlerSuite.java
@@ -60,12 +60,12 @@ public void testRegisterExecutor() {
     RpcResponseCallback callback = mock(RpcResponseCallback.class);
 
     ExecutorShuffleInfo config = new ExecutorShuffleInfo(new String[] {"/a", "/b"}, 16, "sort");
-    byte[] registerMessage = new RegisterExecutor("app0", "exec1", config).toByteArray();
+    ByteBuffer registerMessage = new RegisterExecutor("app0", "exec1", config).toByteBuffer();
     handler.receive(client, registerMessage, callback);
     verify(blockResolver, times(1)).registerExecutor("app0", "exec1", config);
 
-    verify(callback, times(1)).onSuccess((byte[]) any());
-    verify(callback, never()).onFailure((Throwable) any());
+    verify(callback, times(1)).onSuccess(any(ByteBuffer.class));
+    verify(callback, never()).onFailure(any(Throwable.class));
   }
 
   @SuppressWarnings("unchecked")
@@ -77,17 +77,18 @@ public void testOpenShuffleBlocks() {
     ManagedBuffer block1Marker = new NioManagedBuffer(ByteBuffer.wrap(new byte[7]));
     when(blockResolver.getBlockData("app0", "exec1", "b0")).thenReturn(block0Marker);
     when(blockResolver.getBlockData("app0", "exec1", "b1")).thenReturn(block1Marker);
-    byte[] openBlocks = new OpenBlocks("app0", "exec1", new String[] { "b0", "b1" }).toByteArray();
+    ByteBuffer openBlocks = new OpenBlocks("app0", "exec1", new String[] { "b0", "b1" })
+      .toByteBuffer();
     handler.receive(client, openBlocks, callback);
     verify(blockResolver, times(1)).getBlockData("app0", "exec1", "b0");
     verify(blockResolver, times(1)).getBlockData("app0", "exec1", "b1");
 
-    ArgumentCaptor<byte[]> response = ArgumentCaptor.forClass(byte[].class);
+    ArgumentCaptor<ByteBuffer> response = ArgumentCaptor.forClass(ByteBuffer.class);
     verify(callback, times(1)).onSuccess(response.capture());
     verify(callback, never()).onFailure((Throwable) any());
 
     StreamHandle handle =
-      (StreamHandle) BlockTransferMessage.Decoder.fromByteArray(response.getValue());
+      (StreamHandle) BlockTransferMessage.Decoder.fromByteBuffer(response.getValue());
     assertEquals(2, handle.numChunks);
 
     @SuppressWarnings("unchecked")
@@ -104,7 +105,7 @@ public void testOpenShuffleBlocks() {
   public void testBadMessages() {
     RpcResponseCallback callback = mock(RpcResponseCallback.class);
 
-    byte[] unserializableMsg = new byte[] { 0x12, 0x34, 0x56 };
+    ByteBuffer unserializableMsg = ByteBuffer.wrap(new byte[] { 0x12, 0x34, 0x56 });
     try {
       handler.receive(client, unserializableMsg, callback);
       fail("Should have thrown");
@@ -112,7 +113,7 @@ public void testBadMessages() {
       // pass
     }
 
-    byte[] unexpectedMsg = new UploadBlock("a", "e", "b", new byte[1], new byte[2]).toByteArray();
+    ByteBuffer unexpectedMsg = new UploadBlock("a", "e", "b", new byte[1], new byte[2]).toByteBuffer();
     try {
       handler.receive(client, unexpectedMsg, callback);
       fail("Should have thrown");
@@ -120,7 +121,7 @@ public void testBadMessages() {
       // pass
     }
 
-    verify(callback, never()).onSuccess((byte[]) any());
-    verify(callback, never()).onFailure((Throwable) any());
+    verify(callback, never()).onSuccess(any(ByteBuffer.class));
+    verify(callback, never()).onFailure(any(Throwable.class));
   }
 }
diff --git a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/OneForOneBlockFetcherSuite.java b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/OneForOneBlockFetcherSuite.java
index b35a6d685dd02..2590b9ce4c1f1 100644
--- a/network/shuffle/src/test/java/org/apache/spark/network/shuffle/OneForOneBlockFetcherSuite.java
+++ b/network/shuffle/src/test/java/org/apache/spark/network/shuffle/OneForOneBlockFetcherSuite.java
@@ -134,14 +134,14 @@ private BlockFetchingListener fetchBlocks(final LinkedHashMap<String, ManagedBuf
     doAnswer(new Answer<Void>() {
       @Override
       public Void answer(InvocationOnMock invocationOnMock) throws Throwable {
-        BlockTransferMessage message = BlockTransferMessage.Decoder.fromByteArray(
-          (byte[]) invocationOnMock.getArguments()[0]);
+        BlockTransferMessage message = BlockTransferMessage.Decoder.fromByteBuffer(
+          (ByteBuffer) invocationOnMock.getArguments()[0]);
         RpcResponseCallback callback = (RpcResponseCallback) invocationOnMock.getArguments()[1];
-        callback.onSuccess(new StreamHandle(123, blocks.size()).toByteArray());
+        callback.onSuccess(new StreamHandle(123, blocks.size()).toByteBuffer());
         assertEquals(new OpenBlocks("app-id", "exec-id", blockIds), message);
         return null;
       }
-    }).when(client).sendRpc((byte[]) any(), (RpcResponseCallback) any());
+    }).when(client).sendRpc(any(ByteBuffer.class), any(RpcResponseCallback.class));
 
     // Respond to each chunk request with a single buffer from our blocks array.
     final AtomicInteger expectedChunkIndex = new AtomicInteger(0);

From 96bf468c7860be317c20ccacf259910968d2dc83 Mon Sep 17 00:00:00 2001
From: Sean Owen <sowen@cloudera.com>
Date: Mon, 30 Nov 2015 17:33:09 -0800
Subject: [PATCH 514/518] [SPARK-12049][CORE] User JVM shutdown hook can cause
 deadlock at shutdown

Avoid potential deadlock with a user app's shutdown hook thread by more narrowly synchronizing access to 'hooks'

Author: Sean Owen <sowen@cloudera.com>

Closes #10042 from srowen/SPARK-12049.
---
 .../spark/util/ShutdownHookManager.scala      | 33 +++++++++----------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/util/ShutdownHookManager.scala b/core/src/main/scala/org/apache/spark/util/ShutdownHookManager.scala
index 4012dca3ecdf8..620f226a23e15 100644
--- a/core/src/main/scala/org/apache/spark/util/ShutdownHookManager.scala
+++ b/core/src/main/scala/org/apache/spark/util/ShutdownHookManager.scala
@@ -206,7 +206,7 @@ private[spark] object ShutdownHookManager extends Logging {
 private [util] class SparkShutdownHookManager {
 
   private val hooks = new PriorityQueue[SparkShutdownHook]()
-  private var shuttingDown = false
+  @volatile private var shuttingDown = false
 
   /**
    * Install a hook to run at shutdown and run all registered hooks in order. Hadoop 1.x does not
@@ -232,28 +232,27 @@ private [util] class SparkShutdownHookManager {
     }
   }
 
-  def runAll(): Unit = synchronized {
+  def runAll(): Unit = {
     shuttingDown = true
-    while (!hooks.isEmpty()) {
-      Try(Utils.logUncaughtExceptions(hooks.poll().run()))
+    var nextHook: SparkShutdownHook = null
+    while ({ nextHook = hooks.synchronized { hooks.poll() }; nextHook != null }) {
+      Try(Utils.logUncaughtExceptions(nextHook.run()))
     }
   }
 
-  def add(priority: Int, hook: () => Unit): AnyRef = synchronized {
-    checkState()
-    val hookRef = new SparkShutdownHook(priority, hook)
-    hooks.add(hookRef)
-    hookRef
-  }
-
-  def remove(ref: AnyRef): Boolean = synchronized {
-    hooks.remove(ref)
+  def add(priority: Int, hook: () => Unit): AnyRef = {
+    hooks.synchronized {
+      if (shuttingDown) {
+        throw new IllegalStateException("Shutdown hooks cannot be modified during shutdown.")
+      }
+      val hookRef = new SparkShutdownHook(priority, hook)
+      hooks.add(hookRef)
+      hookRef
+    }
   }
 
-  private def checkState(): Unit = {
-    if (shuttingDown) {
-      throw new IllegalStateException("Shutdown hooks cannot be modified during shutdown.")
-    }
+  def remove(ref: AnyRef): Boolean = {
+    hooks.synchronized { hooks.remove(ref) }
   }
 
 }

From f73379be2b0c286957b678a996cb56afc96015eb Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Mon, 30 Nov 2015 17:15:47 -0800
Subject: [PATCH 515/518] [HOTFIX][SPARK-12000] Add missing quotes in Jekyll
 API docs plugin.

I accidentally omitted these as part of #10049.
---
 docs/_plugins/copy_api_dirs.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/_plugins/copy_api_dirs.rb b/docs/_plugins/copy_api_dirs.rb
index f2f3e2e653149..174c202e37918 100644
--- a/docs/_plugins/copy_api_dirs.rb
+++ b/docs/_plugins/copy_api_dirs.rb
@@ -117,7 +117,7 @@
 
   puts "Moving to python/docs directory and building sphinx."
   cd("../python/docs")
-  system(make html) || raise("Python doc generation failed")
+  system("make html") || raise("Python doc generation failed")
 
   puts "Moving back into home dir."
   cd("../../")

From 9693b0d5a55bc1d2da96f04fe2c6de59a8dfcc1b Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@appier.com>
Date: Mon, 30 Nov 2015 20:56:42 -0800
Subject: [PATCH 516/518] [SPARK-12018][SQL] Refactor common subexpression
 elimination code

JIRA: https://issues.apache.org/jira/browse/SPARK-12018

The code of common subexpression elimination can be factored and simplified. Some unnecessary variables can be removed.

Author: Liang-Chi Hsieh <viirya@appier.com>

Closes #10009 from viirya/refactor-subexpr-eliminate.
---
 .../sql/catalyst/expressions/Expression.scala | 10 ++----
 .../expressions/codegen/CodeGenerator.scala   | 34 +++++--------------
 .../codegen/GenerateUnsafeProjection.scala    |  4 +--
 3 files changed, 14 insertions(+), 34 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
index 169435a10ea2c..b55d3653a7158 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala
@@ -94,13 +94,9 @@ abstract class Expression extends TreeNode[Expression] {
   def gen(ctx: CodeGenContext): GeneratedExpressionCode = {
     ctx.subExprEliminationExprs.get(this).map { subExprState =>
       // This expression is repeated meaning the code to evaluated has already been added
-      // as a function, `subExprState.fnName`. Just call that.
-      val code =
-        s"""
-           |/* $this */
-           |${subExprState.fnName}(${ctx.INPUT_ROW});
-         """.stripMargin.trim
-      GeneratedExpressionCode(code, subExprState.code.isNull, subExprState.code.value)
+      // as a function and called in advance. Just use it.
+      val code = s"/* $this */"
+      GeneratedExpressionCode(code, subExprState.isNull, subExprState.value)
     }.getOrElse {
       val isNull = ctx.freshName("isNull")
       val primitive = ctx.freshName("primitive")
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
index 2f3d6aeb86c5b..440c7d2fc1156 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
@@ -104,16 +104,13 @@ class CodeGenContext {
   val equivalentExpressions: EquivalentExpressions = new EquivalentExpressions
 
   // State used for subexpression elimination.
-  case class SubExprEliminationState(
-      isLoaded: String,
-      code: GeneratedExpressionCode,
-      fnName: String)
+  case class SubExprEliminationState(isNull: String, value: String)
 
   // Foreach expression that is participating in subexpression elimination, the state to use.
   val subExprEliminationExprs = mutable.HashMap.empty[Expression, SubExprEliminationState]
 
-  // The collection of isLoaded variables that need to be reset on each row.
-  val subExprIsLoadedVariables = mutable.ArrayBuffer.empty[String]
+  // The collection of sub-exression result resetting methods that need to be called on each row.
+  val subExprResetVariables = mutable.ArrayBuffer.empty[String]
 
   final val JAVA_BOOLEAN = "boolean"
   final val JAVA_BYTE = "byte"
@@ -408,7 +405,6 @@ class CodeGenContext {
     val commonExprs = equivalentExpressions.getAllEquivalentExprs.filter(_.size > 1)
     commonExprs.foreach(e => {
       val expr = e.head
-      val isLoaded = freshName("isLoaded")
       val isNull = freshName("isNull")
       val value = freshName("value")
       val fnName = freshName("evalExpr")
@@ -417,18 +413,12 @@ class CodeGenContext {
       val code = expr.gen(this)
       val fn =
         s"""
-           |private void $fnName(InternalRow ${INPUT_ROW}) {
-           |  if (!$isLoaded) {
-           |    ${code.code.trim}
-           |    $isLoaded = true;
-           |    $isNull = ${code.isNull};
-           |    $value = ${code.value};
-           |  }
+           |private void $fnName(InternalRow $INPUT_ROW) {
+           |  ${code.code.trim}
+           |  $isNull = ${code.isNull};
+           |  $value = ${code.value};
            |}
            """.stripMargin
-      code.code = fn
-      code.isNull = isNull
-      code.value = value
 
       addNewFunction(fnName, fn)
 
@@ -448,18 +438,12 @@ class CodeGenContext {
       //   2. Less code.
       // Currently, we will do this for all non-leaf only expression trees (i.e. expr trees with
       // at least two nodes) as the cost of doing it is expected to be low.
-
-      // Maintain the loaded value and isNull as member variables. This is necessary if the codegen
-      // function is split across multiple functions.
-      // TODO: maintaining this as a local variable probably allows the compiler to do better
-      // optimizations.
-      addMutableState("boolean", isLoaded, s"$isLoaded = false;")
       addMutableState("boolean", isNull, s"$isNull = false;")
       addMutableState(javaType(expr.dataType), value,
         s"$value = ${defaultValue(expr.dataType)};")
-      subExprIsLoadedVariables += isLoaded
 
-      val state = SubExprEliminationState(isLoaded, code, fnName)
+      subExprResetVariables += s"$fnName($INPUT_ROW);"
+      val state = SubExprEliminationState(isNull, value)
       e.foreach(subExprEliminationExprs.put(_, state))
     })
   }
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
index 7b6c9373ebe30..68005afb21d2e 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala
@@ -287,8 +287,8 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
     val holderClass = classOf[BufferHolder].getName
     ctx.addMutableState(holderClass, bufferHolder, s"this.$bufferHolder = new $holderClass();")
 
-    // Reset the isLoaded flag for each row.
-    val subexprReset = ctx.subExprIsLoadedVariables.map { v => s"${v} = false;" }.mkString("\n")
+    // Reset the subexpression values for each row.
+    val subexprReset = ctx.subExprResetVariables.mkString("\n")
 
     val code =
       s"""

From a0af0e351e45a8be47a6f65efd132eaa4a00c9e4 Mon Sep 17 00:00:00 2001
From: Yuhao Yang <hhbyyh@gmail.com>
Date: Tue, 1 Dec 2015 09:26:58 +0000
Subject: [PATCH 517/518] [SPARK-11898][MLLIB] Use broadcast for the global
 tables in Word2Vec

jira: https://issues.apache.org/jira/browse/SPARK-11898
syn0Global and sync1Global in word2vec are quite large objects with size (vocab * vectorSize * 8), yet they are passed to worker using basic task serialization.

Use broadcast can greatly improve the performance. My benchmark shows that, for 1M vocabulary and default vectorSize 100, changing to broadcast can help,

1. decrease the worker memory consumption by 45%.
2. decrease running time by 40%.

This will also help extend the upper limit for Word2Vec.

Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #9878 from hhbyyh/w2vBC.
---
 .../scala/org/apache/spark/mllib/feature/Word2Vec.scala    | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
index a47f27b0afb14..655ac0bb5545b 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala
@@ -316,12 +316,15 @@ class Word2Vec extends Serializable with Logging {
       Array.fill[Float](vocabSize * vectorSize)((initRandom.nextFloat() - 0.5f) / vectorSize)
     val syn1Global = new Array[Float](vocabSize * vectorSize)
     var alpha = learningRate
+
     for (k <- 1 to numIterations) {
+      val bcSyn0Global = sc.broadcast(syn0Global)
+      val bcSyn1Global = sc.broadcast(syn1Global)
       val partial = newSentences.mapPartitionsWithIndex { case (idx, iter) =>
         val random = new XORShiftRandom(seed ^ ((idx + 1) << 16) ^ ((-k - 1) << 8))
         val syn0Modify = new Array[Int](vocabSize)
         val syn1Modify = new Array[Int](vocabSize)
-        val model = iter.foldLeft((syn0Global, syn1Global, 0, 0)) {
+        val model = iter.foldLeft((bcSyn0Global.value, bcSyn1Global.value, 0, 0)) {
           case ((syn0, syn1, lastWordCount, wordCount), sentence) =>
             var lwc = lastWordCount
             var wc = wordCount
@@ -405,6 +408,8 @@ class Word2Vec extends Serializable with Logging {
         }
         i += 1
       }
+      bcSyn0Global.unpersist(false)
+      bcSyn1Global.unpersist(false)
     }
     newSentences.unpersist()
 

From 8a0d88332f39e44365c7cbe3fdb9fac251251d85 Mon Sep 17 00:00:00 2001
From: Xusen Yin <yinxusen@gmail.com>
Date: Tue, 1 Dec 2015 22:53:15 +0800
Subject: [PATCH 518/518] fix minor issues

---
 docs/ml-features.md                                           | 2 --
 .../spark/examples/ml/JavaElementwiseProductExample.java      | 4 ++--
 .../apache/spark/examples/ml/JavaOneHotEncoderExample.java    | 2 +-
 .../java/org/apache/spark/examples/ml/JavaPCAExample.java     | 2 +-
 .../org/apache/spark/examples/ml/JavaTokenizerExample.java    | 4 ++--
 examples/src/main/python/ml/normalizer_example.py             | 3 +--
 examples/src/main/python/ml/onehot_encoder_example.py         | 2 +-
 .../org/apache/spark/examples/ml/VectorSlicerExample.scala    | 2 +-
 8 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/docs/ml-features.md b/docs/ml-features.md
index e29e66e69956f..b0fcee6135ac8 100644
--- a/docs/ml-features.md
+++ b/docs/ml-features.md
@@ -279,7 +279,6 @@ Refer to the [NGram Java docs](api/java/org/apache/spark/ml/feature/NGram.html)
 for more details on the API.
 
 {% include_example java/org/apache/spark/examples/ml/JavaNGramExample.java %}
-{% endhighlight %}
 </div>
 
 <div data-lang="python" markdown="1">
@@ -888,7 +887,6 @@ Refer to the [VectorSlicer Java docs](api/java/org/apache/spark/ml/feature/Vecto
 for more details on the API.
 
 {% include_example java/org/apache/spark/examples/ml/JavaVectorSlicerExample.java %}
-{% endhighlight %}
 </div>
 </div>
 
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaElementwiseProductExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaElementwiseProductExample.java
index 15ecd425929b1..2898accec61b0 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaElementwiseProductExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaElementwiseProductExample.java
@@ -29,6 +29,7 @@
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.ml.feature.ElementwiseProduct;
 import org.apache.spark.mllib.linalg.Vector;
+import org.apache.spark.mllib.linalg.VectorUDT;
 import org.apache.spark.mllib.linalg.Vectors;
 import org.apache.spark.sql.DataFrame;
 import org.apache.spark.sql.Row;
@@ -36,7 +37,6 @@
 import org.apache.spark.sql.types.DataTypes;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
-import org.apache.spark.mllib.linalg.VectorUDT;
 // $example off$
 
 public class JavaElementwiseProductExample {
@@ -53,7 +53,7 @@ public static void main(String[] args) {
     ));
 
     List<StructField> fields = new ArrayList<StructField>(2);
-    fields.add(DataTypes.createStructField("id", new VectorUDT(), false));
+    fields.add(DataTypes.createStructField("id", DataTypes.StringType, false));
     fields.add(DataTypes.createStructField("vector", new VectorUDT(), false));
 
     StructType schema = DataTypes.createStructType(fields);
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaOneHotEncoderExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaOneHotEncoderExample.java
index 854354a9334e5..172a9cc6feb28 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaOneHotEncoderExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaOneHotEncoderExample.java
@@ -21,7 +21,7 @@
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.SQLContext;
 
-// $example off$
+// $example on$
 import java.util.Arrays;
 
 import org.apache.spark.api.java.JavaRDD;
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaPCAExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaPCAExample.java
index dc4ba0583b57f..8282fab084f36 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaPCAExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaPCAExample.java
@@ -21,7 +21,7 @@
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.SQLContext;
 
-// $example off$
+// $example on$
 import java.util.Arrays;
 
 import org.apache.spark.api.java.JavaRDD;
diff --git a/examples/src/main/java/org/apache/spark/examples/ml/JavaTokenizerExample.java b/examples/src/main/java/org/apache/spark/examples/ml/JavaTokenizerExample.java
index 213d8bbd8ea78..617dc3f66e3bf 100644
--- a/examples/src/main/java/org/apache/spark/examples/ml/JavaTokenizerExample.java
+++ b/examples/src/main/java/org/apache/spark/examples/ml/JavaTokenizerExample.java
@@ -21,7 +21,7 @@
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.sql.SQLContext;
 
-// $example off$
+// $example on$
 import java.util.Arrays;
 
 import org.apache.spark.api.java.JavaRDD;
@@ -69,7 +69,7 @@ public static void main(String[] args) {
       .setInputCol("sentence")
       .setOutputCol("words")
       .setPattern("\\W");  // alternatively .setPattern("\\w+").setGaps(false);
-    // example off
+    // $example off$
     jsc.stop();
   }
 }
diff --git a/examples/src/main/python/ml/normalizer_example.py b/examples/src/main/python/ml/normalizer_example.py
index 9238d88656ec9..833d93e976a7e 100644
--- a/examples/src/main/python/ml/normalizer_example.py
+++ b/examples/src/main/python/ml/normalizer_example.py
@@ -28,8 +28,7 @@
     sqlContext = SQLContext(sc)
 
     # $example on$
-    dataFrame = sqlContext.read.format("libsvm")
-        .load("data/mllib/sample_libsvm_data.txt")
+    dataFrame = sqlContext.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
 
     # Normalize each Vector using $L^1$ norm.
     normalizer = Normalizer(inputCol="features", outputCol="normFeatures", p=1.0)
diff --git a/examples/src/main/python/ml/onehot_encoder_example.py b/examples/src/main/python/ml/onehot_encoder_example.py
index 21e9dd6fda26c..7529dfd09213a 100644
--- a/examples/src/main/python/ml/onehot_encoder_example.py
+++ b/examples/src/main/python/ml/onehot_encoder_example.py
@@ -40,7 +40,7 @@
     stringIndexer = StringIndexer(inputCol="category", outputCol="categoryIndex")
     model = stringIndexer.fit(df)
     indexed = model.transform(df)
-    encoder = OneHotEncoder(includeFirst=False, inputCol="categoryIndex", outputCol="categoryVec")
+    encoder = OneHotEncoder(dropLast=False, inputCol="categoryIndex", outputCol="categoryVec")
     encoded = encoder.transform(indexed)
     # $example off$
 
diff --git a/examples/src/main/scala/org/apache/spark/examples/ml/VectorSlicerExample.scala b/examples/src/main/scala/org/apache/spark/examples/ml/VectorSlicerExample.scala
index ac72ae60b9f1b..04f19829eff87 100644
--- a/examples/src/main/scala/org/apache/spark/examples/ml/VectorSlicerExample.scala
+++ b/examples/src/main/scala/org/apache/spark/examples/ml/VectorSlicerExample.scala
@@ -21,12 +21,12 @@ package org.apache.spark.examples.ml
 // $example on$
 import org.apache.spark.ml.attribute.{Attribute, AttributeGroup, NumericAttribute}
 import org.apache.spark.ml.feature.VectorSlicer
+import org.apache.spark.mllib.linalg.Vectors
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.types.StructType
 // $example off$
 import org.apache.spark.sql.SQLContext
 import org.apache.spark.{SparkConf, SparkContext}
-import org.apache.spark.mllib.linalg.Vectors
 
 object VectorSlicerExample {
   def main(args: Array[String]): Unit = {