CODEWITHSUNDEEP
pythonpyspark
Victor Hugo Borges
Victor Hugo Borges

Reputation: 413

Pyspark Pivot with multiple aggregations

I have:

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

I wanna have:

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

Where each value is the mean for the price and for the units.

I could just make two pivots, aggregate by price and units, like:

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

then join it.

But I feel like there's a smarter way to do this. Is there?

Upvotes: 1

Views: 4356

Answers (1)

ags29
ags29

Reputation: 2696

Here is one way (basically melt the DF, then pivot)

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

Upvotes: 1

Related Questions