Pyspark Data Manipulation Tutorial | by Armando Rivero | Towards Data Science

This tutorial is meant for data people with some Python experience that are absolute Spark beginners. It will help you installing Pyspark and launching your first script. You’ll learn about Resilient Distributed Datasets (RDDs) and dataframes, the main data structures in Pyspark. We discuss some foundational concepts cause I think they will spare you confusion and debugging later. You’ll learn to do data transformations and read from files or databases.
The main reason to learn Spark is that you will write code that could run in large clusters and process big data. This tutorial only talks about Pyspark, the Python API, but you should know there are 4 languages supported by Spark APIs: Java, Scala, and R in addition to Python. Since Spark core is programmed in Java and Scala, those APIs are the most complete and native-feeling.
The advantage of Pyspark is that Python has already many libraries for data science that you can plug into the pipeline. That, together with the fact that Python rocks!!! can make Pyspark really productive. For instance, if you like pandas, know you can transform a Pyspark dataframe into a pandas dataframe with a single method call. You could then do stuff to the data, and plot it with matplotlib.
Over time you might find Pyspark nearly as powerful and intuitive as pandas or sklearn and use it instead for most of your work. Just give Pyspark a try and it could become the next big thing in your career.
There are many articles on how to create Spark clusters, configure Pyspark to submit scripts to them and so on. All of this is needed to do high performance computation on Spark. However, in most companies they’ll have data or infrastructure engineers that will maintain the clusters and you will only code and run the scripts. So you probably will never be installing a Spark cluster, but if you ever need to, you can deal with the problem when it becomes a problem. OTOH, what if I told you that you can get started in a few minutes with a single command?
pip install pyspark
This will download and install Pyspark on your machine so you can learn the ropes and then, when you get a cluster, it will be relatively easy to use it.
To support this tutorial I created this GitHub repository:
github.com
Most of the code in the examples is better organized on the tutorial_part_1_data_wrangling.py file.
Before getting up to speed a little gotcha. I have found Pyspark will throw errors if I don’t also set some environment variables at the beginning of my main Python script.
import sys, osenvironment = ['PYSPARK_PYTHON', 'PYSPARK_DRIVER_PYTHON']for var in environment:    os.environ[var] = sys.executable
After that you can create the spark session
from pyspark.sql import SparkSessionsession = SparkSession.builder.getOrCreate()
Trust me now, this is pretty much all you need to get started.
Alright, so we can get into the real matter here. First, you must know that Spark is complex. Not only it includes a lot of code for doing many data processing tasks, but it also promises to be highly efficient and error free while executing code distributed across hundreds of machines. You don’t really need this complexity at the beginning, what you need is to abstract things away and forget the gritty details. Something important though is the distinction between the driver and the workers. The driver is the Python (single) process where you issue the orders for workers. The workers are Spark processes that follow the driver’s orders. You might think of the workers as being “in the cloud” (the Spark cluster). Spark assumes that big data will be distributed among the workers which together have enough memory and processing capacity to deal with it. The driver is not expected have enough resources to hold this amount of data. That’s why you need to explicitly say when you want to move data to the driver. We’ll get to that in a moment but beware that in real applications you should be very careful with the amount of data you bring to the driver.
RDDs
The RDDs (Resilient Distributed Datasets) are one of the most important data structures in Spark, and the basis of dataframes. You can think of them as “distributed” arrays. In many regards they behave like lists, with a few details we’ll discuss bellow.
So, how to create an RDD? The most straightforward way is to “parallelize” a Python array.
rdd = session.sparkContext.parallelize([1,2,3])
To start interacting with your RDD, try things like:
rdd.take(num=2)
This will bring the first 2 values of the RDD to the driver. The count method will return the length of the RDD
rdd.count()
If you want to send all the RDD data to the driver as an array you can use collect
rdd.collect()
Be careful though, as we said before, in real applications this could make the driver collapse, as the RDD could be gigabytes in size. In general prefer to use take passing the amount of rows you want and invoke collect only when you’re sure the RDD or dataframe is not too big.
If you know pandas or R dataframes you’ll have a very good idea of what Spark dataframes stand for. They represent tabular (matrix) data with named columns. Deep inside, they are implemented with an RDD of Row objects, which are somewhat similar to a Python named tuple. The greatest power of dataframes is they make you able to put your SQL thinking right into action. We’ll talk about dataframe manipulation later, but let’s start creating a dataframe so you can play with it.
df = session.createDataFrame(  [[1,2,3], [4,5,6]], [‘column1’, ‘column2’, ‘column3’])
The data matrix comes first and the column names later.
You can try the take, count and collect methods as in the RDD case; take and collect will give you a list of Row objects. But to me the most user friendly display method would be show:
df.show(n=3)
It will print a table representation of the dataframe with the first n rows.
One key difference with Python lists is that RDDs, (and also dataframes), are immutable. Immutable data is often required in concurrent applications and functional languages.
Let’s discuss what immutability means in Python. Say you do this:
a = list(range(10))a.append(11)
Here the interpreter first creates a name “a” that points to a list object and in the second line it modifies that very same object. However, when you do
st = “my string”st += “ is pretty”
the interpreter first creates the name “st” that points to a string object, then it creates a brand new string object (at a different location in memory) with “my string is pretty” and points the name “st” to that new object. This happens because strings are immutable, you can’t modify them in place.
In the same way RDDs and dataframes can’t be modified in place, so when you do
my_rdd.map(lambda x: x*100)
my_rdd stays the same. You will need to do
my_rdd = my_rdd.map(lambda x: x*100)
To get the name “my_rdd” pointed to the transformed RDD object.
A more shocking difference with ordinary Python might confuse you at the beginning. Sometimes you’ll notice that a very heavy operation happens instantly. But later you do something little (like printing the first value of the RDD) and it seems to take forever.
Let me try to explain this problem simplifying a lot. In Spark there’s a distinction between transformations and actions. When you change a dataframe that’s a transformation, however, when you actually consume the data (eg df.show(1)) that’s an action. Transformations are lazy loaded, they don’t run when you call them. They are executed when you consume their results via an action. Then all the needed transformations will be planned, optimized and run together. So when you see operations that look instantaneous, even though they’re heavy, it’s because they are transformations and they’ll be run later.
User Defined Functions let you use Python code to operate on dataframe cells. You create a regular Python function, wrap it in a UDF object and pass it to Spark, it will care of making your function available in all the workers and scheduling its execution to transform the data.
import pyspark.sql.functions as funcsimport pyspark.sql.types as typesdef multiply_by_ten(number):    return number*10.0multiply_udf = funcs.udf(multiply_by_ten, types.DoubleType())transformed_df = df.withColumn(    'multiplied', multiply_udf('column1'))transformed_df.show()
First you create a Python function, it could be a method in an object, that’s also a function. Then you create an UDF object. The annoying part is that you need to define the output type, something we aren’t really used to in Python. To be really effective with UDFs you’ll need to learn those types, specially the composite MapType (like dictionaries) and ArrayType (like lists). The benefit is that then you can pass this UDF to the dataframe, tell it which column it will be operating on, and you’ll get fantastic things done without leaving the comfort of your old Python.
One of the main limitations with UDFs, though, is that although they can take several columns as input, they can’t change the row as a whole. If you want to work with the whole row, you’ll need the RDD map.
It’s a lot like using a UDF, you also pass a regular Python function. But in this case, the function will be receiving a full Row object instead of column values. It will be expected to return a full Row as well. This will give you the ultimate power over your rows, with a couple of caveats. First: Row object are immutable, so you need to create a whole new Row and return it. Second: you need to convert the dataframe to an RDD and back again. Fortunately neither of these problems are hard to overcome.
Let me show you a function that will logarithmically transform all the columns in your dataframe. As a nice touch it also transforms each column name as ‘log(column_name)’. I find the easiest to start by getting the row into a dictionary with row.asDict(). The function has a bit of Python magic like dictionary comprehensions and key-word arguments packing (the double stars). Hope you understand all of it.
import pyspark.sql.types as typesdef take_log_in_all_columns(row: types.Row):     old_row = row.asDict()     new_row = {f'log({column_name})': math.log(value)                 for column_name, value in old_row.items()}     return types.Row(**new_row)
This by itself won’t do anything. You need to execute the map.
logarithmic_dataframe = df.rdd.map(take_log_in_all_columns).toDF()
You’ll notice this is a chained method call. First you call rdd, it will give you the underlying RDD where the dataframe rows are stored. Then you apply map on this RDD, where you pass your function. To close you call toDF() that transforms an RDD of rows into a dataframe. See this Stack Overflow question for further discussion.
Since dataframes represent tables, naturally they are endowed with SQL-like operations. I’ll be referring just a few of them to get you excited, but you can expect to find almost isomorphic functionality. Calling select will return a dataframe with only some of the original columns.
df.select('column1', 'column2')
This call to where will return a dataframe with only the rows where the value for column1 is 3.
df.where('column1 = 3')
This call to join will return a dataframe that is, well..., a join of df and df1 through column1 in the same way INNER JOIN would do in SQL.
df.join(df1, [‘column1’], how=’inner’)
In case you need to perform right or left joins, in our example df is like the left table and df1, the right one. Outer joins are also possible.
But on top of that, in Spark you can execute SQL much more directly. You can create a temporal view out of a dataframe
df.createOrReplaceTempView(“table1”)
And then perform a query on top of that view
df2 = session.sql("SELECT column1 AS f1, column2 as f2 from table1")
These queries will return a new dataframe with the corresponding column names and values.
The column abstractions supports several direct operations like arithmetic, for instance, say you want to add column1 plus the multiplication of colum2 and column3. Then you can
df3 = df.withColumn(    'derived_column', df['column1'] + df['column2'] * df['column3'])
For this example we will need to read from a CSV file. I downloaded the adult.data file from https://archive.ics.uci.edu/ml/datasets/adult
I moved the file to the data folder in my repository and renamed it to adult.data.csv. I copied the column names from adult.names into a list:
ADULT_COLUMN_NAMES = [     "age",     "workclass",     "fnlwgt",     "education",     "education_num",     "marital_status",     "occupation",     "relationship",     "race",     "sex",     "capital_gain",     "capital_loss",     "hours_per_week",     "native_country",     "income" ]
After reading from the file I set the column names using this list, since the file lacks headers.
csv_df = session.read.csv(     'data/adult.data.csv', header=False, inferSchema=True )for new_col, old_col in zip(ADULT_COLUMN_NAMES, csv_df.columns):     csv_df = csv_df.withColumnRenamed(old_col, new_col)
After that you can try the following for some quick descriptive statistics
csv_df.describe().show()
To get an aggregation use groupBy together with the agg method. For instance, the following will get you a dataframe with the work hours average and standard deviation by age group. We also sort the dataframe by age.
work_hours_df = csv_df.groupBy(    'age').agg(    funcs.avg('hours_per_week'),    funcs.stddev_samp('hours_per_week')).sort('age')
It’s very likely that you keep your data in a relational database, handled by a RDBMS like MySQL or PostgreSQL. If that’s the case, don’t worry, Spark has ways of interacting with many kinds of data storage. I bet you can google your way through most IO needs you may have.
To connect to a RDBMS you need a JDBC driver, that would be a jar file Spark can use to talk to the database. I will show you a PostgreSQL example.
First you need to download the driver.
In my GitHub repo I put the JDBC driver in the bin folder, but you can have it in any path. Also, you should know the code in this example is at a different file in the repo. I had to do this because unlike the other file, this one will throw errors, as the database credentials are of course fake. You’d need to replace them with your database’s to make it run.
Now you’ll start Pyspark with an extra step:
session = SparkSession.builder.config(    'spark.jars', 'bin/postgresql-42.2.16.jar').config(    'spark.driver.extraClassPath', 'bin/postgresql-42.2.16.jar').getOrCreate()
This creates a session aware of the driver. You can then read from the database like this (replacing the fake configs with your real ones):
url = f"jdbc:postgresql://your_host_ip:5432/your_database"properties = {'user': 'your_user', 'password': 'your_password'}# read from a table into a dataframedf = session.read.jdbc(    url=url, table='your_table_name', properties=properties)
Then you can create a transformed dataframe any way you want and write the data back to the database (maybe at a different table).
transformed_df.write.jdbc(    url=url, table='new_table', mode='append', properties=properties)
The writing modes according to the documentation are:
append: Append contents of this DataFrame to existing data.
overwrite: Overwrite existing data.
ignore: Silently ignore this operation if data already exists.
error or “errorifexists” (default case): Throw an exception if data already exists.
I hope you got encouraged to start learning Pyspark right away. I mean, it will take a few weeks to get productive, depending on how much time per day you can afford, but after that you could perform a lot of your data processing with Pyspark. In case you can’t choose what you use at work, at least you can learn it in your free time. My intro barely scratches the surface and mostly hopes to give you an enticing bite. It’s not much more difficult than pandas or sklearn.
Become fluent at Spark, and it could open you the gates of big data.
Part 2 will cover basic Classification and Regression.
PySpark Recipes by Raju Kumar Mishra. Apress, 2018.