Ajustar un dataframe en randomForest pyspark

Tengo un DataFrame que se ve así:

 +--------------------+------------------+ | features| labels | +--------------------+------------------+ |[-0.38475, 0.568...]| label1 | |[0.645734, 0.699...]| label2 | | ..... | ... | +--------------------+------------------+ 

Ambas columnas son de tipo String (StringType ()), me gustaría encajar esto en spark ml randomForest. Para hacerlo, necesito convertir las columnas de características en un vector que contiene flotadores. ¿Alguien tiene alguna idea de cómo hacerlo?

Si está utilizando Spark 2.x , creo que esto es lo que necesita:

 from pyspark.sql.functions import udf from pyspark.mllib.linalg import Vectors from pyspark.ml.linalg import VectorUDT from pyspark.ml.feature import StringIndexer df = spark.createDataFrame([("[-0.38475, 0.568]", "label1"), ("[0.645734, 0.699]", "label2")], ("features", "label")) def parse(s): try: return Vectors.parse(s).asML() except: return None parse_ = udf(parse, VectorUDT()) parsed = df.withColumn("features", parse_("features")) indexer = StringIndexer(inputCol="label", outputCol="label_indexed") indexer.fit(parsed).transform(parsed).show() ## +----------------+------+-------------+ ## | features| label|label_indexed| ## +----------------+------+-------------+ ## |[-0.38475,0.568]|label1| 0.0| ## |[0.645734,0.699]|label2| 1.0| ## +----------------+------+-------------+ 

Con Spark 1.6 , no es muy diferente:

 from pyspark.sql.functions import udf from pyspark.ml.feature import StringIndexer from pyspark.mllib.linalg import Vectors, VectorUDT df = sqlContext.createDataFrame([("[-0.38475, 0.568]", "label1"), ("[0.645734, 0.699]", "label2")], ("features", "label")) parse_ = udf(Vectors.parse, VectorUDT()) parsed = df.withColumn("features", parse_("features")) indexer = StringIndexer(inputCol="label", outputCol="label_indexed") indexer.fit(parsed).transform(parsed).show() ## +----------------+------+-------------+ ## | features| label|label_indexed| ## +----------------+------+-------------+ ## |[-0.38475,0.568]|label1| 0.0| ## |[0.645734,0.699]|label2| 1.0| ## +----------------+------+-------------+ 

Vectors tiene una función de parse que puede ayudarlo a lograr lo que está tratando de hacer.