SQL 查询到 Py(spark)

Question

我有以下 SQL 查询，我想将其转换为 pyspark。 我想取两列 pp 和 gender 并在 pyspark 中执行以下操作

%sql             
SELECT pp
,       SUM(CASE WHEN Gender = 'M'
               THEN 1.0 ELSE 0.0 END) /
        COUNT(1) AS gender_score
,       count(1) AS total
FROM df 
WHERE gender in ('M', 'F')
GROUP BY pp
HAVING count(1) > 100

我试过这个

from pyspark.sql.functions import expr 
test_df = df.withColumn("gender_score", expr("CASE WHEN Gender == 'M' THEN  1 WHEN Gender == 'F' THEN  0 ELSE 'other' END AS male_score"))

But I don't know how I can select just 2 columns in pyspark and replicate the above-mentioned SQL query to pyspark

提前致谢

Answer 1

如果您想要 python 解决方案：

import pyspark.sql.functions as F

def df_func(df, grp):
    test_df = df.filter( 
        F.col('gender').isin(['M','F'])
        ).groupBy(grp).agg(
            F.mean(
                F.when(F.col('Gender') == F.lit('M'), 1.0).otherwise(0.0)
            ).alias('gender_score'),
            F.count(F.lit(1)).alias('total')
        ).filter('total > 100')
    return test_df

result = df_func(df, 'pp')

您可以通过创建 dataframe 的临时视图来使用 Spark SQL：

df.createOrReplaceTempView('df')
query = """
SELECT pp
,       SUM(CASE WHEN Gender = 'M'
               THEN 1.0 ELSE 0.0 END) /
        COUNT(1) AS gender_score
,       count(1) AS total
FROM df 
WHERE gender in ('M', 'F')
GROUP BY pp
HAVING count(1) > 100
"""
test_df = spark.sql(query)

Answer 2

groupBy agg() 函数采用多个表达式。 所以你可以做一个总和/计数表达式。

val df = spark.sql(""" with t1(
select 10 pp, 'M' gender union all
select 10 pp, 'F' gender union all
select 10 pp, 'M' gender union all
select 40 pp, 'F' gender 
) select * from t1
""")

df.show(false)

+---+------+
|pp |gender|
+---+------+
|10 |M     |
|10 |F     |
|10 |M     |
|40 |F     |
+---+------+

df.groupBy("pp").agg (
    (sum(expr("CASE WHEN Gender = 'M' THEN  1 ELSE 0 END "))/count(lit(1)) ).alias("gender_score"),
    count(lit(1)).alias("total")
).show(false)

+---+------------------+-----+
|pp |gender_score      |total|
+---+------------------+-----+
|40 |0.0               |1    |
|10 |0.6666666666666666|3    |
+---+------------------+-----+

之后，您可以使用另一个过滤器来实现“拥有”。