Spark 高級篇 - 程序不重啓還能支持動態註冊UDF

昨天有位大哥問小弟一個Spark問題，他們想在不停Spark程序的情況下動態更新UDF的邏輯，他一問我這個問題的時候，本豬心裏一驚，Spark**還能這麼玩?我出於程序員的本能回覆他肯定不行，但今天再回過來頭想了一想，昨天腦子肯定進水了，回覆太膚淺了，既然Spark可以通過編程方式註冊UDF，當然把那位大哥的代碼邏輯使用反射加載進去再調用不就行了？這不就是JVM的優勢麼，怪自己的反射沒學到家，說搞就搞起。

分析過程

我會說這波分析過程很無聊，你還會看麼?

跟着本豬看一個Spark註冊UDF的例子

spark.udf.register(name, (a1: String) => a1.toUpperCase)

點擊register的源碼進去看

一個`A1`:參數類型,`RT`:返回類型
def register[RT: TypeTag, A1: TypeTag](name: String, func: Function1[A1, RT]): UserDefinedFunction = {
    val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
    val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: Nil).toOption
    def builder(e: Seq[Expression]) = if (e.length == 1) {
      ScalaUDF(func, dataType, e, inputTypes.getOrElse(Nil), Some(name), nullable, udfDeterministic = true)
    } else {
      ...
    }
    ...
  }
def register[RT: TypeTag, A1: TypeTag, A2: TypeTag](name: String, func: Function2[A1, A2, RT]): UserDefinedFunction
def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag](name: String, func: Function3[A1, A2, A3, RT]): UserDefinedFunction = {
def register[RT: TypeTag, A1: TypeTag, A2: TypeTag, A3: TypeTag, A4: TypeTag](name: String, func: Function4[A1, A2, A3, A4, RT]): UserDefinedFunction = {
...

1. func是上面方法的重點，既然想要動態UDF邏輯代碼，那我們把Function1這個函數實現不就可以了？再利用JVM反射的技術調用，完美。
2. 順便還看出了在scala-2.10.x版本中case class的元素是不能超過 $\color{#FF0000}{22}$ 個的。

上面的UDF註冊的原型其實是

val udf = new Function1[String,String] {
      override def apply(a1: String): String = {
        a1.toUpperCase
      }
}
spark.udf.register(name, udf)

到這裏我有一個 膚淺 並且 大膽 的想法，我把那位大哥的代碼放到apply方法裏面調用不就行了？

val udf = new Function1[String,String] {
      override def apply(a1: String): String = {
        //method.invoke(instance) //使用反射加載代碼，把大哥動態邏輯方法method拿出來調用。
      }
}

1. 但是還有一些問題要解決，我不能強制我的老大哥只能傳遞一個參數吧，那也太年輕不懂事了，至少讓他可以隨意傳 $\color{#FF0000}{22}$ 參數。
2. 唯一的解決方法，就是要控制Function1到Function22函數的動態生成，找了半天沒發現Function的動態生成，然後還發現Spark也是根據參數長度生成FunctionN的，真**刷新本豬的三觀呀。
3. 既然實現方式找到了，那就簡單了，只要通過反射就能 上知天文,下知地理 。

既然是Spark，肯定要用Scala去寫反射了。

case class ClassInfo(clazz: Class[_], instance: Any, defaultMethod: Method, methods: Map[String, Method], func:String) {
  def invoke[T](args: Object*): T = {
    defaultMethod.invoke(instance, args: _*).asInstanceOf[T]
  }
}
object ClassCreateUtils extends Logging{
  private val clazzs = new util.HashMap[String, ClassInfo]()
  private val classLoader = scala.reflect.runtime.universe.getClass.getClassLoader
  private val toolBox = universe.runtimeMirror(classLoader).mkToolBox()
  def apply(func: String): ClassInfo = this.synchronized {
    var clazz = clazzs.get(func)
    if (clazz == null) {
      val (className, classBody) = wrapClass(func)
      val zz = compile(prepareScala(className, classBody))
      val defaultMethod = zz.getDeclaredMethods.head
      val methods = zz.getDeclaredMethods
      clazz = ClassInfo(
        zz,
        zz.newInstance(),
        defaultMethod,
        methods = methods.map { m => (m.getName, m) }.toMap,
        func
      )
      clazzs.put(func, clazz)
      logInfo(s"dynamic load class => $clazz")
    }
    clazz
  }
  def compile(src: String): Class[_] = {
    val tree = toolBox.parse(src)
    toolBox.compile(tree).apply().asInstanceOf[Class[_]]
  }
  def prepareScala(className: String, classBody: String): String = {
    classBody + "\n" + s"scala.reflect.classTag[$className].runtimeClass"
  }
  def wrapClass(function: String): (String, String) = {
    val className = s"dynamic_class_${UUID.randomUUID().toString.replaceAll("-", "")}"
    val classBody =
      s"""
         |class $className{
         |  $function
         |}
            """.stripMargin
    (className, classBody)
  }
}

上面的代碼是小弟給大佬寫好的，不用大佬親自動手了。

使用方法就灰常簡單了我的大佬們。

val infos = ClassCreateUtils(
      """
        |def apply(name:String)=name.toUpperCase
      """.stripMargin
)
    
println(infos.defaultMethod.invoke(infos.instance,"dounine 本豬會一點點 spark"))
# 輸出結果不用猜也知道是
DOUNINE 本豬會一點點 SPARK
# 也可以手動指定方法
println(infos.methods("apply").invoke(infos.instance,"dounine 本豬會一點點 spark"))

根據反射的方法信息生成FunctionN

object ScalaGenerateFuns {

  def apply(func: String): (AnyRef, Array[DataType], DataType) = {
    val (argumentTypes, returnType) = getFunctionReturnType(func)
    (generateFunction(func, argumentTypes.length), argumentTypes, returnType)
  }

  //獲取方法的參數類型及返回類型
  private def getFunctionReturnType(func: String): (Array[DataType], DataType) = {
    val classInfo = ClassCreateUtils(func)
    val method = classInfo.defaultMethod
    val dataType = JavaTypeInference.inferDataType(method.getReturnType)._1
    (method.getParameterTypes.map(JavaTypeInference.inferDataType).map(_._1), dataType)
  }

  //生成22個Function
  def generateFunction(func: String, argumentsNum: Int): AnyRef = {
    lazy val instance = ClassCreateUtils(func).instance
    lazy val method = ClassCreateUtils(func).methods("apply")

    argumentsNum match {
      case 0 => new (() => Any) with Serializable with Logging {
        override def apply(): Any = {
          try {
            method.invoke(instance)
          } catch {
            case e: Exception =>
              logError(e.getMessage)
          }
        }
      }
      case 1 => new (Object => Any) with Serializable with Logging {
        override def apply(v1: Object): Any = {
          try {
            method.invoke(instance, v1)
          } catch {
            case e: Exception =>
              e.printStackTrace()
              logError(e.getMessage)
              null
          }
        }
      }
      case 2 => new ((Object, Object) => Any) with Serializable with Logging {
        override def apply(v1: Object, v2: Object): Any = {
          try {
            method.invoke(instance, v1, v2)
          } catch {
            case e: Exception =>
              logError(e.getMessage)
              null
          }
        }
      }
      //... 麻煩大佬自己去寫剩下的20個了，這裏裝不下了，不然瀏覽器會崩潰的，然後電腦會重啓的，爲了大佬的電腦着想。
}

前戲我們都做完了，高潮的環節來了。

我們最後再照着register的實現方式，把我們動態Function註冊給Spark

1. val ScalaReflection.Schema(dataType, nullable) = ScalaReflection.schemaFor[RT]
2. val inputTypes = Try(ScalaReflection.schemaFor[A1].dataType :: Nil).toOption
3. def builder(e: Seq[Expression]) = if (e.length == 1) {
  ScalaUDF(func, dataType, e, inputTypes.getOrElse(Nil), 
    Some(name), nullable, udfDeterministic = true)
}
4. functionRegistry.createOrReplaceTempFunction(name, builder)

1. 這句代碼比較好理解，就是獲取RT返回值類型，就是我們的returnType
2. 就是參數類型，對應的修改如下

val inputTypes = Try(argumentTypes.toList).toOption

3. 剛開始看到這個時候，我是一臉???，後來看源碼才發現builder是一種自定類型,源碼如下

type FunctionBuilder = Seq[Expression] => Expression

改造方式如下

def builder(e: Seq[Expression]) = ScalaUDF(rf, returnType, e, inputTypes.getOrElse(Nil), Some(name))

4. 看到這句的時候我以爲簡單了，直接使用spark.sessionState.functionRegistry發現編譯不過，看到private[sql]這個作用域的時候有點崩潰，本來是想用下面的方式註冊的。

val udf = UserDefinedFunction(rf, returnType, inputTypes).withName(name)
spark.udf.register(name, udf)

是小弟我想太多了，另闢捷徑，做了那麼多工作難道就白費了？

發現下面這句代碼，瞬間找到了家的方向。

functionRegistry.registerFunction(new FunctionIdentifier(name), builder)

---

人生巔峯

到此，大豬的分析與編碼已經完成，下面是今天給大哥的解決方案。
方法實現可以通過查詢sql得到，或者接口都渴以。

    val spark = SparkSession
      .builder()
      .appName("test")
      .master("local[*]")
      .getOrCreate()

    val name = "hello"

    val (fun, argumentTypes, returnType) = ScalaSourceUDF(
      """
        |def apply(name:String)=name+" => hi"
        |""".stripMargin)

    val inputTypes = Try(argumentTypes.toList).toOption

    def builder(e: Seq[Expression]) = ScalaUDF(fun, returnType, e, inputTypes.getOrElse(Nil), Some(name))

    spark.sessionState.functionRegistry.registerFunction(new FunctionIdentifier(name), builder)

    val rdd = spark
      .sparkContext
      .parallelize(Array(("dounine", "20")))
      .map(x => Row.fromSeq(Array(x._1, x._2)))

    val types = StructType(
      Array(
        StructField("name", StringType),
        StructField("age", StringType)
      )
    )

    spark.createDataFrame(rdd, types).createTempView("log")

    spark.sql("select hello(name) from log").show(false)

真打臉，昨天還說不行的。