Перепишите LogicalPlan в u sh вниз по udf из совокупности - PullRequest
Новая
       92

Перепишите LogicalPlan в u sh вниз по udf из совокупности

1 голос
/ 21 января 2020

Я определил UDF, который увеличивает входное значение на единицу, названный "в c", это код моего udf 

spark.udf.register("inc", (x: Long) => x + 1)

, это мой тест sql

val df = spark.sql("select sum(inc(vals)) from data")
df.explain(true)
df.show()

это оптимизированный план для этого sql

== Optimized Logical Plan ==
Aggregate [sum(inc(vals#4L)) AS sum(inc(vals))#7L]
+- LocalRelation [vals#4L]

Я хочу переписать план и извлечь "in c" из "суммы", как python udf делает. Итак, это оптимизированный план, который я хотел. 

Aggregate [sum(inc_val#6L) AS sum(inc(vals))#7L]
+- Project [inc(vals#4L) AS inc_val#6L]
   +- LocalRelation [vals#4L]

Я обнаружил, что файл исходного кода "ExtractPythonUDFs. scala" предоставляет аналогичную функцию, которая работает на PythonUDF, но он вставляет новый узел с именем " ArrowEval Python ", это логический план pythonudf. 

== Optimized Logical Plan ==
Aggregate [sum(pythonUDF0#7L) AS sum(inc(vals))#4L]
+- Project [pythonUDF0#7L]
   +- ArrowEvalPython [inc(vals#0L)], [pythonUDF0#7L], 200
      +- Repartition 10, true
         +- RelationV2[vals#0L] parquet file:/tmp/vals.parquet

То, что я хочу вставить, это просто" узел проекта ", я не хочу определять новый узел.

это тестовый код моего проекта 

import org.apache.log4j.{Level, Logger}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.{Expression, NamedExpression, ScalaUDF}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.catalyst.rules.Rule

object RewritePlanTest {

  case class UdfRule(spark: SparkSession) extends Rule[LogicalPlan] {

    def collectUDFs(e: Expression): Seq[Expression] = e match {
      case udf: ScalaUDF => Seq(udf)
      case _ => e.children.flatMap(collectUDFs)
    }

    override def apply(plan: LogicalPlan): LogicalPlan = plan match {
      case agg@Aggregate(g, a, _) if (g.isEmpty && a.length == 1) =>
        val udfs = agg.expressions.flatMap(collectUDFs)
        println("================")
        udfs.foreach(println)
        val test = udfs(0).isInstanceOf[NamedExpression]
        println(s"cast ScalaUDF to NamedExpression = ${test}")
        println("================")
        agg
      case _ => plan
    }
  }


  def main(args: Array[String]): Unit = {
    Logger.getLogger("org").setLevel(Level.WARN)

    val spark = SparkSession
      .builder()
      .master("local[*]")
      .appName("Rewrite plan test")
      .withExtensions(e => e.injectOptimizerRule(UdfRule))
      .getOrCreate()

    val input = Seq(100L, 200L, 300L)
    import spark.implicits._
    input.toDF("vals").createOrReplaceTempView("data")

    spark.udf.register("inc", (x: Long) => x + 1)

    val df = spark.sql("select sum(inc(vals)) from data")
    df.explain(true)
    df.show()
    spark.stop()
  }
}

У меня есть извлечение ScalaUDF из узла Aggregate,

, поскольку аргументы необходимы для Project Узел Seq[NamedExpression]

case class Project(projectList: Seq[NamedExpression], child: LogicalPlan)

, но ему не удалось привести ScalaUDF к NamedExpression,

, поэтому я понятия не имею, как построить узел Project.

Может кто-нибудь дать мне несколько советов?

Спасибо.

1 Ответ

0 голосов
/ 22 января 2020

ОК, наконец, я нашел способ ответить на этот вопрос.

Хотя ScalaUDF не может привести к NamedExpression, но Alias может.

Итак, я создаю Alias из ScalaUDF, затем создайте Project. 

import org.apache.log4j.{Level, Logger}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, ExpectsInputTypes, ExprId, Expression, NamedExpression, ScalaUDF}
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LocalRelation, LogicalPlan, Project, Subquery}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.types.{AbstractDataType, DataType}

import scala.collection.mutable

object RewritePlanTest {

  case class UdfRule(spark: SparkSession) extends Rule[LogicalPlan] {

    def collectUDFs(e: Expression): Seq[Expression] = e match {
      case udf: ScalaUDF => Seq(udf)
      case _ => e.children.flatMap(collectUDFs)
    }

    override def apply(plan: LogicalPlan): LogicalPlan = plan match {
      case agg@Aggregate(g, a, c) if g.isEmpty && a.length == 1 => {
        val udfs = agg.expressions.flatMap(collectUDFs)
        if (udfs.isEmpty) {
          agg
        } else {
          val alias_udf = for (i <- 0 until udfs.size) yield Alias(udfs(i), s"udf${i}")()
          val alias_set = mutable.HashMap[Expression, Attribute]()
          val proj = Project(alias_udf, c)
          alias_set ++= udfs.zip(proj.output)
          val new_agg = agg.withNewChildren(Seq(proj)).transformExpressionsUp {
            case udf: ScalaUDF if alias_set.contains(udf) => alias_set(udf)
          }
          println("====== new agg ======")
          println(new_agg)
          new_agg
        }
      }
      case _ => plan
    }
  }


  def main(args: Array[String]): Unit = {
    Logger.getLogger("org").setLevel(Level.WARN)

    val spark = SparkSession
      .builder()
      .master("local[*]")
      .appName("Rewrite plan test")
      .withExtensions(e => e.injectOptimizerRule(UdfRule))
      .getOrCreate()

    val input = Seq(100L, 200L, 300L)
    import spark.implicits._
    input.toDF("vals").createOrReplaceTempView("data")

    spark.udf.register("inc", (x: Long) => x + 1)

    val df = spark.sql("select sum(inc(vals)) from data where vals > 100")
    //    val plan = df.queryExecution.analyzed
    //    println(plan)
    df.explain(true)
    df.show()

    spark.stop()

  }
}

Этот код выводит нужный мне логический план. 

====== new agg ======
Aggregate [sum(udf0#9L) AS sum(inc(vals))#7L]
+- Project [inc(vals#4L) AS udf0#9L]
   +- LocalRelation [vals#4L]
...