如何从PySpark DataFrame每行提取位置未知的目标值z

现有如下PySpark DataFrame：

item_A  item_B item_C
  x       z      y
  z       x      y
  y       x      z
  z       y      x 

所有值均为字符串，已知x和y的值，需要提取每行中z的值（z不固定在某一列），新增一列仅保留z。

尝试过拼接列后用正则提取，但未成功，尝试的代码如下：

main_pattern = r'x|y'
pattern_full = r'(('+ main_pattern+'),)'
df = df.withColumn("vale_z", regexp_extract("columns_concatenated", pattern_full, 1))

提供三种可行的实现方式：

方法一：条件判断遍历列 #

直接针对每一列排除x和y，提取剩余的z值：

from pyspark.sql.functions import col, when

df = df.withColumn(
    "vale_z",
    when(col("item_A").isin(["x", "y"]) == False, col("item_A"))
    .when(col("item_B").isin(["x", "y"]) == False, col("item_B"))
    .when(col("item_C").isin(["x", "y"]) == False, col("item_C"))
)

逻辑直观，适合列数较少的场景。

方法二：数组排除法 #

将每行的列转为数组，排除已知的x和y，剩余内容即为z：

from pyspark.sql.functions import array, array_except, element_at

df = df.withColumn(
    "vale_z",
    element_at(
        array_except(array(col("item_A"), col("item_B"), col("item_C")), array("x", "y")),
        1
    )
)

灵活性强，列数增加时无需大幅修改代码，前提是每行仅存在一个z值。

方法三：修正正则提取逻辑 #

如果坚持用拼接列+正则的方式，调整正则规则匹配非x/y的内容：

from pyspark.sql.functions import concat_ws, regexp_extract

# 先拼接所有列
df = df.withColumn("columns_concatenated", concat_ws(",", col("item_A"), col("item_B"), col("item_C")))
# 匹配并提取非x/y的字符串
df = df.withColumn(
    "vale_z",
    regexp_extract("columns_concatenated", r'(?!x|y)\w+', 0)
)

利用负向预查(?!x|y)定位非x/y的单词，提取第一个匹配结果即可。

内容的提问来源于stack exchange，提问作者FerCTRO