From 07bd868956e8d63294b2acb0b5d01a7ca2b35866 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Thu, 9 Nov 2017 22:57:04 -0800
Subject: [PATCH 1/9] [SPARK-22397][ML]add multiple columns support to
 QuantileDiscretizer

---
 .../ml/feature/QuantileDiscretizer.scala      | 115 ++++++++++--
 .../ml/feature/QuantileDiscretizerSuite.scala | 168 ++++++++++++++++++
 2 files changed, 267 insertions(+), 16 deletions(-)

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 95e8830283de..423916649857 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
@@ -22,7 +22,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.ml._
 import org.apache.spark.ml.attribute.NominalAttribute
 import org.apache.spark.ml.param._
-import org.apache.spark.ml.param.shared.{HasHandleInvalid, HasInputCol, HasOutputCol}
+import org.apache.spark.ml.param.shared.{HasHandleInvalid, HasInputCol, HasInputCols, HasOutputCol, HasOutputCols}
 import org.apache.spark.ml.util._
 import org.apache.spark.sql.Dataset
 import org.apache.spark.sql.types.StructType
@@ -31,7 +31,7 @@ import org.apache.spark.sql.types.StructType
  * Params for [[QuantileDiscretizer]].
  */
 private[feature] trait QuantileDiscretizerBase extends Params
-  with HasHandleInvalid with HasInputCol with HasOutputCol {
+  with HasHandleInvalid with HasInputCol with HasInputCols with HasOutputCol with HasOutputCols {
 
   /**
    * Number of buckets (quantiles, or categories) into which data points are grouped. Must
@@ -50,10 +50,26 @@ private[feature] trait QuantileDiscretizerBase extends Params
   /** @group getParam */
   def getNumBuckets: Int = getOrDefault(numBuckets)
 
+  /**
+   * Array of number of buckets (quantiles, or categories) into which data points are grouped.
+   *
+   * See also [[handleInvalid]], which can optionally create an additional bucket for NaN values.
+   *
+   * @group param
+   */
+  val numBucketsArray = new IntArrayParam(this, "numBucketsArray", "Array of number of buckets " +
+    "(quantiles, or categories) into which data points are grouped. This is for multiple " +
+    "columns input. If numBucketsArray is not set but numBuckets is set, it means user wants " +
+    "to use the same numBuckets across all columns.")
+
+  /** @group getParam */
+  def getNumBucketsArray: Array[Int] = $(numBucketsArray)
+
   /**
    * Relative error (see documentation for
    * `org.apache.spark.sql.DataFrameStatFunctions.approxQuantile` for description)
    * Must be in the range [0, 1].
+   * Note that in multiple columns case, relative error is applied to all columns.
    * default: 0.001
    * @group param
    */
@@ -68,7 +84,9 @@ private[feature] trait QuantileDiscretizerBase extends Params
   /**
    * Param for how to handle invalid entries. Options are 'skip' (filter out rows with
    * invalid values), 'error' (throw an error), or 'keep' (keep invalid values in a special
-   * additional bucket).
+   * additional bucket). Note that in the multiple columns case, the invalid handling is applied
+   * to all columns. That said for 'error' it will throw an error if any invalids are found in
+   * any column, for 'skip' it will skip rows with any invalids in any columns, etc.
    * Default: "error"
    * @group param
    */
@@ -86,6 +104,10 @@ private[feature] trait QuantileDiscretizerBase extends Params
  * categorical features. The number of bins can be set using the `numBuckets` parameter. It is
  * possible that the number of buckets used will be smaller than this value, for example, if there
  * are too few distinct values of the input to create enough distinct quantiles.
+ * Since 2.3.0,
+ * `QuantileDiscretizer` can also map multiple columns at once. Whether it goes to map a column or
+ * multiple columns, it depends on which parameter of `inputCol` and `inputCols` is set. When both
+ * are set, a log warning will be printed and by default it chooses `inputCol`.
  *
  * NaN handling:
  * null and NaN values will be ignored from the column during `QuantileDiscretizer` fitting. This
@@ -129,34 +151,95 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
   @Since("2.1.0")
   def setHandleInvalid(value: String): this.type = set(handleInvalid, value)
 
+  /** @group setParam */
+  @Since("2.3.0")
+  def setNumBucketsArray(value: Array[Int]): this.type = set(numBucketsArray, value)
+
+  /** @group setParam */
+  @Since("2.3.0")
+  def setInputCols(value: Array[String]): this.type = set(inputCols, value)
+
+  /** @group setParam */
+  @Since("2.3.0")
+  def setOutputCols(value: Array[String]): this.type = set(outputCols, value)
+
+  private[feature] def isQuantileDiscretizeMultipleColumns(): Boolean = {
+    if (isSet(inputCols) && isSet(inputCol)) {
+      logWarning("Both `inputCol` and `inputCols` are set, we ignore `inputCols` and this " +
+        "`QuantileDiscretize` only map one column specified by `inputCol`")
+      false
+    } else if (isSet(inputCols)) {
+      true
+    } else {
+      false
+    }
+  }
+
+  private[feature] def getInOutCols: (Array[String], Array[String]) = {
+    if (!isQuantileDiscretizeMultipleColumns) {
+      (Array($(inputCol)), Array($(outputCol)))
+    } else {
+      require($(inputCols).length == $(outputCols).length,
+        "inputCols number do not match outputCols")
+      ($(inputCols), $(outputCols))
+    }
+  }
+
   @Since("1.6.0")
   override def transformSchema(schema: StructType): StructType = {
-    SchemaUtils.checkNumericType(schema, $(inputCol))
-    val inputFields = schema.fields
-    require(inputFields.forall(_.name != $(outputCol)),
-      s"Output column ${$(outputCol)} already exists.")
-    val attr = NominalAttribute.defaultAttr.withName($(outputCol))
-    val outputFields = inputFields :+ attr.toStructField()
+    val (inputColNames, outputColNames) = getInOutCols
+    val existingFields = schema.fields
+    var outputFields = existingFields
+    inputColNames.zip(outputColNames).map { case (inputColName, outputColName) =>
+      SchemaUtils.checkNumericType(schema, inputColName)
+      require(existingFields.forall(_.name != outputColName),
+        s"Output column ${outputColName} already exists.")
+      val attr = NominalAttribute.defaultAttr.withName(outputColName)
+      outputFields :+= attr.toStructField()
+    }
     StructType(outputFields)
   }
 
   @Since("2.0.0")
   override def fit(dataset: Dataset[_]): Bucketizer = {
     transformSchema(dataset.schema, logging = true)
-    val splits = dataset.stat.approxQuantile($(inputCol),
-      (0.0 to 1.0 by 1.0/$(numBuckets)).toArray, $(relativeError))
+    val bucketizer = new Bucketizer(uid).setHandleInvalid($(handleInvalid))
+    if (isQuantileDiscretizeMultipleColumns) {
+      var bucketArray = Array.empty[Int]
+      if (isSet(numBucketsArray)) {
+        bucketArray = $(numBucketsArray)
+      }
+      else {
+        bucketArray = Array($(numBuckets))
+      }
+      val probabilityArray = bucketArray.toSeq.flatMap { numOfBucket =>
+        (0.0 to 1.0 by 1.0 / numOfBucket)
+      }
+      val splitsArray = dataset.stat.approxQuantile($(inputCols),
+        probabilityArray.sorted.toArray.distinct, $(relativeError))
+      val distinctSplitsArray = splitsArray.toSeq.map { splits =>
+        getDistinctSplits(splits)
+      }
+      bucketizer.setSplitsArray(distinctSplitsArray.toArray)
+      copyValues(bucketizer.setParent(this))
+    }
+    else {
+      val splits = dataset.stat.approxQuantile($(inputCol),
+        (0.0 to 1.0 by 1.0 / $(numBuckets)).toArray, $(relativeError))
+      bucketizer.setSplits(getDistinctSplits(splits))
+      copyValues(bucketizer.setParent(this))
+    }
+  }
+
+  private def getDistinctSplits(splits: Array[Double]): Array[Double] = {
     splits(0) = Double.NegativeInfinity
     splits(splits.length - 1) = Double.PositiveInfinity
-
     val distinctSplits = splits.distinct
     if (splits.length != distinctSplits.length) {
       log.warn(s"Some quantiles were identical. Bucketing to ${distinctSplits.length - 1}" +
         s" buckets as a result.")
     }
-    val bucketizer = new Bucketizer(uid)
-      .setSplits(distinctSplits.sorted)
-      .setHandleInvalid($(handleInvalid))
-    copyValues(bucketizer.setParent(this))
+    distinctSplits.sorted
   }
 
   @Since("1.6.0")
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 f219f775b218..d53d72a7bfdc 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
@@ -146,4 +146,172 @@ class QuantileDiscretizerSuite
     val model = discretizer.fit(df)
     assert(model.hasParent)
   }
+
+  test("Multiple Columns: Test observed number of buckets and their sizes match expected values") {
+    val spark = this.spark
+    import spark.implicits._
+
+    val datasetSize = 100000
+    val numBuckets = 5
+    val data1 = Array.range(1, 100001, 1).map(_.toDouble)
+    val data2 = Array.range(1, 200000, 2).map(_.toDouble)
+    val data = (0 until 100000).map { idx =>
+      (data1(idx), data2(idx))
+    }
+    val df: DataFrame = data.toSeq.toDF("input1", "input2")
+
+    val discretizer = new QuantileDiscretizer()
+      .setInputCols(Array("input1", "input2"))
+      .setOutputCols(Array("result1", "result2"))
+      .setNumBuckets(numBuckets)
+    assert(discretizer.isQuantileDiscretizeMultipleColumns())
+    val result = discretizer.fit(df).transform(df)
+
+    val relativeError = discretizer.getRelativeError
+    val isGoodBucket = udf {
+      (size: Int) => math.abs( size - (datasetSize / numBuckets)) <= (relativeError * datasetSize)
+    }
+
+    for (i <- 1 to 2) {
+      val observedNumBuckets = result.select("result" + i).distinct.count
+      assert(observedNumBuckets === numBuckets,
+        "Observed number of buckets does not equal expected number of buckets.")
+
+      val numGoodBuckets = result.groupBy("result" + i).count.filter(isGoodBucket($"count")).count
+      assert(numGoodBuckets === numBuckets,
+        "Bucket sizes are not within expected relative error tolerance.")
+    }
+  }
+
+  test("Multiple Columns: Test on data with high proportion of duplicated values") {
+    val spark = this.spark
+    import spark.implicits._
+
+    val numBuckets = 5
+    val expectedNumBucket = 3
+    val data1 = Array(1.0, 3.0, 2.0, 1.0, 1.0, 2.0, 3.0, 2.0, 2.0, 2.0, 1.0, 3.0)
+    val data2 = Array(1.0, 2.0, 3.0, 1.0, 1.0, 1.0, 1.0, 3.0, 2.0, 3.0, 1.0, 2.0)
+    val data = (0 until data1.length).map { idx =>
+      (data1(idx), data2(idx))
+    }
+    val df: DataFrame = data.toSeq.toDF("input1", "input2")
+    val discretizer = new QuantileDiscretizer()
+      .setInputCols(Array("input1", "input2"))
+      .setOutputCols(Array("result1", "result2"))
+      .setNumBuckets(numBuckets)
+    assert(discretizer.isQuantileDiscretizeMultipleColumns())
+    val result = discretizer.fit(df).transform(df)
+    for (i <- 1 to 2) {
+      val observedNumBuckets = result.select("result" + i).distinct.count
+      assert(observedNumBuckets == expectedNumBucket,
+        s"Observed number of buckets are not correct." +
+          s" Expected $expectedNumBucket but found ($observedNumBuckets")
+    }
+  }
+
+  test("Multiple Columns: Test transform on data with NaN value") {
+    val spark = this.spark
+    import spark.implicits._
+
+    val numBuckets = 3
+    val validData1 = Array(-0.9, -0.5, -0.3, 0.0, 0.2, 0.5, 0.9, Double.NaN, Double.NaN, Double.NaN)
+    val expectedKeep1 = Array(0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0)
+    val validData2 = Array(0.2, -0.1, 0.3, 0.0, 0.1, 0.3, 0.5, 0.8, Double.NaN, Double.NaN)
+    val expectedKeep2 = Array(1.0, 0.0, 2.0, 0.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0)
+
+    val data = (0 until validData1.length).map { idx =>
+      (validData1(idx), validData2(idx), expectedKeep1(idx), expectedKeep2(idx))
+    }
+    val dataFrame: DataFrame = data.toSeq.toDF("input1", "input2", "expected1", "expected2")
+
+    val discretizer = new QuantileDiscretizer()
+      .setInputCols(Array("input1", "input2"))
+      .setOutputCols(Array("result1", "result2"))
+      .setNumBuckets(numBuckets)
+    assert(discretizer.isQuantileDiscretizeMultipleColumns())
+
+    withClue("QuantileDiscretizer with handleInvalid=error should throw exception for NaN values") {
+      intercept[SparkException] {
+        discretizer.fit(dataFrame).transform(dataFrame).collect()
+      }
+    }
+
+    discretizer.setHandleInvalid("keep")
+    discretizer.fit(dataFrame).transform(dataFrame).
+      select("result1", "expected1", "result2", "expected2")
+      .collect().foreach {
+      case Row(r1: Double, e1: Double, r2: Double, e2: Double) =>
+        assert(r1 === e1,
+          s"The result value is not correct after bucketing. Expected $e1 but found $r1")
+        assert(r2 === e2,
+          s"The result value is not correct after bucketing. Expected $e2 but found $r2")
+    }
+
+    discretizer.setHandleInvalid("skip")
+    val result = discretizer.fit(dataFrame).transform(dataFrame)
+    for (i <- 1 to 2) {
+      val skipResults1: Array[Double] = result.select("result" + i).as[Double].collect()
+      assert(skipResults1.length === 7)
+      assert(skipResults1.forall(_ !== 4.0))
+    }
+  }
+
+  test("Multiple Columns: Test numBucketsArray") {
+    val spark = this.spark
+    import spark.implicits._
+
+    val datasetSize = 20
+    val numBucketsArray: Array[Int] = Array(2, 5, 10)
+    val data1 = Array.range(1, 21, 1).map(_.toDouble)
+    val expected1 = Array (0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 4.0, 4.0, 5.0,
+      5.0, 5.0, 6.0, 6.0, 7.0, 8.0, 8.0, 9.0, 9.0, 9.0)
+    val data2 = Array.range(1, 40, 2).map(_.toDouble)
+    val expected2 = Array (0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 4.0, 4.0, 5.0,
+      5.0, 5.0, 6.0, 6.0, 7.0, 8.0, 8.0, 9.0, 9.0, 9.0)
+    val data3 = Array.range(1, 60, 3).map(_.toDouble)
+    val expected3 = Array (0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 4.0, 4.0, 5.0,
+      5.0, 5.0, 6.0, 6.0, 7.0, 8.0, 8.0, 9.0, 9.0, 9.0)
+    val data = (0 until 20).map { idx =>
+      (data1(idx), data2(idx), data3(idx), expected1(idx), expected2(idx), expected3(idx))
+    }
+    val df: DataFrame =
+      data.toSeq.toDF("input1", "input2", "input3", "expected1", "expected2", "expected3")
+
+    val discretizer = new QuantileDiscretizer()
+      .setInputCols(Array("input1", "input2", "input3"))
+      .setOutputCols(Array("result1", "result2", "result3"))
+      .setNumBucketsArray(numBucketsArray)
+    assert(discretizer.isQuantileDiscretizeMultipleColumns())
+    discretizer.fit(df).transform(df).
+      select("result1", "expected1", "result2", "expected2", "result3", "expected3")
+      .collect().foreach {
+      case Row(r1: Double, e1: Double, r2: Double, e2: Double, r3: Double, e3: Double) =>
+        assert(r1 === e1,
+          s"The result value is not correct after bucketing. Expected $e1 but found $r1")
+        assert(r2 === e2,
+          s"The result value is not correct after bucketing. Expected $e2 but found $r2")
+        assert(r3 === e3,
+          s"The result value is not correct after bucketing. Expected $e3 but found $r3")
+    }
+  }
+
+  test("multiple columns: read/write") {
+    val discretizer = new QuantileDiscretizer()
+      .setInputCols(Array("input1", "input2"))
+      .setOutputCols(Array("result1", "result2"))
+      .setNumBucketsArray(Array(5, 10))
+    assert(discretizer.isQuantileDiscretizeMultipleColumns())
+    testDefaultReadWrite(discretizer)
+  }
+
+  test("Both inputCol and inputCols are set") {
+    val discretizer = new QuantileDiscretizer()
+      .setInputCol("input")
+      .setOutputCol("result")
+      .setNumBuckets(3)
+      .setInputCols(Array("input1", "input2"))
+
+    // When both are set, we ignore `inputCols` and just map the column specified by `inputCol`.
+    assert(discretizer.isQuantileDiscretizeMultipleColumns() == false)
+  }
 }

From 87ee0f3ac3ad91bbe13267baf3f9087567bc2cc8 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Sat, 11 Nov 2017 10:17:54 -0800
Subject: [PATCH 2/9] fix binary compatibility issue

---
 .../org/apache/spark/ml/feature/QuantileDiscretizer.scala    | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

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 423916649857..d7863a361008 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
@@ -31,7 +31,7 @@ import org.apache.spark.sql.types.StructType
  * Params for [[QuantileDiscretizer]].
  */
 private[feature] trait QuantileDiscretizerBase extends Params
-  with HasHandleInvalid with HasInputCol with HasInputCols with HasOutputCol with HasOutputCols {
+  with HasHandleInvalid with HasInputCol with HasOutputCol {
 
   /**
    * Number of buckets (quantiles, or categories) into which data points are grouped. Must
@@ -126,7 +126,8 @@ private[feature] trait QuantileDiscretizerBase extends Params
  */
 @Since("1.6.0")
 final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val uid: String)
-  extends Estimator[Bucketizer] with QuantileDiscretizerBase with DefaultParamsWritable {
+  extends Estimator[Bucketizer] with QuantileDiscretizerBase with DefaultParamsWritable
+    with HasInputCols with HasOutputCols {
 
   @Since("1.6.0")
   def this() = this(Identifiable.randomUID("quantileDiscretizer"))

From 5038e21e9f3d0c80f71308f2fc9167e4a7749e82 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Sat, 18 Nov 2017 21:27:40 -0800
Subject: [PATCH 3/9] address comments

---
 .../ml/feature/QuantileDiscretizer.scala      | 44 ++++++++++++++-----
 .../ml/feature/QuantileDiscretizerSuite.scala | 22 ++++------
 2 files changed, 42 insertions(+), 24 deletions(-)

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 d7863a361008..1db2729951eb 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
@@ -209,27 +209,38 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
       var bucketArray = Array.empty[Int]
       if (isSet(numBucketsArray)) {
         bucketArray = $(numBucketsArray)
-      }
-      else {
+      } else {
         bucketArray = Array($(numBuckets))
       }
+
       val probabilityArray = bucketArray.toSeq.flatMap { numOfBucket =>
         (0.0 to 1.0 by 1.0 / numOfBucket)
-      }
+      }.sorted.toArray.distinct
+
       val splitsArray = dataset.stat.approxQuantile($(inputCols),
-        probabilityArray.sorted.toArray.distinct, $(relativeError))
-      val distinctSplitsArray = splitsArray.toSeq.map { splits =>
-        getDistinctSplits(splits)
+        probabilityArray, $(relativeError))
+
+      var distinctSplitsArray = Seq.empty[Array[Double]]
+      if (bucketArray.length > 1) {
+        var idxForColumn = 0
+        distinctSplitsArray = bucketArray.toSeq.map { numOfBuckets =>
+          val splitArrayForEachColumn =
+            getSplitsForEachColumn(numOfBuckets, probabilityArray, splitsArray, idxForColumn)
+          idxForColumn += 1
+          splitArrayForEachColumn
+        }
+      } else {
+        distinctSplitsArray = splitsArray.toSeq.map { splits =>
+          getDistinctSplits(splits)
+        }
       }
       bucketizer.setSplitsArray(distinctSplitsArray.toArray)
-      copyValues(bucketizer.setParent(this))
-    }
-    else {
+    } else {
       val splits = dataset.stat.approxQuantile($(inputCol),
         (0.0 to 1.0 by 1.0 / $(numBuckets)).toArray, $(relativeError))
       bucketizer.setSplits(getDistinctSplits(splits))
-      copyValues(bucketizer.setParent(this))
     }
+    copyValues(bucketizer.setParent(this))
   }
 
   private def getDistinctSplits(splits: Array[Double]): Array[Double] = {
@@ -243,6 +254,19 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
     distinctSplits.sorted
   }
 
+  private def getSplitsForEachColumn(numOfBuckets: Int,
+                                     probabilityArray: Array[Double],
+                                     splitsArray: Array[Array[Double]],
+                                     idxForColumn: Int): Array[Double] = {
+    val probabilityArrayForEachColumn = (0.0 to 1.0 by 1.0 / numOfBuckets)
+    var splitsArrayForEachColumn = Array.empty[Double]
+    for (i <- 0 to probabilityArrayForEachColumn.length - 1) {
+      val index = probabilityArray.indexOf(probabilityArrayForEachColumn(i))
+      splitsArrayForEachColumn :+= splitsArray(idxForColumn)(index)
+    }
+    getDistinctSplits(splitsArrayForEachColumn)
+  }
+
   @Since("1.6.0")
   override def copy(extra: ParamMap): QuantileDiscretizer = defaultCopy(extra)
 }
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 d53d72a7bfdc..f0a88292033e 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
@@ -155,10 +155,7 @@ class QuantileDiscretizerSuite
     val numBuckets = 5
     val data1 = Array.range(1, 100001, 1).map(_.toDouble)
     val data2 = Array.range(1, 200000, 2).map(_.toDouble)
-    val data = (0 until 100000).map { idx =>
-      (data1(idx), data2(idx))
-    }
-    val df: DataFrame = data.toSeq.toDF("input1", "input2")
+    val df = data1.zip(data2).toSeq.toDF("input1", "input2")
 
     val discretizer = new QuantileDiscretizer()
       .setInputCols(Array("input1", "input2"))
@@ -191,10 +188,7 @@ class QuantileDiscretizerSuite
     val expectedNumBucket = 3
     val data1 = Array(1.0, 3.0, 2.0, 1.0, 1.0, 2.0, 3.0, 2.0, 2.0, 2.0, 1.0, 3.0)
     val data2 = Array(1.0, 2.0, 3.0, 1.0, 1.0, 1.0, 1.0, 3.0, 2.0, 3.0, 1.0, 2.0)
-    val data = (0 until data1.length).map { idx =>
-      (data1(idx), data2(idx))
-    }
-    val df: DataFrame = data.toSeq.toDF("input1", "input2")
+    val df = data1.zip(data2).toSeq.toDF("input1", "input2")
     val discretizer = new QuantileDiscretizer()
       .setInputCols(Array("input1", "input2"))
       .setOutputCols(Array("result1", "result2"))
@@ -222,7 +216,7 @@ class QuantileDiscretizerSuite
     val data = (0 until validData1.length).map { idx =>
       (validData1(idx), validData2(idx), expectedKeep1(idx), expectedKeep2(idx))
     }
-    val dataFrame: DataFrame = data.toSeq.toDF("input1", "input2", "expected1", "expected2")
+    val dataFrame = data.toSeq.toDF("input1", "input2", "expected1", "expected2")
 
     val discretizer = new QuantileDiscretizer()
       .setInputCols(Array("input1", "input2"))
@@ -263,18 +257,18 @@ class QuantileDiscretizerSuite
     val datasetSize = 20
     val numBucketsArray: Array[Int] = Array(2, 5, 10)
     val data1 = Array.range(1, 21, 1).map(_.toDouble)
-    val expected1 = Array (0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 4.0, 4.0, 5.0,
-      5.0, 5.0, 6.0, 6.0, 7.0, 8.0, 8.0, 9.0, 9.0, 9.0)
+    val expected1 = Array (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0,
+      1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0)
     val data2 = Array.range(1, 40, 2).map(_.toDouble)
-    val expected2 = Array (0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 4.0, 4.0, 5.0,
-      5.0, 5.0, 6.0, 6.0, 7.0, 8.0, 8.0, 9.0, 9.0, 9.0)
+    val expected2 = Array (0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0,
+      2.0, 2.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 4.0, 4.0)
     val data3 = Array.range(1, 60, 3).map(_.toDouble)
     val expected3 = Array (0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 4.0, 4.0, 5.0,
       5.0, 5.0, 6.0, 6.0, 7.0, 8.0, 8.0, 9.0, 9.0, 9.0)
     val data = (0 until 20).map { idx =>
       (data1(idx), data2(idx), data3(idx), expected1(idx), expected2(idx), expected3(idx))
     }
-    val df: DataFrame =
+    val df =
       data.toSeq.toDF("input1", "input2", "input3", "expected1", "expected2", "expected3")
 
     val discretizer = new QuantileDiscretizer()

From 97ad483d05740b43221d693c2f42f79609422a23 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Thu, 30 Nov 2017 14:34:23 -0800
Subject: [PATCH 4/9] address comments

---
 .../ml/feature/QuantileDiscretizer.scala      |  43 +++----
 .../ml/feature/QuantileDiscretizerSuite.scala | 106 +++++++++++++++++-
 2 files changed, 120 insertions(+), 29 deletions(-)

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 1db2729951eb..fb8e44e1b50b 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
@@ -52,6 +52,7 @@ private[feature] trait QuantileDiscretizerBase extends Params
 
   /**
    * Array of number of buckets (quantiles, or categories) into which data points are grouped.
+   * Each value must be greater than or equal to 2
    *
    * See also [[handleInvalid]], which can optionally create an additional bucket for NaN values.
    *
@@ -59,8 +60,9 @@ private[feature] trait QuantileDiscretizerBase extends Params
    */
   val numBucketsArray = new IntArrayParam(this, "numBucketsArray", "Array of number of buckets " +
     "(quantiles, or categories) into which data points are grouped. This is for multiple " +
-    "columns input. If numBucketsArray is not set but numBuckets is set, it means user wants " +
-    "to use the same numBuckets across all columns.")
+    "columns input. If transforming multiple columns and numBucketsArray is not set, but " +
+    "numBuckets is set, then numBuckets will be applied across all columns.",
+    (arrayOfNumBuckets: Array[Int]) => arrayOfNumBuckets.forall(ParamValidators.gtEq(2)))
 
   /** @group getParam */
   def getNumBucketsArray: Array[Int] = $(numBucketsArray)
@@ -105,9 +107,11 @@ private[feature] trait QuantileDiscretizerBase extends Params
  * possible that the number of buckets used will be smaller than this value, for example, if there
  * are too few distinct values of the input to create enough distinct quantiles.
  * Since 2.3.0,
- * `QuantileDiscretizer` can also map multiple columns at once. Whether it goes to map a column or
- * multiple columns, it depends on which parameter of `inputCol` and `inputCols` is set. When both
- * are set, a log warning will be printed and by default it chooses `inputCol`.
+ * `QuantileDiscretizer ` can map multiple columns at once by setting the `inputCols` parameter.
+ * Note that when both the `inputCol` and `inputCols` parameters are set, a log warning will be
+ * printed and only `inputCol` will take effect, while `inputCols` will be ignored. To specify
+ * the number of bucketsfor each column , the `numBucketsArray ` parameter can be set, or if the
+ *  number of buckets should be the same across columns, `numBuckets` can be set as a convenience.
  *
  * NaN handling:
  * null and NaN values will be ignored from the column during `QuantileDiscretizer` fitting. This
@@ -167,7 +171,7 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
   private[feature] def isQuantileDiscretizeMultipleColumns(): Boolean = {
     if (isSet(inputCols) && isSet(inputCol)) {
       logWarning("Both `inputCol` and `inputCols` are set, we ignore `inputCols` and this " +
-        "`QuantileDiscretize` only map one column specified by `inputCol`")
+        "`QuantileDiscretizer` will only map one column specified by `inputCol`")
       false
     } else if (isSet(inputCols)) {
       true
@@ -206,35 +210,22 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
     transformSchema(dataset.schema, logging = true)
     val bucketizer = new Bucketizer(uid).setHandleInvalid($(handleInvalid))
     if (isQuantileDiscretizeMultipleColumns) {
-      var bucketArray = Array.empty[Int]
-      if (isSet(numBucketsArray)) {
-        bucketArray = $(numBucketsArray)
+      val bucketSeq = if (isSet(numBucketsArray)) {
+        $(numBucketsArray).toSeq
       } else {
-        bucketArray = Array($(numBuckets))
+        Array.fill($(inputCols).length)($(numBuckets)).toSeq
       }
 
-      val probabilityArray = bucketArray.toSeq.flatMap { numOfBucket =>
+      val probabilityArray = bucketSeq.flatMap { numOfBucket =>
         (0.0 to 1.0 by 1.0 / numOfBucket)
       }.sorted.toArray.distinct
 
       val splitsArray = dataset.stat.approxQuantile($(inputCols),
         probabilityArray, $(relativeError))
-
-      var distinctSplitsArray = Seq.empty[Array[Double]]
-      if (bucketArray.length > 1) {
-        var idxForColumn = 0
-        distinctSplitsArray = bucketArray.toSeq.map { numOfBuckets =>
-          val splitArrayForEachColumn =
-            getSplitsForEachColumn(numOfBuckets, probabilityArray, splitsArray, idxForColumn)
-          idxForColumn += 1
-          splitArrayForEachColumn
-        }
-      } else {
-        distinctSplitsArray = splitsArray.toSeq.map { splits =>
-          getDistinctSplits(splits)
-        }
+      val distinctSplits = bucketSeq.zipWithIndex.map { case (numOfBuckets, index) =>
+        getSplitsForEachColumn(numOfBuckets, probabilityArray, splitsArray, index)
       }
-      bucketizer.setSplitsArray(distinctSplitsArray.toArray)
+      bucketizer.setSplitsArray(distinctSplits.toArray)
     } else {
       val splits = dataset.stat.approxQuantile($(inputCol),
         (0.0 to 1.0 by 1.0 / $(numBuckets)).toArray, $(relativeError))
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 f0a88292033e..3fe81f70ea7a 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,6 +18,7 @@
 package org.apache.spark.ml.feature
 
 import org.apache.spark.{SparkException, SparkFunSuite}
+import org.apache.spark.ml.Pipeline
 import org.apache.spark.ml.util.DefaultReadWriteTest
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.sql._
@@ -216,7 +217,7 @@ class QuantileDiscretizerSuite
     val data = (0 until validData1.length).map { idx =>
       (validData1(idx), validData2(idx), expectedKeep1(idx), expectedKeep2(idx))
     }
-    val dataFrame = data.toSeq.toDF("input1", "input2", "expected1", "expected2")
+    val dataFrame = data.toDF("input1", "input2", "expected1", "expected2")
 
     val discretizer = new QuantileDiscretizer()
       .setInputCols(Array("input1", "input2"))
@@ -254,7 +255,6 @@ class QuantileDiscretizerSuite
     val spark = this.spark
     import spark.implicits._
 
-    val datasetSize = 20
     val numBucketsArray: Array[Int] = Array(2, 5, 10)
     val data1 = Array.range(1, 21, 1).map(_.toDouble)
     val expected1 = Array (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0,
@@ -269,13 +269,15 @@ class QuantileDiscretizerSuite
       (data1(idx), data2(idx), data3(idx), expected1(idx), expected2(idx), expected3(idx))
     }
     val df =
-      data.toSeq.toDF("input1", "input2", "input3", "expected1", "expected2", "expected3")
+      data.toDF("input1", "input2", "input3", "expected1", "expected2", "expected3")
 
     val discretizer = new QuantileDiscretizer()
       .setInputCols(Array("input1", "input2", "input3"))
       .setOutputCols(Array("result1", "result2", "result3"))
       .setNumBucketsArray(numBucketsArray)
+
     assert(discretizer.isQuantileDiscretizeMultipleColumns())
+
     discretizer.fit(df).transform(df).
       select("result1", "expected1", "result2", "expected2", "result3", "expected3")
       .collect().foreach {
@@ -289,6 +291,104 @@ class QuantileDiscretizerSuite
     }
   }
 
+  test("Multiple Columns: Compare single/multiple column(s) QuantileDiscretizer in pipeline") {
+    val spark = this.spark
+    import spark.implicits._
+
+    val numBucketsArray: Array[Int] = Array(2, 5, 10)
+    val data1 = Array.range(1, 21, 1).map(_.toDouble)
+    val data2 = Array.range(1, 40, 2).map(_.toDouble)
+    val data3 = Array.range(1, 60, 3).map(_.toDouble)
+    val data = (0 until 20).map { idx =>
+      (data1(idx), data2(idx), data3(idx))
+    }
+    val df =
+      data.toDF("input1", "input2", "input3")
+
+    val multiColsDiscretizer = new QuantileDiscretizer()
+      .setInputCols(Array("input1", "input2", "input3"))
+      .setOutputCols(Array("result1", "result2", "result3"))
+      .setNumBucketsArray(numBucketsArray)
+    val plForMultiCols = new Pipeline()
+      .setStages(Array(multiColsDiscretizer))
+      .fit(df)
+
+    val discretizerForCol1 = new QuantileDiscretizer()
+      .setInputCol("input1")
+      .setOutputCol("result1")
+      .setNumBuckets(numBucketsArray(0))
+
+    val discretizerForCol2 = new QuantileDiscretizer()
+      .setInputCol("input2")
+      .setOutputCol("result2")
+      .setNumBuckets(numBucketsArray(1))
+
+    val discretizerForCol3 = new QuantileDiscretizer()
+      .setInputCol("input3")
+      .setOutputCol("result3")
+      .setNumBuckets(numBucketsArray(2))
+
+    val plForSingleCol = new Pipeline()
+      .setStages(Array(discretizerForCol1, discretizerForCol2, discretizerForCol3))
+      .fit(df)
+
+    val resultForMultiCols = plForMultiCols.transform(df)
+      .select("result1", "result2", "result3")
+      .collect()
+
+    val resultForSingleCol = plForSingleCol.transform(df)
+      .select("result1", "result2", "result3")
+      .collect()
+
+    resultForSingleCol.zip(resultForMultiCols).foreach {
+      case (rowForSingle, rowForMultiCols) =>
+        assert(rowForSingle.getDouble(0) == rowForMultiCols.getDouble(0) &&
+          rowForSingle.getDouble(1) == rowForMultiCols.getDouble(1) &&
+          rowForSingle.getDouble(2) == rowForMultiCols.getDouble(2))
+    }
+  }
+
+  test("Multiple Columns: Comparing setting numBuckets with setting numBucketsArray" +
+    " explicitly with identical values") {
+    val spark = this.spark
+    import spark.implicits._
+
+    val datasetSize = 20
+    val numBucketsArray: Array[Int] = Array(2, 5, 10)
+    val data1 = Array.range(1, 21, 1).map(_.toDouble)
+    val data2 = Array.range(1, 40, 2).map(_.toDouble)
+    val data3 = Array.range(1, 60, 3).map(_.toDouble)
+    val data = (0 until 20).map { idx =>
+      (data1(idx), data2(idx), data3(idx))
+    }
+    val df =
+      data.toDF("input1", "input2", "input3")
+
+    val discretizerSingleNumBuckets = new QuantileDiscretizer()
+      .setInputCols(Array("input1", "input2", "input3"))
+      .setOutputCols(Array("result1", "result2", "result3"))
+      .setNumBuckets(10)
+
+    val discretizerNumBucketsArray = new QuantileDiscretizer()
+      .setInputCols(Array("input1", "input2", "input3"))
+      .setOutputCols(Array("result1", "result2", "result3"))
+      .setNumBucketsArray(Array(10, 10, 10))
+
+    val result1 = discretizerSingleNumBuckets.fit(df).transform(df)
+      .select("result1", "result2", "result3")
+      .collect()
+    val result2 = discretizerNumBucketsArray.fit(df).transform(df)
+      .select("result1", "result2", "result3")
+      .collect()
+
+    result1.zip(result2).foreach {
+      case (row1, row2) =>
+        assert(row1.getDouble(0) == row2.getDouble(0) &&
+          row1.getDouble(1) == row2.getDouble(1) &&
+          row1.getDouble(2) == row2.getDouble(2))
+    }
+  }
+
   test("multiple columns: read/write") {
     val discretizer = new QuantileDiscretizer()
       .setInputCols(Array("input1", "input2"))

From 445bd84a6e5e81896d5c94ada7035b00e2c22337 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Fri, 8 Dec 2017 18:00:42 -0800
Subject: [PATCH 5/9] address comments

---
 .../ml/feature/QuantileDiscretizer.scala      | 58 +++++++++----------
 .../ml/feature/QuantileDiscretizerSuite.scala | 46 +++++++--------
 2 files changed, 47 insertions(+), 57 deletions(-)

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 fb8e44e1b50b..d4195bac8576 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
@@ -107,11 +107,11 @@ private[feature] trait QuantileDiscretizerBase extends Params
  * possible that the number of buckets used will be smaller than this value, for example, if there
  * are too few distinct values of the input to create enough distinct quantiles.
  * Since 2.3.0,
- * `QuantileDiscretizer ` can map multiple columns at once by setting the `inputCols` parameter.
+ * `QuantileDiscretizer` can map multiple columns at once by setting the `inputCols` parameter.
  * Note that when both the `inputCol` and `inputCols` parameters are set, a log warning will be
  * printed and only `inputCol` will take effect, while `inputCols` will be ignored. To specify
- * the number of bucketsfor each column , the `numBucketsArray ` parameter can be set, or if the
- *  number of buckets should be the same across columns, `numBuckets` can be set as a convenience.
+ * the number of buckets for each column, the `numBucketsArray` parameter can be set, or if the
+ * number of buckets should be the same across columns, `numBuckets` can be set as a convenience.
  *
  * NaN handling:
  * null and NaN values will be ignored from the column during `QuantileDiscretizer` fitting. This
@@ -195,7 +195,7 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
     val (inputColNames, outputColNames) = getInOutCols
     val existingFields = schema.fields
     var outputFields = existingFields
-    inputColNames.zip(outputColNames).map { case (inputColName, outputColName) =>
+    inputColNames.zip(outputColNames).foreach { case (inputColName, outputColName) =>
       SchemaUtils.checkNumericType(schema, inputColName)
       require(existingFields.forall(_.name != outputColName),
         s"Output column ${outputColName} already exists.")
@@ -210,22 +210,31 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
     transformSchema(dataset.schema, logging = true)
     val bucketizer = new Bucketizer(uid).setHandleInvalid($(handleInvalid))
     if (isQuantileDiscretizeMultipleColumns) {
-      val bucketSeq = if (isSet(numBucketsArray)) {
-        $(numBucketsArray).toSeq
+      val splitsArray = if (isSet(numBucketsArray)) {
+        val probArrayPerCol = $(numBucketsArray).map { numOfBuckets =>
+          (0.0 to 1.0 by 1.0 / numOfBuckets).toArray
+        }
+
+        val probabilityArray = probArrayPerCol.flatten.sorted.distinct
+        val splitsArrayRaw = dataset.stat.approxQuantile($(inputCols),
+          probabilityArray, $(relativeError))
+
+        splitsArrayRaw.zip(probArrayPerCol).map { case (splits, probs) =>
+          val probSet = probs.toSet
+          val idxSet = probabilityArray.zipWithIndex.collect {
+            case (p, idx) if probSet(p) =>
+              idx
+          }.toSet
+          splits.zipWithIndex.collect {
+            case (s, idx) if idxSet(idx) =>
+              s
+          }
+        }
       } else {
-        Array.fill($(inputCols).length)($(numBuckets)).toSeq
+        dataset.stat.approxQuantile($(inputCols),
+          (0.0 to 1.0 by 1.0 / $(numBuckets)).toArray, $(relativeError))
       }
-
-      val probabilityArray = bucketSeq.flatMap { numOfBucket =>
-        (0.0 to 1.0 by 1.0 / numOfBucket)
-      }.sorted.toArray.distinct
-
-      val splitsArray = dataset.stat.approxQuantile($(inputCols),
-        probabilityArray, $(relativeError))
-      val distinctSplits = bucketSeq.zipWithIndex.map { case (numOfBuckets, index) =>
-        getSplitsForEachColumn(numOfBuckets, probabilityArray, splitsArray, index)
-      }
-      bucketizer.setSplitsArray(distinctSplits.toArray)
+      bucketizer.setSplitsArray(splitsArray.map(getDistinctSplits))
     } else {
       val splits = dataset.stat.approxQuantile($(inputCol),
         (0.0 to 1.0 by 1.0 / $(numBuckets)).toArray, $(relativeError))
@@ -245,19 +254,6 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
     distinctSplits.sorted
   }
 
-  private def getSplitsForEachColumn(numOfBuckets: Int,
-                                     probabilityArray: Array[Double],
-                                     splitsArray: Array[Array[Double]],
-                                     idxForColumn: Int): Array[Double] = {
-    val probabilityArrayForEachColumn = (0.0 to 1.0 by 1.0 / numOfBuckets)
-    var splitsArrayForEachColumn = Array.empty[Double]
-    for (i <- 0 to probabilityArrayForEachColumn.length - 1) {
-      val index = probabilityArray.indexOf(probabilityArrayForEachColumn(i))
-      splitsArrayForEachColumn :+= splitsArray(idxForColumn)(index)
-    }
-    getDistinctSplits(splitsArrayForEachColumn)
-  }
-
   @Since("1.6.0")
   override def copy(extra: ParamMap): QuantileDiscretizer = defaultCopy(extra)
 }
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 3fe81f70ea7a..c9669d9794d2 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
@@ -211,13 +211,10 @@ class QuantileDiscretizerSuite
     val numBuckets = 3
     val validData1 = Array(-0.9, -0.5, -0.3, 0.0, 0.2, 0.5, 0.9, Double.NaN, Double.NaN, Double.NaN)
     val expectedKeep1 = Array(0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0)
-    val validData2 = Array(0.2, -0.1, 0.3, 0.0, 0.1, 0.3, 0.5, 0.8, Double.NaN, Double.NaN)
-    val expectedKeep2 = Array(1.0, 0.0, 2.0, 0.0, 1.0, 2.0, 2.0, 2.0, 3.0, 3.0)
-
-    val data = (0 until validData1.length).map { idx =>
-      (validData1(idx), validData2(idx), expectedKeep1(idx), expectedKeep2(idx))
-    }
-    val dataFrame = data.toDF("input1", "input2", "expected1", "expected2")
+    val expectedSkip1 = Array(0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 2.0)
+    val validData2 = Array(0.2, -0.1, 0.3, 0.0, 0.1, 0.3, 0.5, Double.NaN, Double.NaN, Double.NaN)
+    val expectedKeep2 = Array(1.0, 0.0, 2.0, 0.0, 1.0, 2.0, 2.0, 3.0, 3.0, 3.0)
+    val expectedSkip2 = Array(1.0, 0.0, 2.0, 0.0, 1.0, 2.0, 2.0)
 
     val discretizer = new QuantileDiscretizer()
       .setInputCols(Array("input1", "input2"))
@@ -226,28 +223,25 @@ class QuantileDiscretizerSuite
     assert(discretizer.isQuantileDiscretizeMultipleColumns())
 
     withClue("QuantileDiscretizer with handleInvalid=error should throw exception for NaN values") {
+      val dataFrame: DataFrame = validData1.zip(validData2).toSeq.toDF("input1", "input2")
       intercept[SparkException] {
         discretizer.fit(dataFrame).transform(dataFrame).collect()
       }
     }
 
-    discretizer.setHandleInvalid("keep")
-    discretizer.fit(dataFrame).transform(dataFrame).
-      select("result1", "expected1", "result2", "expected2")
-      .collect().foreach {
-      case Row(r1: Double, e1: Double, r2: Double, e2: Double) =>
-        assert(r1 === e1,
-          s"The result value is not correct after bucketing. Expected $e1 but found $r1")
-        assert(r2 === e2,
-          s"The result value is not correct after bucketing. Expected $e2 but found $r2")
-    }
-
-    discretizer.setHandleInvalid("skip")
-    val result = discretizer.fit(dataFrame).transform(dataFrame)
-    for (i <- 1 to 2) {
-      val skipResults1: Array[Double] = result.select("result" + i).as[Double].collect()
-      assert(skipResults1.length === 7)
-      assert(skipResults1.forall(_ !== 4.0))
+    List(("keep", expectedKeep1, expectedKeep2), ("skip", expectedSkip1, expectedSkip2)).foreach {
+      case (u, v, w) =>
+        discretizer.setHandleInvalid(u)
+        val dataFrame: DataFrame = validData1.zip(validData2).zip(v).zip(w).map {
+          case (((a, b), c), d) => (a, b, c, d)
+        }.toSeq.toDF("input1", "input2", "expected1", "expected2")
+        dataFrame.show
+        val result = discretizer.fit(dataFrame).transform(dataFrame)
+        result.show
+        result.select("result1", "expected1", "result2", "expected2").collect().foreach {
+          case Row(x: Double, y: Double, z: Double, w: Double) =>
+            assert(x === y && w === z)
+        }
     }
   }
 
@@ -348,8 +342,8 @@ class QuantileDiscretizerSuite
     }
   }
 
-  test("Multiple Columns: Comparing setting numBuckets with setting numBucketsArray" +
-    " explicitly with identical values") {
+  test("Multiple Columns: Comparing setting numBuckets with setting numBucketsArray " +
+    "explicitly with identical values") {
     val spark = this.spark
     import spark.implicits._
 

From 0e5971b3a95ae6f105659f6361c0db5a1c9fb9d8 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Mon, 11 Dec 2017 21:22:26 -0800
Subject: [PATCH 6/9] throw Exception if both inputCol and inputCols are set

---
 .../ml/feature/QuantileDiscretizer.scala      | 27 +++++++------------
 .../ml/feature/QuantileDiscretizerSuite.scala | 17 ++++++------
 2 files changed, 18 insertions(+), 26 deletions(-)

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 d4195bac8576..48067e9f9b38 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
@@ -108,9 +108,9 @@ private[feature] trait QuantileDiscretizerBase extends Params
  * are too few distinct values of the input to create enough distinct quantiles.
  * Since 2.3.0,
  * `QuantileDiscretizer` can map multiple columns at once by setting the `inputCols` parameter.
- * Note that when both the `inputCol` and `inputCols` parameters are set, a log warning will be
- * printed and only `inputCol` will take effect, while `inputCols` will be ignored. To specify
- * the number of buckets for each column, the `numBucketsArray` parameter can be set, or if the
+ * Note that only one of `inputCol` and `inputCols` parameters can be set. If both of the
+ * `inputCol` and `inputCols` parameters are set, an Exception will be thrown. To specify the
+ * number of buckets for each column, the `numBucketsArray` parameter can be set, or if the
  * number of buckets should be the same across columns, `numBuckets` can be set as a convenience.
  *
  * NaN handling:
@@ -168,20 +168,13 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
   @Since("2.3.0")
   def setOutputCols(value: Array[String]): this.type = set(outputCols, value)
 
-  private[feature] def isQuantileDiscretizeMultipleColumns(): Boolean = {
-    if (isSet(inputCols) && isSet(inputCol)) {
-      logWarning("Both `inputCol` and `inputCols` are set, we ignore `inputCols` and this " +
-        "`QuantileDiscretizer` will only map one column specified by `inputCol`")
-      false
-    } else if (isSet(inputCols)) {
-      true
-    } else {
-      false
-    }
-  }
-
   private[feature] def getInOutCols: (Array[String], Array[String]) = {
-    if (!isQuantileDiscretizeMultipleColumns) {
+    require((isSet(inputCol) && isSet(outputCol) && !isSet(inputCols) && !isSet(outputCols)) ||
+      (!isSet(inputCol) && !isSet(outputCol) && isSet(inputCols) && isSet(outputCols)),
+      "Only allow to set either inputCol/outputCol, or inputCols/outputCols"
+    )
+
+    if (isSet(inputCol)) {
       (Array($(inputCol)), Array($(outputCol)))
     } else {
       require($(inputCols).length == $(outputCols).length,
@@ -209,7 +202,7 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
   override def fit(dataset: Dataset[_]): Bucketizer = {
     transformSchema(dataset.schema, logging = true)
     val bucketizer = new Bucketizer(uid).setHandleInvalid($(handleInvalid))
-    if (isQuantileDiscretizeMultipleColumns) {
+    if (isSet(inputCols)) {
       val splitsArray = if (isSet(numBucketsArray)) {
         val probArrayPerCol = $(numBucketsArray).map { numOfBuckets =>
           (0.0 to 1.0 by 1.0 / numOfBuckets).toArray
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 c9669d9794d2..54a887a85fab 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
@@ -162,7 +162,6 @@ class QuantileDiscretizerSuite
       .setInputCols(Array("input1", "input2"))
       .setOutputCols(Array("result1", "result2"))
       .setNumBuckets(numBuckets)
-    assert(discretizer.isQuantileDiscretizeMultipleColumns())
     val result = discretizer.fit(df).transform(df)
 
     val relativeError = discretizer.getRelativeError
@@ -194,7 +193,6 @@ class QuantileDiscretizerSuite
       .setInputCols(Array("input1", "input2"))
       .setOutputCols(Array("result1", "result2"))
       .setNumBuckets(numBuckets)
-    assert(discretizer.isQuantileDiscretizeMultipleColumns())
     val result = discretizer.fit(df).transform(df)
     for (i <- 1 to 2) {
       val observedNumBuckets = result.select("result" + i).distinct.count
@@ -220,7 +218,6 @@ class QuantileDiscretizerSuite
       .setInputCols(Array("input1", "input2"))
       .setOutputCols(Array("result1", "result2"))
       .setNumBuckets(numBuckets)
-    assert(discretizer.isQuantileDiscretizeMultipleColumns())
 
     withClue("QuantileDiscretizer with handleInvalid=error should throw exception for NaN values") {
       val dataFrame: DataFrame = validData1.zip(validData2).toSeq.toDF("input1", "input2")
@@ -270,8 +267,6 @@ class QuantileDiscretizerSuite
       .setOutputCols(Array("result1", "result2", "result3"))
       .setNumBucketsArray(numBucketsArray)
 
-    assert(discretizer.isQuantileDiscretizeMultipleColumns())
-
     discretizer.fit(df).transform(df).
       select("result1", "expected1", "result2", "expected2", "result3", "expected3")
       .collect().foreach {
@@ -388,18 +383,22 @@ class QuantileDiscretizerSuite
       .setInputCols(Array("input1", "input2"))
       .setOutputCols(Array("result1", "result2"))
       .setNumBucketsArray(Array(5, 10))
-    assert(discretizer.isQuantileDiscretizeMultipleColumns())
     testDefaultReadWrite(discretizer)
   }
 
   test("Both inputCol and inputCols are set") {
+    val spark = this.spark
+    import spark.implicits._
     val discretizer = new QuantileDiscretizer()
       .setInputCol("input")
       .setOutputCol("result")
       .setNumBuckets(3)
       .setInputCols(Array("input1", "input2"))
-
-    // When both are set, we ignore `inputCols` and just map the column specified by `inputCol`.
-    assert(discretizer.isQuantileDiscretizeMultipleColumns() == false)
+    val df = sc.parallelize(Array(1.0, 2.0, 3.0, 4.0, 5.0, 6.0))
+      .map(Tuple1.apply).toDF("input")
+    // When both inputCol and inputCols are set, we throw Exception.
+    intercept[Exception] {
+      discretizer.fit(df)
+    }
   }
 }

From a030da1549a7138146588e28e0a25ce6ea9ad156 Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Tue, 12 Dec 2017 09:43:27 -0800
Subject: [PATCH 7/9] Address Comments

---
 .../spark/ml/feature/QuantileDiscretizer.scala     | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

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 48067e9f9b38..1ec5f8cb6139 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
@@ -106,12 +106,11 @@ private[feature] trait QuantileDiscretizerBase extends Params
  * categorical features. The number of bins can be set using the `numBuckets` parameter. It is
  * possible that the number of buckets used will be smaller than this value, for example, if there
  * are too few distinct values of the input to create enough distinct quantiles.
- * Since 2.3.0,
- * `QuantileDiscretizer` can map multiple columns at once by setting the `inputCols` parameter.
- * Note that only one of `inputCol` and `inputCols` parameters can be set. If both of the
- * `inputCol` and `inputCols` parameters are set, an Exception will be thrown. To specify the
- * number of buckets for each column, the `numBucketsArray` parameter can be set, or if the
- * number of buckets should be the same across columns, `numBuckets` can be set as a convenience.
+ * Since 2.3.0, `QuantileDiscretizer` can map multiple columns at once by setting the `inputCols`
+ * parameter. If both of the `inputCol` and `inputCols` parameters are set, an Exception will be
+ * thrown. To specify the number of buckets for each column, the `numBucketsArray` parameter can
+ * be set, or if the number of buckets should be the same across columns, `numBuckets` can be
+ * set as a convenience.
  *
  * NaN handling:
  * null and NaN values will be ignored from the column during `QuantileDiscretizer` fitting. This
@@ -171,7 +170,8 @@ final class QuantileDiscretizer @Since("1.6.0") (@Since("1.6.0") override val ui
   private[feature] def getInOutCols: (Array[String], Array[String]) = {
     require((isSet(inputCol) && isSet(outputCol) && !isSet(inputCols) && !isSet(outputCols)) ||
       (!isSet(inputCol) && !isSet(outputCol) && isSet(inputCols) && isSet(outputCols)),
-      "Only allow to set either inputCol/outputCol, or inputCols/outputCols"
+      "QuantileDiscretizer only supports setting either inputCol/outputCol or" +
+        "inputCols/outputCols."
     )
 
     if (isSet(inputCol)) {

From 99726a15b3369d8119750143753d151201c9334c Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Fri, 15 Dec 2017 11:39:16 -0800
Subject: [PATCH 8/9] Address Comments for test case

---
 .../ml/feature/QuantileDiscretizerSuite.scala | 27 +++++++------------
 1 file changed, 10 insertions(+), 17 deletions(-)

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 54a887a85fab..8cdf4917e84f 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
@@ -232,9 +232,7 @@ class QuantileDiscretizerSuite
         val dataFrame: DataFrame = validData1.zip(validData2).zip(v).zip(w).map {
           case (((a, b), c), d) => (a, b, c, d)
         }.toSeq.toDF("input1", "input2", "expected1", "expected2")
-        dataFrame.show
         val result = discretizer.fit(dataFrame).transform(dataFrame)
-        result.show
         result.select("result1", "expected1", "result2", "expected2").collect().foreach {
           case Row(x: Double, y: Double, z: Double, w: Double) =>
             assert(x === y && w === z)
@@ -342,8 +340,6 @@ class QuantileDiscretizerSuite
     val spark = this.spark
     import spark.implicits._
 
-    val datasetSize = 20
-    val numBucketsArray: Array[Int] = Array(2, 5, 10)
     val data1 = Array.range(1, 21, 1).map(_.toDouble)
     val data2 = Array.range(1, 40, 2).map(_.toDouble)
     val data3 = Array.range(1, 60, 3).map(_.toDouble)
@@ -386,19 +382,16 @@ class QuantileDiscretizerSuite
     testDefaultReadWrite(discretizer)
   }
 
-  test("Both inputCol and inputCols are set") {
-    val spark = this.spark
-    import spark.implicits._
-    val discretizer = new QuantileDiscretizer()
-      .setInputCol("input")
-      .setOutputCol("result")
-      .setNumBuckets(3)
-      .setInputCols(Array("input1", "input2"))
-    val df = sc.parallelize(Array(1.0, 2.0, 3.0, 4.0, 5.0, 6.0))
-      .map(Tuple1.apply).toDF("input")
-    // When both inputCol and inputCols are set, we throw Exception.
-    intercept[Exception] {
-      discretizer.fit(df)
+  test("multiple columns: Both inputCol and inputCols are set") {
+    intercept[IllegalArgumentException] {
+      new QuantileDiscretizer().setInputCol("in").setInputCols(Array("in1", "in2")).getInOutCols
+    }
+  }
+
+  test("multiple columns: Mismatched sizes of inputCols / outputCols") {
+    intercept[IllegalArgumentException] {
+      new QuantileDiscretizer().setInputCols(Array("in1", "in2"))
+        .setOutputCols(Array("out1")).getInOutCols
     }
   }
 }

From 486b68d1de9e9dc480133d8680baebc98f6e572c Mon Sep 17 00:00:00 2001
From: Huaxin Gao <huaxing@us.ibm.com>
Date: Thu, 21 Dec 2017 10:35:01 -0800
Subject: [PATCH 9/9] address comment

---
 .../ml/feature/QuantileDiscretizerSuite.scala | 29 +++++++++++++++----
 1 file changed, 23 insertions(+), 6 deletions(-)

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 8cdf4917e84f..e9a75e931e6a 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
@@ -374,7 +374,7 @@ class QuantileDiscretizerSuite
     }
   }
 
-  test("multiple columns: read/write") {
+  test("Multiple Columns: read/write") {
     val discretizer = new QuantileDiscretizer()
       .setInputCols(Array("input1", "input2"))
       .setOutputCols(Array("result1", "result2"))
@@ -382,16 +382,33 @@ class QuantileDiscretizerSuite
     testDefaultReadWrite(discretizer)
   }
 
-  test("multiple columns: Both inputCol and inputCols are set") {
+  test("Multiple Columns: Both inputCol and inputCols are set") {
+    val spark = this.spark
+    import spark.implicits._
+    val discretizer = new QuantileDiscretizer()
+      .setInputCol("input")
+      .setOutputCol("result")
+      .setNumBuckets(3)
+      .setInputCols(Array("input1", "input2"))
+    val df = sc.parallelize(Array(1.0, 2.0, 3.0, 4.0, 5.0, 6.0))
+      .map(Tuple1.apply).toDF("input")
+    // When both inputCol and inputCols are set, we throw Exception.
     intercept[IllegalArgumentException] {
-      new QuantileDiscretizer().setInputCol("in").setInputCols(Array("in1", "in2")).getInOutCols
+      discretizer.fit(df)
     }
   }
 
-  test("multiple columns: Mismatched sizes of inputCols / outputCols") {
+  test("Multiple Columns: Mismatched sizes of inputCols / outputCols") {
+    val spark = this.spark
+    import spark.implicits._
+    val discretizer = new QuantileDiscretizer()
+      .setInputCols(Array("input"))
+      .setOutputCols(Array("result1", "result2"))
+      .setNumBuckets(3)
+    val df = sc.parallelize(Array(1.0, 2.0, 3.0, 4.0, 5.0, 6.0))
+      .map(Tuple1.apply).toDF("input")
     intercept[IllegalArgumentException] {
-      new QuantileDiscretizer().setInputCols(Array("in1", "in2"))
-        .setOutputCols(Array("out1")).getInOutCols
+      discretizer.fit(df)
     }
   }
 }