Pyspark Pivot с несколькими скоплениями

Question

У меня есть:

d = [(100,1,23,10),(100,2,45,11),(100,3,67,12),(100,4,78,13),(101,1,23,10),(101,2,45,13),(101,3,67,14),(101,4,78,15),(102,1,23,10),(102,2,45,11),(102,3,67,16),(102,4,78,18)]
mydf = spark.createDataFrame(d,['id','day','price','units'])
mydf.show()

+---+---+-----+-----+
| id|day|price|units|
+---+---+-----+-----+
|100|  1|   23|   10|
|100|  2|   45|   11|
|100|  3|   67|   12|
|100|  4|   78|   13|
|101|  1|   23|   10|
|101|  2|   45|   13|
|101|  3|   67|   14|
|101|  4|   78|   15|
|102|  1|   23|   10|
|102|  2|   45|   11|
|102|  3|   67|   16|
|102|  4|   78|   18|
+---+---+-----+-----+

Я хочу иметь:

+---+-----+----+----+----+----+
| id|  ?  |   1|   2|   3|   4|
+---+-----+----+----+----+----+
|100|units|10.0|11.0|12.0|13.0|
|101|units|10.0|13.0|14.0|15.0|
|102|units|10.0|11.0|16.0|18.0|
|100|price|23.0|45.0|67.0|78.0|
|101|price|23.0|45.0|67.0|78.0|
|102|price|23.0|45.0|67.0|78.0|
+---+-----+----+----+----+----+

Где каждое значение является средним для price и для units.

Я мог бы просто сделать два стержня, агрегировать по price и units, например:

mydf.groupby('id').pivot('day').agg(F.mean('units')).show() и mydf.groupby('id').pivot('day').agg(F.mean('price')).show()

, а затем присоединиться к нему.

Но я чувствую, что есть более разумный способ сделать это. Есть ли?

Answer 1

Вот один из способов (в основном плавить DF, а затем поворачиваться)

# First combine price and units into a map column
mydf = mydf.withColumn("price_units", F.create_map(F.lit("price"), "price", F.lit("units"), "units"))

# Now explode to get a melted dataframe
mydf = mydf.select("id", "day", F.explode("price_units").alias("name", "value"))

+---+---+-----+-----+
| id|day| name|value|
+---+---+-----+-----+
|100|  1|price|   23|
|100|  1|units|   10|
|100|  2|price|   45|
|100|  2|units|   11|
|100|  3|price|   67|
etc

# Then pivot
mydf.groupby("id", "name").pivot("day").agg(F.mean("value")).show()

+---+-----+----+----+----+----+
| id| name|   1|   2|   3|   4|
+---+-----+----+----+----+----+
|100|price|23.0|45.0|67.0|78.0|
|101|price|23.0|45.0|67.0|78.0|
|102|units|10.0|11.0|16.0|18.0|
|100|units|10.0|11.0|12.0|13.0|
|101|units|10.0|13.0|14.0|15.0|
|102|price|23.0|45.0|67.0|78.0|
+---+-----+----+----+----+----+