pyspark 中的情况总和

Question

我正在尝试将 hql 脚本转换为 pyspark。我正在努力如何在 groupby 子句之后的聚合中实现 case when 语句的总和。例如。

dataframe1 = dataframe0.groupby(col0).agg(
            SUM(f.when((col1 == 'ABC' | col2 == 'XYZ'), 1).otherwise(0)))

pyspark 中可以吗？我在执行此类语句时遇到错误。谢谢

Answer 1

您可以使用 withColumn 创建一个列，其中包含要求和的值，然后对其进行聚合。例如：

from pyspark.sql import functions as F, types as T

schema = T.StructType([
    T.StructField('key', T.IntegerType(), True),
    T.StructField('col1', T.StringType(), True),
    T.StructField('col2', T.StringType(), True)
])

data = [
    (1, 'ABC', 'DEF'),
    (1, 'DEF', 'XYZ'),
    (1, 'DEF', 'GHI')
]

rdd = sc.parallelize(data)
df = sqlContext.createDataFrame(rdd, schema)



result = df.withColumn('value', F.when((df.col1 == 'ABC') | (df.col2 == 'XYZ'), 1).otherwise(0)) \
           .groupBy('key') \
              .agg(F.sum('value').alias('sum'))

result.show(100, False)

打印出这个结果：

+---+---+
|key|sum|
+---+---+
|1  |2  |
+---+---+