Spark - select multiple columns from array object - scala

Question

I have a dataset with following schema
root
|-- id: string (nullable = true)
|-- name: string (nullable = true)
|-- subEntities: array (nullable = true)
| |-- element: struct (containsNull = true)
| | |-- status: string (nullable = true)
| | |-- subEntityId: long (nullable = true)
| | |-- subEntityName: string (nullable = true)
dataset.select($"id", $"name", $"subEntities.subEntityId", $"subEntities.subEntityName") put subEntityId and subEntityName into separate arrays. How to select multiple columns and put them into single array?

Answer 1

If working on Spark >= 2.4 you can use the transform function to generate an array which contains a subset of the original array's fields:
import org.apache.spark.sql.functions.expr
dataset.withColumn("newArray", expr("transform(subEntities, i -> struct(i.subEntityId, i.subEntityName))"))
// or with select
dataset.select(
$"id",
$"name",
expr("transform(subEntities, i -> struct(i.subEntityId, i.subEntityName))").as("newArray")
)

Answer 2

.withColumn("status",col("subEntities").getField("status"))
.withColumn("subEntityId",col("subEntities").getField("subEntityId"))
To extract value out of your array
Below is working example
import org.apache.spark.sql.functions._
object ExplodeArrauy {
def main(args: Array[String]): Unit = {
val spark = Constant.getSparkSess
import spark.implicits._
val df = List(bean57("1",Array(bean55("aaa",2),bean55("aaa1",21))),
bean57("2",Array(bean55("bbb",3),bean55("bbb3",31)))).toDF
df
.withColumn("status",col("subEntities").getField("status"))
.withColumn("subEntityId",col("subEntities").getField("subEntityId"))
.show()
}
}
case class bean57(id:String,subEntities:Array[bean55])
case class bean55(status: String,subEntityId:Long)