pyspark - PySpark - 从列表中获取字符串位置

Question

我有一个包含 FN 列的数据框和这些列值的子集的列表，例如

    **FN**
    ABC
    DEF
    GHI
    JKL
    MNO
    
    List:
    ["GHI","DEF"]

我想在我的数据框中添加一列，如果列值存在于列表中，我记录列表中的位置，即我的最终 DF

    FN     POS
    ABC
    DEF    1
    GHI    0
    JKL
    MNO

我的代码如下

    from pyspark.sql.functions import udf
    from pyspark.sql.types import StringType
    l = ["GHI","DEF"]

    x = udf(lambda fn, p = l: p.index(fn), StringType())
    
    df = df.withColumn('POS', when(col("FN").isin(l), x(col("FN"))).otherwise(lit('')))

但是在运行时，我得到一个ob aborted due to stage failure带有一系列其他异常的“J”异常，唯一有意义的部分是“ ValueError: 'JKL' is not in list”（JKL 是我的 DF 列中的随机其他列）

如果p.index(fn)我只是输入“fn”而不是“”，我会在新列中得到正确的列值，同样如果我使用“ p.index("DEF")”，我会得到“1”，所以这些单独工作，有什么想法为什么会出现异常？

TIA

---

编辑：我设法通过在 lambda 中执行 if-else 来解决这个问题，这几乎意味着它在 withColumn 语句中的“isin”检查之前执行 lambda。我想知道什么（除了上述是否属实），是否有人对如何以更好的方式实现这一目标有更好的建议？

标签： pyspark

Answer 1

这是我的尝试。我已经为给定的列表制作了一个数据框并加入它们。

from pyspark.sql.functions import *

l = ['GHI','DEF']
m = [(l[i], i) for i in range(0, len(l))]

df2 = spark.createDataFrame(m).toDF('FN', 'POS')
df1 = spark.createDataFrame(['POS','ABC','DEF','GHI','JKL','MNO'], "string").toDF('FN')

df1.join(df2, ['FN'], 'left').show()

+---+----+
| FN| POS|
+---+----+
|JKL|null|
|MNO|null|
|DEF|   1|
|POS|null|
|GHI|   0|
|ABC|null|
+---+----+