Adding new column using other existing columns Spark/Scala - 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|
+--------+-------------------------------------------------+-------+

Answer 1

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|
+--------+-------------------------------------------------+-------+