Как создать новый столбец на основе значений двух столбцов в Spark

Question

Итак, я хочу создать новый столбец в моем фрейме данных, строки которого зависят от значений из двух столбцов, а также содержат условие.

Я пробовал это, но это не работает.

some_value = ...
df = df.withColumn("new_col", col("col1") if col("col2") == some_value else None)

Как правильно это сделать?

Answer 1

Вы можете использовать функцию when.

newval='10002'
df2 = df.withColumn("new_col", when((col("col1") == lit(newval)), "col1").otherwise(None))
df2.show()

Answer 2

Тривиальный пример использования с expr, также можно использовать when:

val df3 = df2.withColumn("new_col", expr("case when c1 = 1 and c2 = 101 then c1 + c2 else 999 end"))

Answer 3

Это сработало для меня,

 new_col_expr = when(col("col2").eqNullSafe(some_value), col("col1")).otherwise(None)
 df = df.withColumn("new_col", new_col_expr)