Dataframe regexp_extract values from string like array

Viewed 1595

My DataFrame looks like as follows:

StudentID          Marks
100                ["20", "25.5", "40.23", "50"]
200                ["30", "20", "25", "40"]
300                ["20", "25", "50", "35"]

I need to extract the marks within the array and create a new DataFrame. However, am unable to extract beyond the second value in the DF (don't know how to select all marks via the regex ([0-9]+)(?:\.[0-9]+){3}.

df1.select(regexp_extract('StudentID', '(\w+)(,)', 1).alias("C1"), 
             regexp_extract('Marks', '([0-9]+)(?:\.[0-9]+){3}', 0).alias("C2"))

Ultimately, need to create a new DataFrame with below format:

StudentID  C1    C2    C3     C4
100        20    25.5  40.23  50
200        30    20    25     40
300        20    25    50     35

Thank you in advance.

2 Answers

You can split the string, and, then use element_at to pull substrings into separate columns:

df1.withColumn("marks_array", split( regexp_replace(col("Marks"), "\\[|\\]|\"", ""), ",")  )
      .withColumn("C1", element_at(col("marks_array"), 1))
      .withColumn("C2", element_at(col("marks_array"), 2))
      .withColumn("C3", element_at(col("marks_array"), 3))
      .withColumn("C4", element_at(col("marks_array"), 4))
      .drop("marks_array", "Marks")
      .show(false)
+---------+---+-----+------+---+
|StudentID|C1 |C2   |C3    |C4 |
+---------+---+-----+------+---+
|100      |20 | 25.5| 40.23| 50|
|200      |30 | 20  | 25   | 40|
|300      |20 | 25  | 50   | 35|
+---------+---+-----+------+---+

You can use from_json to convert the string column Marks to an array of strings. Then get elements of the array to create each column.

However, if you don't know the size of the arrays, you can use transform function to convert it to a map then explode the map and pivot to get the desired output.

transform_expr = """transform(from_json(Marks, 'array<string>'), 
                              (x, i) -> struct(concat('C', i+1), x)
                             )
                 """

df.select(col("*"), explode(map_from_entries(expr(transform_expr)))) \
  .groupBy("StudentID").pivot("key").agg(first("value")) \
  .show()

#+---------+---+----+-----+---+
#|StudentID|C1 |C2  |C3   |C4 |
#+---------+---+----+-----+---+
#|100      |20 |25.5|40.23|50 |
#|200      |30 |20  |25   |40 |
#|300      |20 |25  |50   |35 |
#+---------+---+----+-----+---+

Note: transfrom and map_from_entries functions are only available for Spark 2.4+

Related