A Decent Guide to DataFrames in Spark 3.0 for Beginners | by David Vrba | Towards Data Science

Apache Spark is a distributed engine that provides a couple of APIs for the end-user to build data processing pipelines. The most commonly used API in Apache Spark 3.0 is the DataFrame API that is very popular especially because it is user-friendly, easy to use, very expressive (similarly to SQL), and in 3.0 quite rich and mature. The API is suitable especially for processing data with some structure and using just a few lines of code you can compose a query that carries out several non-trivial transformations.
In Spark, it is quite typical that one thing can be achieved in different ways and this flexibility brings power but sometimes also confusion. If there are more ways, what is the difference between them? Which way is better or more efficient? As we will see, this is the case also with column transformations on a DataFrame and we will try to clarify this to avoid some misconceptions.
In this article, we will explain what a DataFrame is, how you can create it, and how you can work with it to perform various data transformations. We will use PySpark API in the latest stable version which is at the time of writing (January 2021) 3.0.
DataFrame in Spark is an abstraction that allows us to work with distributed data in a nice way. It represents data that has a tabular structure, each record in the dataset is like a row that has some fields, each field has a name and a data type so each field is like a column in a table. It is indeed very similar to a table in a database, we just need to keep in mind that DataFrame represents data that will be distributed on a cluster during the processing. Also, if your data doesn’t really have a tabular structure, for example, each record is just a string with some text message, you can still represent it with a DataFrame that will have just one column.
There are six basic ways how to create a DataFrame:
The most basic way is to transform another DataFrame. For example:
The most basic way is to transform another DataFrame. For example:
# transformation of one DataFrame creates another DataFramedf2 = df1.orderBy('age')
2. You can also create a DataFrame from an RDD. RDD is a low-level data structure in Spark which also represents distributed data, and it was used mainly before Spark 2.x. It is slowly becoming more like an internal API in Spark but you can still use it if you want and in particular, it allows you to create a DataFrame as follows:
df = spark.createDataFrame(rdd, schema)
3. The next and more useful way (especially for prototyping) is to create a DataFrame from a local collection, for example, from a list:
l = [(1, 'Alice'), (2, 'Bob')] # each tuple will become a rowdf = spark.createDataFrame(l, schema=['id', 'name'])
This is useful especially if you want to test your transformations, but you don’t want to run them on the real data yet. If you know the schema, you can create a small DataFrame like this.
4. For prototyping, it is also useful to quickly create a DataFrame that will have a specific number of rows with just a single column id using a sequence:
df = spark.range(10) # creates a DataFrame with one column id
5. The next option is by using SQL. We pass a valid SQL statement as a string argument to the sql() function:
df = spark.sql("show tables") # this creates a DataFrame
6. And finally, the most important option how to create a DataFrame is by reading the data from a source:
The options here are specified as the key-value pairs and there are different keys available based on the format that your data is in. To check all the possible options see the PySpark docs. There are also various alternatives and shortcuts to this piece of code, for example:
df = spark.read.parquet(path, **options)df = spark.read.format('json').schema(schema).load(path, **options)
If our datasource is stored using the Hive metastore, we can access it using its table name as follows:
df = spark.read.table('table_name')# or just:df = spark.table('table_name')
This is actually more elegant because you don’t have to deal with the schema, the file format, or the path to the data, all you need to know is the name of the table and Spark picks all the needed information from the metastore. Notice that apart from Hive table format there are also other table formats such as Delta or Iceberg and we can cover them in some separate articles.
So we showed six possible ways how to create a DataFrame and as you can see all of them (except for the first one) start with the spark object which is an instance of the SparkSession. The SparkSession is created at the beginning of your Spark application and it is your ticket to the world of DataFrames, you can not create a DataFrame without it. So typically at the beginning of your Spark application, you create the object as follows:
from pyspark.sql import SparkSessionspark = (    SparkSession    .builder    .appName('name_of_your_application')    .enableHiveSupport() # if you want to use Hive metastore    .getOrCreate())
There are two types of operations you can call on a DataFrame, namely transformations, and actions. The transformations are lazy which means that they don’t trigger the computation when you call them, but instead, they just build up a query plan under the cover. So when you call for example this:
result = df.dropDuplicates(['user_id']).orderBy('age')
two transformations are called, but no computation has been done, the data itself is still sitting in the storage and waiting for you to process it. Instead, you just created a DataFrame result that has a query plan which is saying that the data will be deduplicated by the user_id column and sorted by the age column. But all of this will happen after you materialize the query and this is when you call some action. So after calling an action which is usually a function by which you ask for some output, Spark will trigger the computation, which itself is composed of multiple steps: the query plan is optimized by the Spark optimizer, a physical plan is generated, then the plan is compiled into RDD DAG (directed acyclic graph) which is next divided into stages and tasks that are executed on the cluster.
If you want to see the query plan, you can just call:
df.explain()
To read more about query plans, feel free to check my other article or watch my talk from the Spark Summit Europe 2019.
The transformations themselves can be divided into two groups, DataFrame transformations, and column transformations. The first group transform the entire DataFrame, for example
df.select(col1, col2, col3)df.filter(col('user_id') == 123)df.orderBy('age')...
The most frequently used DataFrame transformations are probably the following (but it of course depends on the use case):
select(), withColumn() — for projecting columnsfilter() — for filteringorderBy(), sort(), sortWithinPartitions() — for sortingdistinct(), dropDuplicates() — for deduplicationjoin ()— for joining (see my other article about joins in Spark 3.0)groupBy ()— for aggregations
select(), withColumn() — for projecting columns
filter() — for filtering
orderBy(), sort(), sortWithinPartitions() — for sorting
distinct(), dropDuplicates() — for deduplication
join ()— for joining (see my other article about joins in Spark 3.0)
groupBy ()— for aggregations
For the complete list of DataFrame transformations see the class DataFrame in PySpark docs.
On the other hand, the column transformations are used to transform individual columns and they are used either inside the withColumn(), select(), or selectExpr() transformations that allow you to add a new column to the DataFrame:
from pyspark.sql.functions import initcap# capitalize the first letter of the user name and save it to a new # column name_capdf.withColumn('name_cap', initcap('user_name'))
Here the column transformation is achieved with the function initcap() which transforms the string from the user_name column. On the other hand, the withColumn transformation is a DataFrame transformation because it transforms the entire DataFrame — it adds a new column. So the typical use case is that you want to add a new column that is derived from other columns using some transformations.
For the column transformations, it is typical that they can be expressed using multiple different ways, in the example above the same can be achieved with select (or equivalently using expr or selectExpr):
df.select('*', initcap('user_name').alias('name_cap'))df.select('*', expr('initcap(user_name) as name_cap'))df.selectExpr('*', 'initcap(user_name) as name_cap')
Notice the difference between the select and withColumn, the latter projects all columns from the DataFrame and adds one new with a given name, while select only projects those columns that are passed as the argument, so if you want to have there also all columns, you either need to list them all explicitly, or you use the notation with an asterisk as we have above. Also, the resulting column of the transformation inside select will have a default name, so if you want to use a custom name, you have to rename it using alias().
Also, notice that the function expr() allows you to pass in a SQL expression as a string in quotes, similarly, selectExpr() is just a shortcut for using select() + expr(). In the expression, you can use any of the SQL functions that you can find in the SQL docs.
Now you might be confused what is the use case for the expr function and we will explain that shortly. Before we do that, let’s see yet another way how to express column transformations and it is using the methods from the Column class. Consider an example in which we want to do a substring of length 3 taken from the beginning of the word. Using the method substr() from the Column class we can do it as follows:
df.withColumn('new_col', col('user_name').substr(1, 3))
As you can see these methods from the Column class are called on a column object that you can construct either using the col function or using the dot or bracket notations as follows:
# all of these create a column object:col('user_name')df.user_namedf['user_name']
The example above can be alternatively written using the substring() function from the pyspark.sql.functions package as follows:
df.withColumn('new_col', substring('user_name', 1, 3))
Here user_name can be either a string — the name of the column, or it can be a column object, the function is more flexible and can handle both, which is typical for the functions from this package, however, this flexibility has its limitations as you will see in a moment.
You might be wondering what is the difference between using the SQL functions with expr or using directly the functions from pyspark.sql.functions package (that we call here DSL functions). Well, if you check both documentations (SQL, DSL), you will see that many of these functions are really the same (in the sense that they have the same name and provide you the same functionality), but there are some cases in which the SQL functions are more flexible and allow you to achieve what is not easy with the DSL functions.
Let’s consider again the example with the substring, the second two arguments, position, and length of the DSL substring function have to be integers — they can not be column objects. But sometimes there is a use case where you need to take the argument from another column, for example, you want the length of the substring to be dynamic, and you want to use a different value on each row. The DSL substring function doesn’t allow you to do it — the position and length need to be constant, the same on each row.
This is the limitation that we mentioned above and it can be circumvented using the SQL functions inside expr(). Consider this example in which we want to make a substring of one column based on information in another column. Let’s say we want to take the first n characters, where n is stored in another column called len:
l = [('abcd', 2), ('abcd', 1)]df = spark.createDataFrame(l, ['str', 'len'])df.show()+----+----------+| str|    len   |+----+----------+|abcd|         2||abcd|         1|+----+----------+
And let’s now see what options we have to implement the substring of len characters:
# (1) raises error:df.withColumn('new_col', substring('str', 1, 'len')) # (2) raises error:df.withColumn('new_col', substring(col('str'), 1, col('len')))# (3) raises error:df.withColumn('new_col', substring(col('str'), lit(1), col('len')))# (4) raises error:df.withColumn('new_col', col('str').substr(1, 'len'))# (5) raises error:df.withColumn('new_col', col('str').substr(1, col('len')))# (6) this works:df.withColumn('new_col', col('str').substr(lit(1), col('len')))# (7) this works:df.withColumn('new_col', expr('substring(str, 1, len)')).show()+----+---+-------+| str|len|new_col|+----+---+-------+|abcd|  2|     ab||abcd|  1|      a|+----+---+-------+
As you can see we are not able to do it using substring from the pyspark.sql.package, each of the attempts (1–3) fails. But if we use the substring function from the SQL API inside expr it works (attempt 7), the SQL function is more flexible and is able to realize that len is actually a name of a column in the DataFrame and it takes the length from there. We can see that attempt 6 in which we use the substr method also works but we need to make sure that both arguments are columns, that is why we have to use lit(1) for the first argument (the lit function will make a column object from the constant that we pass in). The class Column is however not so rich in terms of methods it provides as compared to the SQL functions, so in some other cases, we will not find a corresponding function here. One such example is discussed in this StackOverflow question, another one is related to array_sort, which I described in my other article here.
To summarize, if you want to create a column transformation, there are three places where you can look for an appropriate function in the documentation:
pyspark.sql.functions (here we call them DSL functions)
Column class (methods of the Column class — they are called on a column object)
SQL functions (from the SQL API to be used inside expr, or selectExpr, here we call them SQL functions)
All of them are equivalent in terms of performance because they are lazy and simply update the query plan under the hood. The SQL functions used with expr() (or selectExpr()) have the advantage to be sometimes more flexible, than their DSL counterparts and allow you to pass all arguments as column names. Very often you don’t need this extra flexibility and which approach you decide to use is usually a matter of style and it is possibly also based on some agreement between the team members on what they prefer to use in their codebase.
There are three more advanced groups of column transformations, which are user-defined functions, higher-order functions, and window functions and we can cover them in some separate articles since they deserve a more detailed exploration. Very briefly explained, the user-defined functions allow us to extend the DataFrame API through a very simple interface and so using native Python (Scala/Java) functions we can implement any custom transformation that is not available in the API yet, and we can then use it as the column transformation. The higher-order functions are quite well supported since Spark 2.4 and they are used to transform and manipulate complex data types such as arrays or maps. The window functions allow us to compute various aggregations or ranking over groups of rows that are defined by a window and a frame.
As explained above, actions are functions in which we ask for some output. These functions trigger the computation and run a job on the Spark cluster. The usual situation is that the action runs one job, however, in some cases, it can run even more jobs (see for example my other article where I explain why the show() function can run more jobs). Some typical actions are:
count() — computes the number of rows in the dataset that is represented by the DataFrame.
show() — prints on the screen 20 records from the dataset that is represented by the DataFrame.
collect() — prints all the records on the screen. This should be used carefully since it literally collects all the data from all executors and brings it to the driver, this is potentially dangerous because if the data is big, it can crash the driver since its memory is limited. This function can be useful if the data is already filtered or aggregated enough so its size is not a problem for the driver.
toPandas() — this is an analogy of collect(), the result however is not a list of records but instead it is a Pandas dataframe.
take(n) — also an analogy of collect(), but it collects only n records. This can be useful if you want to check whether there is some data, or the DataFrame is empty (df.take(1)).
write — creates a DataFrame writer which allows saving data to external storage.
Let’s see a final diagram (with some examples) that describes the operations in Spark:
In this note, we took an introduction path to Spark SQL and DataFrames in particular. We first explained what a DataFrame is, we have seen six ways how it can be created and we spent some time explaining what kind of operations you can do with it. We also pointed out the flexibility of the SQL functions that allow you to bypass the limitations of the corresponding DSL functions when you need to provide multiple arguments as columns from the DataFrame. We explained what is the difference between various approaches to column transformations which is a subject that is often confusing especially for beginners in Spark.