pyspark - 当一列中的值在另一列中时标记数据

Question

我正在尝试根据 2 个条件在我的数据集中创建一个标志，第一个很简单。CheckingCol = CheckingCol2。

第二个更复杂。我有一个名为 TranID 的列和一个名为 RevID 的列。

如果 RevID 位于 TranID AND CheckingCol = CheckingCol2 中，则该标志应返回“是”。否则，标志应返回“否”。

我的数据如下所示：

TranID   RevID   CheckingCol  CheckingCol2
1        2       ABC          ABC
2        1       ABC          ABC
3        6       ABCDE        ABCDE
4        3       ABCDE        ABC
5        7       ABCDE        ABC

预期的结果是：

TranID   RevID   CheckingCol  CheckingCol2  Flag
1        2       ABC          ABC           Yes
2        1       ABC          ABC           Yes
3        6       ABCDE        ABCDE         No
4        3       ABCDE        ABC           No
5        7       ABCDE        ABC           No

我试过使用：

df.withColumn("TotalMatch", when((col("RevID").contains(col("TranID"))) & (col("CheckingColumn") == col("CheckingColumn2")), "Yes").otherwise("No"))

但它没有用，我无法在网上找到任何关于如何做到这一点的信息。

任何帮助都会很棒！

Answer 1

从 TranID 列获取唯一值作为数组，然后使用 isIn() 函数检查该数组中的 RevID

from pyspark.sql import functions as sf
unique_values = df1.agg(sf.collect_set("TranID").alias("uniqueIDs"))
unique_values.show()
+---------------+
|       uniqueIDs|
+---------------+
|[3, 1, 2, 5, 4]|
+---------------+

required_array = unique_values.take(1)[0].uniqueIDs
['3', '1', '2', '5', '4']

df2 = df1.withColumn("Flag", sf.when( (sf.col("RevID").isin(required_array) & (sf.col("CheckingCol") ==sf.col("CheckingCol2")) ) , "Yes").otherwise("No"))

注意：检查 RevID 和 TranID 列中的 null 和 NoneType 值，因为它们会影响结果