Built-in functions

Spark tutorial · PySpark.in

PySpark Functions

PySpark, the Python API for Apache Spark, comes with a rich set of functions to manipulate and analyze data at scale. This tutorial series introduces beginners and intermediate users to some of the most useful PySpark functions, organized into thematic chapters. Each chapter provides an overview of the function group, typical use cases, code examples with DataFrames, and best practices. By the end of the series, you will know how to use PySpark's built-in functions for common transformations, aggregate data effectively, handle array and string data, work with dates, and even extend Spark with your own user-defined functions when necessary.

1. Built-in DataFrame Functions: col, lit, when, regexp_extract, concat

Overview: PySpark provides many built-in functions in the pyspark.sql.functions module for common data transformations. In this section, we'll cover five fundamental functions:

col() – Returns a Column object referencing a DataFrame column by name.
lit() – Creates a Column of a literal (constant) value.
when() – Performs conditional logic (if/else) on columns, similar to SQL CASE WHEN.
regexp_extract() – Extracts substrings using a regular expression with capture groups
concat() – Concatenates multiple columns into one (works on strings, numbers, or arrays).

Typical Use Cases:

Dynamically referencing columns (using col) especially when column names are stored in variables or when building expressions on the fly.
Adding constant values or flags to all rows (using lit for default values or identifiers).
Deriving new columns based on conditions (using when ... .otherwise for categorizing or flagging data).
Parsing and extracting information from text columns (using regexp_extract for fields like IDs, codes, or email addresses embedded in strings).
Combining multiple fields into one (using concat to merge first and last name into full name, or merge address components, etc., often with a lit as a separator).

Using col() and lit()

The col() function is often used to reference DataFrame columns in expressions instead of using string column names. This is useful in complex expressions or to avoid ambiguity. For example, we can use col("age") to refer to the age column. The lit() function creates a literal column – for instance, adding a constant value to each row.

Example: Add a constant column and use col in an expression. Suppose we want to add a new column for country (assume all employees are in USA) and calculate a 10% bonus of their salary:

```

from pyspark.sql.functions import col, lit

from pyspark.sql import SparkSession

from pyspark.sql.functions import col, lit, when, regexp_extract, concat

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

# Sample employee data

data = [

("Alice", "Smith", 34, "alice.smith@xyz.com", "Sales", 80000),

("Bob", "Jones", 23, "bob.jones@xyz.com", "Marketing", 50000),

("Charlie", "Lee", 36, "charlie.lee@xyz.com", "Sales", 120000),

("Diana", "Diaz", 29, "diana.diaz@xyz.com", "Finance", 70000)

]

columns = ["first_name", "last_name", "age", "email", "department", "salary"]

df = spark.createDataFrame(data, columns)

df_with_const = df.withColumn("country", lit("USA"))

df_with_const = df_with_const.withColumn("bonus", col("salary") * lit(0.10))

df_with_const.select("first_name", "salary", "bonus", "country").show()

```

Here we used lit("USA") to add a constant string column, and col("salary") * lit(0.10) to compute 10% of the salary. Using col in expressions ensures the code works even if we later reuse it with different column names (since we could replace the string in col() with a variable).

Best Practice: Using col() is preferred over dataframe["columnName"] notation in many cases because it allows dynamic construction of expressions and is required in some contexts (e.g., inside when or other function calls). It’s also a lightweight operation that doesn’t incur performance cost. Use lit() whenever you need to inject a literal value (number, string, etc.) into a DataFrame computation instead of a column – for example, filtering where a column equals a fixed string, or adding a constant column.

Using when() for Conditional Columns

The when(condition, value) function (often combined with .otherwise(default)) allows you to create new columns based on conditional logic, similar to an if-else or CASE in SQL. This is useful for bucketing or flagging rows.

Example: Classify employees as "Young" or "Adult" based on age.

```

from pyspark.sql.functions import when

from pyspark.sql import SparkSession

from pyspark.sql.functions import col, lit, when, regexp_extract, concat

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

# Sample employee data

data = [

("Alice", "Smith", 34, "alice.smith@xyz.com", "Sales", 80000),

("Bob", "Jones", 23, "bob.jones@xyz.com", "Marketing", 50000),

("Charlie", "Lee", 36, "charlie.lee@xyz.com", "Sales", 120000),

("Diana", "Diaz", 29, "diana.diaz@xyz.com", "Finance", 70000)

]

columns = ["first_name", "last_name", "age", "email", "department", "salary"]

df = spark.createDataFrame(data, columns)

df_with_status = df.withColumn(

"status", when(col("age") < 30, "Young").otherwise("Adult")

)

df_with_status.select("first_name", "age", "status").show()

```

We added a column status that is "Young" if age < 30, otherwise "Adult". You can chain multiple when conditions to handle more categories (just like multiple WHEN ... THEN clauses in SQL), and call .otherwise() at the end for the default case. Using when is much more efficient and readable than writing a Python UDF for simple conditional logic, since it will be executed within Spark’s engine.

Best Practice: Always terminate a series of when() calls with an .otherwise() to handle the false condition (or unexpected values); otherwise, Spark will assign NULL to entries not matching any when condition. The when function is optimized by Spark (unlike a Python if/else in a UDF) and should be preferred for conditional column creation.

Using regexp_extract() for Pattern Matching

The regexp_extract(column, regex, idx) function extracts substrings that match a regex capture group. The idx parameter specifies which capture group to return (0 for the entire match, 1 for the first group, etc.). This is very handy for parsing strings like IDs, codes, or emails.

Example: Extract the domain from the email address. In our DataFrame, the email field contains addresses like "alice.smith@xyz.com". We can extract the domain (the part after '@') using a regex:

```

from pyspark.sql.functions import regexp_extract

from pyspark.sql import SparkSession

from pyspark.sql.functions import col, lit, when, regexp_extract, concat

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

# Sample employee data

data = [

("Alice", "Smith", 34, "alice.smith@xyz.com", "Sales", 80000),

("Bob", "Jones", 23, "bob.jones@xyz.com", "Marketing", 50000),

("Charlie", "Lee", 36, "charlie.lee@xyz.com", "Sales", 120000),

("Diana", "Diaz", 29, "diana.diaz@xyz.com", "Finance", 70000)

]

columns = ["first_name", "last_name", "age", "email", "department", "salary"]

df = spark.createDataFrame(data, columns)

# Extract domain of email after the '@' symbol

df_with_domain = df.withColumn(

"email_domain", regexp_extract(col("email"), "@(\\w+\\.\\w+)", 1)

)

df_with_domain.select("email", "email_domain").show()

```

The regex pattern @(\\w+\\.\\w+) means "an '@' followed by one or more word characters, a literal dot, and one or more word characters". We used capture group 1 (the part inside parentheses) to get everything after the "@". As a result, all emails in this example have the domain "xyz.com". The regexp_extract function returns an empty string if the pattern isn't found; if we wanted a NULL instead, we could handle that with an additional condition. This function is powerful for extracting structured pieces from unstructured text.

Best Practice: Always test your regex patterns to ensure they behave as expected, and be mindful of Spark's double-escape requirements in patterns (notice the double backslash \\ in the pattern above). Also, if you need to extract multiple parts, you may call regexp_extract multiple times with different group indices or use regexp_extract_all (Spark 3.2+) for all occurrences.

Using concat() to Combine Columns

The concat(col1, col2, ..., colN) function concatenates multiple columns into one column. If the inputs are strings, it will produce a string by combining them end-to-end. (For array columns, it will concatenate the arrays, and for numeric types it will also just smoosh them together as string representation, so typically you cast numbers to string first or use formatting.)

Example: Create a full name by concatenating first and last names.

```

from pyspark.sql.functions import concat

from pyspark.sql import SparkSession

from pyspark.sql.functions import col, lit, when, regexp_extract, concat

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

# Sample employee data

data = [

("Alice", "Smith", 34, "alice.smith@xyz.com", "Sales", 80000),

("Bob", "Jones", 23, "bob.jones@xyz.com", "Marketing", 50000),

("Charlie", "Lee", 36, "charlie.lee@xyz.com", "Sales", 120000),

("Diana", "Diaz", 29, "diana.diaz@xyz.com", "Finance", 70000)

]

columns = ["first_name", "last_name", "age", "email", "department", "salary"]

df = spark.createDataFrame(data, columns)

df_with_fullname = df.withColumn(

"full_name", concat(col("first_name"), lit(" "), col("last_name"))

)

df_with_fullname.select("first_name", "last_name", "full_name").show()

```

We used lit(" ") to insert a space between first and last name during concatenation. The result is a new full_name string. As shown, "Charlie Lee" has two spaces in the output because in our input "Charlie" had a trailing space originally (just an example of how concat will not trim inputs). If any of the columns were numeric, we would first cast them to string (or use concat_ws which can handle mixing types by converting to string internally).

Note: concat() will produce NULL if any of the columns being concatenated are NULL. If you want to treat nulls as empty strings, consider using concat_ws (concatenate with separator) or handle nulls via coalesce or na.fill before concatenation.

Built-in Functions Best Practices

Use built-in column functions (col, lit, when, etc.) for transformations whenever possible. They are optimized by Spark's Catalyst engine, whereas UDFs (user-defined functions) are not and can slow down your job.
col("name") is useful for dynamic column references and is required when using columns in expressions (you cannot use a raw string in many function calls). It simply wraps the column name in a Column object and has no performance overhead.
lit(value) is the correct way to include constant values in DataFrame operations – avoid mixing Python literals into expressions without lit, as Spark might not recognize them as columns.
when() with .otherwise() is the go-to for conditional logic on DataFrames. It’s more efficient than writing the equivalent in pure Python. For multiple conditions, you can chain .when() calls (each subsequent .when acts like an "elif").
regexp_extract() is powerful for text processing. Keep your regex patterns efficient and consider using non-capturing groups (?: ) if you have complex regex but only need to extract certain parts. Always specify the correct group index: 0 for full match, 1 for first capturing group, etc.
concat() is handy for string concatenation. Remember to insert separators manually (using lit as shown) or use concat_ws if you want to specify a delimiter. Also, be cautious with nulls as they can result in null output – you may need to handle or fill nulls prior to concatenation.

With these built-in functions, you can cover a wide range of transformations without writing any custom code. In the next sections, we'll explore aggregate functions, array and string utilities, date functions, and finally how to write custom UDFs when needed.

Built-in functions

1. Built-in DataFrame Functions: col, lit, when, regexp_extract, concat

More Spark tutorials