时间戳

首页 > 解决方案 > 如何强制 Spark SQL 进入代码生成模式？

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

问题描述

我正在Expression使用自定义 codegen 编写自定义 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 使用 codegen 模式？

标签： apache-sparkapache-spark-sql

解决方案

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

推荐阅读