大型交叉连接DataFrame UDF聚合求和性能优化咨询

• 90%以上的性能损耗来自普通Python UDF的跨进程开销：当前使用的逐行Python UDF需要在JVM和Python解释器之间逐行做数据序列化/反序列化，每行单独调用scipy、numpy函数，完全无法利用Spark的向量化执行、整阶段代码生成优化，性能比原生JVM计算低2~3个数量级，和C++实现的性能差距完全符合这个开销比例，和DataFrame的创建方式没有直接关系。
• 次要开销来自笛卡尔积的默认分区策略不合理：多轮crossJoin使用默认分区配置时，容易出现分区粒度过小、调度开销暴涨、数据倾斜的问题，维度越高这个问题越明显。
• 额外算力浪费：逐行重复计算正态分布PDF的常量项，没有做前置预计算。

• 优先使用Spark原生内置函数替换Python UDF，实现全JVM端向量化计算
  正态分布PDF可以直接拆解为基础数学表达式，不需要依赖scipy库，所有运算都可以用pyspark.sql.functions下的内置数学函数实现，内置函数会经过Spark Catalyst优化器自动生成向量化字节码，性能可以接近原生C++实现的水平。
  可直接复用的多维度实现代码如下：

  from pyspark.sql import SparkSession
  from pyspark.sql.functions import col, exp, sqrt, log, pow, lit, when
  
  # 常量定义
  eps   = 10.1
  iterx = 0.1
  sens  = 1
  scale = 0.7262317494490057
  N = 2 # 修改为实际维度数即可
  PI = 3.141592653589793
  
  # 预计算全局常量，避免逐行重复计算
  const_pdf = 1/(scale * sqrt(lit(2 * PI)))
  exp_eps = exp(lit(eps))
  iterx_pow = pow(lit(iterx), lit(N))
  
  # 初始化Spark会话，开启自适应执行优化
  spark = SparkSession.builder.master('local[*]').appName('test') \
      .config("spark.sql.adaptive.enabled", "true") \
      .getOrCreate()
  
  # 生成单维度取值序列
  base_dim = spark.range(0, 101, 1) \
      .withColumn('x', lit(-5.0) + col('id') * lit(iterx)) \
      .drop('id')
  # 根据维度数生成N维笛卡尔积
  df = base_dim.withColumnRenamed('x', 'x1')
  for i in range(2, N+1):
      df = df.crossJoin(base_dim.withColumnRenamed('x', f'x{i}'))
  
  # 纯原生函数计算di，无任何UDF
  p1 = lit(1)
  p2 = lit(1)
  for col_name in [f'x{i}' for i in range(1, N+1)]:
      x = col(col_name)
      p1 = p1 * const_pdf * exp(-0.5 * pow((x - lit(0))/lit(scale), 2))
      p2 = p2 * const_pdf * exp(-0.5 * pow((x - lit(sens))/lit(scale), 2))
  p1 = p1 * iterx_pow
  p2 = p2 * iterx_pow
  pl = log(p1 / p2)
  di = when(pl > lit(eps), p1 - exp_eps * p2).otherwise(lit(0.0))
  
  # 直接聚合求和
  dsum = df.agg(di.sum()).collect()[0][0]
  

• 如果后续需要使用无法用SQL表达的复杂Python逻辑，替换普通Python UDF为Pandas向量化UDF（pandas_udf），以批量列式数据为单位在Python和JVM之间传输，结合numpy/pandas的向量化运算，性能比逐行UDF高10~100倍。
• 调整笛卡尔积分区配置，降低调度开销：根据总数据量设置合理的shuffle分区数，单分区数据量控制在128MB_{256MB即可，比如1亿行单精度浮点数据约1.2GB，设置8}16个分区即可，不需要使用默认的200个分区；配合开启的AQE自适应执行，自动合并过小分区、处理倾斜。
• 终极优化：针对当前计算逻辑做数学化简，彻底避免生成超大全量笛卡尔积。当前计算逻辑中p1、p2都是各维度独立概率密度的乘积，对全量笛卡尔积的求和可以直接拆解为各维度单维度求和结果的乘积，计算复杂度从O(101^N)直接降到O(N*101)，哪怕是10个维度的场景，也不需要跑分布式任务，普通单核机器毫秒级就能出结果，完全不会出现数据量随维度指数膨胀的问题。

内容的提问来源于stack exchange，提问作者be_excellent_to_each_other