根据条件计算不同的列值 pyspark

Question

I have a column with 2 possible values: 'users' or 'not_users'

What I want to do is to countDistinct values when those values are 'users'

This is the code I'm using:

output = (df
           .withColumn('week', f.expr('DATE_FORMAT(DATE_SUB(registration_date, 1), "Y-ww")'))
           .groupby('week') 
           .agg(f.countDistinct('customer_id').alias('count_total_users'),
                f.countDistinct('vegetables_customers').alias('count_vegetable_users')   
     
               )
         
          )

display(output)

this is the output (not desired):

Week        count_total_users      count_vegetable_users
2020-40            2345                        2
2020-41            5678                        2
2020-42            3345                        2
2020-43            5689                        2

desired output:

Week        count_total_users      count_vegetable_users
2020-40            2345                        457
2020-41            5678                        1987
2020-42            3345                        2308
2020-43            5689                        4000

This desired output should be the count distinct for 'users' values inside the column it belongs to.

Any clue?

Answer 1

Is df2 the result that you want?

df.show()
+----+-----------+--------------------+
|week|customer_id|vegetables_customers|
+----+-----------+--------------------+
|   1|          1|               users|
|   1|          2|           not_users|
|   1|          3|               users|
|   2|          1|           not_users|
|   2|          2|           not_users|
|   2|          3|               users|
+----+-----------+--------------------+

df2 = df.groupBy('week').agg(
    F.countDistinct('customer_id').alias('count_total_users'),
    F.countDistinct(
        F.when(
            F.col('vegetables_customers') == 'users', 
            F.col('customer_id')
        )
    ).alias('count_vegetable_users')
)

df2.show()
+----+-----------------+---------------------+
|week|count_total_users|count_vegetable_users|
+----+-----------------+---------------------+
|   1|                3|                    2|
|   2|                3|                    1|
+----+-----------------+---------------------+