如何计算pyspark数据帧中值的条件概率？

debugcn 发表于 Dev

雅西

我想在不收集的情况下通过pyspark中的列类型值来计算评级列中的评级（'A'，'B'，'C'）的条件概率。

输入：

    company     model    rating   type
0   ford       mustang     A      coupe
1   chevy      camaro      B      coupe
2   ford       fiesta      C      sedan
3   ford       focus       A      sedan
4   ford       taurus      B      sedan
5   toyota     camry       B      sedan

输出：

    rating   type    conditional_probability
0     A      coupe   0.50   
1     B      coupe   0.33
2     C      sedan   1.00
3     A      sedan   0.50
4     B      sedan   0.66

萨夫万

您可以groupby用来获取单独的ratings以及ratings和types的单独组合中的项目计数，并使用这些值来计算条件概率。

from pyspark.sql import functions as F

ratings_cols = ["company", "model", "rating", "type"]
ratings_values = [
    ("ford", "mustang", "A", "coupe"),
    ("chevy", "camaro", "B", "coupe"),
    ("ford", "fiesta", "C", "sedan"),
    ("ford", "focus", "A", "sedan"),
    ("ford", "taurus", "B", "sedan"),
    ("toyota", "camry", "B", "sedan"),
]
ratings_df = spark.createDataFrame(data=ratings_values, schema=ratings_cols)
ratings_df.show()
# +-------+-------+------+-----+                                                  
# |company|  model|rating| type|
# +-------+-------+------+-----+
# |   ford|mustang|     A|coupe|
# |  chevy| camaro|     B|coupe|
# |   ford| fiesta|     C|sedan|
# |   ford|  focus|     A|sedan|
# |   ford| taurus|     B|sedan|
# | toyota|  camry|     B|sedan|
# +-------+-------+------+-----+

probability_df = (ratings_df.groupby(["rating", "type"])
                            .agg(F.count(F.lit(1)).alias("rating_type_count"))
                            .join(ratings_df.groupby("rating").agg(F.count(F.lit(1)).alias("rating_count")), on="rating")
                            .withColumn("conditional_probability", F.round(F.col("rating_type_count")/F.col("rating_count"), 2))
                            .select(["rating", "type", "conditional_probability"])
                            .sort(["type", "rating"]))

probability_df.show()
# +------+-----+-----------------------+                                          
# |rating| type|conditional_probability|
# +------+-----+-----------------------+
# |     A|coupe|                    0.5|
# |     B|coupe|                   0.33|
# |     A|sedan|                    0.5|
# |     B|sedan|                   0.67|
# |     C|sedan|                    1.0|
# +------+-----+-----------------------+

本文收集自互联网，转载请注明来源。

如有侵权，请联系[email protected] 删除。