How to force Spark SQL into codegen mode?

Viewed 270

I'm writing a custom Spark catalyst Expression with custom codegen, but it seems that Spark (3.0.0) doesn't want to use the generated code, and falls back to interpreted mode.

I create my SparkSession in a pretty standard way, except that I try to force codegen:

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

And then I have this custom Expression with both interpreted mode and codegen defined:

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

which I register into the session's registry:

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

To invoke the function/Expression, I do

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

But instead of the expected exception from the doGenCode function, I got the exception from the nullSafeEval function which should not be run at all.

How do I force Spark to use codegen mode?

1 Answers

Enabling codegen is done via setting spark.sql.codegen to True

Related