[GitHub] [spark] zhengruifeng commented on a change in pull request #28176: [SPARK-31301][ML] Flatten the result dataframe of tests in testChiSquare

[GitBox <gi...@apache.org>]

zhengruifeng commented on a change in pull request #28176: [SPARK-31301][ML] Flatten the result dataframe of tests in testChiSquare
URL: https://github.com/apache/spark/pull/28176#discussion_r406614269
 
 

 ##########
 File path: mllib/src/main/scala/org/apache/spark/ml/stat/ChiSquareTest.scala
 ##########
 @@ -63,40 +55,47 @@ object ChiSquareTest {
    */
   @Since("2.2.0")
   def test(dataset: DataFrame, featuresCol: String, labelCol: String): DataFrame = {
-    val spark = dataset.sparkSession
-    import spark.implicits._
-
-    SchemaUtils.checkColumnType(dataset.schema, featuresCol, new VectorUDT)
-    SchemaUtils.checkNumericType(dataset.schema, labelCol)
-    val rdd = dataset.select(col(labelCol).cast("double"), col(featuresCol)).as[(Double, Vector)]
-      .rdd.map { case (label, features) => OldLabeledPoint(label, OldVectors.fromML(features)) }
-    val testResults = OldStatistics.chiSqTest(rdd)
-    val pValues = Vectors.dense(testResults.map(_.pValue))
-    val degreesOfFreedom = testResults.map(_.degreesOfFreedom)
-    val statistics = Vectors.dense(testResults.map(_.statistic))
-    spark.createDataFrame(Seq(ChiSquareResult(pValues, degreesOfFreedom, statistics)))
+    test(dataset, featuresCol, labelCol, false)
   }
 
   /**
    * @param dataset  DataFrame of categorical labels and categorical features.
    *                 Real-valued features will be treated as categorical for each distinct value.
    * @param featuresCol  Name of features column in dataset, of type `Vector` (`VectorUDT`)
    * @param labelCol  Name of label column in dataset, of any numerical type
-   * @return Array containing the SelectionTestResult for every feature against the label.
+   * @param flatten  If false, the returned DataFrame contains only a single Row, otherwise, one
+   *                 row per feature.
    */
   @Since("3.1.0")
-  def testChiSquare(
-      dataset: Dataset[_],
+  def test(
+      dataset: DataFrame,
       featuresCol: String,
-      labelCol: String): Array[SelectionTestResult] = {
-
+      labelCol: String,
+      flatten: Boolean): DataFrame = {
     SchemaUtils.checkColumnType(dataset.schema, featuresCol, new VectorUDT)
     SchemaUtils.checkNumericType(dataset.schema, labelCol)
-    val input = dataset.select(col(labelCol).cast(DoubleType), col(featuresCol)).rdd
-      .map { case Row(label: Double, features: Vector) =>
-        OldLabeledPoint(label, OldVectors.fromML(features))
-      }
-    val chiTestResult = OldStatistics.chiSqTest(input)
-    chiTestResult.map(r => new ChiSqTestResult(r.pValue, r.degreesOfFreedom, r.statistic))
+
+    val spark = dataset.sparkSession
+    import spark.implicits._
+
+    val data = dataset.select(col(labelCol).cast("double"), col(featuresCol)).rdd
+      .map { case Row(label: Double, vec: Vector) => (label, OldVectors.fromML(vec)) }
+    val resultRDD = OldChiSqTest.computeChiSquared(data)
+
+    if (flatten) {
+      resultRDD.map { case (col, (pValue, degreesOfFreedom, statistic, _)) =>
+        (col, pValue, degreesOfFreedom, statistic)
+      }.toDF("featureIndex", "pValue", "degreesOfFreedom", "statistic")
+    } else {
+      resultRDD.map { case (col, (pValue, degreesOfFreedom, statistic, _)) =>
+        (0, (col, pValue, degreesOfFreedom, statistic))
+      }.groupByKey().map { case (_, seq) =>
 
 Review comment:
   collect results in some executor instead of the driver

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org