PySpark split and explode example

Code description

This code snippet shows you how to define a function to split a string column to an array of strings using Python built-in split function. It then explodes the array element from the split into using PySpark built-in explode function.

Sample output

+----------+-----------------+--------------------+-----+
|  category|            users|         users_array| user|
+----------+-----------------+--------------------+-----+
|Category A|user1,user2,user3|[user1, user2, us...|user1|
|Category A|user1,user2,user3|[user1, user2, us...|user2|
|Category A|user1,user2,user3|[user1, user2, us...|user3|
|Category B|      user3,user4|      [user3, user4]|user3|
|Category B|      user3,user4|      [user3, user4]|user4|
+----------+-----------------+--------------------+-----+

Code snippet

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, ArrayType
from pyspark.sql.functions import udf, explode

appName = "PySpark Example - split"
master = "local"

# Create Spark session
spark = SparkSession.builder \
    .appName(appName) \
    .master(master) \
    .getOrCreate()

# List
data = [{
    'category': 'Category A',
    'users': 'user1,user2,user3'
}, {
    'category': 'Category B',
    'users': 'user3,user4'
}]

# Create DF
df = spark.createDataFrame(data)
df.show()


@udf(ArrayType(StringType(), False))
def udf_split(text: str, separator: str = ','):
    """
    Spit udf
    """
    return text.split(sep=separator)


df = df.withColumn('users_array', udf_split(df['users']))
df.show()

# Explode - convert array to rows

df  = df.withColumn('user', explode(df['users_array']))
df.show()