1. Extract first element of a single vector column:
To get the first element of a vector column, you can use the answer from this SO: discussion Access element of a vector in a Spark DataFrame (Logistic Regression probability vector)
Here's a reproducible example:
>>> from pyspark.sql import functions as f
>>> from pyspark.sql.types import FloatType
>>> df = spark.createDataFrame([{"col1": [0.2], "col2": [0.25]},
{"col1": [0.45], "col2":[0.85]}])
>>> df.show()
+------+------+
| col1| col2|
+------+------+
| [0.2]|[0.25]|
|[0.45]|[0.85]|
+------+------+
>>> firstelement=f.udf(lambda v:float(v[0]),FloatType())
>>> df.withColumn("col1", firstelement("col1")).show()
+----+------+
|col1| col2|
+----+------+
| 0.2|[0.25]|
|0.45|[0.85]|
+----+------+
2. Extract first element of multiple vector columns:
To generalize the above solution to multiple columns, apply a for loop
. Here's an example:
>>> from pyspark.sql import functions as f
>>> from pyspark.sql.types import FloatType
>>> df = spark.createDataFrame([{"col1": [0.2], "col2": [0.25]},
{"col1": [0.45], "col2":[0.85]}])
>>> df.show()
+------+------+
| col1| col2|
+------+------+
| [0.2]|[0.25]|
|[0.45]|[0.85]|
+------+------+
>>> firstelement=f.udf(lambda v:float(v[0]),FloatType())
>>> df = df.select([firstelement(c).alias(c) for c in df.columns])
>>> df.show()
+----+----+
|col1|col2|
+----+----+
| 0.2|0.25|
|0.45|0.85|
+----+----+