2015-09-16 15 views
10

複数の数値列に機能がある場合、pysparkのDataFrameでML(例えばKMeans)を実行する簡潔な方法があるのだろうかと思います。Spark ML/pysparkでプログラムで機能ベクトルを作成

I.e. Irisデータセットのように:

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

私は、特徴ベクトルは、新しい列とコードで繰り返しハードコーディング元列として手動で追加してDataSetを再作成することなく、関数kmeansを使用したいと思います。

ソリューション私は改善したいと思います:

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) 

を、私は何かのようなソリューション、を探しています:あなたはVectorAssemblerを使用することができます

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

答えて

20

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) 

MLパイプラインを使用してk-meansと組み合わせることができます:

from pyspark.ml import Pipeline 

pipeline = Pipeline(stages=[assembler, kmeans_estimator]) 
model = pipeline.fit(df) 
関連する問題