python - 比较两列以在 Spark DataFrame 中创建新列

Question

我有一个包含 2 列的 Spark DataFrame，我正在尝试使用其他两列和 when else 操作创建一个新列。

df_newcol = df.withColumn("Flag", when(col("a") <= lit(ratio1) | col("b") <= lit(ratio1), 1).otherwise(2))

但这会引发错误

ValueError: Cannot convert column into bool: please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions.

我之前在一个列中使用 when 和 else ，而在多列中使用它时，我们是否必须以不同的方式编写逻辑。

谢谢。

Answer 1

您有运算符优先级问题，当比较与逻辑运算符(例如 &)混合时，请确保将比较运算符放在括号中。和 | ，用它固定，你甚至不需要 lit ，标量也应该起作用:

import pyspark.sql.functions as F
df = spark.createDataFrame([[1, 2], [2, 3], [3, 4]], ['a', 'b'])

以下两项都应该有效:

df.withColumn('flag', F.when((F.col("a") <= F.lit(2)) | (F.col("b") <= F.lit(2)), 1).otherwise(2)).show()
+---+---+----+
|  a|  b|flag|
+---+---+----+
|  1|  2|   1|
|  2|  3|   1|
|  3|  4|   2|
+---+---+----+

df.withColumn('flag', F.when((F.col("a") <= 2) | (F.col("b") <= 2), 1).otherwise(2)).show()
+---+---+----+
|  a|  b|flag|
+---+---+----+
|  1|  2|   1|
|  2|  3|   1|
|  3|  4|   2|
+---+---+----+