使用 case when 为原始列中的每个 notNull 分配 1 的 spark df 的每一列添加相应的列

Question

我有一个样本 dataframe：

df = spark.createDataFrame([('name1','id1',1,None,3),('name2','id2',None,2,5)],['NAME','personID','col1','col2','col3'])

我的用例有 15 列

我想做的是使用 case when 和循环，添加与原始列中除前两列之外的每一列相对应的新列。 在这些新列中，如果不是 Null，它将给出 1 的值，否则为 0。

我的目标是获得如下内容：

+--------+--------+--------+-------+-------+-------+------+------+
|Name    | ID     | col1   | col2  |  col3 | col1_N|col2_N|col3_N|
+--------+--------+--------+-------+-------+-------+------+------+
|name1   | id1    | 1      | Null  | 3     | 1     | 0    | 1    |
|name2   | id2    | Null   | 2     | 5     | 0     | 1    | 1    |
+--------+--------+--------+-------+-------+-------+------+------+

前五列是原始列，最后三列将添加来自“col1”、“col2”和“col3”值的相应 1 或 0。

我正在处理的最后一个代码创建了一个新代码，但不保留原始 dataframe 值。

df.select([when(col(c).isNotNull(), 1).otherwise(0).alias(c + '_N') for c in df.columns])

为此我得到：

+-------+-------+-------+------+------+
| Name_N|  ID_N | col1_N|col2_N|col3_N|
+-------+-------+-------+------+------+
| 1     | 1     | 1     | 0    | 1    |
| 1     | 1     | 0     | 1    | 1    |
+-------+-------+-------+------+------+

以上可以接受，但我需要保留 Name 和 ID 列的原始值。

我得到一个 InvalidArgument 如下：

df.select(['*'],[when(col(c).isNotNull(), 1).otherwise(0).alias(c + '_N') for c in df.columns])


TypeError: Invalid argument, not a string or column: ['*'] of type <class 'list'>. For column literals, use 'lit', 'array', 'struct' or 'create_map' function.

我以为先选择全部会给我原始的所有列

更新：不知何故这奏效了，但我只得到最后一列：

for c in df.columns[2:]:
  sdf = df.withColumn(c+'_N', when(col(c).isNotNull(),1).otherwise(0))

但这就是我得到的：

+--------+--------+--------+-------+-------+------+
|Name    | ID     | col1   | col2  |  col3 |col3_N|
+--------+--------+--------+-------+-------+------+
|name1   | id1    | 1      | Null  | 3     | 1    |
|name2   | id2    | Null   | 2     | 5     | 1    |
+--------+--------+--------+-------+-------+------+
``

I only got the last original column

Answer 1

使用如下所示的列表理解将给出预期的结果。

df.select([col(c) if c in ['NAME', 'personID'] else when(col(c).isNotNull(), 1).otherwise(0).alias(f"{c}_N")  for c in df.columns]).show()

+-----+--------+------+------+------+
| NAME|personID|col1_N|col2_N|col3_N|
+-----+--------+------+------+------+
|name1|     id1|     1|     0|     1|
|name2|     id2|     0|     1|     1|
+-----+--------+------+------+------+

Answer 2

只需通过指定一段列来修复您的第一种方法并简化 boolean 条件：

df.select([col(c).isNotNull().cast("integer").alias(c + '_N') for c in df.columns[2:]])