Iterate over Elements of Array in PySpark DataFrame

Iterating over elements of an array column in a PySpark DataFrame can be done in several efficient ways, such as explode()from pyspark.sql.functions transforms each element of an array into a new row, effectively "flattening" the array column.

Since PySpark DataFrames are distributed across a cluster, you don’t typically use traditional Python for loops for array iteration. Instead, PySpark provides built-in SQL functions such as explode(), posexplode(), transform(), and filter() that handle array iteration at scale. For small datasets, you can also collect the data and use a standard Python loop.

In this article, we’ll explore different approaches to iterating through array elements in a PySpark DataFrame, including built-in functions and UDFs. You’ll also find real-world examples with clear code comments to make the concepts easy to follow.

Key Points-

Use explode() → Flattens an array column so that each element becomes a new row.
Use posexplode() → Similar to explode(), but also returns the position (index) of each element.
Use transform() → Applies a function to each element of an array without changing row count.
Use filter() → Keeps only the array elements that satisfy a given condition.
Use aggregate() → Performs custom reductions (like summing, concatenating) over array elements.
Use Python collect() and loops → Works for small datasets but not recommended for large-scale data.
Use UDFs → Define custom logic to transform arrays into maps, new arrays, or complex structures.
Prefer built-in functions → Functions like explode() and transform() are optimized and faster than UDFs.
Choose method based on goal → Use explode() when you need new rows, transform() when modifying within arrays, filter() when reducing elements.
Scalability matters → Always prefer Spark SQL functions over Python loops to leverage distributed computing.

Use explode()to Iterate Array in a PySpark

You can use the explode() function to convert an array column into individual rows. Let’s create a DataFrame with an array-type column and apply the explode() function to transform each element of the array into a separate row.


# Using explode() to iterate array elements
# Create a DataFrame
from pyspark.sql import SparkSession
from pyspark.sql.functions import explode

spark = SparkSession.builder.appName("IterateArrayExample").getOrCreate()

data = [
    ("James", ["Java", "Scala", "C++"]),
    ("Michael", ["Spark", "Java", "C++"]),
    ("Robert", ["CSharp", "VB"])
]

schema = ["name", "languages"]

df = spark.createDataFrame(data, schema)
df.show(truncate=False)

# Explode array into rows
df_exploded = df.withColumn("language", explode("languages"))
df_exploded.show(truncate=False)

Yields below the output.

Use posexplode() to Iterate Array in a PySpark

You can use posexplode() to retrieve both the array element and its position (index). Let's pass an array column of the given DataFrame into this function to return each element along with its corresponding position.


# Using posexplode() to iterate array elements along with its position
from pyspark.sql.functions import posexplode
df.select("name", posexplode("languages")).show()

Yields below the output.

Use transform() to Iterate Array in a PySpark

You can use the transform() function on an array column to modify each element within the array. In this example, we apply the upper() function to convert all array elements from lowercase to uppercase.


# Using transform() to transform the array elements
from pyspark.sql.functions import expr

df1 = df.withColumn("upper_languages", expr("transform(languages, x -> upper(x))"))
df1.show(truncate=False)

Yields below the output.


# Output:
-------+------------------+------------------+
|name   |languages         |upper_languages   |
+-------+------------------+------------------+
|James  |[Java, Scala, C++]|[JAVA, SCALA, C++]|
|Michael|[Spark, Java, C++]|[SPARK, JAVA, C++]|
|Robert |[CSharp, VB]      |[CSHARP, VB]      |
+-------+------------------+------------------+

Use filter() to Filter the Array Elements

You can use the filter() function from pyspark.sql.functions to keep only the elements of an array column that satisfy a given condition. It takes two arguments: the array column and a lambda function that defines the filtering logic.


# Filter the array elements
from pyspark.sql.functions import filter

df.select("name", filter("languages", lambda x: x.contains("a"))).show()

Yields below the output.


+-------+--------------------------------------------------------------------------------------------+
|   name|filter(languages, lambdafunction(contains(namedlambdavariable(), a), namedlambdavariable()))|
+-------+--------------------------------------------------------------------------------------------+
|  James|                                                                               [Java, Scala]|
|Michael|                                                                               [Spark, Java]|
| Robert|                                                                                    [CSharp]|
+-------+--------------------------------------------------------------------------------------------+

Iterate Array Elements using Python

You can use the collect() method to bring the DataFrame to the driver and then iterate over it using regular Python loops. This works fine for small datasets, but it’s not recommended for larger ones since it may cause memory issues.


# Use for loop implement small dataset
rows = df.select("name", "languages").collect()

for row in rows:
    print(f"Student: {row['name']}")
    for lang in row['languages']:
        print(" -", lang)

Yields below the output.


Student: James
 - Java
 - Scala
 - C++
Student: Michael
 - Spark
 - Java
 - C++
Student: Robert
 - CSharp
 - VB

Using UDF to Build a Map from Array Elements

You can use a PySpark UDF (User Defined Function) to transform array elements into a key–value map. In this example, first we define a UDF that takes the array column along with the individual score columns (Java, Scala, C++) and builds a dictionary where each language is mapped to its corresponding score.


# To Build a Map from Array Elements
from pyspark.sql.functions import udf, col
from pyspark.sql.types import MapType, StringType

# Create Spark session
spark = SparkSession.builder.appName("example").getOrCreate()

# Sample data
data = [
    ("James", ["Java", "Scala", "C++"], 5, 8, 3),
    ("Michael", ["Java"], 2, 9, 1)
]

schema = ["name", "languages", "Java", "Scala", "C++"]

df2 = spark.createDataFrame(data, schema=schema)
df2.show(truncate=False)

# Define UDF to create a dictionary mapping language -> score
@udf(MapType(StringType(), StringType()))
def json_map(languages, Java, Scala, Cpp):   # use safe Python var names
    result = {}
    mapping = {"Java": Java, "Scala": Scala, "C++": Cpp}
    for f in languages:
        result[f] = str(mapping.get(f, None))
    return result

# Apply UDF
df2 = df2.withColumn(
    "language_map",
    json_map(col("languages"), col("Java"), col("Scala"), col("C++"))
)
df2.show(truncate=False)

Yields below the output.


# Output:
+-------+------------------+----+-----+---+
|name   |languages         |Java|Scala|C++|
+-------+------------------+----+-----+---+
|James  |[Java, Scala, C++]|5   |8    |3  |
|Michael|[Java]            |2   |9    |1  |
+-------+------------------+----+-----+---+

+-------+------------------+----+-----+---+---------------------------------+
|name   |languages         |Java|Scala|C++|language_map                     |
+-------+------------------+----+-----+---+---------------------------------+
|James  |[Java, Scala, C++]|5   |8    |3  |{Java -> 5, C++ -> 3, Scala -> 8}|
|Michael|[Java]            |2   |9    |1  |{Java -> 2}                      |
+-------+------------------+----+-----+---+---------------------------------+

Frequently Asked Questions of Iterate Array in PySpark

The best way is to use built-in functions like explode(), transform(), or filter(). These functions are optimized for distributed processing and scale better than Python loops.

explode() converts each array element into a separate row, while transform() modifies elements within the array without changing the row structure.

You can use the explode() or posexplode() function to flatten arrays into multiple rows. This is useful when you want each element treated as an individual record.

The transform() function allows you to apply a lambda expression to modify each element directly inside the array column.

Use the filter() function with a lambda condition to keep only the matching elements inside the array.

Conclusion

In this article, we explored different ways to iterate over arrays in PySpark, including exploding arrays into rows, applying transformations, filtering elements, and creating custom mappings with UDFs.

The best approach depends on whether you need to expand arrays into multiple rows or transform them within the same column.

By learning these techniques, you’ll be better equipped to work with array data efficiently in real-world PySpark pipelines.

Happy Learning!!

Iterate over Elements of Array in PySpark DataFrame

Use explode()to Iterate Array in a PySpark

Use posexplode() to Iterate Array in a PySpark

Use transform() to Iterate Array in a PySpark

Use filter() to Filter the Array Elements

Iterate Array Elements using Python

Using UDF to Build a Map from Array Elements

Frequently Asked Questions of Iterate Array in PySpark

Conclusion

Related Articles