python - How to extract an element from a array in pyspark

Question

I have a data frame with following type

col1|col2|col3|col4
xxxx|yyyy|zzzz|[1111],[2222]

I want my output to be following type

col1|col2|col3|col4|col5
xxxx|yyyy|zzzz|1111|2222

My col4 is an array and I want to convert it to a separate column. What needs to be done?

I saw many answers with flatmap but they are increasing a row, I want just the tuple to be put in another column but in the same row

Following is my actual schema:

root
 |-- PRIVATE_IP: string (nullable = true)
 |-- PRIVATE_PORT: integer (nullable = true)
 |-- DESTINATION_IP: string (nullable = true)
 |-- DESTINATION_PORT: integer (nullable = true)
 |-- collect_set(TIMESTAMP): array (nullable = true)
 |    |-- element: string (containsNull = true)

Also can please some one help me with explanation on both dataframes and RDD's

What's the schema of your data frame? Can you show df.printSchema()? — Psidom
Does all cells in the array column have the same number of elements? Always 2? What if another row have three elements in the array? — Psidom
No all the elements have exactly 2 elements. Because the element in the array are a start date and end date. — AnmolDave
Also this is my actual requirment if you can help me with it. stackoverflow.com/questions/45252906/… — AnmolDave

Psidom Psidom · Accepted Answer · 2017-07-22T14:00:54

Create sample data:

from pyspark.sql import Row
x = [Row(col1="xx", col2="yy", col3="zz", col4=[123,234])]
rdd = sc.parallelize([Row(col1="xx", col2="yy", col3="zz", col4=[123,234])])
df = spark.createDataFrame(rdd)
df.show()
#+----+----+----+----------+
#|col1|col2|col3|      col4|
#+----+----+----+----------+
#|  xx|  yy|  zz|[123, 234]|
#+----+----+----+----------+

Use getItem to extract element from the array column as this, in your actual case replace col4 with collect_set(TIMESTAMP):

df = df.withColumn("col5", df["col4"].getItem(1)).withColumn("col4", df["col4"].getItem(0))
df.show()
#+----+----+----+----+----+
#|col1|col2|col3|col4|col5|
#+----+----+----+----+----+
#|  xx|  yy|  zz| 123| 234|
#+----+----+----+----+----+

python - How to extract an element from a array in pyspark

1 Answers