How to unpack list of lists that are in string format?

user_1400724 0
Viewed 649

I have a PySpark dataframe with a column containing a StructField of stringtype which has a dynamic length list of lists.

df.schema: StructType(List(StructField(id,StringType,true),StructField(recs,StringType,true)))

|id     | recs |

|ABC|[66, [["AB", 10]]]
|XYZ|[66, [["XY", 10], ["YZ", 20]]]
|DEF|[66, [["DE", 10], ["EF", 20], ["FG", 30]]]

I am trying to flatten the lists to something like this

|id | like_id
|ABC|AB|
|XYZ|XY|
|XYZ|YZ|
|DEF|DE|
|DEF|EF|
|DEF|FG|

What did I try:

I tried using array expressions, it threw me an error since the recs is of StringType as expected

I was able to process this using json loads and itertools in pandas, but I need this processing to happen in spark as the dataframe is large ~30 million and the resulting would be 10 times.

df["recs"].apply(
        lambda x: [rec_id[0] for rec_id in json.loads(x)[1:][0]]
    )
    for i, row in df.iterrows():
        ....
pyspark apache-spark-sql apache-spark python
3 Answers

IIUC, you can split the String in recs column using a pattern , (?=\[\[)|\]$, find the 2nd element and then use from_json to retrieve the array of arrays:

from pyspark.sql import functions as F

df1 = df.withColumn('recs1', F.split('recs', ', (?=\[\[)|\]$')[1]) \
    .withColumn('recs2', F.from_json('recs1', 'array<array<string>>'))

Where: the split-pattern , (?=\[\[)|\]$ contains two sub-patterns:

  • comma followed by a SPACE which must be followed by two opening brackets , (?=\[\[)
  • closing-bracked at the end of String \]$

Result:

df1.show(truncate=False)
+---+------------------------------------------+------------------------------------+------------------------------+
|id |recs                                      |recs1                               |recs2                         |
+---+------------------------------------------+------------------------------------+------------------------------+
|ABC|[66, [["AB", 10]]]                        |[["AB", 10]]                        |[[AB, 10]]                    |
|XYZ|[66, [["XY", 10], ["YZ", 20]]]            |[["XY", 10], ["YZ", 20]]            |[[XY, 10], [YZ, 20]]          |
|DEF|[66, [["DE", 10], ["EF", 20], ["FG", 30]]]|[["DE", 10], ["EF", 20], ["FG", 30]]|[[DE, 10], [EF, 20], [FG, 30]]|
+---+------------------------------------------+------------------------------------+------------------------------+

Then you can use explode to get the desired result:

df1.selectExpr("id", "explode_outer(recs2) as recs") \
    .selectExpr("id", "recs[0] as like_id") \
    .show()
+---+-------+
| id|like_id|
+---+-------+
|ABC|     AB|
|XYZ|     XY|
|XYZ|     YZ|
|DEF|     DE|
|DEF|     EF|
|DEF|     FG|
+---+-------+

In short, we can write the above code into the following:

df_new = df.selectExpr("id", r"explode_outer(from_json(split(recs, ', (?=\\[\\[)|\\]$')[1], 'array<array<string>>')) as recs") \
    .selectExpr("id", "recs[0] as like_id")
user_95107290

A way to solve it's cleaning the column content and split:

Remember to import:

import pyspark.sql.functions as f

# Remove any non string character
df = df.withColumn('only_ids', f.trim(f.regexp_replace('recs', r'[^a-zA-Z\s:]', '')))

# Change blank spaces to commas
df = df.withColumn('clear_blank_spaces', f.regexp_replace('only_ids', r'\W+', ','))

# Split the values by comma
df = df.withColumn('like_ids', f.split('clear_blank_spaces', ','))

# Just for DEBUG
df.show(truncate=False)

# Explode like_ids to transform array to rows
df = df.select('id', f.explode('like_ids').alias('like_ids'))

# Final result
df.show(truncate=False)

First output:

+---+------------------------------------------+----------+------------------+------------+
|id |recs                                      |only_ids  |clear_blank_spaces|like_ids    |
+---+------------------------------------------+----------+------------------+------------+
|ABC|[66, [["AB", 10]]]                        |AB        |AB                |[AB]        |
|XYZ|[66, [["XY", 10], ["YZ", 20]]]            |XY  YZ    |XY,YZ             |[XY, YZ]    |
|DEF|[66, [["DE", 10], ["EF", 20], ["FG", 30]]]|DE  EF  FG|DE,EF,FG          |[DE, EF, FG]|
+---+------------------------------------------+----------+------------------+------------+

Second output:

+---+--------+
|id |like_ids|
+---+--------+
|ABC|AB      |
|XYZ|XY      |
|XYZ|YZ      |
|DEF|DE      |
|DEF|EF      |
|DEF|FG      |
+---+--------+
user_60802760

If your data really looks this clean, then you can just split by " and get the entries with an odd-numbered index in the resulting array.

import pyspark.sql.functions as F

df2 = df.select(
    'id',
    F.posexplode(F.split('recs', '"')).alias('pos', 'like_id')
).filter('cast(pos % 2 as boolean)').drop('pos')

df2.show()
+---+-------+
| id|like_id|
+---+-------+
|ABC|     AB|
|XYZ|     XY|
|XYZ|     YZ|
|DEF|     DE|
|DEF|     EF|
|DEF|     FG|
+---+-------+
user_141657300
Related