apache-spark - 如何强制 Spark SQL 进入代码生成模式？

Question

我正在编写自定义 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)
  }
}

我注册到 session 的注册表中:

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()

但不是来自 doGenCode 的预期异常函数，我从 nullSafeEval 得到异常根本不应该运行的函数。
如何强制 Spark 使用代码生成模式？

Answer 1

启用代码生成是通过将 spark.sql.codegen 设置为 True 来完成的