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

Question

我正在Expression使用自定义代码生成器编写自定义 Spark 催化剂，但似乎 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()

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

如何强制 Spark 使用代码生成模式？

Answer 1

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