Erstellen Sie Feature-Vektor programmgesteuert in Spark ML/pyspark

Question

Ich frage mich, ob es eine kurze Möglichkeit zum Ausführen von ML (z. B. KMeans) auf einem Datenrahmen in Pyspark gibt, wenn ich die Features in mehreren numerischen Spalten haben.

I.e. wie in dem Iris-Datensatz:

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

Ich mag KMeans verwenden, ohne den Datensatz neu zu erstellen mit dem Merkmalsvektor manuell als neue Spalte und den ursprünglichen Spalten fest einprogrammiert wiederholt im Code hinzugefügt.

Die Lösung würde ich verbessern möchte:

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) 

ich nach einer Lösung gesucht, die etwas ist, wie:

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

Sie VectorAssembler verwenden können:

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) 

Es kann mit k-Mitteln mit ML-Pipeline kombiniert werden:

from pyspark.ml import Pipeline 

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