Kierk
Kierk

Reputation: 498

Conver pyspark column to a list

I thought this would be easy but can't find the answer :-)

How do I convert the name column in to a list. I am hoping I can get isin to work rather than a join against another datframe column. But isin seems to require a list (if I understand correctly).

Create the datframe:

from pyspark import SparkContext, SparkConf, SQLContext
from datetime import datetime
sc = SparkContext().getOrCreate()
sqlContext = SQLContext(sc)

data2 = [
  ('George', datetime(2010, 3, 24, 3, 19, 58), 3),
  ('Sally', datetime(2009, 12, 12, 17, 21, 30), 5),
  ('Frank', datetime(2010, 11, 22, 13, 29, 40), 2),
  ('Paul', datetime(2010, 2, 8, 3, 31, 23), 8),
  ('Jesus', datetime(2009, 1, 1, 4, 19, 47), 2),
  ('Lou', datetime(2010, 3, 2, 4, 33, 51), 3),
]
 
df2 = sqlContext.createDataFrame(data2, ['name', 'trial_start_time', 'purchase_time'])
df2.show(truncate=False)

Should look like:

+------+-------------------+-------------+
|name  |trial_start_time   |purchase_time|
+------+-------------------+-------------+
|George|2010-03-24 07:19:58|3            |
|Sally |2009-12-12 22:21:30|5            |
|Frank |2010-11-22 18:29:40|2            |
|Paul  |2010-02-08 08:31:23|8            |
|Jesus |2009-01-01 09:19:47|2            |
|Lou   |2010-03-02 09:33:51|3            |
+------+-------------------+-------------+

I am not sure if collect is the closest I can come to this.

df2.select("name").collect()

[Row(name='George'),
 Row(name='Sally'),
 Row(name='Frank'),
 Row(name='Paul'),
 Row(name='Jesus'),
 Row(name='Lou')]

Any suggestions on how to output the name column to a list?

It may need to look something like this:

[George, Sally, Frank, Paul, Jesus, Lou]

Upvotes: 1

Views: 92

Answers (1)

notNull
notNull

Reputation: 31540

Use collect_list function and then collect to get list variable.

Example:

from pyspark.sql.functions import *
df2.agg(collect_list(col("name")).alias("name")).show(10,False)
#+----------------------------------------+
#|name                                    |
#+----------------------------------------+
#|[George, Sally, Frank, Paul, Jesus, Lou]|
#+----------------------------------------+

lst=df2.agg(collect_list(col("name"))).collect()[0][0] 
lst
#['George', 'Sally', 'Frank', 'Paul', 'Jesus', 'Lou']

Upvotes: 1

Related Questions