Apply VarianceThreshold feature selection technique on pyspark DataFrame

user_11170350 0
Viewed 440

I am working with pyspark. I converted the features VectorAssembler with following code

from pyspark.ml.feature import VectorAssembler
assembler = VectorAssembler(inputCols=list(df.columns[1::]), outputCol='features')
transformed_data = assembler.transform(df)

After that I scaled the data

from pyspark.ml.feature import StandardScaler
Scalerizer=StandardScaler().setInputCol("features").setOutputCol("Scaled_features")
scaled_data=Scalerizer.fit(transformed_data).transform(transformed_data)

Now I want to do feature selection. But the VarianceThreshold techniuqe is not available. But I know the code of it

var=np.nanvar(df,axis=0)
var=(var>threshold)
df[:,var]

But I have no clue how to implement it on spark dataframe

scikit-learn pyspark python
1 Answers

Is this what you wanted?

import pyspark.sql.functions as F

df = spark.createDataFrame([['2','3'],['4','999']]).toDF('col1', 'col2')
# +----+----+
# |col1|col2|
# +----+----+
# |   2|   3|
# |   4| 999|
# +----+----+


var = df.select([F.variance(col) for col in df.columns]).collect()[0]
# Row(var_samp(col1)=2.0, var_samp(col2)=496008.0)

threshold = 10
df2 = df.select([col for col_id, col in enumerate(df.columns) if var[col_id] < threshold])
# +----+
# |col1|
# +----+
# |   2|
# |   4|
# +----+
user_141657300
Related