将QuantileDiscretizer应用于DataFrame中的所有列

时间:2017-07-26 17:52:35

标签: scala apache-spark dataframe apache-spark-mllib

假设我有一个id为100列的数据帧。我想在每一列上应用QuantileDiscretizer并返回一个新的数据框,其id列与带有离散值的新列相关联。 仅适用于两列的示例:

输入

 id | col1 | col2
----|------|------
 0  | 18.0 | 20.0
----|------|------
 1  | 19.0 | 30.0
----|------|------
 2  | 8.0  | 35.0
----|------|------
 3  | 5.0  | 10.0
----|------|------
 4  | 2.2  | 5.0

输出

 id | col1Disc | col2Disc
----|----------|------
 0  | 2        | 2
----|----------| ------
 1  | 2        | 3
----|----------|------
 2  | 1        | 3
----|----------|------
 3  | 2        | 1
----|----------|------
 4  | 0        | 0

1 个答案:

答案 0 :(得分:0)

您可以使用Pipeline API:

import org.apache.spark.ml.Pipeline

val df = Seq(
  (0, 18.0, 20.0), (1, 19.0, 30.0), (2, 8.0, 35.0), (3, 5.0, 10.0), (4, 2.2, 5.0)
).toDF("id", "col1", "col2")


val pipeline = new Pipeline().setStages(for {
  c <- df.columns
  if c != "id"
} yield new QuantileDiscretizer().setInputCol(c).setOutputCol(s"${c}Disc"))

val result = pipeline.fit(df).transform(df)
result.drop(df.columns.diff(Seq("id")): _*).show


+---+--------+--------+
| id|col1Disc|col2Disc|
+---+--------+--------+
|  0|     1.0|     1.0|
|  1|     1.0|     1.0|
|  2|     1.0|     1.0|
|  3|     0.0|     0.0|
|  4|     0.0|     0.0|
+---+--------+--------+