PySpark 在每個 DataFrame 行上執行純 Python 函數

Question

我有數百萬行的 Spark DataFrame DF1 。 每行最多有 100 列。

col1 | col2 | col3 | ... | colN
--------------------------------
v11  | v12  | v13  | ... | v1N
v21  | v22  | v23  | ... | v2N
...  | ...  | ...  | ... | ...

此外，我還有另一個 DataFrame DF2 ，其中有數百行帶有名稱和正文列。 Name 包含函數名，body 包含純 Python 代碼，返回 true 或 false 的布爾函數。 這些函數在其邏輯中可以引用 DF1 中單行中的任何列。

func_name | func_body
-----------------------------------------------
func1     |   col2 < col45
func2     |   col11.contains("London") and col32*col15 < col21
funcN     |   .... 

我需要將這兩個數據幀 - DF1 與 DF2 連接起來，並將 Df2 中的每個函數應用到 DF1 中的每一行。 每個函數都必須能夠接受來自 DF1 的參數，假設具有鍵/值對的字典數組，這些鍵/值對表示來自 DF1 的相應行的所有列的名稱/值。

我知道如何加入 DF1 和 DF2，而且，我知道 Python 函數的執行不會以分布式方式工作。 暫時沒問題。 這是一個暫時的解決方案。 我只需要將來自 DF1 的所有行分發到工作節點上，並將每個 Python 函數應用於 Apache Spark 應用程序不同任務中的每一行 DF1。 評估eval()它們並傳遞帶有鍵/值對的字典數組，正如我上面提到的。

通常，每個 Python 函數都是一個標記，如果某些函數返回 true，我想將其分配給 DF1 中的行。 例如，這是生成的 DataFrame DF3 ：

col1 | col2 | col3 | ... | colN | tags
--------------------------------------
v11  | v12  | v13  | ... | v1N  | [func1, func76, funcN]
v21  | v22  | v23  | ... | v2N  | [func32]
...  | ...  | ...  | ... | ...  | [..., ..., ..., ..., ...]

PySpark 是否可行，如果可以，請舉例說明如何實現？ 使用DF.columns中的Map作為輸入參數的 UDF 函數是正確的方法還是可以以更簡單的方式完成？ Spark 對一次可以注冊多少 UDF 函數（數量）有任何限制嗎？

Answer 1

您可以使用可以使用expr進行評估的 SQL 表達式來實現這一點。 但是，您將無法加入 2 個數據幀，因為 SQL 表達式不能作為列值進行評估（請參閱此帖子），因此您必須將這些函數收集到一個列表中（因為您只有數百行，它可以放入內存中）。

這是一個可以根據您的要求進行調整的工作示例：

data1 = [(1, "val1", 4, 5, "A", 10), (0, "val2", 7, 8, "B", 20),
         (9, "val3", 8, 1, "C", 30), (10, "val4", 2, 9, "D", 30),
         (20, "val5", 6, 5, "E", 50), (3, "val6", 100, 2, "X", 45)]

df1 = spark.createDataFrame(data1, ["col1", "col2", "col3", "col4", "col5", "col6"])

data2 = [("func1", "col1 + col3 = 5 and col2 like '%al1'"),
         ("func2", "col6 = 30 or col1 * col4 > 20"),
         ("func3", "col5 in ('A', 'B', 'C') and col6 - col1 < 30"),
         ("func4", "col2 like 'val%' and col1 > 0")]

df2 = spark.createDataFrame(data2, ["func_name", "func_body"])

# get functions into a list
functions = df2.collect()

# case/when expression to evaluate the functions
satisfied_expr = [when(expr(f.func_body), lit(f.func_name)) for f in functions]

# add new column tags
df1.withColumn("tags", array(*satisfied_expr)) \
    .withColumn("tags", expr("filter(tags, x -> x is not null)")) \
    .show(truncate=False)

添加數組列tags ， filter函數用於去除不滿足表達式對應的空值。 此功能僅從 Spark 2.4+ 開始可用，對於舊版本，您必須使用和 UDF。

給出：

+----+----+----+----+----+----+---------------------+
|col1|col2|col3|col4|col5|col6|tags                 |
+----+----+----+----+----+----+---------------------+
|1   |val1|4   |5   |A   |10  |[func1, func3, func4]|
|0   |val2|7   |8   |B   |20  |[func3]              |
|9   |val3|8   |1   |C   |30  |[func2, func3, func4]|
|10  |val4|2   |9   |D   |30  |[func2, func4]       |
|20  |val5|6   |5   |E   |50  |[func2, func4]       |
|3   |val6|100 |2   |X   |45  |[func4]              |
+----+----+----+----+----+----+---------------------+