You can do this with rdd.map()
or using DataFrames and udf()
:
RDD
First create a sample dataset:
text = """abelia,fl,nc
abelia x grandiflora,fl,nc
abelmoschus,ct,dc,fl,hi,il,ky,la,md,mi,ms,nc,sc,va,pr,vi"""
rdd = sc.parallelize(map(lambda x: (x,), text.split("\n")))
rdd.toDF(["rawText"]).show(truncate=False)
#+--------------------------------------------------------+
#|rawText |
#+--------------------------------------------------------+
#|abelia,fl,nc |
#|abelia x grandiflora,fl,nc |
#|abelmoschus,ct,dc,fl,hi,il,ky,la,md,mi,ms,nc,sc,va,pr,vi|
#+--------------------------------------------------------+
Now use map()
twice. First to map each record to a list by splitting on ,
. The second maps the split string into a tuple of the form (x[0], x[1:])
:
rdd.map(lambda x: x[0].split(','))\
.map(lambda x: (x[0], x[1:]))\
.toDF(["plant", "states"])\
.show(truncate=False)
#+--------------------+------------------------------------------------------------+
#|plant |states |
#+--------------------+------------------------------------------------------------+
#|abelia |[fl, nc] |
#|abelia x grandiflora|[fl, nc] |
#|abelmoschus |[ct, dc, fl, hi, il, ky, la, md, mi, ms, nc, sc, va, pr, vi]|
#+--------------------+------------------------------------------------------------+
You could also have done this in one call to map()
but I split it in two for readability.
Dataframe
import pyspark.sql.functions as f
df = sqlCtx.createDataFrame(map(lambda x: (x,), text.split("\n")), ["rawText"])
# define udf to split a string on comma and return all
# of the elements except the first one
get_states = f.udf(lambda x: x.split(',')[1:], ArrayType(StringType()))
df.withColumn('plant', f.split('rawText', ',')[0])\
.withColumn('states', get_states('rawText'))\
.select('plant', 'states')\
.show(truncate=False)
#+--------------------+------------------------------------------------------------+
#|plant |states |
#+--------------------+------------------------------------------------------------+
#|abelia |[fl, nc] |
#|abelia x grandiflora|[fl, nc] |
#|abelmoschus |[ct, dc, fl, hi, il, ky, la, md, mi, ms, nc, sc, va, pr, vi]|
#+--------------------+------------------------------------------------------------+