Introduction

The following are some of the most used array functions available in Spark SQL. These functions enable various operations on arrays within Spark SQL DataFrame columns, facilitating array manipulation and analysis.

Note: Some of the following functions are supported in Apache Spark 3.5.0; if any function doesn’t work for you, you are probably using a lower Spark version.

| FUNCTION NAME | SQL ARRAY FUNCTIONS DESCRIPTION | | ------------------ | --------------------------------------------------------------------------------------------------------------------------------------------- | | array() | Creates a new array from the given input columns. | | array_contains() | Returns true if the array contains the given value. | | array_append() | Appends the element to the source array and returns an array containing all elements. The new element/column is added at the end of the array. | | array_insert() | Returns an array after adding the element at the specified position. | | arrays_overlap() | Returns true if a1 and a2 have at least one non-null element in common. | | array_distinct() | Removes duplicate values from the array. | | array_except() | Computes the set difference of the two arrays. | | array_intersect() | Computes the set intersection of the two arrays. | | array_join() | Concatenates the elements of the given array using a delimiter. | | array_max() | Returns the maximum value in the array. | | array_min() | Returns the minimum value in the array. | | array_size() | Returns the size of an array. | | array_position() | Returns the (1-based) index of the first occurrence of the given value in the array. | | array_remove() | Remove all null elements from the input array. | | array_compact() | Remove all null elements from the input array. | | array_prepend() | Returns an array after appending the array at the beginning. | | array_repeat() | Repeat an array n times. | | array_sort() | Sorts the input array in ascending order according to the natural ordering of the array elements. | | array_union() | Computes the union of the two arrays. | | element_at() | Returns element of the array at the given index. | | slice() | Returns sliced array. |

Array Function Examples

Before we use these functions, let’s create a DataFrame with a few array columns. I will use this DataFrame for all my examples below.

from pyspark.sql import SparkSession
from pyspark.sql.functions import array_contains, array_sort, array_join, col, array_union, array_size, element_at, expr

# Create SparkSession
spark = SparkSession.builder \
    .master("local[1]") \
    .appName("AIStudyLabs.com") \
    .getOrCreate()

# Create DataFrame
data = [
    ("James,,Smith", ["Java", "Scala", "C++"], ["Python", "PHP"], "CA"),
    ("Michael,Rose,", ["Spark", "Java", "C++"], ["AWS", "Scala", "Scala"], "NJ"),
    ("Robert,,Williams", ["AWS", "Scala"], None, "NV")
]

columns = ["name", "languages_school", "additional_language", "state"]
df = spark.createDataFrame(data, schema=columns)
df.show(truncate=False)

+-----------------+-------------------+---------------------+-----+
|name             |languages_school   |additional_language  |state|
+-----------------+-------------------+---------------------+-----+
|James,,Smith     |[Java, Scala, C++] |[Python, PHP]        |CA   |
|Michael,Rose,    |[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |
|Robert,,Williams |[AWS, Scala]       |NULL                 |NV   |
+-----------------+-------------------+---------------------+-----+

array_contains()

Function array_contains() in Spark returns true if the array contains the specified value. Returns null value if the array itself is null; otherwise, it returns false. This is primarily used to filter rows from the DataFrame.

The following example returns the DataFrame df3 by including only rows where the list column "languages_school" contains the value "Java".

# Using array_contains
df3 = df.filter(array_contains(col("languages_school"), "Java"))
df3.show(truncate=False)

+-------------+-------------------+---------------------+-----+
|name         |languages_school   |additional_language  |state|
+-------------+-------------------+---------------------+-----+
|James,,Smith |[Java, Scala, C++] |[Python, PHP]        |CA   |
|Michael,Rose,|[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |
+-------------+-------------------+---------------------+-----+

array_sort()

array_sort() function arranges the input array in ascending order. The elements within the array must be sortable. When you have NaN values in an array, the following applies:

• For double/float type, NaN is considered greater than any non-NaN elements.
• Null elements are positioned at the end of the resulting array.

Example

# Using array_sort
df2 = df.withColumn("array_sort", array_sort(col("languages_school")))
df2.show(truncate=False)

+----------------+-------------------+---------------------+-----+-------------------+
|name            |languages_school   |additional_language  |state|array_sort         |
+----------------+-------------------+---------------------+-----+-------------------+
|James,,Smith    |[Java, Scala, C++] |[Python, PHP]        |CA   |[C++, Java, Scala] |
|Michael,Rose,   |[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |[C++, Java, Spark] |
|Robert,,Williams|[AWS, Scala]       |NULL                 |NV   |[AWS, Scala]       |
+----------------+-------------------+---------------------+-----+-------------------+

From the code df2 = df.withColumn("array_sort", array_sort(col("languages_school"))), a new DataFrame df2 is created based on the DataFrame df, where a new column named "array_sort" is added. This new column contains the original values of the "languages_school" column sorted in ascending order within arrays.

array_join()

This function combines all elements of the list/array column using the delimiter. When the nullReplacement parameter is used, the array containing null values is replaced with 'nullReplacement'.

Example

# Using array_join
df4 = df.withColumn("array_join", array_join(col("languages_school"), "|"))
df4.show(truncate=False)

+----------------+-------------------+---------------------+-----+--------------+
|name            |languages_school   |additional_language  |state|array_join    |
+----------------+-------------------+---------------------+-----+--------------+
|James,,Smith    |[Java, Scala, C++] |[Python, PHP]        |CA   |Java|Scala|C++|
|Michael,Rose,   |[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |Spark|Java|C++|
|Robert,,Williams|[AWS, Scala]       |NULL                 |NV   |AWS|Scala     |
+----------------+-------------------+---------------------+-----+--------------+

array_append()

In PySpark, to append elements to an array, you'll use the withColumn combined with concat. Below is how you achieve this operation:

from pyspark.sql.functions import concat

# Using array_append (with concat in PySpark)
df5 = df.withColumn("array_append", expr("concat(languages_school, array('PHP'))"))
df5.show(truncate=False)

+----------------+-------------------+---------------------+-----+----------------------+
|name            |languages_school   |additional_language  |state|array_append          |
+----------------+-------------------+---------------------+-----+----------------------+
|James,,Smith    |[Java, Scala, C++] |[Python, PHP]        |CA   |[Java, Scala, C++, PHP]|
|Michael,Rose,   |[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |[Spark, Java, C++, PHP]|
|Robert,,Williams|[AWS, Scala]       |NULL                 |NV   |[AWS, Scala, PHP]     |
+----------------+-------------------+---------------------+-----+----------------------+

array_union()

Similarly, the array_union function combines the elements from both columns, removing duplicates, and returns an array that contains all unique elements from both input arrays. If there are any null arrays or columns, they are ignored in the union operation.

# Using array_union
df6 = df.withColumn("array_union", array_union(col("languages_school"), col("additional_language")))
df6.show(truncate=False)

+----------------+-------------------+---------------------+-----+-------------------------------+
|name            |languages_school   |additional_language  |state|array_union                    |
+----------------+-------------------+---------------------+-----+-------------------------------+
|James,,Smith    |[Java, Scala, C++] |[Python, PHP]        |CA   |[Java, Scala, C++, Python, PHP]|
|Michael,Rose,   |[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |[Spark, Java, C++, AWS, Scala] |
|Robert,,Williams|[AWS, Scala]       |NULL                 |NV   |[AWS, Scala]                   |
+----------------+-------------------+---------------------+-----+-------------------------------+

array_size()

The array_size() returns the total number of elements in the array column. If your input array column is null, it returns null.

# Using array_size
df7 = df.withColumn("array_size", array_size(col("languages_school")))
df7.show(truncate=False)

+----------------+-------------------+---------------------+-----+----------+
|name            |languages_school   |additional_language  |state|array_size|
+----------------+-------------------+---------------------+-----+----------+
|James,,Smith    |[Java, Scala, C++] |[Python, PHP]        |CA   |3         |
|Michael,Rose,   |[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |3         |
|Robert,,Williams|[AWS, Scala]       |NULL                 |NV   |2         |
+----------------+-------------------+---------------------+-----+----------+

This returns a new DataFrame with a column containing the array size of the column languages_school.

element_at()

element_at() returns an element of the array at the specified index.

# Using element_at
df8 = df.withColumn("element_at", element_at(col("languages_school"), 2))
df8.show(truncate=False)

+----------------+-------------------+---------------------+-----+----------+
|name            |languages_school   |additional_language  |state|element_at|
+----------------+-------------------+---------------------+-----+----------+
|James,,Smith    |[Java, Scala, C++] |[Python, PHP]        |CA   |Scala     |
|Michael,Rose,   |[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |Java      |
|Robert,,Williams|[AWS, Scala]       |NULL                 |NV   |Scala     |
+----------------+-------------------+---------------------+-----+----------+

array_position()

Use array_position() to find the position of the first occurrence of the value in the given array. It returns null if either of the arguments is null.

Note that the position is not zero-based but 1-based index. Returns 0 if the value could not be found in the array.

# Using array_position
df9 = df.withColumn("array_position", expr("array_position(languages_school, 'Scala')"))
df9.show(truncate=False)

+----------------+-------------------+---------------------+-----+--------------+
|name            |languages_school   |additional_language  |state|array_position|
+----------------+-------------------+---------------------+-----+--------------+
|James,,Smith    |[Java, Scala, C++] |[Python, PHP]        |CA   |2             |
|Michael,Rose,   |[Spark, Java, C++] |[AWS, Scala, Scala]  |NJ   |3             |
|Robert,,Williams|[AWS, Scala]       |NULL                 |NV   |2             |
+----------------+-------------------+---------------------+-----+--------------+