pyspark：優化 pandas udf，對 Z6A8064B5DF47945550DZ?53C7 中的每一行執行 sql 查詢

Question

我正在利用 pyspark 和 pandas udfs 來加速對包含約 3,500,000 行的 dataframe 的計算。 本質上，我正在從滿足某些條件的表中加載行。 然后，我按列“some_col”對這些進行分組，該列大致應將數據分成 4 組。

然后我應用一個 function 來計算metric_1和metric_2 ，其中每個指標是我的表中與當前行中的某些值匹配的條目數。 然后在final_result中將最終計算設置為(metric_1 / metric_2)

我的代碼大綱

import numpy as np

def calc_metric(arg):
    query = """
        select some_id
            from table
            where some_col = {arg}
    """.format(arg)

    con = # sql connection initialization (not important)

    df = pd.read_sql(con=con, sql=query)

    return

def calculate_metric(df):
    df[metric_1] = df.apply(calc_metric, args=('A'), axis=1)
    df[metric_2] = df.apply(calc_metric, args=('B'), axis=1)
    df[final_metric] = np.where(df[metric_2] != 0,  
            str(float((df[metric_1] / df[metric_2]))), str(float(0)))
    return

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def calc_metric(df):
    df = calculate_metric(df)
    return df

df = df.groupBy('chosen_groupby_col').apply(calc_metric)

當我在本地運行一行時，整個過程似乎真的很慢。 我知道默認情況下 apache spark 設置spark.sql.shuffle.partitions=200 ，我注意到最后一個任務（實際執行這些操作）何時運行，它真的很慢。

這是因為我每行執行一個 sql 查詢嗎？ 或者 Spark JVM 和pyspark之間的序列化仍然需要很長時間？ 有沒有更優化的方法來做到這一點，其中更少的數據被洗牌並且運行得非常快。 這最終將在 aws emr 集群中運行。 我已經閱讀了很多關於pandas_udf GROUPED_MAP 函數應該如何快速的帖子，但這太慢了......

編輯 1：廣播 SQL 連接當我嘗試廣播 sql 連接時出現錯誤：

Traceback (most recent call last):
  File "/Users/megan/apache-spark/python/lib/pyspark.zip/pyspark/broadcast.py", line 113, in dump
    pickle.dump(value, f, 2)
TypeError: can't pickle _thread._local objects
Traceback (most recent call last):
  File "/Users/megan/apache-spark/python/lib/pyspark.zip/pyspark/broadcast.py", line 113, in dump
TypeError: can't pickle _thread._local objects

Answer 1

問題甚至可能不是每行執行 SQL 查詢，而是每次都打開兩個新連接。 因此，您忽略not important的部分可能是最重要的部分。 您可以嘗試廣播連接。