Python Pyspark-将具有最大值的列转换为单独的1和0项
我在《熊猫》中有一个解决这个问题的工作版本,但我在将它翻译成pyspark时遇到了麻烦 我的输入数据框如下所示:Python Pyspark-将具有最大值的列转换为单独的1和0项,python,apache-spark,pyspark,apache-spark-sql,etl,Python,Apache Spark,Pyspark,Apache Spark Sql,Etl,我在《熊猫》中有一个解决这个问题的工作版本,但我在将它翻译成pyspark时遇到了麻烦 我的输入数据框如下所示: 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
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
目前,我能想到的最多的是如何根据列的值(无论是否为最大值)将列设置为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|
+---+-----+-----+-----+
提前谢谢 假设您当前的结果是
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
it来操作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|
+---+-----+-----+-----+
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|
+---+-----+-----+-----+