Is there an efficent method similar to sql window functions in PySpark? - performance

Question

I am dealing with a huge dataframe containing 3 columns and 5 Bil. rows. The size of the data is 360 GB. For the analysis of the Data I am using the following set-up:
-Jupyternotebooks running on a AWS r4.16xlarge
-PySpark Kernel
The table called customer_sales looks similar to the following example:
+--------------------+----------+-------+
| business_unit_id | customer | sales |
+--------------------+----------+-------+
| 1 + a + 5000 +
| 1 + b + 2000 +
| 1 + c + 3000 +
| 1 + d + 5000 +
| 2 + f + 600 +
| 2 + c + 7000 +
| 3 + j + 200 +
| 3 + k + 800 +
| 3 + c + 4500 +
Now I want to get for each business_unit_id the customer with the highest sales. If there is a draw in sales between several customer I want to get them all. The information should be stored in a table called best_customers_for_each_unit. So for the above illustrated example the table best_customers_for_each_unit looks as follows:
+--------------------+----------+-------+
| business_unit_id | customer | sales |
+--------------------+----------+-------+
| 1 + a + 5000 +
| 1 + d + 5000 +
| 2 + c + 7000 +
| 3 + c + 4500 +
In a second step I want to count how often it happens that a customer is the one with the highest sales in a specific business_unit_id. The output of this query will be:
+----------+-------+
| customer | count |
+----------+-------+
| a + 1 +
| b + 1 +
| c + 2 +
For the first query I used spark.sql with window functions. The used query looks as follows:
best_customers_for_each_unit = spark.sql("""
SELECT
business_unit_id,
customer,
sales
FROM (
SELECT
business_unit_id,
customer,
sales,
dense_rank() OVER (PARTITION BY business_unit_id ORDER BY sales DESC)as rank
FROM customer_sales) tmp
WHERE
rank =1
""")
For the second query I used the following PySpark snipped:
best_customers_for_each_unit.groupBy("customer").count()
My queries do actually work, but it takes ages to only process a minor part of the data. So do you know any efficient way to do such queries with PySpark?
Regards