pyspark：优化对数据帧中的每一行执行 sql 查询的 pandas udf？

Question

我正在利用 pyspark 和 pandas udfs 来加速包含约 3,500,000 行数据帧的计算。本质上，我正在从满足某些条件的表中加载行。然后，我按列“some_col”对这些进行分组，该列大致应将数据分成 4 组。

然后我应用一个计算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 查询吗？ Or is serializing between Spark JVM and pyspark still taking a long time? 有没有更优化的方法来做到这一点，其中更少的数据被洗牌并且运行得非常快。这最终将在 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 可能是最重要的部分。您可以尝试broadcasting 连接。