Skip to content
Tutorial Reference

Python PySpark: How to Add a Column from a List of Values Using a UDF in PySpark

Learn how to map a Python list to a PySpark DataFrame column using row numbering, monotonically_increasing_id(), and a User Defined Function (UDF) with step-by-step examples.

When working with PySpark DataFrames, you may encounter situations where you have a list of values that needs to be added as a new column to an existing DataFrame. Unlike Pandas, PySpark does not allow direct column assignment from a Python list because DataFrames are distributed across a cluster. However, you can achieve this by combining row numbering with a User Defined Function (UDF).

In this guide, you will learn how to map a Python list to a PySpark DataFrame column using monotonically_increasing_id, row_number, and a UDF - with clear examples and explanations.

Understanding the Approach

The strategy involves three steps:

  1. Assign sequential row numbers to the DataFrame using row_number() with monotonically_increasing_id() as the ordering column.
  2. Create a UDF that takes each row number as input and returns the corresponding value from the Python list.
  3. Apply the UDF to create the new column, then drop the temporary row number column.

Adding a String Column from a List

In this example, we create a DataFrame with student data and add a Names column from a Python list.

Step-by-Step Code

from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StringType
from pyspark.sql.functions import row_number, monotonically_increasing_id
from pyspark.sql.window import Window

spark = SparkSession.builder.getOrCreate()

df = spark.createDataFrame(
    [(1, 10000, 400),
     (2, 14000, 500),
     (3, 12000, 800)],
    ["Roll_Number", "Fees", "Fine"]
)

print("Original DataFrame:")
df.show()

Output:

Original DataFrame:
+-----------+-----+----+
|Roll_Number| Fees|Fine|
+-----------+-----+----+
|          1|10000| 400|
|          2|14000| 500|
|          3|12000| 800|
+-----------+-----+----+

Now add the list as a new column:

student_names = ["Anna", "David", "Marty"]

labels_udf = F.udf(
    lambda idx: student_names[idx - 1] if idx and idx <= len(student_names) else None,
    StringType()
)

window_spec = Window.orderBy(F.lit(1))
df_with_id = df.withColumn(
    "row_id",
    row_number().over(window_spec)
)

result_df = df_with_id.withColumn("Names", labels_udf("row_id"))

result_df.drop("row_id").show()

Output:

+-----------+-----+----+-----+
|Roll_Number| Fees|Fine|Names|
+-----------+-----+----+-----+
|          1|10000| 400| Anna|
|          2|14000| 500|David|
|          3|12000| 800|Marty|
+-----------+-----+----+-----+

How It Works

  1. monotonically_increasing_id() generates unique, monotonically increasing IDs across the DataFrame.
  2. row_number().over(Window.orderBy(...)) converts those IDs into consecutive integers starting from 1.
  3. The UDF labels_udf receives each row number and returns student_names[idx - 1] (subtracting 1 because list indexing starts at 0).

Adding an Integer Column from a List

The same technique works for numeric data. Simply change the UDF's return type to IntegerType().

from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import IntegerType
from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

spark = SparkSession.builder.getOrCreate()

df = spark.createDataFrame(
    [(1, "James", 10000),
     (2, "Emma", 14000),
     (3, "Liam", 12000),
     (4, "Olivia", 11000),
     (5, "Noah", 13000)],
    ["Roll_Number", "Name", "Fees"]
)

fine_data = [200, 300, 400, 0, 500]

fine_udf = F.udf(
    lambda idx: fine_data[idx - 1] if idx and idx <= len(fine_data) else None,
    IntegerType()
)

window_spec = Window.orderBy("Roll_Number")

df_with_id = df.withColumn(
    "row_id",
    row_number().over(window_spec)
)

result_df = df_with_id.withColumn("Fine", fine_udf("row_id"))

result_df.drop("row_id").show()

Output:

+-----------+-------+-----+----+
|Roll_Number|   Name| Fees|Fine|
+-----------+-------+-----+----+
|          1|  James|10000| 200|
|          2|   Emma|14000| 300|
|          3|   Liam|12000| 400|
|          4| Olivia|11000|   0|
|          5|   Noah|13000| 500|
+-----------+-------+-----+----+

Common Mistake: Mismatched List and DataFrame Sizes

If your list has fewer elements than the DataFrame has rows, the UDF will throw an IndexError.

Wrong: list has 3 elements but DataFrame has 5 rows

fine_data = [200, 300, 400]  # Only 3 values for 5 rows!

fine_udf = F.udf(lambda idx: fine_data[idx - 1], IntegerType())

df_with_id = df.withColumn(
    "row_id",
    row_number().over(Window.orderBy(monotonically_increasing_id()))
)

result_df = df_with_id.withColumn("Fine", fine_udf("row_id"))
result_df.show()

Error

PythonException: IndexError: list index out of range

Correct: add validation in the UDF

fine_data = [200, 300, 400]

safe_udf = F.udf(
    lambda idx: fine_data[idx - 1] if idx <= len(fine_data) else None,
    IntegerType()
)

result_df = df_with_id.withColumn("Fine", safe_udf("row_id"))
result_df.drop("row_id").show()

Output:

+-----------+-------+-----+----+
|Roll_Number|   Name| Fees|Fine|
+-----------+-------+-----+----+
|          1|  James|10000| 200|
|          2|   Emma|14000| 300|
|          3|   Liam|12000| 400|
|          4| Olivia|11000|null|
|          5|   Noah|13000|null|
+-----------+-------+-----+----+

Alternative Approach: Using createDataFrame and Join

For better performance - especially with large DataFrames - consider converting the list into a DataFrame and joining instead of using a UDF:

from pyspark.sql import SparkSession
from pyspark.sql.functions import row_number, monotonically_increasing_id
from pyspark.sql.window import Window

spark = SparkSession.builder.getOrCreate()

df = spark.createDataFrame(
    [(1, 10000), (2, 14000), (3, 12000)],
    ["Roll_Number", "Fees"]
)

df = df.withColumn(
    "row_id",
    row_number().over(Window.orderBy(monotonically_increasing_id()))
)

names = ["Aman", "Ishita", "Vinayak"]
names_df = spark.createDataFrame(
    [(i + 1, name) for i, name in enumerate(names)],
    ["row_id", "Name"]
)

result_df = df.join(names_df, on="row_id").drop("row_id")
result_df.show()

Output:

+-----------+-----+-----+
|Roll_Number| Fees| Name|
+-----------+-----+-----+
|          1|10000|James|
|          2|14000| Emma|
|          3|12000| Liam|
+-----------+-----+-----+

Summary

Approach Best For Performance
UDF with row_number Quick prototyping, small DataFrames Slower (Python UDF overhead)
Join with list-based DataFrame Production workloads, large DataFrames Faster (native Spark execution)

Key steps to add a column from a list using a UDF:

  1. Generate row numbers using row_number() over monotonically_increasing_id().
  2. Define a UDF that maps each row number to the corresponding list element.
  3. Apply the UDF to create the new column and drop the temporary row number column.
  4. Always validate that the list length matches the DataFrame row count to avoid IndexError.

For performance-critical applications, prefer the join-based approach over UDFs.

Navigation