Elemento de acceso de un vector en un dataframe de chispa (vector de probabilidad de regresión logística)

Entrené un modelo de LogisticRegression en PySpark (paquete ML) y el resultado de la predicción es un DataSrame de PySpark ( cv_predictions ) (ver [1]). La columna de probability (ver [2]) es un tipo de vector (ver [3]).

 [1] type(cv_predictions_prod) pyspark.sql.dataframe.DataFrame [2] cv_predictions_prod.select('probability').show(10, False) +----------------------------------------+ |probability | +----------------------------------------+ |[0.31559134817066054,0.6844086518293395]| |[0.8937864350711228,0.10621356492887715]| |[0.8615878905395029,0.1384121094604972] | |[0.9594427633777901,0.04055723662220989]| |[0.5391547673698157,0.46084523263018434]| |[0.2820729747752462,0.7179270252247538] | |[0.7730465873083118,0.22695341269168817]| |[0.6346585276598942,0.3653414723401058] | |[0.6346585276598942,0.3653414723401058] | |[0.637279255218404,0.362720744781596] | +----------------------------------------+ only showing top 10 rows [3] cv_predictions_prod.printSchema() root ... |-- rawPrediction: vector (nullable = true) |-- probability: vector (nullable = true) |-- prediction: double (nullable = true) 

¿Cómo creo el análisis del vector del dataframe de PySpark, de modo que creo una nueva columna que solo extrae el primer elemento de cada vector de probability ?

Esta pregunta es similar a, pero las soluciones en los enlaces a continuación no funcionaron / no fueron claras para mí:

Cómo acceder a los valores de denseVector en PySpark

¿Cómo acceder al elemento de una columna VectorUDT en un Spark DataFrame?

Actualizar:

Parece que hay un error en la chispa que le impide acceder a elementos individuales en un vector denso durante una statement de selección. Normalmente, debería poder acceder a ellos como si fuera una matriz numpy, pero al intentar ejecutar el código publicado anteriormente, puede obtener el error pyspark.sql.utils.AnalysisException: "Can't extract value from probability#12;"

Entonces, una forma de manejar esto para evitar este error tonto es usar un udf. Similar a la otra pregunta, puede definir un udf de la siguiente manera:

 from pyspark.sql.functions import udf from pyspark.sql.types import FloatType firstelement=udf(lambda v:float(v[0]),FloatType()) cv_predictions_prod.select(firstelement('probability')).show() 

Detrás de escena, esto todavía accede a los elementos del DenseVector como una matriz numpy, pero no lanza el mismo error que antes.


Respuesta original: un vector denso es solo un envoltorio para una matriz numpy. Por lo tanto, puede acceder a los elementos de la misma manera que lo haría a los elementos de una matriz numpy.

Hay varias formas de acceder a elementos individuales de una matriz en un dataframe. Una es llamar explícitamente a la columna cv_predictions_prod['probability'] en su statement de selección. Al llamar explícitamente a la columna, puede realizar operaciones en esa columna, como seleccionar el primer elemento en la matriz. Por ejemplo:

 cv_predictions_prod.select(cv_predictions_prod['probability'][0]).show() 

Debería resolver el problema.