[SPARK-11057] [SQL] Add correlation and covariance matrices

sql/core/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala

@@ -52,13 +52,39 @@ final class DataFrameStatFunctions private[sql](df: DataFrame) {
     StatFunctions.calculateCov(df, Seq(col1, col2))
   }
 
+  /**
+   * Calculate the sample covariance between columns of a DataFrame.
+   *
+   * @return a covariance matrix as a DataFrame
+   *
+   * {{{
+   *    val df = sc.parallelize(0 until 10).toDF("id").withColumn("rand1", rand(seed=10))
+   *      .withColumn("rand2", rand(seed=27))
+   *    val covmatrix = df.stat.cov()
+   *    covmatrix.show()
+   *    +---------+------------------+-------------------+-------------------+
+   *    |FieldName|                id|              rand1|              rand2|
+   *    +---------+------------------+-------------------+-------------------+
+   *    |       id| 9.166666666666666| 0.4131594565676311| 0.7012982830955725|
+   *    |    rand1|0.4131594565676311|0.11982701890603772|0.06500805072758595|
+   *    |    rand2|0.7012982830955725|0.06500805072758595|0.09383550706974164|
+   *    +---------+------------------+-------------------+-------------------+
+   * }}}
+   *
+   * @since 1.6.0
+   */
+  def cov(): DataFrame = {
+    StatFunctions.calculateCov(df)
+  }
+
   /**
    * Calculates the correlation of two columns of a DataFrame. Currently only supports the Pearson
    * Correlation Coefficient. For Spearman Correlation, consider using RDD methods found in
    * MLlib's Statistics.
    *
    * @param col1 the name of the column
    * @param col2 the name of the column to calculate the correlation against
+   * @param method the name of the correlation method
    * @return The Pearson Correlation Coefficient as a Double.
    *
    * {{{
@@ -96,6 +122,63 @@ final class DataFrameStatFunctions private[sql](df: DataFrame) {
     corr(col1, col2, "pearson")
   }
 
+  /**
+   * Calculates the correlation of columns in the DataFrame. Currently only supports the Pearson
+   * Correlation Coefficient. For Spearman Correlation, consider using RDD methods found in
+   * MLlib's Statistics.
+   *
+   * @param method the name of the correlation method
+   * @return The Pearson Correlation matrix as a DataFrame.
+   *
+   * {{{
+   *    val df = sc.parallelize(0 until 10).toDF("id").withColumn("rand1", rand(seed=10))
+   *      .withColumn("rand2", rand(seed=27))
+   *    val corrmatrix = df.stat.corr()
+   *    corrmatrix.show()
+   *    +---------+------------------+------------------+------------------+
+   *    |FieldName|                id|             rand1|             rand2|
+   *    +---------+------------------+------------------+------------------+
+   *    |       id|               1.0|   0.3942163209095|0.7561595709319909|
+   *    |    rand1|   0.3942163209095|               1.0|0.6130644931298477|
+   *    |    rand2|0.7561595709319909|0.6130644931298477|               1.0|
+   *    +---------+------------------+------------------+------------------+
+   * }}}
+   *
+   * @since 1.6.0
+   */
+ def corr(method: String): DataFrame = {
+    require(method == "pearson", "Currently only the calculation of the Pearson Correlation " +
+      "coefficient is supported.")
+    StatFunctions.pearsonCorrelation(df)
+ }
+
+  /**
+   * Calculates the correlation of columns in the DataFrame. Currently only supports the Pearson
+   * Correlation Coefficient. For Spearman Correlation, consider using RDD methods found in
+   * MLlib's Statistics.
+   *
+   * @return The Pearson Correlation matrix as a DataFrame.
+   *
+   * {{{
+   *    val df = sc.parallelize(0 until 10).toDF("id").withColumn("rand1", rand(seed=10))
+   *      .withColumn("rand2", rand(seed=27))
+   *    val corrmatrix = df.stat.corr()
+   *    corrmatrix.show()
+   *    +---------+------------------+------------------+------------------+
+   *    |FieldName|                id|             rand1|             rand2|
+   *    +---------+------------------+------------------+------------------+
+   *    |       id|               1.0|   0.3942163209095|0.7561595709319909|
+   *    |    rand1|   0.3942163209095|               1.0|0.6130644931298477|
+   *    |    rand2|0.7561595709319909|0.6130644931298477|               1.0|
+   *    +---------+------------------+------------------+------------------+
+   * }}}
+   *
+   * @since 1.6.0
+   */
+ def corr(): DataFrame = {
+    corr("pearson")
+ }
+
   /**
    * Computes a pair-wise frequency table of the given columns. Also known as a contingency table.
    * The number of distinct values for each column should be less than 1e4. At most 1e6 non-zero

sql/core/src/main/scala/org/apache/spark/sql/execution/stat/StatFunctions.scala

@@ -19,7 +19,8 @@ package org.apache.spark.sql.execution.stat
 
 import org.apache.spark.Logging
 import org.apache.spark.sql.{Row, Column, DataFrame}
-import org.apache.spark.sql.catalyst.expressions.{GenericMutableRow, Cast}
+import org.apache.spark.sql.catalyst.expressions.{GenericMutableRow, Cast, AttributeReference}
+import scala.collection.mutable.ArrayBuffer
 import org.apache.spark.sql.catalyst.plans.logical.LocalRelation
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types._
@@ -33,6 +34,31 @@ private[sql] object StatFunctions extends Logging {
     counts.Ck / math.sqrt(counts.MkX * counts.MkY)
   }
 
+  /** Calculate the Pearson Correlation matrix for given DataFrame */
+  private[sql] def pearsonCorrelation(df: DataFrame): DataFrame = {
+    val fieldNames = df.schema.fieldNames
+    val dfStructAttrs = ArrayBuffer[AttributeReference](
+              AttributeReference("FieldName", StringType, true)())
+    val rows = fieldNames.map{fname => val countsRow = new GenericMutableRow(fieldNames.length + 1)
+        countsRow.update(0, UTF8String.fromString(fname))
+        countsRow
+    }.toSeq
+    // generates field types of the output DataFrame
+    for(field <- fieldNames) dfStructAttrs += AttributeReference(field, DoubleType, true)()
+
+    // fills the correlation matrix by computing column-by-column correlations
+    for (i <- 0 to fieldNames.length - 1){
+      for (j <- 0 to i){
+          val corr = pearsonCorrelation(df, Seq(fieldNames(i), fieldNames(j)))
+          rows(i).setDouble(j + 1, corr)
+          rows(j).setDouble(i + 1, corr)
+      }
+      rows(i).setDouble(i + 1, 1.0)
+    }
+
+    new DataFrame(df.sqlContext, new LocalRelation(dfStructAttrs, rows))
+  }
+
   /** Helper class to simplify tracking and merging counts. */
   private class CovarianceCounter extends Serializable {
     var xAvg = 0.0 // the mean of all examples seen so far in col1
@@ -102,6 +128,34 @@ private[sql] object StatFunctions extends Logging {
     counts.cov
   }
 
+ /**
+   * Calculate the covariance of two numerical columns of a DataFrame.
+   * @param df The DataFrame
+   * @return the covariance matrix.
+   */
+  private[sql] def calculateCov(df: DataFrame): DataFrame = {
+    val fieldNames = df.schema.fieldNames
+    val dfStructAttrs = ArrayBuffer[AttributeReference](
+             AttributeReference("FieldName", StringType, true)())
+    val rows = fieldNames.map{fname => val countsRow = new GenericMutableRow(fieldNames.length + 1)
+        countsRow.update(0, UTF8String.fromString(fname))
+        countsRow
+    }.toSeq
+    // generates field types of the output DataFrame
+    for(field <- fieldNames) dfStructAttrs += AttributeReference(field, DoubleType, true)()
+
+    // fills the covariance matrix by computing column-by-column covariances
+    for (i <- 0 to fieldNames.length-1){
+      for (j <- 0 to i){
+          val cov = calculateCov(df, Seq(fieldNames(i), fieldNames(j)))
+          rows(i).setDouble(j + 1, cov)
+          rows(j).setDouble(i + 1, cov)
+      }
+    }
+
+    new DataFrame(df.sqlContext, new LocalRelation(dfStructAttrs, rows))
+  }
+
   /** Generate a table of frequencies for the elements of two columns. */
   private[sql] def crossTabulate(df: DataFrame, col1: String, col2: String): DataFrame = {
     val tableName = s"${col1}_$col2"

sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala

@@ -85,6 +85,35 @@ class DataFrameStatSuite extends QueryTest with SharedSQLContext {
     assert(math.abs(corr3 - 0.95723391394758572) < 1e-12)
   }
 
+  test("pearson correlation matrix") {
+    val df1 = Seq.tabulate(10)(i => (i, 2.0 * i, toLetter(i))).toDF("singles", "doubles", "letters")
+
+    intercept[IllegalArgumentException] {
+      df1.stat.corr() // doesn't accept non-numerical dataTypes
+    }
+
+    val df2 = Seq.tabulate(10)(i => (i, 2 * i, i * -1.0)).toDF("a", "b", "c")
+    val results = df2.stat.corr()
+
+    val row1 = results.where($"FieldName" === "a").collect()(0)
+    assert(row1.getString(0) == "a")
+    assert(row1.getDouble(1) == 1.0)
+    assert(row1.getDouble(2) == 1.0)
+    assert(row1.getDouble(3) == -1.0)
+
+    val row2 = results.where($"FieldName" === "b").collect()(0)
+    assert(row2.getString(0) == "b")
+    assert(row2.getDouble(1) == 1.0)
+    assert(row2.getDouble(2) == 1.0)
+    assert(row2.getDouble(3) == -1.0)
+
+    val row3 = results.where($"FieldName" === "c").collect()(0)
+    assert(row3.getString(0) == "c")
+    assert(row3.getDouble(1) == -1.0)
+    assert(row3.getDouble(2) == -1.0)
+    assert(row3.getDouble(3) == 1.0)
+  }
+
   test("covariance") {
     val df = Seq.tabulate(10)(i => (i, 2.0 * i, toLetter(i))).toDF("singles", "doubles", "letters")
 
@@ -98,6 +127,27 @@ class DataFrameStatSuite extends QueryTest with SharedSQLContext {
     assert(math.abs(decimalRes) < 1e-12)
   }
 
+  test("covariance matrix") {
+    val df1 = Seq.tabulate(10)(i => (i, 2.0 * i, toLetter(i))).toDF("singles", "doubles", "letters")
+
+    intercept[IllegalArgumentException] {
+      df1.stat.cov() // doesn't accept non-numerical dataTypes
+    }
+
+    val df2 = Seq.tabulate(10)(i => (i, 2.0 * i)).toDF("singles", "doubles")
+    val results = df2.stat.cov()
+
+    val row1 = results.where($"FieldName" === "singles").collect()(0)
+    assert(row1.getString(0) == "singles")
+    assert(row1.getDouble(1) == 9.166666666666666)
+    assert(row1.getDouble(2) == 18.333333333333332)
+
+    val row2 = results.where($"FieldName" === "doubles").collect()(0)
+    assert(row2.getString(0) == "doubles")
+    assert(row2.getDouble(1) == row1.getDouble(2))
+    assert(row2.getDouble(2) == 36.666666666666664)
+  }
+
   test("crosstab") {
     val rng = new Random()
     val data = Seq.tabulate(25)(i => (rng.nextInt(5), rng.nextInt(10)))