In [1]:
# create entry points to spark
try:
    sc.stop()
except:
    pass
from pyspark import SparkContext, SparkConf
from pyspark.sql import SparkSession
sc=SparkContext()
spark = SparkSession(sparkContext=sc)

Example DataFrame


In [2]:
mtcars = spark.read.csv(path='../../data/mtcars.csv',
                        sep=',',
                        encoding='UTF-8',
                        comment=None,
                        header=True, 
                        inferSchema=True)

DataFrame to RDD

A DataFrame can be easily converted to an RDD by calling the pyspark.sql.DataFrame.rdd() function. Each element in the returned RDD is an pyspark.sql.Row object. An Row is a list of key-value pairs.


In [3]:
mtcars.rdd.take(2)


Out[3]:
[Row(_c0='Mazda RX4', mpg=21.0, cyl=6, disp=160.0, hp=110, drat=3.9, wt=2.62, qsec=16.46, vs=0, am=1, gear=4, carb=4),
 Row(_c0='Mazda RX4 Wag', mpg=21.0, cyl=6, disp=160.0, hp=110, drat=3.9, wt=2.875, qsec=17.02, vs=0, am=1, gear=4, carb=4)]

With an RDD object, we can apply a set of mapping functions, such as map, mapValues, flatMap, flatMapValues and a lot of other methods that come from RDD.


In [4]:
mtcars_map = mtcars.rdd.map(lambda x: (x['_c0'], x['mpg']))
mtcars_map.take(5)


Out[4]:
[('Mazda RX4', 21.0),
 ('Mazda RX4 Wag', 21.0),
 ('Datsun 710', 22.8),
 ('Hornet 4 Drive', 21.4),
 ('Hornet Sportabout', 18.7)]

In [5]:
mtcars_mapvalues = mtcars_map.mapValues(lambda x: [x, x * 10])
mtcars_mapvalues.take(5)


Out[5]:
[('Mazda RX4', [21.0, 210.0]),
 ('Mazda RX4 Wag', [21.0, 210.0]),
 ('Datsun 710', [22.8, 228.0]),
 ('Hornet 4 Drive', [21.4, 214.0]),
 ('Hornet Sportabout', [18.7, 187.0])]

RDD to DataFrame

To convert an RDD to a DataFrame, we can use the SparkSession.createDataFrame() function. Every element in the RDD has be to an Row object.

Create an RDD


In [6]:
rdd_raw = sc.textFile('../../data/mtcars.csv')
rdd_raw.take(5)


Out[6]:
[',mpg,cyl,disp,hp,drat,wt,qsec,vs,am,gear,carb',
 'Mazda RX4,21,6,160,110,3.9,2.62,16.46,0,1,4,4',
 'Mazda RX4 Wag,21,6,160,110,3.9,2.875,17.02,0,1,4,4',
 'Datsun 710,22.8,4,108,93,3.85,2.32,18.61,1,1,4,1',
 'Hornet 4 Drive,21.4,6,258,110,3.08,3.215,19.44,1,0,3,1']

Save the first row to a variable


In [7]:
header = rdd_raw.map(lambda x: x.split(',')).filter(lambda x: x[1] == 'mpg').collect()[0]
header[0] = 'model'
header


Out[7]:
['model',
 'mpg',
 'cyl',
 'disp',
 'hp',
 'drat',
 'wt',
 'qsec',
 'vs',
 'am',
 'gear',
 'carb']

Save the rest to a new RDD


In [8]:
rdd = rdd_raw.map(lambda x: x.split(',')).filter(lambda x: x[1] != 'mpg')
rdd.take(2)


Out[8]:
[['Mazda RX4',
  '21',
  '6',
  '160',
  '110',
  '3.9',
  '2.62',
  '16.46',
  '0',
  '1',
  '4',
  '4'],
 ['Mazda RX4 Wag',
  '21',
  '6',
  '160',
  '110',
  '3.9',
  '2.875',
  '17.02',
  '0',
  '1',
  '4',
  '4']]

Convert RDD elements to RDD Row objects

First we define a function which takes a list of column names and a list of values and create a Row of key-value pairs. Since keys in an Row object are variable names, we can’t simply pass a dictionary to the Row() function. We can think of a dictionary as an argument list and use the ** to unpack the argument list.

See an example.


In [9]:
from pyspark.sql import Row
my_dict = dict(zip(['a', 'b', 'c'], range(1, 4)))
Row(**my_dict)


Out[9]:
Row(a=1, b=2, c=3)

Let’s define the function.


In [10]:
def list_to_row(keys, values):
    row_dict = dict(zip(keys, values))
    return Row(**row_dict)

Convert elements to RDD objects


In [11]:
rdd_rows = rdd.map(lambda x: list_to_row(header, x))
rdd_rows.take(3)


Out[11]:
[Row(am='1', carb='4', cyl='6', disp='160', drat='3.9', gear='4', hp='110', model='Mazda RX4', mpg='21', qsec='16.46', vs='0', wt='2.62'),
 Row(am='1', carb='4', cyl='6', disp='160', drat='3.9', gear='4', hp='110', model='Mazda RX4 Wag', mpg='21', qsec='17.02', vs='0', wt='2.875'),
 Row(am='1', carb='1', cyl='4', disp='108', drat='3.85', gear='4', hp='93', model='Datsun 710', mpg='22.8', qsec='18.61', vs='1', wt='2.32')]

Now we can convert the RDD to a DataFrame.


In [12]:
df = spark.createDataFrame(rdd_rows)
df.show(5)


+---+----+---+----+----+----+---+-----------------+----+-----+---+-----+
| am|carb|cyl|disp|drat|gear| hp|            model| mpg| qsec| vs|   wt|
+---+----+---+----+----+----+---+-----------------+----+-----+---+-----+
|  1|   4|  6| 160| 3.9|   4|110|        Mazda RX4|  21|16.46|  0| 2.62|
|  1|   4|  6| 160| 3.9|   4|110|    Mazda RX4 Wag|  21|17.02|  0|2.875|
|  1|   1|  4| 108|3.85|   4| 93|       Datsun 710|22.8|18.61|  1| 2.32|
|  0|   1|  6| 258|3.08|   3|110|   Hornet 4 Drive|21.4|19.44|  1|3.215|
|  0|   2|  8| 360|3.15|   3|175|Hornet Sportabout|18.7|17.02|  0| 3.44|
+---+----+---+----+----+----+---+-----------------+----+-----+---+-----+
only showing top 5 rows


In [ ]: