Adding new column using other existing columns Spark/Scala

Question

I want to add a new column using other existing columns. This must be released on conditions. This is an example of my Dataframe :

val data = Seq(("WHT20177", "CTHT WO/MTR# : WHT20212/BTI0426; WHT20177/BTH0393"),
                ("WHT55637", "CTHT WO/MTR# : WHT50747/BTI2699; WHT55637/SQL1239"))

val dataFrame = data.toDF("prev_wo", "ref_wo")

+--------+-------------------------------------------------+
|prev_wo |ref_wo                                           |
+--------+-------------------------------------------------+
|WHT20177|CTHT WO/MTR# : WHT20212/BTI0426; WHT20177/BTH0393|
|WHT55637|CTHT WO/MTR# : WHT50747/BTI2699; WHT55637/SQL1239|
+--------+-------------------------------------------------+

The column "ref_wo" must contain "prev_wo", in that case only, I must put the following element in a new column which I shall name "col1".

For the first line, the value to extract is "BTH0393", for the second line, the value to extract is "SQL1239" I am trying this in Spark Scala using two different methods. The first one reacts only to the first line, and the second one reacts only to the second line.

First method :

def addNewColumn(df: DataFrame): DataFrame = {

  val prev_wo = dataFrame.select("prev_wo").collectAsList().get(0).mkString(",")

  val regex_extract = ("(?<=" + prev_wo + "\/)(.{7})").r

  df
    .withColumn("col1",
      when($"ref_wo".contains(col("prev_wo")),
        regexp_extract(col("ref_wo"), regex_extract.toString(), 1))
        .otherwise(null)
    )

}

val new_dataFrame = dataFrame
  .transform(addNewColumn)

OUTPUT :

+--------+-------------------------------------------------+-------+
|prev_wo |ref_wo                                           |col1   |
+--------+-------------------------------------------------+-------+
|WHT20177|CTHT WO/MTR# : WHT20212/BTI0426; WHT20177/BTH0393|BTH0393|
|WHT55637|CTHT WO/MTR# : WHT50747/BTI2699; WHT55637/SQL1239|       |
+--------+-------------------------------------------------+-------+

Second method :

def addColumn(df: DataFrame): DataFrame = {
  var out = df

  df.collect().foreach(row => {

    val prev_wo = row.getValuesMap(Seq("prev_wo")).get("prev_wo").getOrElse("")

    val regex_extract = ("(?<=" + prev_wo + "\/)(.{7})").r

    out = out
      .withColumn("col1",
        when($"ref_wo".contains(col("prev_wo")),
          regexp_extract(col("ref_wo"), regex_extract.toString(), 1))
          .otherwise(null)
      )
  })

  out
}

val new_dataFrame = dataFrame
  .transform(addColumn)

OUTPUT

+--------+-------------------------------------------------+-------+
|prev_wo |ref_wo                                           |col1   |
+--------+-------------------------------------------------+-------+
|WHT20177|CTHT WO/MTR# : WHT20212/BTI0426; WHT20177/BTH0393|       |
|WHT55637|CTHT WO/MTR# : WHT50747/BTI2699; WHT55637/SQL1239|SQL1239|
+--------+-------------------------------------------------+-------+

user_14165730 · Answer

You can use regexp_extract with a pattern dynamically generated from prev_wo:

dataFrame.withColumn("col1", expr("regexp_extract(ref_wo, concat(prev_wo, '/(.{7})'), 1)")).show(false)
+--------+-------------------------------------------------+-------+
|prev_wo |ref_wo                                           |col1   |
+--------+-------------------------------------------------+-------+
|WHT20177|CTHT WO/MTR# : WHT20212/BTI0426; WHT20177/BTH0393|BTH0393|
|WHT55637|CTHT WO/MTR# : WHT50747/BTI2699; WHT55637/SQL1239|SQL1239|
+--------+-------------------------------------------------+-------+

Adding new column using other existing columns Spark/Scala

1 Answers