Pfthb

Looking to explode a nested array w/ Spark into batches. The column below is a nested array from an XML files. Now attempting to write the time series data into batches in order to write over to a NoSQL database. For example:

+-------+-----------------------+
| ID | Example |
+-------+-----------------------+
| A| [[1,2],[3,4],[5,6]] |
+-------+-----------------------+

Output with batches of size 2

+-------+-----------------------+
| ID | Example |
+-------+-----------------------+
| A| [[1,2],[3,4]] |
+-------+-----------------------+
| A| [[5,6]] |
+-------+-----------------------+

For Spark v 2.1+

You can take advantage of pyspark.sql.functions.posexplode() to explode your column along with the index it appears in your array and then divide the resultant position by n to create groups.

For example, here is the output of using posexplode() on your DataFrame:

import pyspark.sql.functions as f
df.select('ID', f.posexplode('Example')).show()
#+---+---+------+
#| ID|pos| col|
#+---+---+------+
#| A| 0|[1, 2]|
#| A| 1|[3, 4]|
#| A| 2|[5, 6]|
#+---+---+------+

Notice that we get two columns: pos and col instead of just one. Since we want groups of n, we can simply divide the pos by n and take the floor to get groups.

n = 2
df.select('ID', f.posexplode('Example'))
 .withColumn("group", f.floor(f.col("pos")/n))
 .show(truncate=False)
#+---+---+------+-----+
#|ID |pos|col |group|
#+---+---+------+-----+
#|A |0 |[1, 2]|0 |
#|A |1 |[3, 4]|0 |
#|A |2 |[5, 6]|1 |
#+---+---+------+-----+

Now group by the "ID" and the "group" and use pyspark.sql.functions.collect_list() to get your desired output.

df.select('ID', f.posexplode('Example'))
 .withColumn("group", f.floor(f.col("pos")/n))
 .groupBy("ID", "group")
 .agg(f.collect_list("col").alias("Example"))
 .sort("group")
 .drop("group")
 .show(truncate=False)
#+---+----------------------------------------+
#|ID |Example |
#+---+----------------------------------------+
#|A |[WrappedArray(1, 2), WrappedArray(3, 4)]|
#|A |[WrappedArray(5, 6)] |
#+---+----------------------------------------+

You'll see that I also sorted by the "group" column and dropped it, but this is optional depending on your needs.

For Older Versions of Spark

There are some other methods for Spark versions below 2.1. All of these methods produce the same output as above.

1. Using udf

You can use a udf to break your array into groups. For example:

def get_groups(array, n):
 return filter(lambda x: x, [array[i*n:(i+1)*n] for i in range(len(array))])

get_groups_of_2 = f.udf(
 lambda x: get_groups(x, 2),
 ArrayType(ArrayType(ArrayType(IntegerType())))
)

df.select("ID", f.explode(get_groups_of_2("Example")).alias("Example"))
 .show(truncate=False)

The get_groups() function will take an array and return an array of groups of n elements.

2. Using rdd

Another option is to serialize to rdd and use the get_groups() function inside of a call to map(). Then convert back to a DataFrame. You'll have to specify the schema for this conversion to work properly.

n = 2

schema = StructType(
 [
 StructField("ID", StringType()),
 StructField("Example", ArrayType(ArrayType(ArrayType(IntegerType()))))
 ]
)

df.rdd.map(lambda x: (x["ID"], get_groups(x["Example"], n=n)))
 .toDF(schema)
 .select("ID", f.explode("Example").alias("Example"))
 .show(truncate=False)