Select Columns that Satisfy a Condition in PySpark

29 Jun, 2021
In this article, we are going to select columns in the dataframe based on the condition using the where() function in Pyspark. 
Let’s create a sample dataframe with employee data.
Python3
# importing moduleimport pyspark  # importing sparksession from pyspark.sql modulefrom pyspark.sql import SparkSession  # creating sparksession and giving an app namespark = SparkSession.builder.appName('sparkdf').getOrCreate()  # list  of employee datadata = [[1, "sravan", "company 1"], [2, "ojaswi", "company 1"],        [3, "rohith", "company 2"], [4, "sridevi", "company 1"],         [1, "sravan", "company 1"], [4, "sridevi", "company 1"]]  # specify column namescolumns = ['ID', 'NAME', 'Company']  # creating a dataframe from the lists of datadataframe = spark.createDataFrame(data, columns)  # display dataframedataframe.show()
Output:
This method is used to return the dataframe based on the given condition. It can take a condition and returns the dataframe
Syntax:
where(dataframe.column condition)
Here dataframe is the input dataframeThe column  is the column name where we have to raise a condition
Here dataframe is the input dataframe
The column  is the column name where we have to raise a condition
After applying the where clause, we will select the data from the dataframe
Syntax:
dataframe.select('column_name').where(dataframe.column condition)
Here dataframe is the input dataframeThe column is the column name where we have to raise a condition
Here dataframe is the input dataframe
The column is the column name where we have to raise a condition
Example 1: Python program to return ID based on condition
Python3
# importing moduleimport pyspark  # importing sparksession from pyspark.sql modulefrom pyspark.sql import SparkSession  # creating sparksession and giving an app namespark = SparkSession.builder.appName('sparkdf').getOrCreate()  # list  of employee datadata = [[1, "sravan", "company 1"], [2, "ojaswi", "company 1"],         [3, "rohith", "company 2"], [4, "sridevi", "company 1"],         [1, "sravan", "company 1"], [4, "sridevi", "company 1"]]  # specify column namescolumns = ['ID', 'NAME', 'Company']  # creating a dataframe from the lists of datadataframe = spark.createDataFrame(data, columns)  # select ID where ID less than 3dataframe.select('ID').where(dataframe.ID < 3).show()
Output:
Example 2: Python program to select ID and name  where ID =4.
Python3
# importing moduleimport pyspark  # importing sparksession from pyspark.sql modulefrom pyspark.sql import SparkSession  # creating sparksession and giving an app namespark = SparkSession.builder.appName('sparkdf').getOrCreate()  # list  of employee datadata = [[1, "sravan", "company 1"], [2, "ojaswi", "company 1"],        [3, "rohith", "company 2"], [4, "sridevi", "company 1"],         [1, "sravan", "company 1"], [4, "sridevi", "company 1"]]  # specify column namescolumns = ['ID', 'NAME', 'Company']  # creating a dataframe from the lists of datadataframe = spark.createDataFrame(data, columns)  # select ID and name  where ID =4dataframe.select(['ID', 'NAME']).where(dataframe.ID == 4).show()
Output:
Example 3: Python program to select all column based on condition
Python3
# importing moduleimport pyspark  # importing sparksession from pyspark.sql modulefrom pyspark.sql import SparkSession  # creating sparksession and giving an app namespark = SparkSession.builder.appName('sparkdf').getOrCreate()  # list  of employee datadata = [[1, "sravan", "company 1"], [2, "ojaswi", "company 1"],        [3, "rohith", "company 2"], [4, "sridevi", "company 1"],         [1, "sravan", "company 1"], [4, "sridevi", "company 1"]]  # specify column namescolumns = ['ID', 'NAME', 'Company']  # creating a dataframe from the lists of datadataframe = spark.createDataFrame(data, columns)  # select all columns e  where name = sridevidataframe.select(['ID', 'NAME', 'Company']).where(    dataframe.NAME == 'sridevi').show()
Output:
Picked
Python-Pyspark
Python
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.