Pyspark - 将具有最大值的列转换为单独的 1 和 0 条目答案

【问题标题】：Pyspark - Transform columns with maximum values into separate 1 and 0 entriesPyspark - 将具有最大值的列转换为单独的 1 和 0 条目
【发布时间】：2021-03-29 14:14:56
【问题描述】：

我在 pandas 中有一个解决这个问题的工作版本，但我无法将其转换为 pyspark。

我的输入 DataFrame 如下所示：

test_df = pd.DataFrame({
    'id': [1],
    'cat_1': [2],
    'cat_2': [2],
    'cat_3': [1]
})
test_df_spark = spark.createDataFrame(test_df)
test_df_spark.show()

+---+-----+-----+-----+
| id|cat_1|cat_2|cat_3|
+---+-----+-----+-----+
|  1|    2|    2|    1| <- non-maximum
+---+-----+-----+-----+
         ^     ^
         |     |
     maximum maximum

我想：

在 cat_1、cat_2、cat_3 中获取具有最大值的列（1 个或更多）。在示例中，它们是 cat_1 和 cat_2。
这些列的值应为 1。其余非最大列将设置为 0。
应将具有 1 个值的列分成单独的行。

产生的 DataFrame 应如下所示：

+---+-----+-----+-----+
| id|cat_1|cat_2|cat_3|
+---+-----+-----+-----+
|  1|    1|    0|    0|
|  1|    0|    1|    0|
+---+-----+-----+-----+

目前，我能弄清楚的最多的是如何根据列的值（无论是否为最大值）将列设置为 1 或 0，但我仍然失踪如何生成附加条目：

columns = ['cat_1', 'cat_2', 'cat_3']
(
    test_df_spark
    .withColumn(
        'max_value',
        F.greatest(
            *columns
        )
    )
    .select(
        'id',
        *[F.when(F.col(c) == F.col('max_value'), F.lit(1)).otherwise(F.lit(0)).alias(c) for c in columns]
    )
    .show()
)

+---+-----+-----+-----+
| id|cat_1|cat_2|cat_3|
+---+-----+-----+-----+
|  1|    1|    1|    0|
+---+-----+-----+-----+

提前致谢！

【问题讨论】：

标签： python apache-spark pyspark apache-spark-sql etl

【解决方案1】：

假设你当前的结果是df1:

columns = ['cat_1', 'cat_2', 'cat_3']
df1 = (
    test_df_spark
    .withColumn(
        'max_value',
        F.greatest(
            *columns
        )
    )
    .select(
        'id',
        *[F.when(F.col(c) == F.col('max_value'), F.lit(1)).otherwise(F.lit(0)).alias(c) for c in columns]
    )
)

您可以通过创建结构数组和inline 来操纵df1 以获得所需的结果：

df2 = df1.select(
    'id', 
    F.array(*[
        F.when(
            F.col(c1) == 1, 
            F.struct(*[
                F.lit(1).alias(c2) if i1 == i2 else F.lit(0).alias(c2) 
                for i2, c2 in enumerate(columns)
            ])
        ) 
        for i1, c1 in enumerate(columns)
    ]).alias('cat')
).selectExpr(
    'id', 
    'inline(filter(cat, x -> x is not null))'
)

df2.show()
+---+-----+-----+-----+
| id|cat_1|cat_2|cat_3|
+---+-----+-----+-----+
|  1|    1|    0|    0|
|  1|    0|    1|    0|
+---+-----+-----+-----+

【讨论】：