Adding a Vectors Column to a pyspark DataFrame

How do I add a Vectors.dense column to a pyspark dataframe?

import pandas as pd
from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.ml.linalg import DenseVector

py_df = pd.DataFrame.from_dict({"time": [59., 115., 156., 421.], "event": [1, 1, 1, 0]})

sc = SparkContext(master="local")
sqlCtx = SQLContext(sc)
sdf = sqlCtx.createDataFrame(py_df)
sdf.withColumn("features", DenseVector(1))

Gives an error in file anaconda3/lib/python3.6/site-packages/pyspark/sql/dataframe.py, line 1848:

AssertionError: col should be Column

It doesn't like the DenseVector type as a column. Essentially, I have a pandas dataframe that I'd like to transform to a pyspark dataframe and add a column of the type Vectors.dense. Is there another way of doing this?

1 answer

  • answered 2018-04-14 15:14 user9613318

    Constant Vectors cannot be added as literal. You have to use udf:

    from pyspark.sql.functions import udf
    from pyspark.ml.linalg import VectorUDT
    
    one = udf(lambda: DenseVector([1]), VectorUDT())
    sdf.withColumn("features", one()).show()
    

    But I am not sure why you need that at all. If you want to transform existing columns into Vectors use appropriate pyspark.ml tools, like VectorAssembler - Encode and assemble multiple features in PySpark

    from pyspark.ml.feature import VectorAssembler
    
    VectorAssembler(inputCols=["time"], outputCol="features").transform(sdf)