在 Scala 中转换 spark 数据帧的模式

Question

Input data-frame:

{ 
  "C_1" : "A",
  "C_2" : "B",
  "C_3" : [ 
            {
              "ID" : "ID1",
              "C3_C2" : "V1",
              "C3_C3" : "V2"
            },
            {
              "ID" : "ID2",
              "C3_C2" : "V3",
              "C3_C3" : "V4"
            },
            {
              "ID" : "ID3",
              "C3_C2" : "V4",
              "C3_C3" : "V5"
            },
            ..
          ]
}

Desired Output:

{ 
  "C_1" : "A",
  "C_2" : "B",
  "ID1" : {
              "C3_C2" : "V2",
              "C3_C3" : "V3"
          },
  "ID2" : {
              "C3_C2" : "V2",
              "C3_C3" : "V3"
          },
  "ID3" : {
              "C3_C2" : "V4",
              "C3_C3" : "V5"
          },
  ..
}

C_3 is an array of n structs with each item having a unique ID . The new data-frame is expected to convert the n structs in C_3 into separate columns and name the columns as per the value of ID .

I am new to Spark & Scala. Any thought on how of achieve this transformation will be very helpful.

Thanks!

Answer 1

You can explode the structs, then pivot by ID:

val df2 = df.selectExpr("C_1","C_2","inline(C_3)")
            .groupBy("C_1","C_2")
            .pivot("ID")
            .agg(first(struct("C3_C2","C3_C3")))

df2.show
+---+---+--------+--------+--------+
|C_1|C_2|     ID1|     ID2|     ID3|
+---+---+--------+--------+--------+
|  A|  B|[V1, V2]|[V3, V4]|[V4, V5]|
+---+---+--------+--------+--------+

df2.printSchema
root
 |-- C_1: string (nullable = true)
 |-- C_2: string (nullable = true)
 |-- ID1: struct (nullable = true)
 |    |-- C3_C2: string (nullable = true)
 |    |-- C3_C3: string (nullable = true)
 |-- ID2: struct (nullable = true)
 |    |-- C3_C2: string (nullable = true)
 |    |-- C3_C3: string (nullable = true)
 |-- ID3: struct (nullable = true)
 |    |-- C3_C2: string (nullable = true)
 |    |-- C3_C3: string (nullable = true)

Answer 2

[posting my hacky solution for reference].
@mck's answer is probably the neat way to do it but it wasn't sufficient for my use-case. My data-frame had a lot of columns and using all of them on group-by then pivot was an expensive operation.
For my use-case, the IDs in C_3 were unique and known values, hence this is the assumption made in this solution. I was able to achieve the transformation as follows:

case class C3_Interm_Struct(C3_C2: String, C3_C3: String)
case class C3_Out_Struct(ID1: C3_Interm_Struct, ID2: C3_Interm_Struct) //Manually add fields as required

val custom_flat_func = udf((a: Seq[Row]) =>
    {
        var _id1: C3_Interm_Struct= null;
        var _id2: C3_Interm_Struct = null;
        for(item<-a)
        {
            val intermData = C3_Interm_Struct(item(1).toString, item(2).toString)

            if(item(0).equals("ID1")) {
                _id1 = intermData
            }
            else if(item(0).equals("ID2")) {
                _id2 = intermData
            }
            else if()//Manual expansion
                ..
        }

        Seq(C3_Out_Struct(_id1, _id2)) //Return type has to be Seq
    }
)

val flatDf = df.withColumn("C_3", custom_flat_func($"C_3")).selectExpr("C_1", "C_2", "inline(C_3)") //Expand the Seq which has only 1 Row
flatDf.first.prettyJson

Output:

{ 
  "C_1" : "A",
  "C_2" : "B",
  "ID1" : {
              "C3_C2" : "V2",
              "C3_C3" : "V3"
          },
  "ID2" : {
              "C3_C2" : "V2",
              "C3_C3" : "V3"
          }
}

udf s are generally slow but this is much faster than pivot with group-by .

[There may be more efficient solutions possible, I am not aware of them at the time of writing this]