Get Min and Max from values of another column after a Groupby in PySpark

user_10037461 0
Viewed 2281

I´m trying to get the min and max values from a column´s values after doing a groupby in two other columns in pyspark.

The dataset looks like:

|      country      |     company    |   value   |
|-------------------|----------------|-----------|
|          arg      |       hh       |   3       |
|          arg      |       hh       |   2       |
|          arg      |       go       |   4       |
|          arg      |       go       |   3       |
|          bra      |       go       |   1       |
|          bra      |       go       |   2       |
|          bra      |       hh       |   3       |
|          bra      |       hh       |   2       |

My current implementation is this one:

from pyspark.sql.functions import col, first, min, max

new_df = df.groupBy("country", "company").agg(first("value").alias("value"), 
                                              min("value").alias("min_value"), 
                                              max("value").alias("max_value")
                                             )

But the result I´m getting is not correct, since I get this:

|      country      |     company    |   value   |   min_value   |   max_value   |
|-------------------|----------------|-----------|---------------|---------------|
|          arg      |       hh       |   3       |   3           |   3           |
|          arg      |       hh       |   2       |   2           |   2           |
|          arg      |       go       |   4       |   4           |   4           |
|          arg      |       go       |   3       |   3           |   3           |
|          bra      |       go       |   1       |   1           |   1           |
|          bra      |       go       |   2       |   2           |   2           |
|          bra      |       hh       |   3       |   3           |   3           |
|          bra      |       hh       |   2       |   2           |   2           |

And I wish to get something like:

|      country      |     company    |   value   |   min_value   |   max_value   |
|-------------------|----------------|-----------|---------------|---------------|
|          arg      |       hh       |   3       |   2           |   3           |
|          arg      |       hh       |   2       |   2           |   3           |
|          arg      |       go       |   4       |   3           |   4           |
|          arg      |       go       |   3       |   3           |   4           |
|          bra      |       go       |   1       |   1           |   2           |
|          bra      |       go       |   2       |   1           |   2           |
|          bra      |       hh       |   3       |   2           |   3           |
|          bra      |       hh       |   2       |   2           |   3           |
pyspark dataframe python
1 Answers

Do a join with the grouped dataframe

from pyspark.sql.functions import min, max

df.join(df.groupby('country', 'company').agg(min('value').alias('min_value'), 
                                             max('value').alias('max_value')),
                                             on=['country', 'company'])

which is the (unordered) result you are looking for

+-------+-------+-----+---------+---------+
|country|company|value|min_value|max_value|
+-------+-------+-----+---------+---------+
|    bra|     go|    1|        1|        2|
|    bra|     go|    2|        1|        2|
|    bra|     hh|    3|        2|        3|
|    bra|     hh|    2|        2|        3|
|    arg|     hh|    3|        2|        3|
|    arg|     hh|    2|        2|        3|
|    arg|     go|    4|        3|        4|
|    arg|     go|    3|        3|        4|
+-------+-------+-----+---------+---------+
user_95343900
Related