如何强制 Spark SQL 进入代码生成模式？答案

【问题标题】：How to force Spark SQL into codegen mode?如何强制 Spark SQL 进入代码生成模式？
【发布时间】：2021-09-29 09:55:38
【问题描述】：

我正在使用自定义 codegen 编写自定义 Spark 催化剂 Expression，但 Spark (3.0.0) 似乎不想使用生成的代码，并退回到解释模式。

我以非常标准的方式创建我的 SparkSession，除了我尝试强制代码生成：

val spark = SparkSession.builder()
  .appName("test-spark")
  .master("local[5]")
  .config("spark.sql.codegen.factoryMode", "CODEGEN_ONLY")
  .config("spark.sql.codegen.fallback", "false")
  .getOrCreate()

然后我有这个自定义的Expression，同时定义了解释模式和代码生成：

case class IsTrimmedExpr(child: Expression) extends UnaryExpression with ExpectsInputTypes {
  override def inputTypes: Seq[DataType] = Seq(StringType)
  override lazy val dataType: DataType = BooleanType

  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
    throw new RuntimeException("expected code gen")
    nullSafeCodeGen(ctx, ev, input => s"($input.trim().equals($input))")
  }

  override protected def nullSafeEval(input: Any): Any = {
    throw new RuntimeException("should not eval")
    val str = input.asInstanceOf[org.apache.spark.unsafe.types.UTF8String]
    str.trim.equals(str)
  }
}

我在会话的注册表中注册：

spark.sessionState.functionRegistry.registerFunction(
  FunctionIdentifier("is_trimmed"), {
    case Seq(s) => IsTrimmedExpr(s)
  }
)

要调用函数/表达式，我会这样做

val df = Seq("   abc", "def", "56 ", "  123  ", "what is a trim").toDF("word")
df.selectExpr("word", "is_trimmed(word)").show()

但是，我从 nullSafeEval 函数中得到了一个根本不应该运行的异常，而不是来自 doGenCode 函数的预期异常。

如何强制 Spark 使用 codegen 模式？

【问题讨论】：

标签： apache-spark apache-spark-sql

【解决方案1】：

通过将 spark.sql.codegen 设置为 True 来启用 codegen

【讨论】：

这是很久以前的 Spark 的默认设置。