Crear un vector de características mediante progtwigción en Spark ML / pyspark

Me pregunto si hay una forma concisa de ejecutar ML (por ejemplo, KMeans) en un DataFrame en pyspark si tengo las características en varias columnas numéricas.

Es decir, como en el conjunto de datos Iris :

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

Me gustaría usar KMeans sin recrear el DataSet con el vector de características agregado manualmente como una nueva columna y las columnas originales codificadas repetidamente en el código.

La solución que me gustaría mejorar:

 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) 

Estoy buscando una solución, que es algo como:

 feature_cols = ["a1", "a2", "a3", "a4"] prediction_col_name = "prediction"   

Puedes usar 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) 

Se puede combinar con k-means utilizando ML Pipeline:

 from pyspark.ml import Pipeline pipeline = Pipeline(stages=[assembler, kmeans_estimator]) model = pipeline.fit(df)