VarianceThreshold Function For Data Cleansing - scala

Question

I have the following function that I want to use to see how many features are selected based on different Threshold values for the variance.
import org.apache.spark.ml.feature.VarianceThresholdSelector
def varianceThreshold(df: DataFrame, thresholds: Seq[Threshold]): Seq[(Threshold, DataFrame)] = {
thresholds.map(threshold => {
val selector = new VarianceThresholdSelector()
.setVarianceThreshold(threshold)
.setFeaturesCol("features")
.setOutputCol("selectedFeatures")
(threshold, selector.fit(df).transform(df))
})
}
So far so good. I have a DataFrame that looks like this:
Now my question is if col2 is the predictor variable, i.e., the value that I'm trying to predict, how can I then have all the other columns grouped so that I can pass that as a feature as such. For example., I came across this example from the Spark documentation:
import org.apache.spark.ml.feature.VarianceThresholdSelector
import org.apache.spark.ml.linalg.Vectors
val data = Seq(
(1, Vectors.dense(6.0, 7.0, 0.0, 7.0, 6.0, 0.0)),
(2, Vectors.dense(0.0, 9.0, 6.0, 0.0, 5.0, 9.0)),
(3, Vectors.dense(0.0, 9.0, 3.0, 0.0, 5.0, 5.0)),
(4, Vectors.dense(0.0, 9.0, 8.0, 5.0, 6.0, 4.0)),
(5, Vectors.dense(8.0, 9.0, 6.0, 5.0, 4.0, 4.0)),
(6, Vectors.dense(8.0, 9.0, 6.0, 0.0, 0.0, 0.0))
)
val df = spark.createDataset(data).toDF("id", "features")
val selector = new VarianceThresholdSelector()
.setVarianceThreshold(8.0)
.setFeaturesCol("features")
.setOutputCol("selectedFeatures")
val result = selector.fit(df).transform(df)
println(s"Output: Features with variance lower than" +
s" ${selector.getVarianceThreshold} are removed.")
result.show()
So for my example what will be the featureCol or rather how can I get my individual columns as a featuresCol array?

Answer 1

Here is what I did to get the effect that I wanted:
type Threshold = Double
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.feature.VarianceThresholdSelector
def varianceThreshold(df: DataFrame, thresholds: Seq[Threshold]): Seq[(Threshold, DataFrame)] = {
val assembler = new VectorAssembler()
.setInputCols(df.columns.tail)
.setOutputCol("features")
val output = assembler.transform(df)
thresholds.map(threshold => {
val selector = new VarianceThresholdSelector()
.setVarianceThreshold(threshold)
.setFeaturesCol("features")
.setOutputCol("selectedFeatures")
(threshold, selector.fit(output).transform(output))
})
}