Creare il vettore di funzionalità a livello di programmazione in Spark ML/pyspark

Question

Mi chiedo se esiste un modo conciso per eseguire ML (ad esempio KMean) su un DataFrame in pyspark se ho le funzionalità in più colonne numeriche.

I.e. come nel Iris dati.Infatti

(a1=5.1, a2=3.5, a3=1.4, a4=0.2, id=u'id_1', label=u'Iris-setosa', binomial_label=1) 

Vorrei utilizzare K-Means senza ricreare DataSet con il vettore di caratteristiche aggiunto manualmente come una nuova colonna e le colonne originali hardcoded ripetutamente nel codice.

La soluzione mi piacerebbe migliorare:

from pyspark.mllib.linalg import Vectors 
from pyspark.sql.types import Row 
from pyspark.ml.clustering import KMeans, KMeansModel 

iris = sqlContext.read.parquet("/opt/data/iris.parquet") 
iris.first() 
# Row(a1=5.1, a2=3.5, a3=1.4, a4=0.2, id=u'id_1', label=u'Iris-setosa', binomial_label=1) 

df = iris.map(lambda r: Row(
        id = r.id, 
        a1 = r.a1, 
        a2 = r.a2, 
        a3 = r.a3, 
        a4 = r.a4, 
        label = r.label, 
        binomial_label=r.binomial_label, 
        features = Vectors.dense(r.a1, r.a2, r.a3, r.a4)) 
        ).toDF() 


kmeans_estimator = KMeans()\ 
    .setFeaturesCol("features")\ 
    .setPredictionCol("prediction")\ 
kmeans_transformer = kmeans_estimator.fit(df) 

predicted_df = kmeans_transformer.transform(df).drop("features") 
predicted_df.first() 
# Row(a1=5.1, a2=3.5, a3=1.4, a4=0.2, binomial_label=1, id=u'id_1', label=u'Iris-setosa', prediction=1) 

Sto cercando una soluzione, che è qualcosa di simile:

feature_cols = ["a1", "a2", "a3", "a4"] 
prediction_col_name = "prediction" 
<dataframe independent code for KMeans> 
<New dataframe is created, extended with the `prediction` column.> 

Answer 1

È possibile utilizzare VectorAssembler:

from pyspark.ml.feature import VectorAssembler 

ignore = ['id', 'label', 'binomial_label'] 
assembler = VectorAssembler(
    inputCols=[x for x in df.columns if x not in ignore], 
    outputCol='features') 

assembler.transform(df) 

Può essere combinato con k-means utilizzando ML Pipeline:

from pyspark.ml import Pipeline 

pipeline = Pipeline(stages=[assembler, kmeans_estimator]) 
model = pipeline.fit(df)