Question

我想执行一些条件分支以避免计算不必要的节点，但是我注意到，如果条件语句中的源列是UDF，则无论如何都将解析否则：

 @pandas_udf("double", PandasUDFType.SCALAR)
 def udf_that_throws_exception(*cols):
   raise Exception('Error')

 @pandas_udf("int", PandasUDFType.SCALAR)
 def simple_mul_udf(*cols):
   result = cols[0]
   for c in cols[1:]:
     result *= c
   return result

 df = spark.range(0,5)

 df = df.withColumn('A', lit(1))
 df = df.withColumn('B', lit(2))

 df = df.withColumn('udf', simple_mul('A','B'))
 df = df.withColumn('sql', expr('A*B'))

 df = df.withColumn('res', when(df.sql < 100, lit(1)).otherwise(udf_that_throws(lit(0))))

上面的代码按预期工作，在这种情况下该语句始终为true，因此永远不会调用引发异常的UDF。

但是，如果我将条件更改为使用 df.udf ，那么突然之间就会调用否则UDF，即使条件结果未更改，我也会得到异常。

我认为我可以通过从条件中删除UDF来对其进行混淆，但是无论如何，都会出现相同的结果

  df = df.withColumn('cond', when(df.udf < 100, lit(1)).otherwise(lit(0)))
  df = df.withColumn('res', when(df.cond == lit(1), lit(1)).otherwise(udf_that_throws_exception(lit(0))))

我认为这与Spark优化的方式有关，这很好，但是我正在寻找任何方法来实现此目的而又不会产生成本。有什么想法吗？

修改根据要求获取更多信息。我们正在编写一个可以接受任意模型的处理引擎，并且代码生成图形。在此过程中，我们会根据运行时的值状态做出决策。我们大量使用熊猫UDF。因此，想象一下这样一种情况：我们在图中有多个路径，并且根据运行时的某些条件，我们希望遵循这些路径之一，而其余所有路径都保持不变。

我想将此逻辑编码到图中，这样就不必在代码中收集和分支了。

我提供的示例代码仅用于演示目的。我面临的问题是，如果IF语句中使用的列是UDF，或者，如果它是从UDF派生的，则即使未实际使用过OTHERWISE条件，也总是执行该条件。如果我不介意IF / ELSE是廉价的操作（例如文字），但是如果UDF列（可能在两边）导致大的聚合或实际上被丢弃的其他长度的过程会怎样呢？

Answer 1

在PySpark中，UDF是预先计算的，因此您会得到这种次优的行为。您还可以从查询计划中看到它：

== Physical Plan ==
*(2) Project [id#753L, 1 AS A#755, 2 AS B#758, pythonUDF1#776 AS udf#763, CASE WHEN (pythonUDF1#776 < 100) THEN 1.0 ELSE pythonUDF2#777 END AS res#769]
+- ArrowEvalPython [simple_mul_udf(1, 2), simple_mul_udf(1, 2), udf_that_throws_exception(0)], [id#753L, pythonUDF0#775, pythonUDF1#776, pythonUDF2#777]
   +- *(1) Range (0, 5, step=1, splits=8)

ArrowEvalPython运算符负责计算UDF，此后将在Project运算符中评估条件。

在条件（最佳行为）下调用df.sql时，您会得到不同的行为的原因是，这是一种特殊情况，其中此列中的值是恒定的（两列{{1} }和A不变），Spark优化器可以事先对其进行评估（在查询计划处理过程中的驱动程序中，在集群上执行实际作业之前），因此它知道B分支无需评估。如果此otherwise列中的值是动态的（例如，像sql列中那样），则该行为将再次变为次优，因为Spark事先不知道id部分永远都不应采用地方。

如果要避免这种次优的行为（即使不需要，也可以在otherwise中调用udf），一种可能的解决方案是在udf中评估这种情况，例如：

otherwise

如何在条件语句中停止Spark解析UDF列

1 个答案: