如何使用Quasiquotes编写Scalafix规则替换unionAll为union

Spark 2.4中使用unionAll方法合并DataFrame/Dataset，示例代码如下：

import org.apache.spark.sql.{DataFrame, Dataset}

object UnionRewrite {
  def inSource(
    df1: DataFrame,
    df2: DataFrame,
    df3: DataFrame,
    ds1: Dataset[String],
    ds2: Dataset[String]
  ): Unit = {
    val res1 = df1.unionAll(df2)
    val res2 = df1.unionAll(df2).unionAll(df3)
    val res3 = Seq(df1, df2, df3).reduce(_ unionAll _)
    val res4 = ds1.unionAll(ds2)
    val res5 = Seq(ds1, ds2).reduce(_ unionAll _)
  }
}

而Spark 3.+版本中unionAll已被弃用，需替换为union，等效代码如下：

import org.apache.spark.sql.{DataFrame, Dataset}

object UnionRewrite {
  def inSource(
    df1: DataFrame,
    df2: DataFrame,
    df3: DataFrame,
    ds1: Dataset[String],
    ds2: Dataset[String]
  ): Unit = {
    val res1 = df1.union(df2)
    val res2 = df1.union(df2).union(df3)
    val res3 = Seq(df1, df2, df3).reduce(_ union _)
    val res4 = ds1.union(ds2)
    val res5 = Seq(ds1, ds2).reduce(_ union _)
  }
}

现有非Quasiquotes实现的Scalafix规则 #

目前已实现不依赖Quasiquotes的Scalafix规则，代码如下：

override def fix(implicit doc: SemanticDocument): Patch = {
  def matchOnTree(t: Tree): Patch = {
    t.collect {
      case Term.Apply(
          Term.Select(_, deprecated @ Term.Name(name)),
          _
          ) if config.deprecatedMethod.contains(name) =>
        Patch.replaceTree(
          deprecated,
          config.deprecatedMethod(name)
        )
      case Term.Apply(
          Term.Select(_, _ @Term.Name(name)),
          List(
            Term.AnonymousFunction(
              Term.ApplyInfix(
                _,
                deprecatedAnm @ Term.Name(nameAnm),
                _,
                _
              )
            )
          )
          ) if "reduce".contains(name) && config.deprecatedMethod.contains(nameAnm) =>
        Patch.replaceTree(
          deprecatedAnm,
          config.deprecatedMethod(nameAnm)
        )
    }.asPatch
  }

  matchOnTree(doc.tree)
}

使用Quasiquotes改写后的Scalafix规则 #

Quasiquotes能让AST模式匹配代码更简洁、可读性更强，改写后的规则代码如下：

override def fix(implicit doc: SemanticDocument): Patch = {
  def matchOnTree(t: Tree): Patch = {
    t.collect {
      // 匹配直接调用的情况：df.unionAll(df2)
      case q"$target.unionAll($arg)" =>
        Patch.replaceTree(q"unionAll", "union")
      // 匹配reduce中的匿名函数完整写法：(a, b) => a.unionAll(b)
      case q"$seq.reduce { (a, b) => a.unionAll(b) }" =>
        Patch.replaceTree(q"unionAll", "union")
      // 匹配reduce的简写形式：_.unionAll(_)
      case q"$seq.reduce(_.unionAll(_))" =>
        Patch.replaceTree(q"unionAll", "union")
    }.asPatch
  }

  matchOnTree(doc.tree)
}

代码说明： #

• 第一个case匹配直接调用unionAll的场景，直接替换方法名为union。
• 后两个case分别覆盖reduce方法中匿名函数的完整写法和简写形式，精准定位并替换目标方法名。
• 相比非Quasiquotes实现，Quasiquotes用贴近Scala语法的模式匹配，代码结构更直观，后续维护和扩展更便捷。

内容的提问来源于stack exchange，提问作者mvasyliv