pyspark if else

How to use conditional statements in PySpark Azure Databricks?

by

Are you looking to find out how to use conditional statements of PySpark DataFrame using Azure Databricks cloud or maybe you are looking for a solution, to get values based on if else condition in PySpark Databricks using the when() and otherwise() functions? If you are looking for any of these problem solutions, you have landed on the correct page. I will also help you how to use PySpark when() and otherwise() functions with multiple examples in Azure Databricks. I will explain it by taking a practical example. So please don’t waste time let’s start with a step-by-step guide to understand how to use the when() and otherwise() functions in PySpark.

In this blog, I will teach you the following with practical examples:

  • Syntax of when() and otherwise()
  • Using when() and Otherwise() on DataFrame
  • Multiple conditions using when()

Syntax:

The Pyspark when() function is a SQL function used to return a value of column type based on a condition.

The Pyspark otherwise() function is a column function used to return a value for matched condition. If otherwise() function is not invoked, None is returned for unmatched conditions.

Syntax:

when().otherwise()

What is the syntax of the when() and otherwise() functions in PySpark Azure Databricks?

The syntax is as follows:

when(condition, matched_value).otherwise(unmatched_value)
Parameter NameRequiredDescription
conditionYesIt represents the condition of if else statement.
matched_valueYesIt represents the matched condition return value.
unmatched_value (Any)OptionalIt represents the unmatched condition return value.
Table 1: when() and otherwise() Method in PySpark Databricks Parameter list with Details

Apache Spark Official Documentation Link: when() and otherwise()

Create a simple DataFrame

Let’s understand the use of the when() and otherwise() functions with various examples. Let’s start by creating a DataFrame.

Gentle reminder:

In Databricks,

  • sparkSession made available as spark
  • sparkContext made available as sc

In case, you want to create it manually, use the below code.

from pyspark.sql.session import SparkSession

spark = SparkSession.builder 
    .master("local[*]") 
    .appName("azurelib.com") 
    .getOrCreate()

sc = spark.sparkContext

a) Create manual PySpark DataFrame

data = [
    (27,190,43,11.9),
    (28,187,28,8.0),
    (29,163,72,27.1),
    (30,179,68,21.2),
    (31,153,54,23.1),
    (32,178,23,7.3),
    (33,195,29,7.6),
    (34,160,59,23.0),
    (35,157,69,28.0),
    (36,189,59,16.5)
]

df = spark.createDataFrame(data, schema=["id","height","weight","bmi"])
df.printSchema()
df.show(5, truncate=False)

"""
root
 |-- id: long (nullable = true)
 |-- height: long (nullable = true)
 |-- weight: long (nullable = true)
 |-- bmi: double (nullable = true)

+---+------+------+----+
|id |height|weight|bmi |
+---+------+------+----+
|27 |190   |43    |11.9|
|28 |187   |28    |8.0 |
|29 |163   |72    |27.1|
|30 |179   |68    |21.2|
|31 |153   |54    |23.1|
+---+------+------+----+
"""

b) Creating a DataFrame by reading files

Download and use the below source file.

Download Sample Bmi.csv
# replace the file_path with the source file location which you have downloaded.

df_2 = spark.read.format("csv").option("header", True).load(file_path)
df_2.printSchema()

"""
root
 |-- id: long (nullable = true)
 |-- height: long (nullable = true)
 |-- weight: long (nullable = true)
 |-- bmi: double (nullable = true)
"""

Note: Here, I will be using the manually created DataFrame.

Before diving deep, let’s see what BMI range value represents:

  • Under weight : BMI < 18.5
  • Normal : BMI >= 18.5 and BMI < 25
  • Over weight : BMI >= 25

How to use if else condition in PySpark Azure Databricks?

Let’s see how to use if else similar statements in PySpark using Azure Databricks.

Example 1:

# Using select()

from pyspark.sql.functions import when, col

df.select("*", when(col("bmi") < 18.5, "under-weight").alias("status")).show(5)

"""
Output:

+---+------+------+----+------------+
| id|height|weight| bmi|      status|
+---+------+------+----+------------+
| 27|   190|    43|11.9|under-weight|
| 28|   187|    28| 8.0|under-weight|
| 29|   163|    72|27.1|        null|
| 30|   179|    68|21.2|        null|
| 31|   153|    54|23.1|        null|
+---+------+------+----+------------+

"""

Note: Since, if hasn’t specified what to happen when the condition is unmatched the function return null as a value.

Example 2:

# Using withColumn()

from pyspark.sql.functions import when, col

df \
.withColumn("status", when(col("bmi") < 18.5, "under-weight") \
.otherwise("unknown").alias("status")) \
.show(5)

"""
Output:

+---+------+------+----+------------+
| id|height|weight| bmi|      status|
+---+------+------+----+------------+
| 27|   190|    43|11.9|under-weight|
| 28|   187|    28| 8.0|under-weight|
| 29|   163|    72|27.1|     unknown|
| 30|   179|    68|21.2|     unknown|
| 31|   153|    54|23.1|     unknown|
+---+------+------+----+------------+

"""

How to use if else statement in PySpark Azure Databricks using SQL expression?

Let’s see how to use SQL CASE WHEN statements in PySpark using Azure Databricks.

Example 1:

# Using select():

from pyspark.sql.functions import expr

df.select("*", expr("""
                    CASE
                    WHEN bmi >= 25 THEN 'over-weight'
                    ELSE 'unknown'
                    END as Status
                    """)
         ).show(5)

"""
Output:

+---+------+------+----+-----------+
| id|height|weight| bmi|     Status|
+---+------+------+----+-----------+
| 27|   190|    43|11.9|    unknown|
| 28|   187|    28| 8.0|    unknown|
| 29|   163|    72|27.1|over-weight|
| 30|   179|    68|21.2|    unknown|
| 31|   153|    54|23.1|    unknown|
+---+------+------+----+-----------+

"""

Example 2:

# Using withColumn():

from pyspark.sql.functions import expr

df.withColumn("status", expr("""
                            CASE
                            WHEN bmi >= 25 THEN 'over-weight'
                            ELSE 'unknown'
                            END AS status
                            """)
             ).show(5)

"""
Output:

+---+------+------+----+-----------+
| id|height|weight| bmi|     status|
+---+------+------+----+-----------+
| 27|   190|    43|11.9|    unknown|
| 28|   187|    28| 8.0|    unknown|
| 29|   163|    72|27.1|over-weight|
| 30|   179|    68|21.2|    unknown|
| 31|   153|    54|23.1|    unknown|
+---+------+------+----+-----------+


"""

How to use multiple if conditions in PySpark Azure Databricks?

Let’s see how to use multiple if similar statements in PySpark using Azure Databricks.

Example 1:

# Using withColumn

from pyspark.sql.functions import when, col, expr

df.select("*", 
           when(col("bmi") < 18.5, "under-weight") \
          .when(col("bmi") >= 25, "over-weight") \
          .otherwise("normal").alias("status")          
         ).show(5)

"""
Output:

+---+------+------+----+------------+
| id|height|weight| bmi|      status|
+---+------+------+----+------------+
| 27|   190|    43|11.9|under-weight|
| 28|   187|    28| 8.0|under-weight|
| 29|   163|    72|27.1| over-weight|
| 30|   179|    68|21.2|      normal|
| 31|   153|    54|23.1|      normal|
+---+------+------+----+------------+

"""

Example 2:

# Using SQL expression

from pyspark.sql.functions import when, col, expr

df.select("*", expr("""
                    CASE
                    WHEN bmi < 18.5 THEN 'under-weight'
                    WHEN bmi >= 25 THEN 'over-weight'
                    ELSE 'normal'
                    END AS status
                    """)         
         ).show(5)

"""
Output:

+---+------+------+----+------------+
| id|height|weight| bmi|      status|
+---+------+------+----+------------+
| 27|   190|    43|11.9|under-weight|
| 28|   187|    28| 8.0|under-weight|
| 29|   163|    72|27.1| over-weight|
| 30|   179|    68|21.2|      normal|
| 31|   153|    54|23.1|      normal|
+---+------+------+----+------------+

"""

I have attached the complete code used in this blog in notebook format to this GitHub link. You can download and import this notebook in databricks, jupyter notebook, etc.

GitHub Link

When should you use conditional statements in PySpark Azure Databricks?

These could be the possible reasons:

  1. To perform the if else statement in PySpark
  2. To add or customize columns based on values

Real World Use Case Scenarios for using conditional statements in PySpark Azure Databricks?

Assume you have given a people dataset. The dataset includes people’s names and ages. You want to create a new column known as “can_vote” and you have been asked to check whether the person is above 18, if he is above 18 represent he/she as “eligible” else “not eligible”. For example (‘Berne’, 23), (‘Suresh’, 16) into (‘Berne’, 23, ‘eligible’), (‘Suresh’, 16, ‘not eligible’).

What are the alternatives to when() and otherwise() functions in PySpark Azure Databricks?

There are multiple alternatives to the when() and otherwise() functions, which is as follows:

  • when() and otherwise() functions
  • PySpark SQL CASE WHEN operations

Final Thoughts

In this article, we have learned about the PySpark when() and otherwise() methods of DataFrame in Azure Databricks along with the examples explained clearly. I have also covered different scenarios with practical examples that could be possible. I hope the information that was provided helped in gaining knowledge.

Please share your comments and suggestions in the comment section below and I will try to answer all your queries as time permits.

Join Telegram Group
Download Free Azure Videos
Connect on LinkedIn

Arud Seka Berne S
Arud Seka Berne S

As a big data engineer, I design and build scalable data processing systems and integrate them with various data sources and databases. I have a strong background in Python and am proficient in big data technologies such as Hadoop, Hive, Spark, Databricks, and Azure. My interest lies in working with large datasets and deriving actionable insights to support informed business decisions.