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

Question

我有一个包含FN列的dataframe和这些列值的子集的列表。

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

我想在我的dataframe中添加一个列，如果列值存在于列表中，则记录列表中的位置，即结束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('')))

但是在运行时，我会得到一个"Job aborted due to stage failure“异常和一系列其他异常，唯一有意义的部分是"ValueError: 'JKL' is not in list”(JKL是DF列中的随机其他列)。

如果我只输入"fn“而不是"p.index(fn)”，那么在我的新列中得到正确的列值，类似地，如果我使用"p.index("DEF")"，我会得到"1“--因此这些都是有效的，有什么例外吗？

提亚

编辑:我在lambda中做了一个if- the，这几乎意味着它是在withColumn语句中的"isin“检查之前执行lambda。我想知道(除了上述是否属实)，是否有人有更好的建议，如何更好地做到这一点？

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|
+---+----+