Some TRICKS about PySpark

1.  DataFrame中如果存在某列类型为vector，如何提取vector中的元素

• pyspark > 3.0.0

from pyspark.ml.functions import vector_to_array

(df.withColumn("xs", vector_to_array("vector")))
    .select(["word"] + [col("xs")[i] for i in range(3)]))

## +-------+-----+-----+-----+
## |   word|xs[0]|xs[1]|xs[2]|
## +-------+-----+-----+-----+
## | assert|  1.0|  2.0|  3.0|
## |require|  0.0|  2.0|  0.0|
## +-------+-----+-----+-----+

• pyspark < 3.0.0
  • 将DataFrame转换为RDD，不再赘述
  • create an UDF

from pyspark.sql.functions import udf, col
from pyspark.sql.types import ArrayType, DoubleType

def to_array(col):
    def to_array_(v):
        return v.toArray().tolist()
    # Important: asNondeterministic requires Spark 2.3 or later
    # It can be safely remove