How to extract the elements of an Array[String] that starts with particular value in Scala?

Viewed 1613

I have a Scala dataframe which has this schema:

filter_msg.printSchema()
root
|-- value: array (nullable = true)
|    |-- element: string (containsNull = true)

and data sample:

|[SD:GK, 3.16.0, OS:Linux, (x86_64), AID:176]|

I want to extract the values from this array string where element starts with say SD and get the value and then If its OS get the value. The problem is the position in the array string is not always same it keeps varying so I cant use

filter_msg.select($"value".getItem(1).as("SD"))

The output should give me a dataframe:

Output=GK | Linux
Output.printSchema()
String,String
3 Answers

You can convert to rdd and extract the values as below

// If you can conferm the data are always in same order
filter_msg.rdd.map(_.getAs[mutable.WrappedArray[String]](0))
  .map(row => {
    val sd = row(0).split(":").tail.head
    val os = row(2).split(":").tail.head
    (sd, os)
  } )
  .toDF("sd", "os")

Or you can use as @SleightX mentioned

filter_msg.rdd.map(_.getAs[mutable.WrappedArray[String]](0))
  .map(row => {
    val sd = row.filter(_.startsWith("SD:")).head.split(":").tail.head
    val os = row.filter(_.startsWith("OS:")).head.split(":").tail.head
    (sd, os)
  } )
  .toDF("sd", "os")

Output:

+---+-----+
|sd |os   |
+---+-----+
|GK |Linux|
+---+-----+

Here is another approach using regex and the regexp_extract function:

import org.apache.spark.sql.functions.{concat_ws, regexp_extract}

val df = Seq(
  Seq("SD:GK", "3.16.0", "OS:Linux", "(x86_64)", "AID:176")
).toDF

df.withColumn("to_str", concat_ws(",", $"value")) //concatenate array items into one string i.e: SD:GK,3.16.0,OS:Linux,(x86_64),AID:176
  .select(
    regexp_extract($"to_str", "SD:(\\w+),", 1) as "SD", //extract SD 
    regexp_extract($"to_str", "OS:(\\w+),", 1) as "OS"  //extract OS
  ).show(false)

// Output
// +---+-----+
// |SD |OS   |
// +---+-----+
// |GK |Linux|
// +---+-----+

UDF can be used:

val df = Seq(Array("SD:GK", "3.16.0", "OS:Linux", "(x86_64)", "AID:176")).toDF("value")

val extractArrayValues = (prefix: String, values: Seq[String]) =>
  values.filter(_.startsWith(prefix + ":")).map(_.split(":")(1)).headOption

val extractUDF = udf(extractArrayValues)
val result = df.select(
  extractUDF(lit("SD"), $"value").alias("SD"),
  extractUDF(lit("OS"), $"value").alias("OS")
)

Result is:

+---+-----+
|SD |OS   |
+---+-----+
|GK |Linux|
+---+-----+
Related