In this article, we will learn how to select columns in PySpark dataframe.
Function used:
In PySpark we can select columns using the select() function. The select() function allows us to select single or multiple columns in different formats.
Syntax: dataframe_name.select( columns_names )
Note: We are specifying our path to spark directory using the findspark.init() function in order to enable our program to find the location of apache spark in our local machine. Ignore this line if you are running the program on cloud. Suppose we have our spark folder in c drive by name of spark so the function would look something like: findspark.init(‘c:/spark’). Not specifying the path sometimes may lead to py4j.protocol.Py4JError error when running the program locally.
Example 1: Select single or multiple columns
We can select single or multiple columns using the select() function by specifying the particular column name. Here we are using our custom dataset thus we need to specify our schema along with it in order to create the dataset.
Python3
# select single and multiple columnsfrom pyspark.sql.types import StructType, StructField, StringType, IntegerType, FloatTypefrom pyspark.sql import SparkSessionimport findsparkfindspark.init('c:/spark')# Initialize our datadata2 = [("Pulkit", 12, "CS32", 82, "Programming"), ("Ritika", 20, "CS32", 94, "Writing"), ("Atirikt", 4, "BB21", 78, None), ("Reshav", 18, None, 56, None) ]# Start spark sessionspark = SparkSession.builder.appName("Student_Info").getOrCreate()# Define schemaschema = StructType([ StructField("Name", StringType(), True), StructField("Roll Number", IntegerType(), True), StructField("Class ID", StringType(), True), StructField("Marks", IntegerType(), True), StructField("Extracurricular", StringType(), True)])# read the dataframedf = spark.createDataFrame(data=data2, schema=schema)# slelct columnsdf.select("Name", "Marks").show()# stop the sessionspark.stop() |
Output:
Note: There are a lot of ways to specify the column names to the select() function. Here we used “column_name” to specify the column. Other ways include (All the examples as shown with reference to the above code):
- df.select(df.Name,df.Marks)
- df.select(df[“Name”],df[“Marks”])
- We can use col() function from pyspark.sql.functions module to specify the particular columns
Python3
from pyspark.sql.functions import coldf.select(col("Name"),col("Marks")).show() |
Note: All the above methods will yield the same output as above
Example 2: Select columns using indexing
Indexing provides an easy way of accessing columns inside a dataframe. Indexing starts from 0 and has total n-1 numbers representing each column with 0 as first and n-1 as last nth column. We can use df.columns to access all the columns and use indexing to pass in the required columns inside a select function. Here is how the code will look like. We are using our custom dataset thus we need to specify our schema along with it in order to create the dataset.
Python3
# select sparkfrom pyspark.sql.types import StructType, StructField, StringType, IntegerType, FloatTypefrom pyspark.sql import SparkSessionimport findsparkfindspark.init('c:/spark')# Initialize our datadata2 = [("Pulkit", 12, "CS32", 82, "Programming"), ("Ritika", 20, "CS32", 94, "Writing"), ("Atirikt", 4, "BB21", 78, None), ("Reshav", 18, None, 56, None) ]# Start spark sessionspark = SparkSession.builder.appName("Student_Info").getOrCreate()# Define schemaschema = StructType([ StructField("Name", StringType(), True), StructField("Roll Number", IntegerType(), True), StructField("Class ID", StringType(), True), StructField("Marks", IntegerType(), True), StructField("Extracurricular", StringType(), True)])# read the dataframedf = spark.createDataFrame(data=data2, schema=schema)# select the columnsdf.select(df.columns[:4]).show()# stop sessionspark.stop() |
Output:
Example 3: Access nested columns of a dataframe
While creating a dataframe there might be a table where we have nested columns like, in a column name “Marks” we may have sub-columns of Internal or external marks, or we may have separate columns for the first middle, and last names in a column under the name. In order to access the nested columns inside a dataframe using the select() function, we can specify the sub-column with the associated column. Here we are using our custom dataset thus we need to specify our schema along with it in order to create the dataset.
Python3
# findsparkfrom pyspark.sql.types import StructType, StructField, StringType, IntegerType, FloatTypefrom pyspark.sql import SparkSessionimport findsparkfindspark.init('c:/spark')# initialize the datadata = [ (("Pulkit", "Dhingra"), 12, "CS32", 82, "Programming"), (("Ritika", "Pandey"), 20, "CS32", 94, "Writing"), (("Atirikt", "Sans"), 4, "BB21", 78, None), (("Reshav", None), 18, None, 56, None)]# start spark sessionspark = SparkSession.builder.appName("Student_Info").getOrCreate()# initialize the schema of the dataschema = StructType([ StructField('name', StructType([ StructField('firstname', StringType(), True), StructField('lastname', StringType(), True) ])), StructField("Roll Number", IntegerType(), True), StructField("Class ID", StringType(), True), StructField("Marks", IntegerType(), True), StructField("Extracurricular", StringType(), True)])# create a dataframedf2 = spark.createDataFrame(data=data, schema=schema)# display the schemadf2.printSchema()# select operationdf2.select("name.firstname", "name.lastname").show(truncate=False)# stop sessionspark.stop() |
Output:
Here we can see we have a dataset of following schema
We have a column name with sub columns as firstname and lastname. Now as we performed the select operation we have an output like
