PySpark – Adding a Column from a list of values using a UDF
A data frame that is similar to a relational table in Spark SQL, and can be created using various functions in SparkSession is known as a Pyspark data frame. There occur various circumstances in which you get data in the list format but you need it in the form of a column in the data frame. If a similar situation has occurred with you, then you can do it easily by assigning increasing IDs to the data frame and then adding the values in a column. Read the article further to know more about it in detail.
PySpark – Adding a Column from a list of values using a UDF
Example 1:
In the example, we have created a data frame with three columns ‘Roll_Number‘, ‘Fees‘, and ‘Fine‘ as follows:
Once created, we assigned continuously increasing IDs to the data frame using the monotonically_increasing_id function. Also, we defined a list of values, i.e., student_names which need to be added as a column to a data frame. Then, with the UDF on increasing Id’s, we assigned values of the list as a column to the data frame and finally displayed the data frame after dropping the increasing Id’s column.
Python3
# PySpark - Adding a Column from a # list of values using a UDF # Import the libraries SparkSession, # functions, IntegerType, # StringType, row_number, # monotonically_increasing_id and Window from pyspark.sql import SparkSession, functions as F from pyspark.sql.types import IntegerType, StringType from pyspark.sql.functions import row_number, monotonically_increasing_id from pyspark.sql.window import Window # Create a spark session using getOrCreate() function spark_session = SparkSession.builder.getOrCreate() # Create a user defined function to assign student names # according to the row index by subtracting row number from 1 labels_udf = F.udf(lambda indx: student_names[indx-1], StringType()) # Create a data frame with three columns 'Roll_Number,' 'Fees' and 'Fine' df = spark_session.createDataFrame( [(1, 10000, 400), (2, 14000, 500), (3, 12000, 800)], ['Roll_Number', 'Fees', 'Fine']) # Define a list of elements student_names = ['Aman', 'Ishita', 'Vinayak'] # Create a column with continuous increasing Id's df = df.withColumn("num_id", row_number().over( Window.orderBy(monotonically_increasing_id()))) # Create a new column by calling the user defined function new_df = df.withColumn('Names', labels_udf('num_id')) # Delete the increasing Id's column and display the data frame new_df.drop('num_id').show() |
Output:
Example 2:
In this example, we have used a data set (link), i.e., basically, a 5*5 data set as follows:
Then, we assigned continuously increasing IDs to the data frame using the monotonically_increasing_id function. Also, we defined a list of values, i.e., fine_data which needs to be added as a column to the data frame. Then, with the UDF on increasing Id’s, we assigned values of the list as a column to the data frame and finally displayed the data frame after dropping the increasing Id’s column.
Python3
# PySpark - Adding a Column from a list of values using a UDF # Import the libraries SparkSession, functions, IntegerType, # StringType, row_number, monotonically_increasing_id and Window from pyspark.sql import SparkSession, functions as F from pyspark.sql.types import IntegerType, StringType from pyspark.sql.functions import row_number, monotonically_increasing_id from pyspark.sql.window import Window # Create a spark session using getOrCreate() function spark_session = SparkSession.builder.getOrCreate() # Create a user defined function to assign student names # according to the row index by subtracting row number from 1 labels_udf = F.udf(lambda indx: fine_data[indx-1], IntegerType()) # Read the CSV file df = csv_file = spark_session.read.csv( '/content/student_data.csv', sep=',', inferSchema=True, header=True) # Define a list of elements fine_data = [200, 300, 400, 0, 500] # Create a column with continuous increasing Id's df = df.withColumn("num_id", row_number().over( Window.orderBy(monotonically_increasing_id()))) # Create a new column by calling the user defined function new_df = df.withColumn('fine', labels_udf('num_id')) # Delete the increasing Id's column and display the data frame new_df.drop('num_id').show() |
Output:
